[bugfix]: synchronize post-FSDP LoRA replicas

fastvideo/tests/train/utils/test_lora.py

@@ -0,0 +1,127 @@
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import pytest
+from torch import nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import DTensor
+
+from fastvideo.train.utils.lora import (
+    _make_replicated_lora_parameter,
+)
+
+
+def _two_rank_replicated_lora_worker(
+    rank,
+    init_method,
+    device_type,
+):
+    backend = "nccl" if device_type == "cuda" else "gloo"
+    if device_type == "cuda":
+        torch.cuda.set_device(rank)
+        device = torch.device("cuda", rank)
+    else:
+        device = torch.device("cpu")
+
+    dist.init_process_group(
+        backend,
+        init_method=init_method,
+        rank=rank,
+        world_size=2,
+    )
+    try:
+        mesh = init_device_mesh(
+            device_type,
+            (1, 2),
+            mesh_dim_names=("replicate", "shard"),
+        )
+        local_initial_value = (
+            torch.tensor([1.0, 2.0], device=device)
+            if rank == 0 else torch.tensor([10.0, 20.0], device=device)
+        )
+        parameter = _make_replicated_lora_parameter(
+            nn.Parameter(local_initial_value),
+            mesh,
+        )
+
+        # A Replicate placement requires identical values. Creation broadcasts
+        # the first mesh rank's value instead of merely assigning the label.
+        assert torch.equal(
+            parameter.to_local(),
+            torch.tensor([1.0, 2.0], device=device),
+        )
+
+        optimizer = torch.optim.SGD([parameter], lr=0.1)
+        coefficient = 1.0 if rank == 0 else 3.0
+        loss = (parameter.to_local() * coefficient).sum()
+        loss.backward()
+
+        # Rank-local gradients [1, 1] and [3, 3] must be averaged before
+        # gradient clipping and the optimizer step.
+        assert isinstance(parameter.grad, DTensor)
+        assert torch.equal(
+            parameter.grad.to_local(),
+            torch.tensor([2.0, 2.0], device=device),
+        )
+
+        # Repeating backward without zero_grad exercises the real gradient
+        # accumulation path. The previously synchronized gradient must remain
+        # intact while the new rank-local contribution is averaged.
+        second_coefficient = 2.0 if rank == 0 else 4.0
+        second_loss = (parameter.to_local() * second_coefficient).sum()
+        second_loss.backward()
+        assert torch.equal(
+            parameter.grad.to_local(),
+            torch.tensor([5.0, 5.0], device=device),
+        )
+
+        optimizer.step()
+        assert torch.equal(
+            parameter.to_local(),
+            torch.tensor([0.5, 1.5], device=device),
+        )
+
+        gathered = [
+            torch.empty_like(parameter.to_local())
+            for _ in range(2)
+        ]
+        dist.all_gather(gathered, parameter.to_local())
+        assert torch.equal(gathered[0], gathered[1])
+    finally:
+        dist.destroy_process_group()
+
+
+def test_replicated_lora_parameter_stays_consistent_across_ranks(
+    tmp_path,
+):
+    if dist.is_initialized():
+        pytest.skip("requires ownership of the default process group")
+
+    rendezvous = (tmp_path / "lora-gradient-rendezvous").resolve().as_uri()
+    mp.spawn(
+        _two_rank_replicated_lora_worker,
+        args=(rendezvous, "cpu"),
+        nprocs=2,
+        join=True,
+    )
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2,
+    reason="requires two CUDA devices",
+)
+def test_replicated_lora_parameter_stays_consistent_on_cuda(
+    tmp_path,
+):
+    if dist.is_initialized():
+        pytest.skip("requires ownership of the default process group")
+
+    rendezvous = (
+        tmp_path / "lora-gradient-cuda-rendezvous"
+    ).resolve().as_uri()
+    mp.spawn(
+        _two_rank_replicated_lora_worker,
+        args=(rendezvous, "cuda"),
+        nprocs=2,
+        join=True,
+    )

fastvideo/train/utils/lora.py

@@ -125,13 +125,71 @@ def _is_excluded_layer(
     return any(excluded in module_name for excluded in excluded_modules)
 
 
+def _register_replicated_gradient_sync(parameter: nn.Parameter, ) -> None:
+    """Average an unmanaged replicated parameter's gradient over its mesh.
+
+    LoRA parameters are attached after ``fully_shard``, so FSDP does not
+    register gradient-reduction hooks for them. ``DTensor.to_local()`` keeps
+    autograd connectivity but preserves the parameter's ``Replicate``
+    placement without inserting a collective. Average the rank-local
+    gradients over every replicated mesh dimension before gradient clipping
+    and the optimizer step.
+    """
+
+    if not isinstance(parameter, DTensor):
+        return
+
+    replicated_dims = [
+        mesh_dim for mesh_dim, placement in enumerate(parameter.placements)
+        if isinstance(placement, Replicate) and parameter.device_mesh.size(mesh_dim) > 1
+    ]
+    if not replicated_dims:
+        return
+
+    def sync_gradient(param: torch.Tensor) -> None:
+        grad = param.grad
+        if grad is None:
+            return
+        local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
+        for mesh_dim in replicated_dims:
+            dist.all_reduce(
+                local_grad,
+                group=parameter.device_mesh.get_group(mesh_dim),
+            )
+            local_grad.div_(parameter.device_mesh.size(mesh_dim))
+
+    parameter.register_post_accumulate_grad_hook(sync_gradient)
+
+
+def _make_replicated_lora_parameter(
+    parameter: nn.Parameter,
+    mesh: DeviceMesh,
+) -> nn.Parameter:
+    """Create a synchronized replicated DTensor for a late-added LoRA weight."""
+
+    placements = [Replicate()] * mesh.ndim
+    replicated = DTensor.from_local(
+        parameter.detach(),
+        device_mesh=mesh,
+        placements=placements,
+        run_check=True,
+    )
+    replicated_parameter = nn.Parameter(
+        replicated,
+        requires_grad=parameter.requires_grad,
+    )
+    _register_replicated_gradient_sync(replicated_parameter)
+    return replicated_parameter
+
+
 def _replicate_lora_parameters(transformer: torch.nn.Module, ) -> None:
     """Wrap LoRA params in replicated DTensors when distributed is active.
 
     The training loaders shard the base transformer with FSDP/HSDP before the
     model plugin sees it. Newly-added LoRA parameters therefore need to be
     explicit replicated DTensors so optimizers/checkpointing can treat them the
-    same way across ranks.
+    same way across ranks. Replicated values are broadcast during creation,
+    and their rank-local gradients are averaged before the optimizer step.
 
     The mesh is reused from the FSDP-wrapped base_layer parameters rather than
     rebuilt via ``init_device_mesh`` — building a parallel mesh with a different
@@ -166,8 +224,6 @@ def _replicate_lora_parameters(transformer: torch.nn.Module, ) -> None:
     if mesh is None:
         return
 
-    placements = [Replicate()] * mesh.ndim
-
     for module in transformer.modules():
         if not isinstance(module, BaseLayerWithLoRA):
             continue
@@ -181,12 +237,11 @@ def _replicate_lora_parameters(transformer: torch.nn.Module, ) -> None:
             param.requires_grad_(True)
             if isinstance(param, DTensor):
                 continue
-            replicated = DTensor.from_local(
-                param.detach(),
-                device_mesh=mesh,
-                placements=placements,
+            setattr(
+                module,
+                attr_name,
+                _make_replicated_lora_parameter(param, mesh),
             )
-            setattr(module, attr_name, nn.Parameter(replicated))
 
 
 def enable_lora_training(