diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index a549c77f0..1458ca52b 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -277,12 +277,14 @@ class BaseLayerWithLoRA(nn.Module): current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()) - # Apply LoRA with alpha scaling - lora_delta = ( - self.slice_lora_b_weights(self.lora_B.to(data)) @ self.slice_lora_a_weights(self.lora_A.to(data))) + # Apply LoRA with alpha scaling. An adapter stored wider than the base (an FP32 adapter on a BF16 DiT) + # is multiplied and added at its own precision and rounded once, as a merge-then-cast reference does. + accumulate_dtype = torch.promote_types(data.dtype, self.lora_A.dtype) + lora_delta = (self.slice_lora_b_weights(self.lora_B.to(accumulate_dtype)) @ self.slice_lora_a_weights( + self.lora_A.to(accumulate_dtype))) scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) lora_delta *= scale * self.lora_strength - data += lora_delta + data = (data.to(accumulate_dtype) + lora_delta).to(data.dtype) self.base_layer.weight.data = data.to(current_device, non_blocking=True) self.merged = True diff --git a/fastvideo/tests/layers/test_lora_merge_precision.py b/fastvideo/tests/layers/test_lora_merge_precision.py new file mode 100644 index 000000000..dea0e3442 --- /dev/null +++ b/fastvideo/tests/layers/test_lora_merge_precision.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contract: merging an adapter stored wider than the base rounds once, after a wide add.""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.lora import linear as lora_linear +from fastvideo.layers.lora.linear import get_lora_layer + +IN, OUT, RANK = 32, 24, 16 + + +@pytest.fixture(autouse=True) +def _cpu_merge(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(torch._dynamo.config, "disable", True) + monkeypatch.setattr(lora_linear, "get_local_torch_device", lambda: torch.device("cpu")) + + +def _merged(adapter_dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + torch.manual_seed(0) + layer = ReplicatedLinear(IN, OUT, bias=False, params_dtype=torch.bfloat16) + with torch.no_grad(): + layer.weight.copy_(torch.randn(OUT, IN)) + base = layer.weight.detach().clone() + wrapped = get_lora_layer(layer) + a, b = torch.randn(RANK, IN, dtype=adapter_dtype), torch.randn(OUT, RANK, dtype=adapter_dtype) * 0.1 + wrapped.set_lora_weights(a, b, lora_alpha=RANK, strength=1.0) + wrapped.merge_lora_weights() + return wrapped.base_layer.weight.detach(), base, a, b + + +def test_fp32_adapter_is_added_in_fp32_and_rounded_once() -> None: + merged, base, a, b = _merged(torch.float32) + expected = (base.float() + b @ a).to(torch.bfloat16) + assert merged.dtype == torch.bfloat16 + assert torch.equal(merged, expected) + + +def test_bf16_adapter_keeps_the_bf16_merge() -> None: + merged, base, a, b = _merged(torch.bfloat16) + assert torch.equal(merged, base + b @ a)