From 32c7a18fb11bf425863fccc2d0b9db8e00edf496 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 30 Oct 2026 11:35:24 -0700 Subject: [PATCH] [bugfix]: merge a wider-than-base LoRA at its own precision under FSDP too The FSDP (DTensor) merge path still cast the adapter to the BF16 base before multiplying, so an FP32 adapter merged differently with FSDP on. Both paths now share one helper, which also moves the factors to the weight's device. --- fastvideo/layers/lora/linear.py | 30 +++++++++---------- .../tests/layers/test_lora_merge_precision.py | 17 +++++++++++ 2 files changed, 31 insertions(+), 16 deletions(-) diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index 1458ca52b..31b21d531 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -254,13 +254,7 @@ class BaseLayerWithLoRA(nn.Module): current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()).full_tensor() - # 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)) - 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 - unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) + unsharded_base_layer.weight = nn.Parameter(self._merged_weight(data).to(current_device)) if isinstance(getattr(self.base_layer, "bias", None), DTensor): unsharded_base_layer.bias = nn.Parameter( self.base_layer.bias.to(get_local_torch_device(), @@ -277,18 +271,22 @@ 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. 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 = (data.to(accumulate_dtype) + lora_delta).to(data.dtype) - self.base_layer.weight.data = data.to(current_device, non_blocking=True) + self.base_layer.weight.data = self._merged_weight(data).to(current_device, non_blocking=True) self.merged = True + def _merged_weight(self, weight: torch.Tensor) -> torch.Tensor: + """``weight`` plus the scaled adapter delta, in ``weight``'s dtype. + + 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; with FSDP or without. + """ + accumulate_dtype = torch.promote_types(weight.dtype, self.lora_A.dtype) + lora_B = self.slice_lora_b_weights(self.lora_B.to(device=weight.device, dtype=accumulate_dtype)) + lora_A = self.slice_lora_a_weights(self.lora_A.to(device=weight.device, dtype=accumulate_dtype)) + scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) + return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * (scale * self.lora_strength)).to(weight.dtype) + @torch.no_grad() # @torch.compile(dynamic=True) def unmerge_lora_weights(self, *, if_merged_only: bool = True) -> None: diff --git a/fastvideo/tests/layers/test_lora_merge_precision.py b/fastvideo/tests/layers/test_lora_merge_precision.py index dea0e3442..6bbf39ca0 100644 --- a/fastvideo/tests/layers/test_lora_merge_precision.py +++ b/fastvideo/tests/layers/test_lora_merge_precision.py @@ -42,3 +42,20 @@ def test_fp32_adapter_is_added_in_fp32_and_rounded_once() -> None: def test_bf16_adapter_keeps_the_bf16_merge() -> None: merged, base, a, b = _merged(torch.bfloat16) assert torch.equal(merged, base + b @ a) + + +def test_merge_helper_used_with_and_without_fsdp_rounds_once() -> None: + torch.manual_seed(1) + layer = ReplicatedLinear(IN, OUT, bias=False, params_dtype=torch.bfloat16) + with torch.no_grad(): + layer.weight.copy_(torch.randn(OUT, IN)) + wrapped = get_lora_layer(layer) + a, b = torch.randn(RANK, IN), torch.randn(OUT, RANK) * 0.1 + wrapped.lora_A, wrapped.lora_B = torch.nn.Parameter(a), torch.nn.Parameter(b) + wrapped.lora_rank, wrapped.lora_alpha, wrapped.lora_strength = RANK, RANK, 0.5 + weight = layer.weight.detach().clone() + + merged = wrapped._merged_weight(weight) + + assert merged.dtype == torch.bfloat16 + assert torch.equal(merged, (weight.float() + 0.5 * (b @ a)).to(torch.bfloat16))