"""The training step loop and the three step bodies it dispatches to.
``run_training_loop`` owns everything that repeats: the step body, NaN
handling, optimizer/scheduler advance, freeze and phase transitions,
metrics, eval, checkpointing, and shutdown. The step body itself is a
plain callable (``StepFn``) chosen once by :func:`select_step_fn`, so a
model family with a different forward contract is a new function rather
than another branch inside the loop.
"""
from __future__ import annotations
import contextlib
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
import torch
import torch.distributed as dist
from kempnerforge.checkpoint.manager import CheckpointManager
from kempnerforge.config.job import JobConfig
from kempnerforge.distributed.utils import clip_grad_norm_
from kempnerforge.metrics.logger import get_logger
from kempnerforge.metrics.tracker import MetricsTracker
from kempnerforge.model.vlm import inner_transformer
from kempnerforge.profiling.profiler import print_profiler_summary
from kempnerforge.resilience.health import NaNDetector, check_nccl_health, params_are_finite
from kempnerforge.resilience.signal_handler import ShutdownHandler
from kempnerforge.training.data_pipeline import DataPipeline, PhaseState, advance_phases
from kempnerforge.training.eval import run_eval
from kempnerforge.training.freeze import apply_freeze_specs, freeze_meta_at_step
from kempnerforge.training.grad import maybe_no_sync
from kempnerforge.training.hooks import HookRunner, StepContext
from kempnerforge.training.optimizer import module_learning_rates
from kempnerforge.training.runtime import PipelineBundle, RuntimeContext, pp_group
if TYPE_CHECKING:
from kempnerforge.model.transformer import Transformer
logger = get_logger(__name__)
LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
_WORKER_KILL_TIMEOUT_S = 5.0
[docs]
class BatchStream:
"""The training dataloader plus its live iterator.
Wraps the epoch-boundary restart and the "no dataloader configured"
case so the step bodies don't each re-implement them. The loader is read
off the pipeline on every access, so swapping ``pipeline.dataloader``
takes effect on the next step instead of leaving a stale snapshot.
"""
[docs]
def __init__(self, pipeline: DataPipeline) -> None:
self.pipeline = pipeline
self._iter: Any = None
self._source: Any = None
@property
def dataloader(self) -> Any:
return self.pipeline.dataloader
@property
def has_data(self) -> bool:
return self.dataloader is not None
[docs]
def ensure_started(self) -> None:
"""Materialize the iterator, if it isn't already, for the loader in use.
``next_batch`` does this itself; the loop calls it up front so building
an iterator (which spawns dataloader workers) stays outside the region
``MetricsTracker`` times.
"""
loader = self.dataloader
if loader is not None and (self._iter is None or self._source is not loader):
self._iter = iter(loader)
self._source = loader
[docs]
def reset(self) -> None:
"""Drop the iterator so the next step takes a fresh one from the loader.
A ``StatefulDataLoader`` re-applies its recorded skip, so this picks up
inside the current epoch rather than restarting it.
"""
self._iter = None
[docs]
def close(self) -> None:
"""Shut the loader's worker processes down explicitly.
Workers ignore SIGTERM (see ``ignore_shutdown_signals_in_worker``) so a
preemption signal cannot kill them out from under an in-flight step.
The same shield means ``Process.terminate()``, which is SIGTERM, cannot
reap them either: left to the interpreter's exit hook, the process would
block forever joining daemonic children that refuse to die. Draining
them through the loader's own sentinel path avoids that.
Rank-local and idempotent — it issues no collectives, so it is safe on
the exception path in ``train_loop``'s ``finally``.
"""
self._iter = None
self._source = None
loader = self.pipeline.dataloader
if loader is None:
return
inner = getattr(loader, "_dataloader", loader)
iterator = getattr(inner, "_iterator", None)
if iterator is not None and hasattr(iterator, "_shutdown_workers"):
iterator._shutdown_workers()
for worker in getattr(iterator, "_workers", ()):
if worker.is_alive():
worker.kill()
worker.join(timeout=_WORKER_KILL_TIMEOUT_S)
if hasattr(inner, "_iterator"):
inner._iterator = None
[docs]
def next_batch(self) -> dict[str, torch.Tensor]:
if self.dataloader is None:
raise RuntimeError("BatchStream has no dataloader; check has_data first")
self.ensure_started()
try:
return next(self._iter)
except StopIteration:
pass
self._iter = iter(self.dataloader)
self._source = self.dataloader
try:
return next(self._iter)
except StopIteration:
raise RuntimeError(
"dataloader yielded no batches for this rank: the dataset is "
"smaller than batch_size x dp_size, a streaming shard is "
"exhausted, or a mixture phase weighted this rank's only "
"source to zero"
) from None
[docs]
@dataclass
class StepResult:
"""What one training step reports back to the loop.
The per-dataset fields stay empty for step bodies that don't produce
them; ``text_tokens`` stays None, which a VLM run treats as an error
rather than logging a zero. ``tokens`` is what this rank ran when the
batch decides it (packed VLM rows), counted as row slots like the unpacked
``batch_size × seq_len``, image slots and padding included; None means
the configured batch per micro-step.
"""
loss: float
grad_norm: float
text_tokens: int | None = None
dataset_token_counts: dict[str, int] = field(default_factory=dict)
dataset_loss_sums: dict[str, float] = field(default_factory=dict)
dataset_loss_counts: dict[str, int] = field(default_factory=dict)
tokens: int | None = None
StepFn = Callable[["TrainingSession", int], StepResult]
[docs]
@dataclass
class TrainingSession:
"""Everything one training run needs after the build phase."""
config: JobConfig
runtime: RuntimeContext
model: torch.nn.Module
optimizer: torch.optim.Optimizer
scheduler: Any
loss_fn: LossFn
step_fn: StepFn
data: DataPipeline
phases: PhaseState
checkpointer: CheckpointManager
tracker: MetricsTracker
hooks: HookRunner
nan_detector: NaNDetector
shutdown_handler: ShutdownHandler
pipeline: PipelineBundle | None = None
eval_dataloader: Any | None = None
profiler: Any | None = None
autocast_dtype: torch.dtype | None = None
_batch_stream: BatchStream | None = field(default=None, init=False, repr=False)
@property
def batches(self) -> BatchStream:
"""The live batch iterator. Rebuilt if ``data`` is replaced wholesale."""
if self._batch_stream is None or self._batch_stream.pipeline is not self.data:
self._batch_stream = BatchStream(self.data)
return self._batch_stream
[docs]
def autocast(self) -> contextlib.AbstractContextManager:
"""``torch.autocast`` to ``autocast_dtype``, or a no-op when it is None."""
if self.autocast_dtype is None:
return contextlib.nullcontext()
return torch.autocast(device_type=self.runtime.device.type, dtype=self.autocast_dtype)
[docs]
def clip_grads(model: torch.nn.Module, max_norm: float, pp_group: Any = None) -> float:
grad_norm = clip_grad_norm_(model, max_norm, pp_group=pp_group)
return grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
[docs]
def add_moe_losses(loss: torch.Tensor, model: torch.nn.Module, mc: Any) -> torch.Tensor:
"""MoE auxiliary + router z-loss terms (only called when ``mc.is_moe``)."""
aux_loss = inner_transformer(model).get_moe_aux_loss() # type: ignore[attr-defined]
loss = loss + mc.moe_aux_loss_weight * aux_loss
if mc.moe_router_z_loss_weight > 0:
z = inner_transformer(model).get_moe_router_z_loss() # type: ignore[attr-defined]
loss = loss + mc.moe_router_z_loss_weight * z
return loss
def _expert_balance(routers: list[tuple[str | None, torch.Tensor]]) -> dict[str, float]:
"""``moe/expert_balance`` over every router, plus one value per MoT modality.
Each router's expert counts are divided by that router's own total before
taking min over max, so routers that see different token counts (a MoT
model's image and text routers) compare as loads. Every flat MoE router sees
the same tokens, so for a text model this equals min over max of the raw
counts. A router whose counts sum to zero is skipped.
Args:
routers: ``(modality, counts)`` per router; ``modality`` is ``None`` for
a flat MoE layer.
Returns:
``moe/expert_balance``, and ``moe/expert_balance/<modality>`` for each
MoT modality; empty when no router saw a token.
"""
loads: list[tuple[str | None, torch.Tensor]] = []
for modality, counts in routers:
total = counts.float().sum()
if total.item() > 0:
loads.append((modality, counts.float() / total))
if not loads:
return {}
groups = {"moe/expert_balance": [load for _, load in loads]}
for modality in sorted({m for m, _ in loads if m is not None}):
groups[f"moe/expert_balance/{modality}"] = [load for m, load in loads if m == modality]
balance = {}
for key, group in groups.items():
values = torch.cat(group)
balance[key] = (values.min() / values.max()).item()
return balance
[docs]
def pipeline_step(session: TrainingSession, step: int) -> StepResult:
"""PP step: collect microbatches, hand them to the schedule, broadcast loss.
``step`` is unused — MoE (the only step-dependent branch) is rejected
with PP in ``JobConfig.validate``.
"""
config = session.config
tc, mc = config.train, config.model
device = session.runtime.device
model = session.model
batches = session.batches
pipeline = session.pipeline
if pipeline is None:
raise RuntimeError("pipeline_step requires TrainingSession.pipeline")
input_ids_list, labels_list, doc_ids_list = [], [], []
for _ in range(tc.grad_accum_steps):
if batches.has_data:
batch = batches.next_batch()
input_ids_list.append(batch["input_ids"].to(device))
labels_list.append(batch["labels"].to(device))
if "doc_ids" in batch:
doc_ids_list.append(batch["doc_ids"].to(device))
else:
input_ids_list.append(
torch.randint(0, mc.vocab_size, (tc.batch_size, tc.seq_len), device=device)
)
labels_list.append(
torch.randint(0, mc.vocab_size, (tc.batch_size, tc.seq_len), device=device)
)
full_input = torch.cat(input_ids_list, dim=0)
full_labels = torch.cat(labels_list, dim=0)
pp_kwargs = {"doc_ids": torch.cat(doc_ids_list, dim=0)} if doc_ids_list else {}
is_first = pipeline.rank == 0
is_last = pipeline.rank == pipeline.size - 1
pp_losses: list[torch.Tensor] = []
with session.autocast():
if is_first:
pipeline.schedule.step(full_input, target=full_labels, losses=pp_losses, **pp_kwargs)
elif is_last:
pipeline.schedule.step(target=full_labels, losses=pp_losses, **pp_kwargs)
else:
pipeline.schedule.step(**pp_kwargs)
if is_last and pp_losses:
avg_loss = sum(loss.item() for loss in pp_losses) / len(pp_losses)
else:
avg_loss = 0.0
group = pp_group(session.runtime)
grad_norm_val = clip_grads(model, tc.grad_clip_norm, group)
loss_tensor = torch.tensor([avg_loss], device=device)
dist.broadcast(loss_tensor, group_src=pipeline.size - 1, group=group)
return StepResult(loss=loss_tensor.item(), grad_norm=grad_norm_val)
_PACKED_BATCH_KEYS = (
"doc_ids",
"position_ids",
"attention_group_ids",
"image_rows",
"image_starts",
"text_rows",
"text_positions",
)
[docs]
def vlm_step(session: TrainingSession, step: int) -> StepResult:
"""VLM step (no PP): pixel_values + input_ids forward, with a text-token count.
A packed batch (one with ``image_rows``) also carries ``doc_ids``,
``position_ids``, the image placement (``image_starts``, ``image_tokens``)
and the text positions, all passed through to the wrapper, with
``attention_group_ids`` when a collator emits them. Its row count varies,
so the step reports the tokens it ran and marks the leading dims dynamic
for ``torch.compile``.
"""
config = session.config
tc, mc = config.train, config.model
device = session.runtime.device
model = session.model
batches = session.batches
total_loss = 0.0
total_text_tokens = 0
packed_tokens: int | None = None
for micro_step in range(tc.grad_accum_steps):
if not batches.has_data:
raise RuntimeError(
"VLM training requires a real dataloader; synthetic fallback "
"(randint) does not produce pixel_values"
)
batch = batches.next_batch()
pixel_values = batch["pixel_values"].to(device)
input_ids = batch["input_ids"].to(device)
labels = batch["labels"].to(device)
frame_mask = batch["frame_mask"].to(device) if "frame_mask" in batch else None
extras: dict[str, Any] = {
key: batch[key].to(device) for key in _PACKED_BATCH_KEYS if key in batch
}
if "image_tokens" in batch:
extras["image_tokens"] = int(batch["image_tokens"])
if "image_rows" in extras:
packed_tokens = (packed_tokens or 0) + input_ids.numel()
if tc.compile_model:
_mark_rows_dynamic(pixel_values, input_ids, labels, *extras.values())
with maybe_no_sync(model, micro_step, tc.grad_accum_steps):
if mc.is_moe:
inner_transformer(model).set_moe_step(step, tc.max_steps) # type: ignore[attr-defined]
with session.autocast():
logits, labels_out = model(
pixel_values, input_ids, labels, frame_mask=frame_mask, **extras
)
loss = session.loss_fn(logits, labels_out)
total_text_tokens += int((labels_out != -100).sum().item())
if mc.is_moe:
loss = add_moe_losses(loss, model, mc)
scaled_loss = loss / tc.grad_accum_steps
scaled_loss.backward()
total_loss += loss.item()
return StepResult(
loss=total_loss / tc.grad_accum_steps,
grad_norm=clip_grads(model, tc.grad_clip_norm),
text_tokens=total_text_tokens,
tokens=packed_tokens,
)
def _mark_rows_dynamic(*tensors: Any) -> None:
"""Mark each tensor's leading dim dynamic so a new packed row count does not recompile."""
from torch._dynamo import maybe_mark_dynamic
for tensor in tensors:
if isinstance(tensor, torch.Tensor):
maybe_mark_dynamic(tensor, 0)
def _sum_over_data_ranks(session: TrainingSession, value: int) -> int:
"""Sum a per-rank count over the data-parallel ranks; tp and pp peers hold the same batch."""
if not dist.is_initialized():
return value * session.data.dp_size
t = torch.tensor([value], device=session.runtime.device, dtype=torch.long)
dist.all_reduce(t, op=dist.ReduceOp.SUM)
return int(t.item()) * session.data.dp_size // session.runtime.world_size
def _tokens_in_step(session: TrainingSession, result: StepResult) -> int:
"""Tokens the step ran across the data-parallel ranks, from the step's count or the config."""
tc = session.config.train
if result.tokens is None:
return tc.batch_size * tc.seq_len * tc.grad_accum_steps * session.data.dp_size
return _sum_over_data_ranks(session, result.tokens)
[docs]
def text_step(session: TrainingSession, step: int) -> StepResult:
"""Standard step (no PP, text-only), with optional per-dataset mixture metrics."""
config = session.config
tc, mc = config.train, config.model
device = session.runtime.device
model = session.model
batches = session.batches
mixture_dataset = session.data.mixture_dataset
total_loss = 0.0
ds_token_counts: dict[str, int] = {}
ds_loss_sums: dict[str, float] = {}
ds_loss_counts: dict[str, int] = {}
for micro_step in range(tc.grad_accum_steps):
batch: dict[str, torch.Tensor] = {}
if batches.has_data:
batch = batches.next_batch()
input_ids = batch["input_ids"].to(device)
labels = batch["labels"].to(device)
doc_ids = batch["doc_ids"].to(device) if "doc_ids" in batch else None
else:
input_ids = torch.randint(0, mc.vocab_size, (tc.batch_size, tc.seq_len), device=device)
labels = torch.randint(0, mc.vocab_size, (tc.batch_size, tc.seq_len), device=device)
doc_ids = None
with maybe_no_sync(model, micro_step, tc.grad_accum_steps):
if mc.is_moe:
inner_transformer(model).set_moe_step(step, tc.max_steps) # type: ignore[attr-defined]
with session.autocast():
logits = model(input_ids, doc_ids=doc_ids)
loss = session.loss_fn(logits, labels)
if mixture_dataset is not None and "dataset_idx" in batch:
ds_idx = batch["dataset_idx"]
with torch.no_grad():
for i, name in enumerate(mixture_dataset.dataset_names):
mask = ds_idx == i
count = int(mask.sum().item())
if count > 0:
ds_token_counts[name] = (
ds_token_counts.get(name, 0) + count * tc.seq_len
)
ds_l = torch.nn.functional.cross_entropy(
logits[mask].reshape(-1, logits.size(-1)),
labels[mask].reshape(-1),
ignore_index=-100,
).item()
ds_loss_sums[name] = ds_loss_sums.get(name, 0) + ds_l
ds_loss_counts[name] = ds_loss_counts.get(name, 0) + 1
if mc.is_moe:
loss = add_moe_losses(loss, model, mc)
scaled_loss = loss / tc.grad_accum_steps
scaled_loss.backward()
total_loss += loss.item()
return StepResult(
loss=total_loss / tc.grad_accum_steps,
grad_norm=clip_grads(model, tc.grad_clip_norm),
dataset_token_counts=ds_token_counts,
dataset_loss_sums=ds_loss_sums,
dataset_loss_counts=ds_loss_counts,
)
[docs]
def select_step_fn(config: JobConfig) -> StepFn:
"""Pick the step body for this job: PP wins, then VLM, then text-only."""
if config.distributed.pp > 1:
return pipeline_step
if config.is_vlm:
return vlm_step
return text_step
def _log_periodic_metrics(session: TrainingSession, result: StepResult, step: int) -> None:
"""Per-module LR / MoE / VLM-text-token / per-dataset series, on the main metrics interval."""
config = session.config
mc = config.model
tracker = session.tracker
model = session.model
module_lrs = module_learning_rates(session.optimizer)
if module_lrs:
tracker.log_eval(module_lrs, step)
if mc.is_moe:
inner = cast("Transformer", inner_transformer(model))
moe_metrics = {"moe/aux_loss": inner.get_moe_aux_loss().item()}
moe_metrics["moe/router_z_loss"] = inner.get_moe_router_z_loss().item()
routers: list[tuple[str | None, torch.Tensor]] = [
(None, counts) for counts in inner.get_expert_counts().values()
]
for per_modality in getattr(inner, "get_mot_expert_counts", dict)().values():
routers.extend(per_modality.items())
moe_metrics.update(_expert_balance(routers))
tracker.log_eval(moe_metrics, step)
if config.is_vlm:
if result.text_tokens is None:
raise RuntimeError(
"VLM run: the step body reported no text-token count, so "
"data/text_tokens_trained would be logged as zero"
)
global_text_tokens = _sum_over_data_ranks(session, result.text_tokens)
tracker.log_eval({"data/text_tokens_trained": float(global_text_tokens)}, step)
if session.data.mixture_dataset is not None and result.dataset_loss_sums:
ds_metrics: dict[str, float] = {}
for name in result.dataset_loss_sums:
ds_metrics[f"loss/{name}"] = (
result.dataset_loss_sums[name] / result.dataset_loss_counts[name]
)
for name, count in result.dataset_token_counts.items():
ds_metrics[f"data/{name}/tokens"] = float(count)
tracker.log_eval(ds_metrics, step)
def _advance_lr_schedule(session: TrainingSession) -> None:
"""Step the scheduler, then scale the new LR by the active phase's ``lr_scale``.
Runs once per step, a step skipped for a non-finite loss included, so
the schedule stays aligned with the step counter.
"""
session.scheduler.step()
if session.phases.lr_scale != 1.0:
for pg in session.optimizer.param_groups:
pg["lr"] *= session.phases.lr_scale
def _apply_freeze_stages(session: TrainingSession, step: int) -> None:
"""FreezeStage hook: apply any stage whose start_step matches this boundary.
AdamW + set_to_none=True default skips frozen params entirely (no SGD
step, no weight decay), so mutating requires_grad mid-training is a
clean no-op for newly-frozen params and re-enables gradient flow for
newly-unfrozen ones.
Async-save fence: drain any in-flight save FIRST so that its
metadata.json — which records the pre-transition spec — lands before we
flip requires_grad. Otherwise a save started at step S-1 could write
metadata after the transition, attaching the post-transition spec to
pre-transition shards.
"""
vlm_cfg = session.config.vlm
if not session.config.is_vlm or vlm_cfg is None or not vlm_cfg.freeze_schedule:
return
pending_stages = [s for s in vlm_cfg.freeze_schedule if s.start_step == step]
if not pending_stages:
return
session.checkpointer.flush_pending_save()
for stage in pending_stages:
flipped = apply_freeze_specs(session.model, stage.specs, vlm_cfg.module_patterns)
logger.info(f"FreezeStage at step={step}: applied {flipped}")
def _apply_step_transitions(session: TrainingSession, step: int) -> None:
"""Apply the freeze stages and data phases due at ``step``.
Runs right after the step counter advances, on a step skipped for a
non-finite loss too, so a transition lands on its configured step. A
phase that fires drops the data iterator, so the next batch is drawn
with the new sampler weights.
"""
_apply_freeze_stages(session, step)
if advance_phases(session.phases, session.data, step):
session.batches.reset()
def _shutdown_if_requested(
session: TrainingSession,
step: int,
tokens_seen: int,
*,
saved_this_step: bool,
nan_skipped: bool = False,
) -> bool:
"""Write the emergency checkpoint and finish when a shutdown is pending.
Returns True when the loop must stop. With ``saved_this_step``, the
scheduled save of ``step`` is the emergency checkpoint, so it is drained
with ``wait()`` rather than written twice. A ``nan_skipped`` step is
saved only if every rank's parameters are finite: after a bad update
(an optimizer step that left non-finite weights although its loss and
gradient norm were finite) they are not, and their checkpoint would
become ``latest`` and count toward ``keep_last_n``, which can prune the
last good one.
"""
if not session.shutdown_handler.should_shutdown():
return False
ckpt_mgr = session.checkpointer
if saved_this_step:
logger.warning(
f"Shutdown requested at step {step}; "
"its scheduled checkpoint is the emergency checkpoint"
)
ckpt_mgr.wait()
elif nan_skipped and not params_are_finite(session.model, session.runtime.device):
logger.error(
f"Shutdown requested at step {step}, skipped for a NaN/Inf loss, but the "
"parameters are not finite: saving no emergency checkpoint, so these weights "
"neither become `latest` nor count toward `checkpoint.keep_last_n`"
)
else:
logger.warning(f"Shutdown requested at step {step} — saving emergency checkpoint")
ckpt_mgr.save(
step=step,
tokens_seen=tokens_seen,
scheduler=session.scheduler,
dataloader=session.data.dataloader,
extra=checkpoint_extra(session.config, step, session.phases),
)
session.shutdown_handler.finish()
return True
[docs]
def run_training_loop(
session: TrainingSession, *, step: int = 0, tokens_seen: int = 0
) -> tuple[int, int]:
"""Run the step loop from ``step`` to ``train.max_steps``.
Returns the final ``(step, tokens_seen)``. Drains any pending async
checkpoint before returning so the caller can tear down the process
group safely. ``session.data`` is read once and must stay fixed for the
duration; only its ``dataloader`` may be swapped mid-run.
"""
config = session.config
tc = config.train
eval_config = config.eval
runtime = session.runtime
model = session.model
optimizer = session.optimizer
scheduler = session.scheduler
ckpt_mgr = session.checkpointer
tracker = session.tracker
hooks = session.hooks
data = session.data
phases = session.phases
prof = session.profiler
logger.info(
f"Starting training: step={step}, max_steps={tc.max_steps}, "
f"batch_size={tc.batch_size}, grad_accum={tc.grad_accum_steps}, "
f"world_size={runtime.world_size}"
)
if phases.phases:
logger.info(f"Phase scheduling: {len(phases.phases)} phase(s) configured")
model.train()
hooks.on_train_begin(config)
if prof is not None:
prof.start()
if step == 0 and config.checkpoint.is_dynamic_milestone(0):
ckpt_mgr.save(
step=0,
tokens_seen=0,
scheduler=scheduler,
dataloader=data.dataloader,
extra=checkpoint_extra(config, 0, phases),
)
hooks.on_checkpoint_save(0, config.checkpoint.dir)
completed_normally = False
ckpt_extra: dict = {}
try:
while step < tc.max_steps:
session.batches.ensure_started()
tracker.start_step()
result = session.step_fn(session, step)
tokens_in_step = _tokens_in_step(session, result)
if not session.nan_detector.check_loss(result.loss, step, grad_norm=result.grad_norm):
optimizer.zero_grad()
if session.nan_detector.should_rollback:
logger.error("Too many consecutive NaNs — stopping")
break
_advance_lr_schedule(session)
step += 1
_apply_step_transitions(session, step)
if _shutdown_if_requested(
session, step, tokens_seen, saved_this_step=False, nan_skipped=True
):
break
continue
optimizer.step()
_advance_lr_schedule(session)
optimizer.zero_grad()
step += 1
tokens_seen += tokens_in_step
_apply_step_transitions(session, step)
current_lr = optimizer.param_groups[0]["lr"]
step_metrics = tracker.end_step(
step=step,
loss=result.loss,
grad_norm=result.grad_norm,
lr=current_lr,
tokens_in_step=tokens_in_step,
)
hooks.on_step_end(
StepContext(
step=step,
loss=result.loss,
grad_norm=result.grad_norm,
lr=current_lr,
tokens_seen=tokens_seen,
model=model,
optimizer=optimizer,
)
)
if step_metrics is not None:
_log_periodic_metrics(session, result, step)
if (
tc.nccl_health_check_interval > 0
and step % tc.nccl_health_check_interval == 0
and not check_nccl_health()
):
logger.error(f"NCCL health check failed at step {step} — stopping")
break
if (
eval_config.enabled
and session.eval_dataloader is not None
and step % eval_config.interval == 0
):
pipeline = session.pipeline
with session.autocast():
eval_metrics = run_eval(
model,
session.eval_dataloader,
session.loss_fn,
runtime.device,
eval_config.steps,
pp_schedule=pipeline.schedule if pipeline is not None else None,
pp_rank=pipeline.rank if pipeline is not None else None,
pp_size=pipeline.size if pipeline is not None else None,
pp_group=pp_group(runtime) if pipeline is not None else None,
pp_n_microbatches=tc.grad_accum_steps if pipeline is not None else None,
)
tracker.log_eval(eval_metrics, step)
hooks.on_eval_end(eval_metrics, step)
if prof is not None:
prof.step()
ckpt_extra = checkpoint_extra(config, step, phases)
saved_this_step = config.checkpoint.should_save(step)
if saved_this_step:
ckpt_mgr.save(
step=step,
tokens_seen=tokens_seen,
scheduler=scheduler,
dataloader=data.dataloader,
extra=ckpt_extra,
)
hooks.on_checkpoint_save(step, config.checkpoint.dir)
if _shutdown_if_requested(session, step, tokens_seen, saved_this_step=saved_this_step):
break
if step >= tc.max_steps:
completed_normally = True
finally:
if prof is not None:
prof.stop()
if runtime.rank == 0:
print_profiler_summary(prof, trace_dir=config.profiling.trace_dir)
session.batches.close()
if completed_normally and not config.checkpoint.should_save(step):
ckpt_mgr.save(
step=step,
tokens_seen=tokens_seen,
scheduler=scheduler,
dataloader=data.dataloader,
extra=ckpt_extra,
)
hooks.on_checkpoint_save(step, config.checkpoint.dir)
ckpt_mgr.wait()
logger.info(f"Training complete: {step} steps, {tokens_seen:,} tokens")
hooks.on_train_end(step, tokens_seen)
return step, tokens_seen