[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

@@ -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:

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))