diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index 44d4a7a38..a549c77f0 100644 --- a/fastvideo/layers/lora/linear.py +++ b/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 diff --git a/fastvideo/tests/layers/test_lora_side_branch.py b/fastvideo/tests/layers/test_lora_side_branch.py index ca2f093f6..04a804df4 100644 --- a/fastvideo/tests/layers/test_lora_side_branch.py +++ b/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)