[bugfix]: synchronize post-FSDP LoRA replicas (#1666)
Co-authored-by: Suckl <Suckl@users.noreply.github.com> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Suckl
co-authored by
Suckl
SolitaryThinker
Claude Opus 5.5
Suckl
SolitaryThinker
Claude Opus 5.5
parent
d6e13f9cc8
commit
a4e7322ba6
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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():
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 "
|
||||
|
||||
Reference in New Issue
Block a user