964 lines
37 KiB
Python
964 lines
37 KiB
Python
"""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
|