[bugfix]: compile only the LoRA delta, not the wrapped base layer

BaseLayerWithLoRA.forward was decorated with torch.compile, so wrapping a linear also compiled its base layer.
Compiling a packed NVFP4/FP8 linear on its own changes its numerics, so a wrapped layer at strength 0 could differ
from the unwrapped layer. forward now adds no compiled region: the base layer runs exactly as it would unwrapped
(eagerly, or traced with a caller's compiled block), and only the adapter delta is compiled.

fastvideo/layers/lora/linear.py

@@ -155,14 +155,10 @@ class BaseLayerWithLoRA(nn.Module):
         """The wrapped layer's bias, or ``None`` when it has none."""
         return getattr(self.base_layer, "bias", None)
 
-    @torch.compile()
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        lora_A = self.lora_A
-        lora_B = self.lora_B
-        if isinstance(self.lora_B, DTensor):
-            lora_B = self.lora_B.to_local()
-            lora_A = self.lora_A.to_local()
-
+        # The wrapper adds no compiled region around the base layer, so it runs exactly as it would unwrapped: eagerly,
+        # or inside a caller's compiled block. Compiling a packed NVFP4/FP8 linear on its own changes its numerics, and
+        # a wrapped layer with no adapter or strength 0 must equal the unwrapped layer exactly.
         if self.lora_A is None or (self.side_branch and self._strength_is_zero):
             # No adapter, or strength 0 on a side branch, is exactly the unadapted layer, down to the output dtype.
             # A merge-mode wrapper reaches this branch while its first adapter is still loading or after a load
@@ -170,25 +166,32 @@ class BaseLayerWithLoRA(nn.Module):
             # not depend on an adapter that is not there.
             return self.base_layer(x)
 
-        if not self.merged and not self.disable_lora:
-            lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
-            lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
-            alpha_over_rank = 1.0
-            if self.lora_alpha != self.lora_rank:
-                alpha_over_rank = float(self.lora_alpha) / float(self.lora_rank)  # type: ignore[arg-type]
-            out, output_bias = self.base_layer(x)
-            if self.side_branch:
-                # Keep chunk temporaries opaque to Inductor so fullgraph capture cannot retain a full-size delta.
-                _add_low_rank_inplace(out, x, lora_A_sliced, lora_B_sliced, self._side_branch_strength, alpha_over_rank)
-                return out, output_bias
-            scale = float(self.lora_strength) * alpha_over_rank
-            delta = x @ lora_A_sliced.transpose(0, 1) @ lora_B_sliced.transpose(0, 1)
-            if scale != 1.0:
-                delta = delta * scale
-            return out + delta, output_bias
-        else:
-            out, output_bias = self.base_layer(x)
+        out, output_bias = self.base_layer(x)
+        if self.merged or self.disable_lora:
             return out.to(x), output_bias
+        return self._add_adapter(out, x), output_bias
+
+    @torch.compile()
+    def _add_adapter(self, out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        lora_A = self.lora_A
+        lora_B = self.lora_B
+        if isinstance(lora_B, DTensor):
+            lora_B = lora_B.to_local()
+            lora_A = lora_A.to_local()
+        lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
+        lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
+        alpha_over_rank = 1.0
+        if self.lora_alpha != self.lora_rank:
+            alpha_over_rank = float(self.lora_alpha) / float(self.lora_rank)  # type: ignore[arg-type]
+        if self.side_branch:
+            # Keep chunk temporaries opaque to Inductor so fullgraph capture cannot retain a full-size delta.
+            _add_low_rank_inplace(out, x, lora_A_sliced, lora_B_sliced, self._side_branch_strength, alpha_over_rank)
+            return out
+        scale = float(self.lora_strength) * alpha_over_rank
+        delta = x @ lora_A_sliced.transpose(0, 1) @ lora_B_sliced.transpose(0, 1)
+        if scale != 1.0:
+            delta = delta * scale
+        return out + delta
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A

fastvideo/tests/layers/test_lora_side_branch.py

@@ -68,6 +68,75 @@ def test_a_wrapped_layer_without_an_adapter_runs_the_base_layer() -> None:
         assert torch.equal(wrapped(x)[0], base(x)[0])
 
 
+@pytest.mark.parametrize("packed", [False, True])
+def test_wrapped_layer_at_strength_zero_is_bitwise_the_base_layer(packed: bool) -> None:
+    base = _packed_linear() if packed else _dense_linear()
+    x = _x()
+    with torch.no_grad():
+        # Taken before wrapping, so a merge that wrongly rewrites the shared base weight cannot also change it.
+        expected = base(x)[0].clone()
+    layer = get_lora_layer(base)
+    layer.set_lora_weights(*_factors(), lora_alpha=8, strength=0.0)
+    assert layer.lora_A is not None and layer.lora_strength == 0.0
+    with torch.no_grad():
+        assert torch.equal(layer(x)[0], expected)
+
+
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+@pytest.mark.parametrize("adapter", [False, True])
+def test_wrapping_does_not_change_a_compiled_block(monkeypatch: pytest.MonkeyPatch, backend: str,
+                                                   adapter: bool) -> None:
+    """Under a caller's compile the base layer is traced with the block, as it is unwrapped; the wrapper adds no
+    region of its own, so the compiled block is unchanged with no adapter or at strength 0."""
+    if backend == "inductor" and not torch.cuda.is_available():
+        pytest.skip("Inductor execution requires a CUDA GPU for this regression")
+    device = torch.device("cuda" if backend == "inductor" else "cpu")
+    monkeypatch.setattr(torch._dynamo.config, "disable", False)
+    torch._dynamo.reset()
+
+    class Block(torch.nn.Module):
+
+        def __init__(self, proj: torch.nn.Module) -> None:
+            super().__init__()
+            self.proj = proj
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.proj(x)[0] + x[..., :OUT]
+
+    x = _x().to(device)
+    try:
+        with torch.no_grad():
+            expected = torch.compile(Block(_packed_linear().to(device)), backend=backend, fullgraph=True)(x)
+            torch._dynamo.reset()
+            layer = get_lora_layer(_packed_linear().to(device))
+            if adapter:
+                layer.set_lora_weights(*(factor.to(device) for factor in _factors()), lora_alpha=8, strength=0.0)
+            actual = torch.compile(Block(layer), backend=backend, fullgraph=True)(x)
+        assert torch.equal(actual, expected)
+    finally:
+        torch._dynamo.reset()
+
+
+def test_wrapper_adds_no_compiled_region_around_the_base_layer(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Compiling a packed base layer on its own changes its numerics, so an uncompiled caller runs it eagerly."""
+    monkeypatch.setattr(torch._dynamo.config, "disable", False)
+    torch._dynamo.reset()
+    base = _packed_linear()
+    layer = get_lora_layer(base)
+    layer.set_lora_weights(*_factors(), lora_alpha=8, strength=0.5)
+    seen: list[bool] = []
+    original = base.forward
+
+    def spy(x: torch.Tensor):
+        seen.append(torch.compiler.is_compiling())
+        return original(x)
+
+    monkeypatch.setattr(base, "forward", spy)
+    with torch.no_grad():
+        layer(_x())
+    assert seen == [False]
+
+
 def _factors(seed: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
     generator = torch.Generator().manual_seed(seed)
     return torch.randn(RANK, IN, generator=generator), torch.randn(OUT, RANK, generator=generator)