From 899383b9be3b2c96e5773ce87f761f0a5a0f64c9 Mon Sep 17 00:00:00 2001 From: Suckl Date: Thu, 30 Jul 2026 16:26:15 +0800 Subject: [PATCH] [bugfix]: synchronize post-FSDP LoRA replicas --- fastvideo/tests/train/utils/test_lora.py | 127 +++++++++++++++++++++++ fastvideo/train/utils/lora.py | 71 +++++++++++-- 2 files changed, 190 insertions(+), 8 deletions(-) create mode 100644 fastvideo/tests/train/utils/test_lora.py diff --git a/fastvideo/tests/train/utils/test_lora.py b/fastvideo/tests/train/utils/test_lora.py new file mode 100644 index 0000000000..ec719bdde1 --- /dev/null +++ b/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, + ) diff --git a/fastvideo/train/utils/lora.py b/fastvideo/train/utils/lora.py index d1dcf36387..7cd59d76b8 100644 --- a/fastvideo/train/utils/lora.py +++ b/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(