Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,34 @@
|
||||
"""Distributed trainer package exports."""
|
||||
|
||||
from wall_x.trainer.fsdp_trainer.base_trainer import (
|
||||
BaseDistributedTrainer,
|
||||
DistributedTrainer,
|
||||
all_gather,
|
||||
all_reduce,
|
||||
barrier,
|
||||
cleanup_distributed,
|
||||
get_local_rank,
|
||||
get_rank,
|
||||
get_world_size,
|
||||
is_main_process,
|
||||
setup_distributed,
|
||||
)
|
||||
from wall_x.trainer.fsdp_trainer.fsdp_trainer import FSDPTrainer
|
||||
|
||||
__all__ = [
|
||||
# Base classes and utilities
|
||||
"BaseDistributedTrainer",
|
||||
"DistributedTrainer",
|
||||
"setup_distributed",
|
||||
"cleanup_distributed",
|
||||
"get_rank",
|
||||
"get_world_size",
|
||||
"get_local_rank",
|
||||
"is_main_process",
|
||||
"barrier",
|
||||
"all_reduce",
|
||||
"all_gather",
|
||||
# Trainers
|
||||
# "DDPTrainer",
|
||||
"FSDPTrainer",
|
||||
]
|
||||
@@ -0,0 +1,430 @@
|
||||
"""
|
||||
Base Trainer for Pure PyTorch Distributed Training (DDP/FSDP)
|
||||
No accelerate dependency - can be launched directly with torchrun
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import timedelta
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from wall_x.config.hyperparams_config import AdamWConfig
|
||||
from wall_x.config.train_config import TrainConfig
|
||||
from wall_x.trainer.trainer_utils import seed_all
|
||||
from wall_x.utils.logger import DistributedLogger
|
||||
from wall_x.utils.timers import Timers
|
||||
|
||||
|
||||
def setup_distributed():
|
||||
"""Initialize distributed training environment.
|
||||
|
||||
NCCL watchdog default is 10 min - some dataset backends (e.g. lerobot's
|
||||
per-rank sequential LeRobotDataset indexing) can exceed that during the
|
||||
first build, causing barrier #2 to time out on the waiting rank. Allow
|
||||
the timeout to be raised via ``WALLX_DIST_TIMEOUT_MINUTES`` (default 30).
|
||||
"""
|
||||
if not dist.is_initialized():
|
||||
timeout_min = int(os.environ.get("WALLX_DIST_TIMEOUT_MINUTES", "30"))
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
timeout=timedelta(minutes=timeout_min),
|
||||
)
|
||||
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", 0))
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
||||
return local_rank
|
||||
|
||||
|
||||
def cleanup_distributed():
|
||||
"""Clean up distributed training environment"""
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def get_rank() -> int:
|
||||
"""Get current process rank"""
|
||||
if dist.is_initialized():
|
||||
return dist.get_rank()
|
||||
return 0
|
||||
|
||||
|
||||
def get_world_size() -> int:
|
||||
"""Get total number of processes"""
|
||||
if dist.is_initialized():
|
||||
return dist.get_world_size()
|
||||
return 1
|
||||
|
||||
|
||||
def get_local_rank() -> int:
|
||||
"""Get local rank (GPU index on current node)"""
|
||||
return int(os.environ.get("LOCAL_RANK", 0))
|
||||
|
||||
|
||||
def is_main_process() -> bool:
|
||||
"""Check if current process is the main process (rank 0)"""
|
||||
return get_rank() == 0
|
||||
|
||||
|
||||
# Trainer-level collectives should use an explicit global process group.
|
||||
# FSDP's mesh groups may be 1D or 2D and have sharding-specific semantics,
|
||||
# while trainer barriers/metric reductions need all-rank semantics.
|
||||
_trainer_process_group = None
|
||||
|
||||
|
||||
def set_trainer_process_group(pg):
|
||||
global _trainer_process_group
|
||||
_trainer_process_group = pg
|
||||
|
||||
|
||||
def barrier():
|
||||
"""Synchronize all processes"""
|
||||
if dist.is_initialized():
|
||||
dist.barrier(group=_trainer_process_group)
|
||||
|
||||
|
||||
def all_reduce(tensor: torch.Tensor, op=dist.ReduceOp.SUM) -> torch.Tensor:
|
||||
"""All-reduce tensor across all processes"""
|
||||
if dist.is_initialized():
|
||||
dist.all_reduce(tensor, op=op, group=_trainer_process_group)
|
||||
return tensor
|
||||
|
||||
|
||||
def all_gather(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""All-gather tensor from all processes"""
|
||||
if not dist.is_initialized():
|
||||
return tensor
|
||||
|
||||
group = _trainer_process_group
|
||||
world_size = dist.get_world_size(group=group)
|
||||
gathered = [torch.zeros_like(tensor) for _ in range(world_size)]
|
||||
dist.all_gather(gathered, tensor, group=group)
|
||||
return torch.stack(gathered)
|
||||
|
||||
|
||||
class BaseDistributedTrainer(ABC):
|
||||
"""
|
||||
Abstract base class for distributed training without accelerate.
|
||||
Supports both DDP and FSDP through subclasses.
|
||||
|
||||
Launch with torchrun:
|
||||
torchrun --nproc_per_node=8 train.py --config config.yaml
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def train_loop(self, epoch: int):
|
||||
"""Training loop for one epoch"""
|
||||
raise NotImplementedError("train_loop must be implemented")
|
||||
|
||||
@abstractmethod
|
||||
def val_loop(self):
|
||||
"""Validation loop"""
|
||||
raise NotImplementedError("val_loop must be implemented")
|
||||
|
||||
@abstractmethod
|
||||
def save_checkpoint(self, epoch: int, step: int = 0):
|
||||
"""Save model checkpoint"""
|
||||
raise NotImplementedError("save_checkpoint must be implemented")
|
||||
|
||||
@abstractmethod
|
||||
def load_model(self):
|
||||
"""Load model"""
|
||||
raise NotImplementedError("load_model must be implemented")
|
||||
|
||||
@abstractmethod
|
||||
def load_dataset(self):
|
||||
"""Load dataset"""
|
||||
raise NotImplementedError("load_dataset must be implemented")
|
||||
|
||||
@abstractmethod
|
||||
def backward(self, loss: torch.Tensor):
|
||||
"""Perform backward pass"""
|
||||
raise NotImplementedError("backward must be implemented")
|
||||
|
||||
@abstractmethod
|
||||
def clip_grad_norm(self, max_norm: float) -> torch.Tensor:
|
||||
"""Clip gradient norm"""
|
||||
raise NotImplementedError("clip_grad_norm must be implemented")
|
||||
|
||||
|
||||
class DistributedTrainer(BaseDistributedTrainer):
|
||||
"""
|
||||
Concrete base trainer with common functionality for DDP/FSDP.
|
||||
Subclasses should implement wrap_model, backward, clip_grad_norm, etc.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_config: TrainConfig,
|
||||
wandb_run: Optional[Any] = None,
|
||||
):
|
||||
self.cfg = train_config
|
||||
|
||||
# wandb Run object (metric recorder) - None on non-main ranks and
|
||||
# when use_wandb is disabled. Distinct from self.logger below.
|
||||
self.wandb_run = wandb_run
|
||||
|
||||
# Initialize distributed environment
|
||||
self.local_rank = setup_distributed()
|
||||
self.rank = get_rank()
|
||||
self.world_size = get_world_size()
|
||||
self.device = torch.device(f"cuda:{self.local_rank}")
|
||||
|
||||
# Dataset config: optional backend-provided trainer config.
|
||||
from wall_x.trainer.adapters.base_adapter import load_trainer_data_config
|
||||
|
||||
self.data_config = load_trainer_data_config(self.cfg)
|
||||
|
||||
# Text logger
|
||||
self.logger = DistributedLogger(
|
||||
name=self.__class__.__name__,
|
||||
save_path=self.cfg.checkpoint.save_path,
|
||||
)
|
||||
|
||||
# Training state
|
||||
self.seed = self.cfg.hyperparams.seed
|
||||
self.logger.info(f"seed {self.seed}")
|
||||
|
||||
seed_all(self.seed)
|
||||
self.start_epoch = 0
|
||||
self.global_step = 0
|
||||
self.micro_step = 0
|
||||
self.num_epoch = self.cfg.hyperparams.num_epoch
|
||||
self.dataset_config_path = getattr(self.cfg.data, "dataset_config_path", None)
|
||||
self.initial_step = 0
|
||||
|
||||
# Optimizer hyperparameters (betas / weight_decay / eps)
|
||||
# used by build_optimizer for both AdamW and Muon paths.
|
||||
opt = self.cfg.hyperparams.optimizer
|
||||
if isinstance(opt, AdamWConfig):
|
||||
self.adamw_betas = tuple(opt.betas)
|
||||
self.adamw_weight_decay = opt.weight_decay
|
||||
self.adamw_eps = opt.eps
|
||||
elif getattr(opt, "optimizer_type", None) == "muon":
|
||||
self.adamw_betas = (opt.beta_1, opt.beta_2)
|
||||
self.adamw_weight_decay = opt.weight_decay
|
||||
self.adamw_eps = opt.eps
|
||||
else:
|
||||
self.adamw_betas = (0.9, 0.98)
|
||||
self.adamw_weight_decay = 1e-8
|
||||
self.adamw_eps = 1e-8
|
||||
|
||||
# Performance options
|
||||
self.nvtx = self.cfg.debug.nvtx
|
||||
self.timers = Timers(log_level=0, log_option="minmax")
|
||||
|
||||
# Training hyperparameters
|
||||
self.max_grad_norm = opt.max_grad_norm
|
||||
self.grad_accum_steps = self.cfg.hyperparams.gradient_accumulation_steps
|
||||
self.log_interval = self.cfg.logging.log_interval
|
||||
self.log_stats_buffer = []
|
||||
|
||||
# Model and optimizer (to be set by subclass)
|
||||
self.model = None
|
||||
self.optimizer = None
|
||||
self.lr_scheduler = None
|
||||
|
||||
# Mixed precision
|
||||
dist_cfg = self.cfg.distributed
|
||||
self.use_amp = dist_cfg.use_amp
|
||||
self.amp_dtype = torch.bfloat16 if dist_cfg.bf16 else torch.float16
|
||||
self.grad_scaler = None
|
||||
if self.use_amp and self.amp_dtype == torch.float16:
|
||||
self.grad_scaler = torch.amp.GradScaler("cuda")
|
||||
|
||||
def log(
|
||||
self, message: str, level: int = logging.INFO, main_process_only: bool = True
|
||||
):
|
||||
"""Log a text message via self.logger.
|
||||
|
||||
Kept for backward compat with call sites that still use ``self.log(...)``
|
||||
(new code should call ``self.logger.info/warning/error(...)`` directly).
|
||||
Under the new DistributedLogger the ``main_process_only`` flag is
|
||||
redundant for stdout (stdout is rank-0-only by construction); it still
|
||||
controls whether non-zero ranks write the message to their file log.
|
||||
"""
|
||||
if main_process_only and not is_main_process():
|
||||
return
|
||||
self.logger.log(message, level=level, main_process_only=main_process_only)
|
||||
|
||||
def fit(self):
|
||||
"""Main training loop"""
|
||||
barrier()
|
||||
|
||||
# Optional: validate before training
|
||||
if self.cfg.checkpoint.validate_first:
|
||||
self.val_loop()
|
||||
barrier()
|
||||
|
||||
if self.nvtx:
|
||||
torch.cuda.cudart().cudaProfilerStart()
|
||||
|
||||
# Profiler setup
|
||||
enable_profiling = self.cfg.debug.profile
|
||||
profiler = (
|
||||
self._setup_profiler() if enable_profiling else contextlib.nullcontext()
|
||||
)
|
||||
|
||||
for epoch in range(self.start_epoch, self.num_epoch):
|
||||
# Training
|
||||
self.train_loop(epoch, profiler=profiler)
|
||||
barrier()
|
||||
num_training_steps = getattr(self, "num_training_steps", 0)
|
||||
if num_training_steps > 0 and self.global_step >= num_training_steps:
|
||||
break
|
||||
|
||||
# Save checkpoint
|
||||
if (epoch + 1) % self.cfg.logging.epoch_save_interval == 0:
|
||||
self.save_checkpoint(epoch)
|
||||
|
||||
# Validation
|
||||
self.val_loop()
|
||||
barrier()
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if self.nvtx:
|
||||
torch.cuda.cudart().cudaProfilerStop()
|
||||
|
||||
def _setup_profiler(self):
|
||||
"""Setup PyTorch profiler"""
|
||||
return torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
schedule=torch.profiler.schedule(
|
||||
wait=self.cfg.debug.profile_wait_iters,
|
||||
warmup=self.cfg.debug.profile_warmup_iters,
|
||||
active=self.cfg.debug.profile_active_iters,
|
||||
),
|
||||
on_trace_ready=torch.profiler.tensorboard_trace_handler(
|
||||
self.cfg.debug.profile_save_path,
|
||||
worker_name=f"worker{self.rank}",
|
||||
),
|
||||
record_shapes=True,
|
||||
profile_memory=True,
|
||||
with_stack=True,
|
||||
)
|
||||
|
||||
def gather_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Gather tensor from all processes"""
|
||||
return all_gather(tensor)
|
||||
|
||||
def reduce_tensor(self, tensor: torch.Tensor, average: bool = True) -> torch.Tensor:
|
||||
"""Reduce tensor across all processes"""
|
||||
tensor = tensor.clone()
|
||||
all_reduce(tensor)
|
||||
if average:
|
||||
tensor = tensor / self.world_size
|
||||
return tensor
|
||||
|
||||
def sync_gradients(self) -> bool:
|
||||
"""Check if gradients should be synchronized (for gradient accumulation)"""
|
||||
return (self.micro_step + 1) % self.grad_accum_steps == 0
|
||||
|
||||
def optimizer_zero_grad(self):
|
||||
"""Zero optimizer gradients"""
|
||||
self.optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
def optimizer_step(self):
|
||||
"""Perform optimizer step"""
|
||||
if self.grad_scaler is not None:
|
||||
self.grad_scaler.step(self.optimizer)
|
||||
self.grad_scaler.update()
|
||||
else:
|
||||
self.optimizer.step()
|
||||
|
||||
def lr_scheduler_step(self):
|
||||
"""Perform learning rate scheduler step"""
|
||||
if self.lr_scheduler is not None:
|
||||
self.lr_scheduler.step()
|
||||
|
||||
def get_lr(self) -> float:
|
||||
"""Get current learning rate"""
|
||||
if self.lr_scheduler is not None:
|
||||
return self.lr_scheduler.get_last_lr()[0]
|
||||
return self.optimizer.param_groups[0]["lr"]
|
||||
|
||||
def autocast_context(self):
|
||||
"""Get autocast context for mixed precision"""
|
||||
if self.use_amp:
|
||||
return torch.amp.autocast("cuda", dtype=self.amp_dtype)
|
||||
return contextlib.nullcontext()
|
||||
|
||||
def training_log(
|
||||
self,
|
||||
current_epoch: int,
|
||||
total_epoch: int,
|
||||
current_iter: int,
|
||||
total_iter: int,
|
||||
loss: torch.Tensor,
|
||||
lr: float,
|
||||
time_per_step: float,
|
||||
show_time_details: bool = False,
|
||||
):
|
||||
"""Log training progress"""
|
||||
if not is_main_process():
|
||||
return
|
||||
|
||||
# Prefer the cross-rank reduced loss (train_loss) from
|
||||
# _current_step_stats so the console line matches what wall_x_2509's
|
||||
# fsdp_trainer prints and what wandb sees. Falls back to the raw
|
||||
# rank-0 local tensor when stats haven't been populated yet.
|
||||
_smoothed = getattr(self, "_current_step_stats", None) or {}
|
||||
_loss_to_print = _smoothed.get("train_loss", float(loss))
|
||||
|
||||
log_string = ""
|
||||
log_string += f" epoch {current_epoch:3d}/{total_epoch:3d} |"
|
||||
log_string += f" iter {current_iter:6d}/{total_iter:6d} |"
|
||||
log_string += f" loss {_loss_to_print:.10f} |"
|
||||
log_string += f" lr {lr:.6f} |"
|
||||
log_string += f" time {time_per_step:.4f}s |"
|
||||
|
||||
self.log(log_string)
|
||||
|
||||
if show_time_details:
|
||||
timers_to_log = [
|
||||
"interval-time",
|
||||
"data-load",
|
||||
"forward-compute",
|
||||
"backward-compute",
|
||||
"optimizer",
|
||||
]
|
||||
self.timers.log(timers_to_log, normalizer=1)
|
||||
|
||||
def true_gather(self, value):
|
||||
"""Gather values across all processes and compute mean over non-None values."""
|
||||
device = next(self.model.parameters()).device
|
||||
|
||||
# CRITICAL: must clone() here, not just detach().
|
||||
# detach() only disconnects the computation graph but still shares the
|
||||
# same underlying storage. dist.all_reduce() is **in-place** - it
|
||||
# overwrites that storage with the sum across all ranks. If the caller
|
||||
# passes a tensor that is also referenced elsewhere (e.g. outputs["loss"]
|
||||
# and outputs["flow_loss"] pointing to the same scalar_loss), the
|
||||
# in-place all_reduce silently mutates the original, causing downstream
|
||||
# readers (like training_log) to see a value multiplied by world_size.
|
||||
value_to_gather = (
|
||||
value.detach().clone()
|
||||
if value is not None
|
||||
else torch.tensor(0.0, device=device)
|
||||
)
|
||||
count_to_gather = torch.tensor(1.0 if value is not None else 0.0, device=device)
|
||||
|
||||
# all_reduce is cheaper than all_gather because only the sum is needed.
|
||||
dist.all_reduce(value_to_gather, op=dist.ReduceOp.SUM)
|
||||
dist.all_reduce(count_to_gather, op=dist.ReduceOp.SUM)
|
||||
|
||||
total_processes = count_to_gather.item()
|
||||
if total_processes > 0:
|
||||
return value_to_gather.item() / total_processes
|
||||
else:
|
||||
return None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,400 @@
|
||||
"""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,
|
||||
)
|
||||
@@ -0,0 +1,963 @@
|
||||
"""FSDP trainer implementation."""
|
||||
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
from wall_x.model.core.action.normalizer import create_normalizers
|
||||
from wall_x.trainer.adapters import ADAPTER_REGISTRY, format_adapter_error
|
||||
from wall_x.trainer.fsdp_trainer import checkpoint_io as _ckpt_io
|
||||
from wall_x.trainer.fsdp_trainer.base_trainer import (
|
||||
DistributedTrainer,
|
||||
barrier,
|
||||
is_main_process,
|
||||
)
|
||||
from wall_x.trainer.fsdp_trainer.distribution_strategy import build_strategy
|
||||
from wall_x.trainer.fsdp_trainer.metrics_logger import MetricsLogger
|
||||
from wall_x.trainer.optimizer import get_optimizer
|
||||
from wall_x.trainer.scheduler.scheduler import get_scheduler
|
||||
from wall_x.trainer.utils import move_batch_to_device
|
||||
|
||||
|
||||
class FSDPTrainer(DistributedTrainer):
|
||||
"""
|
||||
FSDP-based Trainer using pure PyTorch FSDP2 (``fully_shard``).
|
||||
|
||||
Features:
|
||||
- No accelerate dependency
|
||||
- Direct torchrun launch
|
||||
- Mixed precision training support (bfloat16/float16)
|
||||
- Gradient accumulation
|
||||
- Full/Sharded state dict saving
|
||||
- CPU offload support
|
||||
- Activation checkpointing
|
||||
|
||||
Launch:
|
||||
torchrun --nproc_per_node=8 --master_port=29500 train_fsdp.py --config config.yaml
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_config,
|
||||
wandb_run: Optional[Any] = None,
|
||||
):
|
||||
if train_config is None:
|
||||
raise ValueError("train_config is required")
|
||||
|
||||
self.model_type = train_config.model_type
|
||||
super().__init__(train_config, wandb_run)
|
||||
|
||||
if self.model_type not in ADAPTER_REGISTRY:
|
||||
raise ValueError(format_adapter_error(self.model_type))
|
||||
adapter_cls = ADAPTER_REGISTRY[self.model_type]
|
||||
self.adapter = adapter_cls(
|
||||
cfg=self.cfg,
|
||||
logger=self.logger,
|
||||
model_type=self.model_type,
|
||||
)
|
||||
|
||||
self.strategy = build_strategy(
|
||||
dataclasses.asdict(self.cfg.distributed),
|
||||
device=self.device,
|
||||
local_rank=self.local_rank,
|
||||
use_dmuon=self.cfg.hyperparams.optimizer.optimizer_type == "dmuon",
|
||||
)
|
||||
|
||||
# Model config
|
||||
self.action_dim = self.cfg.action_dim
|
||||
self.use_selective_recompute = self.cfg.distributed.use_selective_recompute
|
||||
self.show_time_details = self.cfg.debug.show_time_details
|
||||
|
||||
# num_training_steps: read from any scheduler that exposes it
|
||||
# (CosineSchedulerConfig currently; future schedulers may add it).
|
||||
# Used in train_loop to trigger loss_guard_should_stop - works with
|
||||
# constant scheduler too once the field is set.
|
||||
sched = self.cfg.hyperparams.scheduler
|
||||
self.num_training_steps = int(getattr(sched, "num_training_steps", 0) or 0)
|
||||
self.metrics_logger = MetricsLogger(
|
||||
wandb_run=self.wandb_run,
|
||||
log_interval=self.log_interval,
|
||||
smooth_window=self.cfg.logging.loss_log_smooth_window,
|
||||
)
|
||||
|
||||
# Loss tracking
|
||||
self.base_l1_loss = None
|
||||
self.base_l1_loss_detail = {}
|
||||
|
||||
self.load_normalizer()
|
||||
self.load_processor()
|
||||
self.load_model()
|
||||
|
||||
# DDP needs to see correct requires_grad at wrap time, so freeze first.
|
||||
self._freeze_params_if_needed(self.model)
|
||||
|
||||
# Frozen-submodule prefixes for checkpoint filtering. Must be computed
|
||||
# before FSDP wrapping, while child modules still have stable names.
|
||||
self._frozen_prefixes = self._compute_frozen_prefixes()
|
||||
|
||||
# Original shapes captured before FSDP flattening.
|
||||
self.model._orig_param_shapes = {
|
||||
name: p.shape for name, p in self.model.named_parameters()
|
||||
}
|
||||
|
||||
if self._resume_from_single_file():
|
||||
self.load_state_dict(
|
||||
self.model,
|
||||
{"ckpt": self.cfg.checkpoint.resume_from},
|
||||
)
|
||||
|
||||
self._wrap_model(self.model)
|
||||
self._create_optimizer()
|
||||
self._create_scheduler()
|
||||
|
||||
if self._resume_from_training_checkpoint():
|
||||
self.resume_from_checkpoint()
|
||||
|
||||
self.load_dataset()
|
||||
|
||||
self.adapter.init_validation(self.normalizer_action, self.normalizer_propri)
|
||||
|
||||
def load_normalizer(self):
|
||||
norm_cfg = self.cfg.data.normalizer_config or {}
|
||||
custom_stats_path = norm_cfg.get("customized_action_statistic_dof")
|
||||
if not custom_stats_path and self.cfg.data.dataset_type == "lerobot":
|
||||
from wall_x.data.backends.lerobot.build import load_lerobot_normalizers
|
||||
|
||||
loaded = load_lerobot_normalizers(self.cfg)
|
||||
if loaded is not None:
|
||||
self.normalizer_action, self.normalizer_propri = (
|
||||
loaded[0],
|
||||
loaded[1],
|
||||
)
|
||||
self._action_statistic_dof = None
|
||||
self.logger.info(
|
||||
"Loaded LeRobot normalizers from %s with dataset key %s",
|
||||
loaded[2],
|
||||
loaded[3],
|
||||
)
|
||||
return
|
||||
|
||||
merged = {
|
||||
"dof_config": self.cfg.task.dof_config,
|
||||
"agent_pos_config": self.cfg.task.agent_pos_config,
|
||||
"customized_action_statistic_dof": custom_stats_path,
|
||||
"min_key": norm_cfg.get("min_key", "min"),
|
||||
"delta_key": norm_cfg.get("delta_key", "delta"),
|
||||
}
|
||||
self.normalizer_action, self.normalizer_propri, self._action_statistic_dof = (
|
||||
create_normalizers(merged)
|
||||
)
|
||||
|
||||
def backward(self, loss: torch.Tensor):
|
||||
"""Perform backward pass."""
|
||||
if self.grad_scaler is not None:
|
||||
self.grad_scaler.scale(loss).backward()
|
||||
else:
|
||||
loss.backward()
|
||||
|
||||
def clip_grad_norm(self, max_norm: float) -> torch.Tensor:
|
||||
"""Unscale and clip gradient norm via the active distribution strategy."""
|
||||
if self.grad_scaler is not None:
|
||||
self.grad_scaler.unscale_(self.optimizer)
|
||||
|
||||
total_norm = self.strategy.clip_grad_norm(
|
||||
self.model,
|
||||
max_norm,
|
||||
optimizer=self.optimizer,
|
||||
)
|
||||
self._dedicated_param_grad_clip_stats = getattr(
|
||||
self.strategy,
|
||||
"last_grad_clip_stats",
|
||||
None,
|
||||
)
|
||||
return total_norm
|
||||
|
||||
def load_model(self):
|
||||
"""Load and prepare model, optimizer, and scheduler"""
|
||||
self.model_config = self.adapter.build_model_config()
|
||||
self.model = self.adapter.create_model(
|
||||
self.processor, self.tokenizer_mixin, self.model_config
|
||||
)
|
||||
type(self.adapter).log_attention_implementation(self.logger, self.model)
|
||||
self.adapter.load_weights(
|
||||
self.model,
|
||||
self.normalizer_action,
|
||||
self.normalizer_propri,
|
||||
processor=self.processor,
|
||||
)
|
||||
|
||||
def load_processor(self):
|
||||
"""Load processor and tokenizers"""
|
||||
self.adapter.normalizer_action = self.normalizer_action
|
||||
self.adapter.normalizer_propri = self.normalizer_propri
|
||||
processors_dict = self.adapter.load_processor(self._action_statistic_dof)
|
||||
self.processor = processors_dict["processor"]
|
||||
self.train_action_tokenizer = processors_dict["train_action_tokenizer"]
|
||||
self.val_action_tokenizer = processors_dict["val_action_tokenizer"]
|
||||
self.action_mapper = processors_dict["action_mapper"]
|
||||
self.tokenizer_mixin = processors_dict.get("tokenizer_mixin")
|
||||
|
||||
def _freeze_params_if_needed(self, model: torch.nn.Module):
|
||||
"""Freeze non-action parameters when train_action_expert_only is set.
|
||||
|
||||
Must be called BEFORE wrapping with DDP/FSDP so the wrapper sees the
|
||||
correct requires_grad flags and does not expect gradients for frozen params.
|
||||
"""
|
||||
from wall_x.trainer.optimizer.utils import resolve_lr_group_configs
|
||||
|
||||
opt = self.cfg.hyperparams.optimizer
|
||||
if not opt.train_action_expert_only:
|
||||
return
|
||||
lr_groups = resolve_lr_group_configs(
|
||||
opt, self.adapter.default_action_lr_keywords
|
||||
)
|
||||
if not lr_groups:
|
||||
self.log(
|
||||
"WARNING: train_action_expert_only is True but no optimizer LR "
|
||||
"group is set. No parameters will be frozen.",
|
||||
level=logging.WARNING,
|
||||
)
|
||||
return
|
||||
|
||||
frozen_count = 0
|
||||
grouped_count = 0
|
||||
for name, param in model.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
matches_group = any(
|
||||
any(keyword in name for keyword in group.include) for group in lr_groups
|
||||
)
|
||||
if matches_group:
|
||||
grouped_count += 1
|
||||
else:
|
||||
param.requires_grad = False
|
||||
frozen_count += 1
|
||||
|
||||
if grouped_count == 0:
|
||||
raise ValueError(
|
||||
"No grouped params found for train_action_expert_only. "
|
||||
"Please check optimizer.lr_groups or legacy action_lr_keywords."
|
||||
)
|
||||
|
||||
self.log(
|
||||
f"*** train_action_expert_only: frozen {frozen_count} base params, "
|
||||
f"keeping {grouped_count} grouped params trainable ***"
|
||||
)
|
||||
|
||||
def _wrap_model(self, model: torch.nn.Module):
|
||||
"""Prepare model dtype/placement via adapter, then wrap via strategy.
|
||||
|
||||
If the strategy created an explicit trainer process group, register
|
||||
it so trainer-level barriers and metric collectives keep all-rank
|
||||
semantics independent of the FSDP mesh topology.
|
||||
"""
|
||||
self.adapter.convert_to_mix_precision_hint(
|
||||
model,
|
||||
device=self.device,
|
||||
use_fsdp=self.cfg.distributed.use_fsdp,
|
||||
log_fn=self.log,
|
||||
)
|
||||
self.model = self.strategy.wrap(model)
|
||||
|
||||
trainer_pg = getattr(self.strategy, "trainer_process_group", None)
|
||||
if trainer_pg is not None:
|
||||
from wall_x.trainer.fsdp_trainer.base_trainer import (
|
||||
set_trainer_process_group,
|
||||
)
|
||||
|
||||
set_trainer_process_group(trainer_pg)
|
||||
self.log("[FSDP2] routing trainer collectives through trainer PG")
|
||||
|
||||
def _create_optimizer(self):
|
||||
"""Create optimizer for FSDP wrapped model."""
|
||||
from wall_x.trainer.optimizer.utils import (
|
||||
build_lr_param_groups,
|
||||
resolve_lr_group_configs,
|
||||
uses_legacy_action_lr_groups,
|
||||
)
|
||||
|
||||
opt_cfg = self.cfg.hyperparams.optimizer
|
||||
|
||||
param_groups = None
|
||||
lr_group_configs = resolve_lr_group_configs(
|
||||
opt_cfg, self.adapter.default_action_lr_keywords
|
||||
)
|
||||
if lr_group_configs:
|
||||
if opt_cfg.optimizer_type not in ("adamw", "muon", "dmuon"):
|
||||
raise ValueError(
|
||||
"optimizer.lr_groups are only supported with adamw, muon, "
|
||||
"or dmuon"
|
||||
)
|
||||
base_group_name = (
|
||||
"base_lr_group" if uses_legacy_action_lr_groups(opt_cfg) else "base"
|
||||
)
|
||||
param_groups = build_lr_param_groups(
|
||||
self.model,
|
||||
opt_cfg,
|
||||
lr_group_configs,
|
||||
base_group_name=base_group_name,
|
||||
)
|
||||
summary = ", ".join(
|
||||
f"{group.name}={group.lr} include={group.include}"
|
||||
for group in lr_group_configs
|
||||
)
|
||||
self.log(
|
||||
f"setting optimizer LR groups: base={opt_cfg.learning_rate}; "
|
||||
f"{summary} ({opt_cfg.optimizer_type})"
|
||||
)
|
||||
|
||||
optimizer_kwargs = {
|
||||
"opt_cfg": opt_cfg,
|
||||
"param_groups": param_groups,
|
||||
}
|
||||
if opt_cfg.optimizer_type == "dmuon":
|
||||
optimizer_kwargs["log_fn"] = self.log
|
||||
|
||||
self.optimizer = get_optimizer(
|
||||
opt_cfg.optimizer_type,
|
||||
self.model,
|
||||
**optimizer_kwargs,
|
||||
)
|
||||
|
||||
def _create_scheduler(self):
|
||||
"""Create learning rate scheduler"""
|
||||
from wall_x.config.hyperparams_config import (
|
||||
ConstantSchedulerConfig,
|
||||
CosineSchedulerConfig,
|
||||
)
|
||||
|
||||
sched = self.cfg.hyperparams.scheduler
|
||||
lr = self.cfg.hyperparams.optimizer.learning_rate
|
||||
|
||||
if isinstance(sched, ConstantSchedulerConfig):
|
||||
self.lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
|
||||
self.optimizer, factor=1.0, total_iters=0
|
||||
)
|
||||
elif isinstance(sched, CosineSchedulerConfig):
|
||||
if sched.num_training_steps <= 0:
|
||||
raise ValueError(
|
||||
"num_training_steps must be > 0 for cosine scheduler. "
|
||||
"Please set it explicitly in the config."
|
||||
)
|
||||
min_lr = sched.min_lr if sched.min_lr is not None else 0.1 * lr
|
||||
self.lr_scheduler = get_scheduler(
|
||||
optimizer=self.optimizer,
|
||||
lr_scheduler_type="cosine",
|
||||
num_warmup_steps=sched.num_warmup_steps,
|
||||
num_training_steps=sched.num_training_steps,
|
||||
peak_lr=lr,
|
||||
end_lr=min_lr,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported scheduler type: {type(sched).__name__}")
|
||||
|
||||
def load_dataset(self):
|
||||
"""Load training and validation datasets"""
|
||||
torch.cuda.empty_cache()
|
||||
barrier()
|
||||
# Dispatch to adapter with unified signature (supports model-specific kwargs)
|
||||
extra_dataset_kwargs = dict(
|
||||
model_config=self.model_config,
|
||||
tokenizer_mixin=self.tokenizer_mixin,
|
||||
normalizer_action=self.normalizer_action,
|
||||
normalizer_propri=self.normalizer_propri,
|
||||
)
|
||||
|
||||
resume_path = self.cfg.checkpoint.resume_from
|
||||
if resume_path is not None and os.path.isdir(resume_path):
|
||||
# Step and epoch restore are generic. Backend-owned episode
|
||||
# metadata is optional and only loaded when the adapter exposes
|
||||
# the hook.
|
||||
is_incomplete_epoch = False
|
||||
if hasattr(self.adapter, "_load_episode_indices"):
|
||||
resume_indices = self.adapter._load_episode_indices(
|
||||
resume_path, self.rank
|
||||
)
|
||||
if resume_indices is not None:
|
||||
is_incomplete_epoch = resume_indices["is_incomplete_epoch"]
|
||||
extra_dataset_kwargs["resume_indices"] = resume_indices
|
||||
|
||||
train_state = self.adapter.load_step_and_epoch(
|
||||
resume_path, is_incomplete_epoch
|
||||
)
|
||||
self.global_step = train_state["global_step"]
|
||||
self.start_epoch = train_state["start_epoch"]
|
||||
self.log(
|
||||
f"global_step: {self.global_step}, start_epoch: "
|
||||
f"{self.start_epoch} from {resume_path}"
|
||||
)
|
||||
# global_step records completed steps, so resume from the next step
|
||||
# to avoid triggering _should_save_checkpoint() immediately.
|
||||
self.global_step += 1
|
||||
(
|
||||
self.dataset,
|
||||
self.train_dataloader,
|
||||
self.train_num,
|
||||
) = self.adapter.load_dataset(
|
||||
self.data_config,
|
||||
self.processor,
|
||||
self.rank,
|
||||
self.world_size,
|
||||
**extra_dataset_kwargs,
|
||||
)
|
||||
|
||||
def train_loop(self, epoch: int, profiler=contextlib.nullcontext()):
|
||||
"""Execute training for a single epoch"""
|
||||
self.model.train()
|
||||
|
||||
# VGDynamicRobotDataset (v1 path) doesn't expose set_epoch; other
|
||||
# dataset backends (v2, lerobot) do. Dispatch conditionally.
|
||||
if hasattr(self.dataset, "set_epoch"):
|
||||
self.dataset.set_epoch(epoch)
|
||||
total = len(self.train_dataloader)
|
||||
# Drop last
|
||||
stop_step = total - total % self.grad_accum_steps
|
||||
|
||||
# Disable automatic GC to prevent random ~200ms stalls during compute
|
||||
gc_interval = self.cfg.logging.gc_interval_steps
|
||||
gc.disable()
|
||||
|
||||
t0 = time.time()
|
||||
with profiler:
|
||||
self.timers("interval-time", log_level=0).start(barrier=True)
|
||||
self.timers("data-load", log_level=0).start(barrier=True)
|
||||
|
||||
for i, batch in enumerate(self.train_dataloader, self.initial_step):
|
||||
self.timers("data-load").stop()
|
||||
|
||||
# Save first batch for offline profiling (rank 0 only)
|
||||
debug_batch_path = self.cfg.debug.save_debug_batch_path
|
||||
if debug_batch_path and i == self.initial_step and self.rank == 0:
|
||||
os.makedirs(os.path.dirname(debug_batch_path) or ".", exist_ok=True)
|
||||
torch.save({"batch": batch}, debug_batch_path)
|
||||
self.log(f"Saved debug batch to {debug_batch_path}")
|
||||
|
||||
# Move batch to device
|
||||
batch = move_batch_to_device(batch, self.device)
|
||||
|
||||
if self.cfg.debug.enable_mfu_profile and self.global_step == 0:
|
||||
self.log(
|
||||
"[MFU] Forward FLOPs profiling is not available in the public package.",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
# Forward pass with autocast
|
||||
self.timers("forward-compute", log_level=0).start(barrier=False)
|
||||
with self.autocast_context():
|
||||
outputs = self.adapter.forward(
|
||||
self.model, batch, global_step=self.global_step
|
||||
)
|
||||
self.timers("forward-compute").stop()
|
||||
|
||||
loss = self.adapter.extract_loss(outputs)
|
||||
|
||||
# Check for NaN loss -- replace with 0 instead of `continue`
|
||||
# to avoid skipping collective ops (backward, all_reduce, etc.)
|
||||
# which would cause NCCL deadlock across ranks.
|
||||
nan_loss = torch.isnan(loss)
|
||||
if nan_loss:
|
||||
self.log(
|
||||
f"Warning: nan in loss at epoch: {epoch}, step: {i}",
|
||||
level=logging.WARNING,
|
||||
)
|
||||
loss = torch.zeros_like(loss)
|
||||
|
||||
# Backward pass
|
||||
self.timers("backward-compute", log_level=0).start(barrier=True)
|
||||
|
||||
context = (
|
||||
contextlib.nullcontext()
|
||||
if self.sync_gradients()
|
||||
else self.strategy.no_sync(self.model)
|
||||
)
|
||||
with context:
|
||||
scaled_loss = loss / self.grad_accum_steps
|
||||
self.backward(scaled_loss)
|
||||
self.timers("backward-compute").stop()
|
||||
|
||||
# Gradient sync and optimizer step
|
||||
if self.sync_gradients():
|
||||
self.timers("optimizer", log_level=0).start(barrier=True)
|
||||
|
||||
# Per-component grad norms (before clipping)
|
||||
self.timers("optimizer-grad-norms", log_level=0).start(
|
||||
barrier=False
|
||||
)
|
||||
self._component_grad_norms = {}
|
||||
self.adapter.collect_grad_norms(
|
||||
self.model,
|
||||
self._component_grad_norms,
|
||||
device=self.device,
|
||||
reduce_tensor_fn=self.reduce_tensor,
|
||||
params_sharded=self.strategy.params_sharded,
|
||||
)
|
||||
self.timers("optimizer-grad-norms").stop()
|
||||
|
||||
# Clip gradients
|
||||
self.timers("optimizer-clip", log_level=0).start(barrier=False)
|
||||
if self.cfg.hyperparams.optimizer.enable_grad_clip:
|
||||
total_norm = self.clip_grad_norm(self.max_grad_norm)
|
||||
else:
|
||||
self._dedicated_param_grad_clip_stats = None
|
||||
total_norm = 0
|
||||
self.timers("optimizer-clip").stop()
|
||||
|
||||
self.timers("optimizer-step", log_level=0).start(barrier=False)
|
||||
self.optimizer_step()
|
||||
self.timers("optimizer-step").stop()
|
||||
|
||||
self.timers("optimizer-zero-grad", log_level=0).start(barrier=False)
|
||||
self.optimizer_zero_grad()
|
||||
self.timers("optimizer-zero-grad").stop()
|
||||
self.timers("optimizer").stop()
|
||||
|
||||
# Scheduler step
|
||||
self.lr_scheduler_step()
|
||||
|
||||
# Logging
|
||||
self.timers("logging", log_level=0).start(barrier=True)
|
||||
_t_metrics_start = time.time()
|
||||
self._log_training_metrics(
|
||||
epoch, i, total, loss, total_norm, outputs
|
||||
)
|
||||
_t_metrics_ms = (time.time() - _t_metrics_start) * 1000
|
||||
|
||||
t1 = time.time()
|
||||
self.training_log(
|
||||
epoch,
|
||||
self.num_epoch,
|
||||
i,
|
||||
total,
|
||||
loss,
|
||||
self.get_lr(),
|
||||
t1 - t0,
|
||||
self.show_time_details,
|
||||
)
|
||||
self.timers("logging").stop()
|
||||
t0 = time.time()
|
||||
|
||||
# Optional: log per-step breakdown to pinpoint spikes (e.g. param_norms every 100 steps)
|
||||
if (
|
||||
self.show_time_details
|
||||
and _t_metrics_ms > 5000
|
||||
and is_main_process()
|
||||
):
|
||||
self.log(
|
||||
f"[Step time breakdown] _log_training_metrics took {_t_metrics_ms:.0f} ms at global_step={self.global_step}",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
# Checkpoint saving (FSDP full state_dict can take 10-20s every save_interval steps)
|
||||
if self._should_save_checkpoint():
|
||||
_t_save_start = time.time()
|
||||
self.save_checkpoint(epoch, self.global_step)
|
||||
if is_main_process():
|
||||
self.log(
|
||||
f"[Step time breakdown] save_checkpoint took {(time.time() - _t_save_start):.1f} s at global_step={self.global_step}",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
# Validation
|
||||
if self._should_validate():
|
||||
self.val_loop()
|
||||
|
||||
self.global_step += 1
|
||||
self.micro_step = 0
|
||||
|
||||
# Manual GC outside timing window to avoid random stalls (can add 1-5s every gc_interval_steps)
|
||||
if self.global_step % gc_interval == 0:
|
||||
_t_gc_start = time.time()
|
||||
gc.collect()
|
||||
if is_main_process() and self.show_time_details:
|
||||
self.log(
|
||||
f"[Step time breakdown] gc.collect took {(time.time() - _t_gc_start):.1f} s at global_step={self.global_step}",
|
||||
level=logging.INFO,
|
||||
)
|
||||
else:
|
||||
self.micro_step += 1
|
||||
|
||||
del batch
|
||||
self.timers("interval-time").stop()
|
||||
|
||||
# Drop last
|
||||
if i == stop_step:
|
||||
break
|
||||
|
||||
if not isinstance(profiler, contextlib.nullcontext):
|
||||
profiler.step()
|
||||
|
||||
if (
|
||||
self.num_training_steps > 0
|
||||
and self.global_step >= self.num_training_steps
|
||||
):
|
||||
break
|
||||
|
||||
# Setup timers for next iteration
|
||||
if i < total - 1:
|
||||
self.timers("interval-time", log_level=0).start(barrier=True)
|
||||
self.timers("data-load", log_level=0).start(barrier=True)
|
||||
|
||||
# Re-enable automatic GC after training loop
|
||||
gc.enable()
|
||||
|
||||
# Reset dataloader for next epoch
|
||||
self.train_dataloader = self.dataset.get_train_dataloader()
|
||||
|
||||
def _log_training_metrics(self, epoch, step, total, loss, total_norm, outputs):
|
||||
"""Collect per-step stats, delegate buffering + wandb emission."""
|
||||
lr = self.get_lr()
|
||||
train_loss = self.reduce_tensor(loss.detach()).item()
|
||||
|
||||
step_stats = {
|
||||
"lr": lr,
|
||||
"train_loss": train_loss,
|
||||
"grad_norm": (
|
||||
total_norm.item() if torch.is_tensor(total_norm) else float(total_norm)
|
||||
),
|
||||
}
|
||||
for idx, group in enumerate(self.optimizer.param_groups):
|
||||
group_name = group.get("group_name", f"group_{idx}")
|
||||
step_stats[f"lr_group/{group_name}"] = float(group["lr"])
|
||||
|
||||
# Model-family-specific auxiliary losses / accuracies.
|
||||
self.adapter.collect_output_stats(
|
||||
outputs,
|
||||
step_stats,
|
||||
reduce_tensor_fn=self.reduce_tensor,
|
||||
true_gather_fn=self.true_gather,
|
||||
tokenizer_mixin=self.tokenizer_mixin,
|
||||
)
|
||||
|
||||
# Per-component grad norms (captured before clip & optimizer.zero_grad).
|
||||
if hasattr(self, "_component_grad_norms"):
|
||||
step_stats.update(self._component_grad_norms)
|
||||
|
||||
dedicated_clip_stats = getattr(self, "_dedicated_param_grad_clip_stats", None)
|
||||
if dedicated_clip_stats is not None:
|
||||
step_stats.update(
|
||||
{
|
||||
"muon_grad_norm": dedicated_clip_stats["total_norm"],
|
||||
"muon_grad_clip_coef": dedicated_clip_stats["clip_coef"],
|
||||
"muon_grad_clipped": float(dedicated_clip_stats["clipped"]),
|
||||
}
|
||||
)
|
||||
|
||||
if self.global_step % 100 == 0:
|
||||
self.adapter.collect_param_norms(
|
||||
self.model,
|
||||
step_stats,
|
||||
device=self.device,
|
||||
reduce_tensor_fn=self.reduce_tensor,
|
||||
params_sharded=self.strategy.params_sharded,
|
||||
)
|
||||
|
||||
self._current_step_raw_stats = step_stats
|
||||
|
||||
# Display-smoothing rolling window (DZ-style)
|
||||
self._current_step_stats = self.metrics_logger.smooth(step_stats)
|
||||
|
||||
def training_log(
|
||||
self,
|
||||
current_epoch,
|
||||
total_epoch,
|
||||
current_train_iter,
|
||||
total_train_iter,
|
||||
loss,
|
||||
lr,
|
||||
time_per_step,
|
||||
show_time_details=False,
|
||||
):
|
||||
# timers.log() contains all_gather - must run on ALL ranks before the
|
||||
# is_main_process() guard to avoid NCCL deadlock.
|
||||
if show_time_details:
|
||||
self.timers.log(
|
||||
[
|
||||
"interval-time",
|
||||
"data-load",
|
||||
"forward-compute",
|
||||
"backward-compute",
|
||||
"optimizer",
|
||||
"optimizer-grad-norms",
|
||||
"optimizer-clip",
|
||||
"optimizer-step",
|
||||
"optimizer-zero-grad",
|
||||
"logging",
|
||||
],
|
||||
normalizer=1,
|
||||
)
|
||||
|
||||
main = is_main_process()
|
||||
stats = getattr(self, "_current_step_stats", None) or {}
|
||||
|
||||
# MFU is computed after step time is known, then merged into both the
|
||||
# console stats and the MetricsLogger buffer so wandb records it.
|
||||
mfu_info = None
|
||||
if self.cfg.debug.enable_mfu and hasattr(self.adapter, "compute_mfu"):
|
||||
unwrapped = (
|
||||
self.model.module if hasattr(self.model, "module") else self.model
|
||||
)
|
||||
mfu_info = self.adapter.compute_mfu(unwrapped, time_per_step)
|
||||
if mfu_info is not None:
|
||||
self.log(
|
||||
"[MFU-rank] rank={} mfu={:.2f}% step={:.3f}s "
|
||||
"flops_step={:.3f}T flops_fwd={:.3f}T "
|
||||
"seq={} latent={} source={}".format(
|
||||
getattr(self, "rank", 0),
|
||||
mfu_info["mfu"] * 100.0,
|
||||
time_per_step,
|
||||
mfu_info["flops_per_step_T"],
|
||||
mfu_info.get("flops_forward_T", 0.0) or 0.0,
|
||||
mfu_info.get("seq_dims"),
|
||||
mfu_info.get("latent_dims"),
|
||||
mfu_info.get("latent_source"),
|
||||
),
|
||||
main_process_only=False,
|
||||
)
|
||||
if main and mfu_info is not None and stats is not None:
|
||||
mfu_stats = {
|
||||
"mfu": mfu_info["mfu"],
|
||||
"flops_per_step_T": mfu_info["flops_per_step_T"],
|
||||
"flops_forward_T": mfu_info.get("flops_forward_T"),
|
||||
}
|
||||
profile_fwd_flops = mfu_info.get("profile_fwd_flops_T")
|
||||
if profile_fwd_flops is not None:
|
||||
mfu_stats["profile_fwd_flops_T"] = profile_fwd_flops
|
||||
stats.update(mfu_stats)
|
||||
raw_stats = getattr(self, "_current_step_raw_stats", None)
|
||||
if raw_stats is not None:
|
||||
raw_stats.update(mfu_stats)
|
||||
|
||||
if main:
|
||||
raw_stats = getattr(self, "_current_step_raw_stats", None)
|
||||
if raw_stats is not None:
|
||||
self.metrics_logger.record_step(raw_stats, is_main=True)
|
||||
avg_stats = self.metrics_logger.flush_if_due(self.global_step)
|
||||
if avg_stats is not None:
|
||||
self.logger.info(
|
||||
f"[FSDP Train] Step {self.global_step}: {avg_stats}"
|
||||
)
|
||||
|
||||
if not main:
|
||||
return
|
||||
|
||||
loss_to_print = stats.get(
|
||||
"train_loss", loss.item() if torch.is_tensor(loss) else float(loss)
|
||||
)
|
||||
fields = self.adapter.console_fields(tokenizer_mixin=self.tokenizer_mixin)
|
||||
|
||||
line = self.metrics_logger.format_training_line(
|
||||
epoch=current_epoch,
|
||||
total_epoch=total_epoch,
|
||||
current_iter=current_train_iter,
|
||||
total_iter=total_train_iter,
|
||||
loss=loss_to_print,
|
||||
lr=lr,
|
||||
time_per_step=time_per_step,
|
||||
stats=stats,
|
||||
fields=fields,
|
||||
mfu_info=mfu_info,
|
||||
)
|
||||
self.log(line)
|
||||
|
||||
def _should_save_checkpoint(self) -> bool:
|
||||
"""Check if checkpoint should be saved"""
|
||||
log = self.cfg.logging
|
||||
return (
|
||||
self.global_step >= log.ignore_until_interval
|
||||
and self.global_step % log.save_interval == 0
|
||||
and self.global_step != 0
|
||||
)
|
||||
|
||||
def _should_validate(self) -> bool:
|
||||
"""Check if validation should be performed"""
|
||||
log = self.cfg.logging
|
||||
return (
|
||||
self.global_step >= log.ignore_until_interval
|
||||
and self.global_step % log.val_interval == 0
|
||||
and self.global_step > 0
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def val_loop(self):
|
||||
"""Delegate the full validation run to the adapter.
|
||||
|
||||
Skips silently when the dataset has no val split - adapter's
|
||||
run_validation iterates ``val_dataloader`` with ``tqdm(...,
|
||||
total=len(...))`` and cannot accept None. v2 returns None here
|
||||
when the YAML only declares a train split.
|
||||
"""
|
||||
self.val_dataloader = self.dataset.get_val_dataloader()
|
||||
if self.val_dataloader is None:
|
||||
if is_main_process():
|
||||
self.logger.info("No val split configured, skipping validation.")
|
||||
return
|
||||
save_path = self.cfg.checkpoint.save_path
|
||||
self.adapter.run_validation(
|
||||
model=self.model,
|
||||
val_dataloader=self.val_dataloader,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
device=self.device,
|
||||
autocast_context=self.autocast_context,
|
||||
reduce_fn=self.reduce_tensor,
|
||||
gather_fn=self.true_gather,
|
||||
logger=self.wandb_run if is_main_process() else None,
|
||||
global_step=self.global_step,
|
||||
output_path=os.path.join(save_path, f"val_rank_{self.rank}"),
|
||||
tokenizer_mixin=self.tokenizer_mixin,
|
||||
)
|
||||
barrier()
|
||||
|
||||
def _compute_frozen_prefixes(self) -> Optional[Tuple[str, ...]]:
|
||||
"""Return state-dict key prefixes for submodules that are fully frozen.
|
||||
|
||||
For models with a ``pipe`` container, any child not present in
|
||||
``cfg.model.trainable_models`` is treated as frozen. The matching key
|
||||
prefixes let checkpoint_io drop those entries from the saved file.
|
||||
Empty / missing ``trainable_models`` disables filtering and preserves
|
||||
"save everything" behavior.
|
||||
"""
|
||||
if not hasattr(self.model, "pipe"):
|
||||
return None
|
||||
trainable_models = getattr(self.cfg.model, "trainable_models", None)
|
||||
if trainable_models is None:
|
||||
return None
|
||||
if isinstance(trainable_models, str):
|
||||
trainable_set = {
|
||||
s.strip() for s in trainable_models.split(",") if s.strip()
|
||||
}
|
||||
else:
|
||||
trainable_set = set(trainable_models)
|
||||
if not trainable_set:
|
||||
self.log(
|
||||
"WARNING: trainable_models is empty; not filtering frozen "
|
||||
"entries from checkpoint to avoid saving an empty file."
|
||||
)
|
||||
return None
|
||||
frozen_prefixes = tuple(
|
||||
f"pipe.{name}."
|
||||
for name, _ in self.model.pipe.named_children()
|
||||
if name not in trainable_set
|
||||
)
|
||||
if frozen_prefixes:
|
||||
self.log(
|
||||
f"Frozen submodule prefixes excluded from checkpoints: "
|
||||
f"{list(frozen_prefixes)}"
|
||||
)
|
||||
return frozen_prefixes
|
||||
return None
|
||||
|
||||
def save_checkpoint(self, epoch: int, step: int = 0):
|
||||
"""Save model checkpoint via checkpoint_io (dispatches on model wrapper type)."""
|
||||
save_path = self.cfg.checkpoint.save_path
|
||||
ckpt_path = f"{save_path}/{epoch}_{step}" if step else f"{save_path}/{epoch}"
|
||||
_ckpt_io.save_checkpoint(
|
||||
ckpt_path=ckpt_path,
|
||||
model=self.model,
|
||||
optimizer=self.optimizer,
|
||||
lr_scheduler=self.lr_scheduler,
|
||||
config=dataclasses.asdict(self.cfg),
|
||||
rank=self.rank,
|
||||
is_main=is_main_process(),
|
||||
epoch=epoch,
|
||||
global_step=self.global_step,
|
||||
seed=self.seed,
|
||||
normalizer_action=self.normalizer_action,
|
||||
normalizer_propri=self.normalizer_propri,
|
||||
dataset=getattr(self, "dataset", None) if step != 0 else None,
|
||||
grad_scaler=self.grad_scaler,
|
||||
log_fn=self.log,
|
||||
frozen_prefixes=self._frozen_prefixes,
|
||||
)
|
||||
_ckpt_io.finalize_save()
|
||||
self.log(f"Saved checkpoint to {ckpt_path}")
|
||||
|
||||
def load_state_dict(self, model, resume_config):
|
||||
"""Load state dict with fused-weight conversion + try_harder support."""
|
||||
return _ckpt_io.load_weights(
|
||||
model=model,
|
||||
resume_config=resume_config,
|
||||
model_class=self._checkpoint_model_class(),
|
||||
log_fn=self.log,
|
||||
)
|
||||
|
||||
def _checkpoint_model_class(self):
|
||||
model_class = getattr(self, "ModelClass", None)
|
||||
if model_class is None and hasattr(type(self.adapter), "model_class"):
|
||||
model_class = type(self.adapter).model_class()
|
||||
return model_class
|
||||
|
||||
def _resume_from_single_file(self) -> bool:
|
||||
path = self.cfg.checkpoint.resume_from
|
||||
return bool(path) and str(path).endswith((".safetensors", ".pth"))
|
||||
|
||||
def _resume_from_training_checkpoint(self) -> bool:
|
||||
path = self.cfg.checkpoint.resume_from
|
||||
return bool(path) and not self._resume_from_single_file()
|
||||
|
||||
def resume_from_checkpoint(self):
|
||||
"""Resume training from checkpoint via checkpoint_io."""
|
||||
_ckpt_io.resume_from_checkpoint(
|
||||
model=self.model,
|
||||
optimizer=self.optimizer,
|
||||
lr_scheduler=self.lr_scheduler,
|
||||
resume_config={"ckpt": self.cfg.checkpoint.resume_from},
|
||||
rank=self.rank,
|
||||
grad_scaler=self.grad_scaler,
|
||||
model_class=self._checkpoint_model_class(),
|
||||
log_fn=self.log,
|
||||
)
|
||||
barrier()
|
||||
|
||||
def predict_action_loop(
|
||||
self,
|
||||
current_step=0,
|
||||
max_iteration=None,
|
||||
prediction_type="flow_action",
|
||||
mode="acc",
|
||||
):
|
||||
"""Delegate action prediction to the adapter."""
|
||||
del mode # unused in current dispatch; kept for signature compatibility
|
||||
return self.adapter.predict(
|
||||
prediction_type,
|
||||
model=self._unwrapped_model_for_inference(),
|
||||
val_dataloader=self.dataset.get_val_dataloader(),
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
device=self.device,
|
||||
processor=self.processor,
|
||||
tokenizer_mixin=self.tokenizer_mixin,
|
||||
logger=self.wandb_run,
|
||||
current_step=current_step,
|
||||
max_iteration=max_iteration,
|
||||
)
|
||||
|
||||
def predict_text_loop(self, current_step=0, max_samples=None, save_dir="./"):
|
||||
"""Delegate text prediction to the adapter."""
|
||||
return self.adapter.predict(
|
||||
"text",
|
||||
model=self._unwrapped_model_for_inference(),
|
||||
val_dataloader=self.dataset.get_val_dataloader(),
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
device=self.device,
|
||||
processor=self.processor,
|
||||
tokenizer_mixin=self.tokenizer_mixin,
|
||||
logger=self.wandb_run,
|
||||
current_step=current_step,
|
||||
max_samples=max_samples,
|
||||
save_dir=save_dir,
|
||||
)
|
||||
|
||||
def _unwrapped_model_for_inference(self):
|
||||
"""Return the underlying module for inference: unwrap DDP, pass FSDP through."""
|
||||
if isinstance(self.model, DDP):
|
||||
return self.model.module
|
||||
return self.model
|
||||
@@ -0,0 +1,3 @@
|
||||
from wall_x.trainer.fsdp_trainer.metrics.norms import DistributedNormAccumulator
|
||||
|
||||
__all__ = ["DistributedNormAccumulator"]
|
||||
@@ -0,0 +1,123 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from typing import Callable, DefaultDict, Optional
|
||||
|
||||
import torch
|
||||
from torch.distributed.tensor import DTensor, Replicate
|
||||
|
||||
|
||||
class DistributedNormAccumulator:
|
||||
"""Accumulate L2 norms without leaking parallel-layout branches to adapters."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
device: torch.device,
|
||||
reduce_tensor_fn: Optional[Callable[[torch.Tensor, bool], torch.Tensor]],
|
||||
):
|
||||
self.device = device
|
||||
self.reduce_tensor_fn = reduce_tensor_fn
|
||||
self._regular: DefaultDict[str, list[torch.Tensor]] = defaultdict(list)
|
||||
self._dtensor: DefaultDict[str, list[torch.Tensor]] = defaultdict(list)
|
||||
self._owner_only: DefaultDict[str, list[torch.Tensor]] = defaultdict(list)
|
||||
|
||||
@staticmethod
|
||||
def _is_dtensor(tensor: torch.Tensor) -> bool:
|
||||
return isinstance(tensor, DTensor)
|
||||
|
||||
def add(self, group: Optional[str], tensor: Optional[torch.Tensor]) -> None:
|
||||
if group is None or tensor is None or tensor.numel() == 0:
|
||||
return
|
||||
tensor = tensor.detach()
|
||||
if self._is_dtensor(tensor):
|
||||
self._dtensor[group].append(tensor)
|
||||
else:
|
||||
self._regular[group].append(tensor)
|
||||
|
||||
def add_owner_only(
|
||||
self, group: Optional[str], tensor: Optional[torch.Tensor]
|
||||
) -> None:
|
||||
if group is None or tensor is None or tensor.numel() == 0:
|
||||
return
|
||||
self._owner_only[group].append(tensor.detach())
|
||||
|
||||
def squared_norm(self, group: str) -> torch.Tensor:
|
||||
with torch.no_grad():
|
||||
total = torch.zeros((), dtype=torch.float32, device=self.device)
|
||||
total = total + self._regular_squared_norm(self._regular[group])
|
||||
total = total + self._owner_only_squared_norm(self._owner_only[group])
|
||||
total = total + self._dtensor_squared_norm(self._dtensor[group])
|
||||
return total
|
||||
|
||||
def norm(self, group: str) -> torch.Tensor:
|
||||
return self.squared_norm(group).sqrt()
|
||||
|
||||
def _regular_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor:
|
||||
if not tensors:
|
||||
return torch.zeros((), dtype=torch.float32, device=self.device)
|
||||
per = self._foreach_norm(tensors)
|
||||
return self._sum_norms_squared(per)
|
||||
|
||||
def _owner_only_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor:
|
||||
local_sq = self._regular_squared_norm(tensors)
|
||||
if self.reduce_tensor_fn is None:
|
||||
return local_sq
|
||||
return self.reduce_tensor_fn(local_sq, average=False)
|
||||
|
||||
def _dtensor_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor:
|
||||
if not tensors:
|
||||
return torch.zeros((), dtype=torch.float32, device=self.device)
|
||||
|
||||
per = self._foreach_norm(tensors)
|
||||
|
||||
# Group by (mesh, placements) so each unique layout costs only one
|
||||
# collective (redistribute).
|
||||
groups: DefaultDict[object, list[torch.Tensor]] = defaultdict(list)
|
||||
for norm in per:
|
||||
if self._is_dtensor(norm):
|
||||
key = (id(norm.device_mesh), tuple(norm.placements))
|
||||
else:
|
||||
key = None
|
||||
groups[key].append(norm)
|
||||
|
||||
parts: list[torch.Tensor] = []
|
||||
for group_tensors in groups.values():
|
||||
try:
|
||||
stacked = torch.stack(group_tensors)
|
||||
local = self._to_local_replicated(stacked).to(
|
||||
device=self.device, dtype=torch.float32
|
||||
)
|
||||
parts.append(local.pow(2).sum())
|
||||
except (RuntimeError, NotImplementedError):
|
||||
# Stack/redistribute not supported for this layout in the
|
||||
# current PyTorch version; fall back to per-tensor for safety.
|
||||
for norm in group_tensors:
|
||||
local = self._to_local_replicated(norm).to(
|
||||
device=self.device, dtype=torch.float32
|
||||
)
|
||||
parts.append(local.pow(2).sum())
|
||||
|
||||
return torch.stack(parts).sum()
|
||||
|
||||
def _foreach_norm(self, tensors: list[torch.Tensor]) -> list[torch.Tensor]:
|
||||
try:
|
||||
return torch._foreach_norm(tensors, 2.0)
|
||||
except (RuntimeError, NotImplementedError):
|
||||
return [torch.linalg.vector_norm(t, 2.0) for t in tensors]
|
||||
|
||||
def _sum_norms_squared(self, norms: list[torch.Tensor]) -> torch.Tensor:
|
||||
if not norms:
|
||||
return torch.zeros((), dtype=torch.float32, device=self.device)
|
||||
local_norms = [
|
||||
self._to_local_replicated(norm).to(device=self.device, dtype=torch.float32)
|
||||
for norm in norms
|
||||
]
|
||||
return torch.stack(local_norms).pow(2).sum()
|
||||
|
||||
def _to_local_replicated(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
if not self._is_dtensor(tensor):
|
||||
return tensor.to(device=self.device)
|
||||
mesh = tensor.device_mesh
|
||||
tensor = tensor.redistribute(placements=[Replicate()] * mesh.ndim)
|
||||
return tensor.to_local().to(device=self.device)
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Training metric buffering and wandb emission helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from typing import Any, Deque, Dict, List, Optional, Tuple
|
||||
|
||||
|
||||
class MetricsLogger:
|
||||
"""Buffer per-step stats and emit averaged batches to wandb.
|
||||
|
||||
Also formats the per-step console line via ``format_training_line``,
|
||||
which consumes an ``(stat_key, pretty_label, fmt_spec)`` triple list
|
||||
from ``adapter.console_fields()``.
|
||||
|
||||
``smooth_window`` controls a separate rolling-mean buffer used purely
|
||||
for display smoothing (console + tqdm). It is independent of the
|
||||
``log_interval`` wandb-flush buffer; the smoothed dict is returned by
|
||||
:meth:`smooth` for the caller to assign to ``_current_step_stats``.
|
||||
1 (default) preserves historical per-step values.
|
||||
"""
|
||||
|
||||
def __init__(self, *, wandb_run, log_interval: int, smooth_window: int = 1):
|
||||
self._wandb_run = wandb_run
|
||||
self._log_interval = log_interval
|
||||
self._buffer: List[Dict[str, Any]] = []
|
||||
self._smooth_window = max(1, int(smooth_window))
|
||||
self._smooth_buffers: Dict[str, Deque[float]] = {}
|
||||
|
||||
def record_step(self, step_stats: Dict[str, Any], *, is_main: bool) -> None:
|
||||
"""Push a per-step stats dict into the rank-0 buffer (no-op off rank-0)."""
|
||||
if is_main:
|
||||
self._buffer.append(step_stats)
|
||||
|
||||
def flush_if_due(self, global_step: int) -> Optional[Dict[str, Any]]:
|
||||
"""At log_interval boundaries, average + emit; return avg dict or None."""
|
||||
if global_step % self._log_interval != 0 or not self._buffer:
|
||||
return None
|
||||
avg = self._average(self._buffer)
|
||||
if self._wandb_run is not None and hasattr(self._wandb_run, "log"):
|
||||
self._wandb_run.log(avg, step=global_step)
|
||||
self._buffer = []
|
||||
return avg
|
||||
|
||||
def smooth(self, stats: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Return *stats* with numeric entries replaced by their rolling mean.
|
||||
|
||||
Uses one ``deque(maxlen=smooth_window)`` per metric key. Non-numeric
|
||||
entries pass through unchanged. When ``smooth_window <= 1`` this is
|
||||
the identity (historical behavior).
|
||||
"""
|
||||
if self._smooth_window <= 1:
|
||||
return stats
|
||||
out: Dict[str, Any] = {}
|
||||
for key, val in stats.items():
|
||||
if not isinstance(val, (int, float)):
|
||||
out[key] = val
|
||||
continue
|
||||
buf = self._smooth_buffers.get(key)
|
||||
if buf is None or buf.maxlen != self._smooth_window:
|
||||
buf = deque(maxlen=self._smooth_window)
|
||||
self._smooth_buffers[key] = buf
|
||||
buf.append(float(val))
|
||||
out[key] = sum(buf) / len(buf)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _average(buf: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
all_keys: set = set()
|
||||
for stats in buf:
|
||||
all_keys.update(stats.keys())
|
||||
avg: Dict[str, Any] = {}
|
||||
for key in all_keys:
|
||||
values = [s[key] for s in buf if key in s]
|
||||
if values:
|
||||
avg[key] = sum(values) / len(values)
|
||||
return avg
|
||||
|
||||
def format_training_line(
|
||||
self,
|
||||
*,
|
||||
epoch: int,
|
||||
total_epoch: int,
|
||||
current_iter: int,
|
||||
total_iter: int,
|
||||
loss: float,
|
||||
lr: float,
|
||||
time_per_step: float,
|
||||
stats: Dict[str, Any],
|
||||
fields: List[Tuple[str, str, str]],
|
||||
mfu_info: Optional[Dict[str, Any]] = None,
|
||||
) -> str:
|
||||
"""Render a line matching pre-refactor training_log output exactly.
|
||||
|
||||
Layout: " epoch E/T | iter I/T | loss L | [adapter fields] | lr L |
|
||||
time_current_backward_step Ts | [MFU P% |]"
|
||||
"""
|
||||
parts = [
|
||||
" epoch {:3d}/{:3d} |".format(epoch, total_epoch),
|
||||
" iter {:6d}/{:6d} |".format(current_iter, total_iter),
|
||||
" loss {:.6f} |".format(loss),
|
||||
]
|
||||
for key, label, fmt in fields:
|
||||
value = stats.get(key)
|
||||
if value is not None:
|
||||
parts.append(" {} {:{fmt}} |".format(label, value, fmt=fmt))
|
||||
parts.append(" lr {:.6f} |".format(lr))
|
||||
parts.append(" time_current_backward_step {:.6f}s |".format(time_per_step))
|
||||
if mfu_info is not None:
|
||||
parts.append(" MFU {:.2f}% |".format(mfu_info["mfu"] * 100))
|
||||
return "".join(parts)
|
||||
@@ -0,0 +1,181 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
FSDP Training Entry Point
|
||||
Launch with torchrun:
|
||||
torchrun --nproc_per_node=8 --master_port=29500 train_fsdp.py --config config.yaml
|
||||
"""
|
||||
import argparse
|
||||
import dataclasses
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
import torch
|
||||
import wandb
|
||||
|
||||
from wall_x.config.loader import load_config as load_typed_config
|
||||
from wall_x.trainer.fsdp_trainer import (
|
||||
FSDPTrainer,
|
||||
cleanup_distributed,
|
||||
is_main_process,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TeeOutput:
|
||||
"""Write output to both the terminal and a file."""
|
||||
|
||||
def __init__(self, file_path, mode="w"):
|
||||
self.terminal = sys.stdout
|
||||
self.log = open(file_path, mode)
|
||||
|
||||
def write(self, message):
|
||||
self.terminal.write(message)
|
||||
self.log.write(message)
|
||||
self.log.flush()
|
||||
|
||||
def flush(self):
|
||||
self.terminal.flush()
|
||||
self.log.flush()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="FSDP Training Script")
|
||||
parser.add_argument(
|
||||
"--config",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to training config YAML file",
|
||||
)
|
||||
# FSDP specific overrides
|
||||
parser.add_argument(
|
||||
"--fsdp_sharding_strategy",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["full_shard", "shard_grad_op", "no_shard", "hybrid_shard"],
|
||||
help="FSDP sharding strategy (overrides config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable debug mode: reduces buffer size, sets log/save path to debug.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--wandb_offline",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Whether to run wandb in offline mode (overrides config).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--visualize",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Whether to visualize samples during training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log_to_file",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Whether to redirect stdout and stderr to log files.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_config(config_path: str, cli_args=None):
|
||||
"""Load configuration from YAML file into TrainConfig."""
|
||||
return load_typed_config(config_path, cli_args=cli_args)
|
||||
|
||||
|
||||
def setup_logger(cfg):
|
||||
"""Setup wandb logger if enabled"""
|
||||
log = cfg.logging
|
||||
if log.use_wandb and is_main_process():
|
||||
logger.info(
|
||||
"rank %s is initializing wandb , is main process %s",
|
||||
torch.distributed.get_rank(),
|
||||
is_main_process(),
|
||||
)
|
||||
wandb_run = wandb.init(
|
||||
project=log.log_project,
|
||||
name=log.log_name,
|
||||
entity=log.log_entity,
|
||||
config=dataclasses.asdict(cfg),
|
||||
save_code=False,
|
||||
force=False,
|
||||
mode="offline" if log.wandb_offline else "online",
|
||||
)
|
||||
logger.info("Wandb Initialized")
|
||||
return wandb_run
|
||||
return None
|
||||
|
||||
|
||||
def print_fsdp_config(cfg):
|
||||
"""Print FSDP configuration"""
|
||||
if is_main_process():
|
||||
dist = cfg.distributed
|
||||
logger.info("%s", "=" * 60)
|
||||
logger.info("FSDP Configuration:")
|
||||
logger.info(" use_fsdp: %s", dist.use_fsdp)
|
||||
logger.info("%s", "=" * 60)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
# Redirect logs to files only on the main process to avoid write races.
|
||||
if args.log_to_file:
|
||||
import yaml as _yaml
|
||||
|
||||
with open(args.config, "r") as f:
|
||||
_tmp = _yaml.safe_load(f)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
# support both old (save_path) and new (checkpoint.save_path) schema
|
||||
log_dir = _tmp.get("save_path") or _tmp.get("checkpoint", {}).get(
|
||||
"save_path", "./ckpt"
|
||||
)
|
||||
if args.debug:
|
||||
log_dir = "./ckpt/debug"
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
log_file = os.path.join(log_dir, f"training_log_{timestamp}.log")
|
||||
sys.stdout = TeeOutput(log_file, mode="w")
|
||||
sys.stderr = TeeOutput(log_file.replace(".log", "_stderr.log"), mode="w")
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
|
||||
logger.info("\n%s", "=" * 80)
|
||||
logger.info("LOG TO FILE MODE: All output will be saved to:")
|
||||
logger.info(" STDOUT: %s", log_file)
|
||||
logger.info(" STDERR: %s", log_file.replace(".log", "_stderr.log"))
|
||||
logger.info("%s\n", "=" * 80)
|
||||
else:
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
|
||||
|
||||
torch.cuda.init()
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", 0))
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.cuda.set_device(device)
|
||||
torch.distributed.init_process_group("nccl", device_id=device)
|
||||
cfg = load_config(args.config, cli_args=args)
|
||||
wandb_run = setup_logger(cfg)
|
||||
print_fsdp_config(cfg)
|
||||
|
||||
try:
|
||||
trainer = FSDPTrainer(
|
||||
train_config=cfg,
|
||||
wandb_run=wandb_run,
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Training failed with error: %s", e)
|
||||
raise
|
||||
|
||||
finally:
|
||||
cleanup_distributed()
|
||||
if wandb_run is not None:
|
||||
wandb_run.finish()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user