diff --git a/fastvideo/tests/train/callbacks/test_grad_clip.py b/fastvideo/tests/train/callbacks/test_grad_clip.py index 3fee9f63e..5f5950a50 100644 --- a/fastvideo/tests/train/callbacks/test_grad_clip.py +++ b/fastvideo/tests/train/callbacks/test_grad_clip.py @@ -71,9 +71,10 @@ class TestGradNormClipCallback: method = _Method(targets={"m": m}) cb.on_before_optimizer_step(method=method, iteration=0) - # No clipping applied; ``get_grad_clip_targets`` not consulted. + # No clipping applied, but ``get_grad_clip_targets`` is consulted so + # replicated LoRA gradients are synchronized even when clipping is off. assert _grad_norm(m) == before - assert method.iter_seen is None + assert method.iter_seen == 0 def test_large_grads_get_clipped(self) -> None: m = _make_module(grad_value=10.0) diff --git a/fastvideo/tests/train/utils/test_lora.py b/fastvideo/tests/train/utils/test_lora.py new file mode 100644 index 000000000..7acb1efa8 --- /dev/null +++ b/fastvideo/tests/train/utils/test_lora.py @@ -0,0 +1,167 @@ +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, + sync_replicated_lora_gradients, +) + + +class _LoraHolder(nn.Module): + """Minimal stand-in for a LoRA-wrapped layer (only lora_A is exercised).""" + + def __init__(self, parameter): + super().__init__() + self.lora_A = parameter + + +def _sync(parameter): + """Run the production post-backward pass over a bare parameter.""" + + return sync_replicated_lora_gradients(_LoraHolder(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. + _sync(parameter) + 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() + _sync(parameter) + 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]) + + # Rank 1's loss skips the parameter entirely, so its autograd hook + # never fires. The pass must still join the collective on both ranks + # (otherwise rank 0 blocks forever) and zero-fill the absent gradient, + # so the replicas keep stepping together. + parameter.grad = None + if rank == 0: + (parameter.to_local() * 4.0).sum().backward() + else: + (torch.zeros((), device=device, requires_grad=True) * 1.0).sum().backward() + if rank == 0: + assert parameter.grad is not None + else: + assert parameter.grad is None + _sync(parameter) + assert torch.equal( + parameter.grad.to_local(), + torch.tensor([2.0, 2.0], device=device), + ) + + optimizer.step() + 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/callbacks/grad_clip.py b/fastvideo/train/callbacks/grad_clip.py index e15761062..ec883c7a5 100644 --- a/fastvideo/train/callbacks/grad_clip.py +++ b/fastvideo/train/callbacks/grad_clip.py @@ -12,6 +12,8 @@ from typing import TYPE_CHECKING from fastvideo.logger import init_logger from fastvideo.train.callbacks.callback import Callback +from fastvideo.train.utils.lora import ( + sync_replicated_lora_gradients, ) from fastvideo.train.utils.optimizer import ( clip_grad_norm_if_needed, ) @@ -42,11 +44,19 @@ class GradNormClipCallback(Callback): method: TrainingMethod, iteration: int = 0, ) -> None: + # LoRA parameters are attached after fully_shard, so FSDP never reduces + # them; average their replicated gradients on every rank before the + # norm is measured (a rank that skipped a layer still joins the + # collective, otherwise the norm -- and the clip coefficient -- would + # differ per rank). + targets = method.get_grad_clip_targets(iteration) + for module in targets.values(): + sync_replicated_lora_gradients(module) + max_norm = self._max_grad_norm if max_norm <= 0.0: return - targets = method.get_grad_clip_targets(iteration) tracker = getattr(method, "tracker", None) for name, module in targets.items(): diff --git a/fastvideo/train/utils/lora.py b/fastvideo/train/utils/lora.py index d1dcf3638..f2aadce93 100644 --- a/fastvideo/train/utils/lora.py +++ b/fastvideo/train/utils/lora.py @@ -125,13 +125,90 @@ def _is_excluded_layer( return any(excluded in module_name for excluded in excluded_modules) +def _replicated_dims(parameter: DTensor) -> list[int]: + return [ + mesh_dim for mesh_dim, placement in enumerate(parameter.placements) + if isinstance(placement, Replicate) and parameter.device_mesh.size(mesh_dim) > 1 + ] + + +def sync_replicated_lora_gradients(transformer: torch.nn.Module) -> int: + """Average every replicated LoRA gradient over its mesh, on every rank. + + A per-parameter ``register_post_accumulate_grad_hook`` is not enough: the + hook only fires on ranks where that parameter actually received a + gradient, so a rank whose loss skipped a layer issues fewer collectives + than its peers -- the ranks then pair the wrong tensors together, or hang + outright when the totals differ. This pass runs after backward and before + clipping/the optimizer step, walks the LoRA parameters in a stable order, + and joins the collective on every rank, zero-filling the gradient where the + rank contributed none (a rank that skipped the layer contributes exactly + zero to the average). + + Returns the number of parameters synchronized. + """ + + if not dist.is_available() or not dist.is_initialized(): + return 0 + + synced = 0 + for module in transformer.modules(): + for attr_name in ("lora_A", "lora_B"): + parameter = getattr(module, attr_name, None) + if not isinstance(parameter, DTensor) or not parameter.requires_grad: + continue + replicated_dims = _replicated_dims(parameter) + if not replicated_dims: + continue + grad = parameter.grad + if grad is None: + local_grad = torch.zeros_like(parameter.to_local()) + parameter.grad = DTensor.from_local( + local_grad, + device_mesh=parameter.device_mesh, + placements=parameter.placements, + run_check=False, + ) + else: + 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)) + synced += 1 + return synced + + +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, + ) + 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 +243,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 +256,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( diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 7eb3847fa..a026435de 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -986,6 +986,42 @@ def _get_total_norm( total_norm = torch.linalg.vector_norm(torch.stack([norm.to(first_device) for norm in norms]), norm_type) + # The per-tensor norms above come from ``to_local()``, so for DTensor grads + # they describe only this rank's shard. Reduce across the mesh so every + # rank derives the same clip coefficient: sum the squared norms over sharded + # dims (each rank holds a disjoint slice) and average them over replicated + # dims (every rank holds the same values). + mesh_dims: dict[tuple[int, int], tuple[torch.distributed.device_mesh.DeviceMesh, int, bool]] = {} + for tensor in tensors: + if not isinstance(tensor, torch.distributed.tensor.DTensor): + continue + for mesh_dim, placement in enumerate(tensor.placements): + if tensor.device_mesh.size(mesh_dim) <= 1: + continue + key = (id(tensor.device_mesh), mesh_dim) + replicate = isinstance(placement, torch.distributed.tensor.Replicate) + previous = mesh_dims.get(key) + mesh_dims[key] = (tensor.device_mesh, mesh_dim, replicate if previous is None else previous[2] and replicate) + + if mesh_dims and dist.is_available() and dist.is_initialized(): + if norm_type == math.inf: + # Max commutes with both reductions: exact across sharded dims + # (disjoint slices) and idempotent across replicated dims. The + # sum-of-powers reduction below is only valid for finite norms. + for mesh, mesh_dim, _replicate in mesh_dims.values(): + dist.all_reduce( + total_norm, + op=dist.ReduceOp.MAX, + group=mesh.get_group(mesh_dim), + ) + else: + accumulated = total_norm.pow(norm_type) + for mesh, mesh_dim, replicate in mesh_dims.values(): + dist.all_reduce(accumulated, group=mesh.get_group(mesh_dim)) + if replicate: + accumulated.div_(mesh.size(mesh_dim)) + total_norm = accumulated.pow(1.0 / norm_type) + if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()): raise RuntimeError(f"The total norm of order {norm_type} for gradients from " "`parameters` is non-finite, so it cannot be clipped. To disable "