401 lines
14 KiB
Python
401 lines
14 KiB
Python
"""Distribution strategies for model wrapping and gradient coordination."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from contextlib import contextmanager
|
|
from importlib import import_module
|
|
from typing import NamedTuple, Optional
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
|
|
from torch.distributed.fsdp import (
|
|
MixedPrecisionPolicy,
|
|
OffloadPolicy,
|
|
)
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FSDP2Layout(NamedTuple):
|
|
"""Per-call FSDP2 wrap configuration produced by ``FSDPStrategy``."""
|
|
|
|
mesh: DeviceMesh
|
|
dp_process_group: Optional[dist.ProcessGroup]
|
|
shard_process_group: Optional[dist.ProcessGroup]
|
|
replicate_process_group: Optional[dist.ProcessGroup]
|
|
trainer_process_group: Optional[dist.ProcessGroup]
|
|
mp_policy: Optional[MixedPrecisionPolicy]
|
|
offload_policy: Optional[OffloadPolicy]
|
|
reshard_after_forward: bool
|
|
|
|
|
|
class DistributionStrategy(ABC):
|
|
"""Common interface for FSDP2 / DDP wrapping + gradient operations."""
|
|
|
|
_cfg: dict
|
|
last_grad_clip_stats: Optional[dict] = None
|
|
|
|
def _maybe_enable_grad_ckpt(self, model: torch.nn.Module) -> None:
|
|
"""Call ``gradient_checkpointing_enable`` if the model exposes it."""
|
|
if not self._cfg.get("use_gradient_checkpointing", False):
|
|
return
|
|
method = getattr(model, "gradient_checkpointing_enable", None)
|
|
if method is None:
|
|
_logger.warning(
|
|
"use_gradient_checkpointing=true but %s has no "
|
|
"gradient_checkpointing_enable method. If this model "
|
|
"routes recomputation through its config the request may "
|
|
"still be honored; otherwise "
|
|
"gradient checkpointing is effectively disabled for "
|
|
"this run.",
|
|
type(model).__name__,
|
|
)
|
|
return
|
|
method()
|
|
|
|
@abstractmethod
|
|
def wrap(self, model: torch.nn.Module) -> torch.nn.Module:
|
|
"""Wrap *model* with the active distribution strategy."""
|
|
|
|
@abstractmethod
|
|
def clip_grad_norm(
|
|
self,
|
|
model: torch.nn.Module,
|
|
max_norm: float,
|
|
*,
|
|
optimizer=None,
|
|
) -> torch.Tensor:
|
|
"""Clip gradient L2 norm; return total norm before clipping."""
|
|
|
|
@abstractmethod
|
|
def no_sync(self, model: torch.nn.Module):
|
|
"""Context manager that disables cross-rank gradient sync.
|
|
|
|
Used by the trainer for gradient accumulation micro-batches.
|
|
FSDP2 toggles ``set_requires_gradient_sync``; DDP yields
|
|
``model.no_sync()``.
|
|
"""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def params_sharded(self) -> bool:
|
|
"""Whether parameters are actually sharded across ranks.
|
|
|
|
FSDP2 always shards (every ``fully_shard`` unit is sharded across
|
|
the mesh). DDP never shards. Used by adapters to decide whether
|
|
per-rank norm contributions need an all-reduce.
|
|
"""
|
|
|
|
|
|
class FSDPStrategy(DistributionStrategy):
|
|
def __init__(self, config: dict):
|
|
self._cfg = config
|
|
self._layout: Optional[FSDP2Layout] = None
|
|
self.last_grad_clip_stats = None
|
|
|
|
@property
|
|
def trainer_process_group(self) -> Optional[dist.ProcessGroup]:
|
|
"""Process group for trainer-level collectives.
|
|
|
|
Returns `dist.group.WORLD` (the implicit all-rank group created
|
|
by `init_process_group`) so that trainer barriers / metric
|
|
reductions behave consistently regardless of FSDP mesh shape
|
|
(1D for full_shard / shard_grad_op, 2D for hybrid_shard /
|
|
_hybrid_shard_zero2). No extra NCCL communicator is allocated.
|
|
"""
|
|
return self._layout.trainer_process_group if self._layout else None
|
|
|
|
@staticmethod
|
|
def _build_trainer_process_group(
|
|
world_size: int,
|
|
) -> Optional[dist.ProcessGroup]:
|
|
if not dist.is_initialized() or world_size <= 1:
|
|
return None
|
|
return dist.group.WORLD
|
|
|
|
def _build_fsdp2_layout(self) -> FSDP2Layout:
|
|
"""Construct the FSDP2 wrap layout from yaml config.
|
|
|
|
Builds the device mesh (1D for full_shard / shard_grad_op, 2D
|
|
for hybrid_shard / _hybrid_shard_zero2), explicit FSDP mesh process
|
|
groups, a trainer-level global process group, the matching
|
|
mixed-precision policy, offload policy, and reshard_after_forward flag.
|
|
"""
|
|
name = self._cfg.get("fsdp_sharding_strategy", "full_shard")
|
|
world_size = dist.get_world_size() if dist.is_initialized() else 1
|
|
dp_pg: Optional[dist.ProcessGroup] = None
|
|
shard_pg: Optional[dist.ProcessGroup] = None
|
|
replicate_pg: Optional[dist.ProcessGroup] = None
|
|
|
|
# 2D mesh for HSDP variants; 1D for full_shard / shard_grad_op.
|
|
if name in ("hybrid_shard", "_hybrid_shard_zero2"):
|
|
replicate_size = self._cfg.get("fsdp_hsdp_replicate_size") or int(
|
|
os.environ.get("LOCAL_WORLD_SIZE", "0")
|
|
)
|
|
if replicate_size <= 0:
|
|
# Fallback: assume one replica per node, equal-sized shard groups.
|
|
# Caller should set fsdp_hsdp_replicate_size explicitly when this
|
|
# default is wrong (e.g. uneven node sizes).
|
|
replicate_size = max(1, world_size // 8) if world_size >= 8 else 1
|
|
if world_size % replicate_size != 0:
|
|
raise ValueError(
|
|
f"world_size ({world_size}) not divisible by HSDP replicate "
|
|
f"size ({replicate_size}); set fsdp_hsdp_replicate_size."
|
|
)
|
|
shard_size = world_size // replicate_size
|
|
mesh = init_device_mesh(
|
|
"cuda",
|
|
(replicate_size, shard_size),
|
|
mesh_dim_names=("replicate", "shard"),
|
|
)
|
|
replicate_pg = mesh.get_group("replicate")
|
|
shard_pg = mesh.get_group("shard")
|
|
else:
|
|
mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",))
|
|
dp_pg = mesh.get_group("dp")
|
|
|
|
trainer_pg = self._build_trainer_process_group(world_size)
|
|
|
|
reshard_after_forward = name not in (
|
|
"shard_grad_op",
|
|
"_hybrid_shard_zero2",
|
|
)
|
|
|
|
offload_policy: Optional[OffloadPolicy] = (
|
|
OffloadPolicy(pin_memory=True)
|
|
if self._cfg.get("fsdp_cpu_offload", False)
|
|
else None
|
|
)
|
|
|
|
mp_policy: Optional[MixedPrecisionPolicy]
|
|
if self._cfg.get("use_mixed_precision", True):
|
|
dtype = torch.bfloat16 if self._cfg.get("bf16", True) else torch.float16
|
|
reduce_dtype = (
|
|
torch.float32
|
|
if self._cfg.get("fsdp_reduce_dtype", "bf16") == "fp32"
|
|
else dtype
|
|
)
|
|
mp_policy = MixedPrecisionPolicy(
|
|
param_dtype=dtype,
|
|
reduce_dtype=reduce_dtype,
|
|
cast_forward_inputs=False,
|
|
)
|
|
else:
|
|
mp_policy = None
|
|
|
|
assert mp_policy is not None, (
|
|
"FSDP2 requires a MixedPrecisionPolicy in this torch version. "
|
|
"Set distributed.use_mixed_precision=true."
|
|
)
|
|
|
|
return FSDP2Layout(
|
|
mesh=mesh,
|
|
dp_process_group=dp_pg,
|
|
shard_process_group=shard_pg,
|
|
replicate_process_group=replicate_pg,
|
|
trainer_process_group=trainer_pg,
|
|
mp_policy=mp_policy,
|
|
offload_policy=offload_policy,
|
|
reshard_after_forward=reshard_after_forward,
|
|
)
|
|
|
|
def _wrap_fsdp(
|
|
self,
|
|
model: torch.nn.Module,
|
|
*,
|
|
use_dmuon: bool = False,
|
|
) -> torch.nn.Module:
|
|
if not hasattr(model, "convert_to_fsdp"):
|
|
raise NotImplementedError(
|
|
f"Model {model.__class__.__name__} has no convert_to_fsdp method"
|
|
)
|
|
self._maybe_enable_grad_ckpt(model)
|
|
layout = self._build_fsdp2_layout()
|
|
self._layout = layout
|
|
wrapped = model.convert_to_fsdp(
|
|
mesh=layout.mesh,
|
|
mp_policy=layout.mp_policy,
|
|
offload_policy=layout.offload_policy,
|
|
reshard_after_forward=layout.reshard_after_forward,
|
|
use_dmuon=use_dmuon,
|
|
)
|
|
torch.cuda.empty_cache()
|
|
return wrapped
|
|
|
|
def wrap(self, model: torch.nn.Module) -> torch.nn.Module:
|
|
return self._wrap_fsdp(model, use_dmuon=False)
|
|
|
|
def clip_grad_norm(
|
|
self,
|
|
model: torch.nn.Module,
|
|
max_norm: float,
|
|
*,
|
|
optimizer=None,
|
|
) -> torch.Tensor:
|
|
del optimizer
|
|
self.last_grad_clip_stats = None
|
|
# Separate DTensor (FSDP2) from regular tensor params - PyTorch's
|
|
# foreach_mul_ in clip_grad_norm_ can't mix them.
|
|
try:
|
|
from torch.distributed.tensor import DTensor
|
|
except ImportError:
|
|
DTensor = None
|
|
dtensor_params = []
|
|
regular_params = []
|
|
for p in model.parameters():
|
|
if p.grad is None:
|
|
continue
|
|
if DTensor is not None and isinstance(p, DTensor):
|
|
dtensor_params.append(p)
|
|
else:
|
|
regular_params.append(p)
|
|
|
|
device = next(iter(p for p in model.parameters())).device
|
|
total_sq = torch.tensor(0.0, device=device)
|
|
if dtensor_params:
|
|
n = torch.nn.utils.clip_grad_norm_(dtensor_params, max_norm)
|
|
total_sq = total_sq + n.to(device).pow(2)
|
|
if regular_params:
|
|
n = torch.nn.utils.clip_grad_norm_(regular_params, max_norm)
|
|
total_sq = total_sq + n.to(device).pow(2)
|
|
return total_sq.sqrt()
|
|
|
|
@contextmanager
|
|
def no_sync(self, model: torch.nn.Module):
|
|
toggled = hasattr(model, "set_requires_gradient_sync")
|
|
if toggled:
|
|
model.set_requires_gradient_sync(False)
|
|
try:
|
|
yield
|
|
finally:
|
|
if toggled:
|
|
model.set_requires_gradient_sync(True)
|
|
|
|
@property
|
|
def params_sharded(self) -> bool:
|
|
return True
|
|
|
|
|
|
class DMuonFSDPStrategy(FSDPStrategy):
|
|
"""FSDP2 plus the DMuon dedicated-parameter runtime."""
|
|
|
|
@staticmethod
|
|
def _runtime():
|
|
return import_module("dmuon")
|
|
|
|
def wrap(self, model: torch.nn.Module) -> torch.nn.Module:
|
|
return self._wrap_fsdp(model, use_dmuon=True)
|
|
|
|
def clip_grad_norm(
|
|
self,
|
|
model: torch.nn.Module,
|
|
max_norm: float,
|
|
*,
|
|
optimizer=None,
|
|
) -> torch.Tensor:
|
|
total_norm = super().clip_grad_norm(model, max_norm, optimizer=optimizer)
|
|
self.last_grad_clip_stats = None
|
|
if optimizer is None:
|
|
return total_norm
|
|
|
|
clip_stats = self._runtime().clip_grad_norm_(optimizer, max_norm)
|
|
self.last_grad_clip_stats = clip_stats.as_dict()
|
|
device = next(iter(p for p in model.parameters())).device
|
|
total_norm_t = (
|
|
total_norm.to(device)
|
|
if torch.is_tensor(total_norm)
|
|
else torch.tensor(float(total_norm), device=device)
|
|
)
|
|
dedicated_norm_t = torch.tensor(float(clip_stats.total_norm), device=device)
|
|
return (total_norm_t.pow(2) + dedicated_norm_t.pow(2)).sqrt()
|
|
|
|
@contextmanager
|
|
def no_sync(self, model: torch.nn.Module):
|
|
with self._runtime().no_sync(model):
|
|
yield
|
|
|
|
|
|
class DDPStrategy(DistributionStrategy):
|
|
def __init__(
|
|
self,
|
|
config: dict,
|
|
*,
|
|
device: torch.device,
|
|
local_rank: int,
|
|
):
|
|
self._cfg = config
|
|
self._device = device
|
|
self._local_rank = local_rank
|
|
self.find_unused_parameters = config.get("find_unused_parameters", False)
|
|
self.broadcast_buffers = config.get("broadcast_buffers", True)
|
|
self.bucket_cap_mb = config.get("bucket_cap_mb", 25)
|
|
self.last_grad_clip_stats = None
|
|
|
|
def wrap(self, model: torch.nn.Module) -> torch.nn.Module:
|
|
# Adapter.convert_to_mix_precision_hint already placed the model on
|
|
# self._device; this .to() is a defensive no-op (same as the original
|
|
# wrap_ddp_model).
|
|
model = model.to(self._device)
|
|
self._maybe_enable_grad_ckpt(model)
|
|
return DDP(
|
|
model,
|
|
device_ids=[self._local_rank],
|
|
output_device=self._local_rank,
|
|
find_unused_parameters=self.find_unused_parameters,
|
|
broadcast_buffers=self.broadcast_buffers,
|
|
bucket_cap_mb=self.bucket_cap_mb,
|
|
)
|
|
|
|
def clip_grad_norm(
|
|
self,
|
|
model: torch.nn.Module,
|
|
max_norm: float,
|
|
*,
|
|
optimizer=None,
|
|
) -> torch.Tensor:
|
|
del optimizer
|
|
self.last_grad_clip_stats = None
|
|
return torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
|
|
|
|
@contextmanager
|
|
def no_sync(self, model: torch.nn.Module):
|
|
with model.no_sync():
|
|
yield
|
|
|
|
@property
|
|
def params_sharded(self) -> bool:
|
|
return False
|
|
|
|
|
|
def build_strategy(
|
|
config: dict,
|
|
*,
|
|
device: torch.device,
|
|
local_rank: int,
|
|
use_dmuon: bool = False,
|
|
) -> DistributionStrategy:
|
|
"""Pick FSDP vs DDP based on ``use_fsdp`` flag.
|
|
|
|
``use_gradient_checkpointing`` is read from the config dict by
|
|
each strategy's ``_maybe_enable_grad_ckpt``; not a separate kwarg here.
|
|
"""
|
|
if config.get("use_fsdp", False):
|
|
if use_dmuon:
|
|
return DMuonFSDPStrategy(config)
|
|
return FSDPStrategy(config)
|
|
if use_dmuon:
|
|
raise NotImplementedError(
|
|
"DMuon dedicated-parameter training requires "
|
|
"distributed.use_fsdp=true in this trainer."
|
|
)
|
|
return DDPStrategy(
|
|
config,
|
|
device=device,
|
|
local_rank=local_rank,
|
|
)
|