3467f8befb
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
141 lines
4.7 KiB
Python
141 lines
4.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Distributed model strategies for the modular trainer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from typing import Any, TypeVar
|
|
from collections.abc import Generator
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch.nn.parallel import DistributedDataParallel
|
|
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_ModelT = TypeVar("_ModelT", bound=torch.nn.Module)
|
|
SUPPORTED_DISTRIBUTED_STRATEGIES = frozenset({"fsdp", "ddp"})
|
|
|
|
|
|
def normalize_distributed_strategy(value: str | None) -> str:
|
|
"""Normalize and validate a modular-training strategy name."""
|
|
strategy = str(value or "fsdp").strip().lower()
|
|
if strategy not in SUPPORTED_DISTRIBUTED_STRATEGIES:
|
|
choices = ", ".join(sorted(SUPPORTED_DISTRIBUTED_STRATEGIES))
|
|
raise ValueError(f"training.distributed.strategy must be one of {{{choices}}}, "
|
|
f"got {strategy!r}")
|
|
return strategy
|
|
|
|
|
|
def is_ddp_strategy(training_config: Any) -> bool:
|
|
distributed = getattr(training_config, "distributed", None)
|
|
return normalize_distributed_strategy(getattr(distributed, "strategy", "fsdp")) == "ddp"
|
|
|
|
|
|
@contextmanager
|
|
def _fork_model_rng(
|
|
device: torch.device,
|
|
seed: int,
|
|
) -> Generator[None, None, None]:
|
|
fork_devices: list[int] = []
|
|
if device.type == "cuda":
|
|
fork_devices = [device.index if device.index is not None else torch.cuda.current_device()]
|
|
with torch.random.fork_rng(devices=fork_devices):
|
|
torch.manual_seed(int(seed))
|
|
yield
|
|
|
|
|
|
def build_replicated_model_from_scratch(
|
|
model_cls: type[_ModelT],
|
|
init_params: dict[str, Any],
|
|
*,
|
|
device: torch.device,
|
|
default_dtype: torch.dtype,
|
|
seed: int,
|
|
) -> _ModelT:
|
|
"""Initialize one deterministic, full model replica on the local device.
|
|
|
|
Model plugins call :func:`wrap_module_ddp` only after applying their
|
|
trainable/frozen parameter policy. This is required because DDP builds its
|
|
gradient reducer from the parameters that require gradients at wrap time.
|
|
"""
|
|
old_dtype = torch.get_default_dtype()
|
|
try:
|
|
torch.set_default_dtype(default_dtype)
|
|
# Official MMAudio constructs on CPU and then calls ``.cuda()`` before
|
|
# DDP. Preserve that RNG/initialization order for scratch parity.
|
|
with _fork_model_rng(device, seed), torch.device("cpu"):
|
|
model = model_cls(**init_params)
|
|
finally:
|
|
torch.set_default_dtype(old_dtype)
|
|
return model.to(device=device)
|
|
|
|
|
|
class DelegatingDistributedDataParallel(DistributedDataParallel):
|
|
"""DDP wrapper that preserves a model plugin's public attributes.
|
|
|
|
FastVideo model plugins access architecture attributes and helper methods
|
|
such as ``latent_seq_len`` and ``normalize`` through ``transformer``.
|
|
Native PyTorch DDP exposes these on ``module`` only; this adapter delegates
|
|
unknown attributes while retaining the standard DDP forward/reducer.
|
|
"""
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
try:
|
|
return super().__getattr__(name)
|
|
except AttributeError:
|
|
module = super().__getattr__("module")
|
|
return getattr(module, name)
|
|
|
|
|
|
def wrap_module_ddp(
|
|
module: _ModelT,
|
|
*,
|
|
device: torch.device,
|
|
broadcast_buffers: bool = False,
|
|
) -> DelegatingDistributedDataParallel:
|
|
"""Wrap a full model replica like official MMAudio's DDP runner."""
|
|
if not dist.is_available() or not dist.is_initialized():
|
|
raise RuntimeError("DDP strategy requires an initialized process group")
|
|
if not any(parameter.requires_grad for parameter in module.parameters()):
|
|
raise ValueError("DDP strategy requires at least one trainable parameter")
|
|
|
|
kwargs: dict[str, Any] = {
|
|
"broadcast_buffers": bool(broadcast_buffers),
|
|
}
|
|
if device.type == "cuda":
|
|
if device.index is None:
|
|
raise ValueError("CUDA DDP requires a concrete local device index")
|
|
kwargs.update({
|
|
"device_ids": [device.index],
|
|
"output_device": device.index,
|
|
})
|
|
wrapped = DelegatingDistributedDataParallel(module, **kwargs)
|
|
logger.info(
|
|
"Wrapped %s with native DDP on %s (broadcast_buffers=%s)",
|
|
type(module).__name__,
|
|
device,
|
|
broadcast_buffers,
|
|
)
|
|
return wrapped
|
|
|
|
|
|
def unwrap_ddp_module(module: _ModelT) -> torch.nn.Module:
|
|
"""Return the underlying module for DDP, otherwise return the input."""
|
|
if isinstance(module, DistributedDataParallel):
|
|
return module.module
|
|
return module
|
|
|
|
|
|
__all__ = [
|
|
"DelegatingDistributedDataParallel",
|
|
"SUPPORTED_DISTRIBUTED_STRATEGIES",
|
|
"build_replicated_model_from_scratch",
|
|
"is_ddp_strategy",
|
|
"normalize_distributed_strategy",
|
|
"unwrap_ddp_module",
|
|
"wrap_module_ddp",
|
|
]
|