[https://nvbugs/6179761][fix] Save LTX-2 BF16 weights to speed up perf (#14639)

Signed-off-by: yibinl-nvidia <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py

@@ -39,13 +39,46 @@
 
 STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
 _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
+# Baseline BF16 peak memory ~75 GiB, saving BF16 weights snopshot total ~108 GiB.
+_BF16_WEIGHTS_SNAPSHOT_FREE_MEMORY_THRESHOLD_GIB = 115.0
 
 
 # ---------------------------------------------------------------------------
 # LoRA helpers
 # ---------------------------------------------------------------------------
 
 
+def _get_free_gpu_memory_gib() -> Optional[float]:
+    """Return free memory on the current CUDA device, or ``None`` if unavailable."""
+    if not torch.cuda.is_available():
+        return None
+
+    try:
+        free_bytes, _ = torch.cuda.mem_get_info()
+    except (RuntimeError, OSError) as exc:
+        logger.warning(f"Unable to query CUDA free memory for BF16 weight snapshots: {exc}")
+        return None
+
+    return free_bytes / (1024**3)
+
+
+def _should_save_bf16_weights(
+    threshold_gib: float = _BF16_WEIGHTS_SNAPSHOT_FREE_MEMORY_THRESHOLD_GIB,
+) -> bool:
+    free_gib = _get_free_gpu_memory_gib()
+    if free_gib is None:
+        logger.debug("BF16 weight snapshots disabled: CUDA free memory is unavailable")
+        return False
+
+    save_state = free_gib > threshold_gib
+    relation = ">" if save_state else "<="
+    logger.debug(
+        f"BF16 weight snapshots {'enabled' if save_state else 'disabled'}: "
+        f"free GPU memory {free_gib:.2f} GiB {relation} {threshold_gib:.2f} GiB threshold"
+    )
+    return save_state
+
+
 def _load_lora_deltas(
     lora_path: str,
     transformer: torch.nn.Module,
@@ -338,11 +371,12 @@ def _apply_lora_deltas(
     module: torch.nn.Module,
     deltas: Dict[str, torch.Tensor],
     sign: float = 1.0,
+    save_bf16_weights: bool = False,
 ) -> tuple:
     """Add (sign=+1) or remove (sign=-1) pre-computed LoRA deltas.
 
-    For standard (BF16/FP16) weights the delta is added directly and
-    later removed by subtracting the same delta.
+    For BF16 weights the delta is added directly and later removed either
+    by restoring an optional saved snapshot or by subtracting the same delta.
     For FP8-quantized weights (same shape, float8 dtype), we
     dequantize → apply → requantize.  FP4 weights are handled through
     the packed-FP4 branch because the current static and dynamic NVFP4
@@ -355,13 +389,15 @@ def _apply_lora_deltas(
     state afterwards.
 
     Returns ``(applied_count, saved_lora_state, snapshot_required_count)`` where
-    *saved_lora_state* maps each touched quantized parameter name to its
-    original tensor.  For packed FP4 it also stores the parent
+    *saved_lora_state* maps each touched snapshotted parameter name to its
+    original tensor.  BF16 weights are snapshotted only when
+    *save_bf16_weights* is true.  This does not change FP8 or FP4 LoRA
+    handling: those paths always snapshot quantized state.  For packed FP4 it
+    also stores the parent
     ``quant_method`` so that stage 2 can run with BF16 weights and then
     restore the exact original FP4 state without another quantization round
-    trip.  Dense BF16/FP16/FP32 weights are not stored in
-    *saved_lora_state*.  *snapshot_required_count* is the number of
-    quantized weights that must be restored from saved snapshots.
+    trip.  *snapshot_required_count* is the number of weights that must be
+    restored from saved snapshots.
     """
     applied = 0
     snapshot_required = 0
@@ -424,8 +460,11 @@ def _apply_lora_deltas(
                 ws_param.data.copy_(new_scale)
                 snapshot_required += 1
             else:
-                # BF16/FP16/FP32: direct in-place addition. These are
-                # restored by subtracting the same delta after stage 2.
+                if save_bf16_weights and param.dtype == torch.bfloat16:
+                    saved_state[param_name] = param.data.clone()
+                    snapshot_required += 1
+                # BF16: direct in-place addition, then restore by snapshot copy
+                # when memory allows. Other dense weights restore by subtraction.
                 param.data.add_(
                     delta.to(param.device, param.dtype),
                     alpha=sign,
@@ -808,11 +847,16 @@ def forward(
         # ================================================================
         # For FP4 models (static-packed or dynamic), stage 2 always runs in
         # BF16: the quant_method is swapped to UnquantizedLinearMethod inside
-        # _apply_lora_deltas and restored afterwards.  BF16 models are unaffected.
+        # _apply_lora_deltas and restored afterwards. FP8 handling is also
+        # unchanged and always restores from saved quantized state. BF16 weights
+        # save snapshots only when enough free GPU memory is available;
+        # otherwise they fall back to on-the-fly LoRA subtraction.
+        save_bf16_weights = _should_save_bf16_weights()
         n, saved_lora_state, snapshot_required = _apply_lora_deltas(
             self.transformer,
             self._distilled_lora_deltas,
             sign=1.0,
+            save_bf16_weights=save_bf16_weights,
         )
         logger.info(f"Merged distilled LoRA ({n} params) for stage 2 (BF16 weights)")
 
@@ -840,8 +884,7 @@ def forward(
             self.transformer.set_ulysses_enabled(True)
             if snapshot_required and not saved_lora_state:
                 raise RuntimeError(
-                    "Quantized LoRA state was not saved; cannot safely restore "
-                    "FP8/FP4 stage 2 weights."
+                    "LoRA state was not saved; cannot safely restore stage 2 weights."
                 )
 
             snapshot_restored = 0
@@ -851,8 +894,8 @@ def forward(
                 _restore_lora_state(self.transformer, saved_lora_state)
                 snapshot_restored = _count_saved_lora_weight_tensors(saved_lora_state)
 
-            # Dense BF16/FP16/FP32 weights are not snapshotted; restore them by
-            # subtracting the LoRA deltas directly.
+            # BF16 weights that were not snapshotted, plus any other dense
+            # floating-point weights, are restored by subtracting LoRA deltas.
             dense_restored = _subtract_dense_lora_deltas(
                 self.transformer,
                 self._distilled_lora_deltas,

tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py

@@ -537,6 +537,58 @@ def test_latent_shape_matches_batch(self):
 class TestTwoStageLoRAHelpers:
     """Test LoRA delta loading and application without checkpoints."""
 
+    def test_bf16_weight_snapshot_gate_uses_cuda_free_memory(self, monkeypatch):
+        from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import (
+            _should_save_bf16_weights,
+        )
+
+        gib = 1024**3
+        monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+        monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (116 * gib, 180 * gib))
+        assert _should_save_bf16_weights()
+
+        monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (115 * gib, 180 * gib))
+        assert not _should_save_bf16_weights()
+
+        def _raise_mem_query_error():
+            raise RuntimeError("mem_get_info failed")
+
+        monkeypatch.setattr(torch.cuda, "mem_get_info", _raise_mem_query_error)
+        assert not _should_save_bf16_weights()
+
+        monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+        assert not _should_save_bf16_weights()
+
+    def test_bf16_weight_snapshot_saved_when_requested(self):
+        from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import (
+            _apply_lora_deltas,
+            _restore_lora_state,
+            _subtract_dense_lora_deltas,
+        )
+
+        linear = torch.nn.Linear(32, 32, bias=False).bfloat16()
+        original_weight = linear.weight.data.clone()
+        deltas = {"weight": torch.randn(32, 32) * 0.1}
+
+        applied, saved_state, snapshot_required = _apply_lora_deltas(
+            linear,
+            deltas,
+            sign=1.0,
+            save_bf16_weights=True,
+        )
+
+        assert applied == 1
+        assert snapshot_required == 1
+        assert "weight" in saved_state
+        assert torch.allclose(saved_state["weight"], original_weight)
+        assert not torch.allclose(linear.weight.data, original_weight)
+
+        removed = _subtract_dense_lora_deltas(linear, deltas, saved_state)
+        assert removed == 0
+
+        _restore_lora_state(linear, saved_state)
+        assert torch.allclose(linear.weight.data, original_weight)
+
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_apply_and_remove_deltas_bf16(self):
         """Merge then unmerge in BF16 should leave weights approximately unchanged."""
@@ -554,7 +606,7 @@ def test_apply_and_remove_deltas_bf16(self):
 
         applied, saved_state, snapshot_required = _apply_lora_deltas(linear, deltas, sign=1.0)
         assert applied == 1, "Expected one parameter to be modified"
-        assert saved_state == {}, "Dense BF16 weights should not be snapshotted"
+        assert saved_state == {}, "BF16 weights should not be snapshotted by default"
         assert snapshot_required == 0
         assert not torch.allclose(linear.weight.data, original_weight), (
             "Weights should have changed after applying delta"
@@ -597,16 +649,16 @@ def test_multi_round_drift_bounded(self):
         rounds = 10
         for _ in range(rounds):
             _, saved_state, snapshot_required = _apply_lora_deltas(model, deltas, sign=1.0)
-            assert saved_state == {}, "Dense BF16 weights should not be snapshotted"
+            assert saved_state == {}, "BF16 weights should not be snapshotted by default"
             assert snapshot_required == 0
             _subtract_dense_lora_deltas(model, deltas, saved_state)
 
         drift = (model.weight.data.float() - original.float()).abs().max().item()
         assert drift < 0.1, f"bf16 drift after {rounds} rounds too large: {drift:.2e}"
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    def test_dense_lora_state_not_saved_and_subtract_restores(self):
-        """Dense weights are restored by subtraction without snapshot storage."""
+    def test_fp32_state_not_saved_and_subtract_restores(self):
+        """FP32 weights restore by subtraction, even when BF16 snapshots are enabled."""
         from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import (
             _apply_lora_deltas,
             _subtract_dense_lora_deltas,
@@ -617,16 +669,21 @@ def test_dense_lora_state_not_saved_and_subtract_restores(self):
         original_weight = linear.weight.data.clone()
 
         deltas = {"weight": torch.randn(32, 32, device=device) * 0.1}
-        _, saved_state, snapshot_required = _apply_lora_deltas(linear, deltas, sign=1.0)
+        _, saved_state, snapshot_required = _apply_lora_deltas(
+            linear,
+            deltas,
+            sign=1.0,
+            save_bf16_weights=True,
+        )
 
-        assert saved_state == {}, "Dense FP32 weights should not be snapshotted"
+        assert saved_state == {}, "FP32 weights should not use the BF16 snapshot path"
         assert snapshot_required == 0
         assert not torch.allclose(linear.weight.data, original_weight)
 
         removed = _subtract_dense_lora_deltas(linear, deltas, saved_state)
         assert removed == 1
         assert torch.allclose(linear.weight.data, original_weight), (
-            "Dense LoRA subtraction should restore weights"
+            "FP32 LoRA subtraction should restore weights"
         )
 
 
@@ -1111,7 +1168,7 @@ def test_two_stage_lora_deltas_match_transformer(self, ltx2_two_stage_assets_exi
             print(f"\n[Two-Stage] LoRA apply rate: {match_rate:.1f}% ({applied}/{total})")
             assert match_rate > 99.0, f"Expected >99% LoRA match rate, got {match_rate:.1f}%"
 
-            assert saved_state == {}, "BF16 checkpoint should not snapshot dense weights"
+            assert saved_state == {}, "BF16 checkpoint should not snapshot weights by default"
             assert snapshot_required == 0
 
             # Verify dense unmerge by subtraction