Data Flow

The path a batch of tokens takes from the dataloader to a committed gradient update. Follow this page with kempnerforge/training/loop.py open in another tab — the structure below maps one-to-one onto the training loop body.

One-slide view

   ┌─ MemoryMappedDataset or MixtureDataset
   │
   ├─ DistributedSampler / MixtureSampler (rank-partitioned indices)
   │
   ├─ StatefulDataLoader ──► batch = {input_ids, labels, doc_ids?}
   │
   ▼
for micro_step in range(grad_accum_steps):       ── maybe_no_sync sets
    ├─ model.set_moe_step(step, max_steps)          DDP/FSDP grad-sync off
    ├─ logits = model(input_ids, doc_ids=doc_ids)   on all but last micro
    ├─ loss   = loss_fn(logits, labels)
    ├─ [+ moe_aux_loss_weight * model.get_moe_aux_loss()]
    └─ (loss / grad_accum_steps).backward()       ── gradient accumulation

grad_norm = clip_grad_norm_(model, grad_clip_norm)
if NaN: zero_grad, maybe stop, skip the update; still run scheduler.step(),
        freeze stages / phases and the shutdown check
optimizer.step() ; scheduler.step() ; optimizer.zero_grad()

tracker.end_step(step, loss, grad_norm, lr, tokens_in_step)
hook_runner.on_step_end(StepContext(...))
[every N steps] eval, NCCL health, profiler.step, ckpt_mgr.save
if shutdown_handler.should_shutdown(): save (unless saved above) + break

Startup, once

Before the loop starts, run_training initializes the collaborators:

  • init_distributed(config.distributed, seed=...) from kempnerforge.distributed.setup — reads RANK / LOCAL_RANK / WORLD_SIZE (torchrun) or SLURM_PROCID / SLURM_NTASKS (multi-node srun), calls dist.init_process_group, builds the DeviceMesh, seeds torch.

  • ShutdownHandler from kempnerforge.resilience.signal_handler — installs SIGTERM / SIGUSR1 handlers for cooperative SLURM preemption.

  • NaNDetector from kempnerforge.resilience.health — tracks consecutive NaN steps; action="warn" by default. After max_consecutive=10 in a row it sets should_rollback and the loop stops; rolling back is manual.

  • Loss function from the registry ("cross_entropy" or "chunked_cross_entropy").

  • Model via build_parallel_model — applies the full parallelism stack (see Parallelism order).

  • Optimizer and scheduler from their registries.

  • CheckpointManager from kempnerforge.checkpoint.manager — owns DCP async save/load, latest-symlink, step_N directories.

  • resolve_resume_path(config.checkpoint.dir) from kempnerforge.resilience.elastic — follows the latest symlink or picks the newest step_N whose save completed. If non-None, ckpt_mgr.load(resume_path) restores model, optimizer, scheduler, dataloader position, and RNG state before the loop starts.

  • MetricsTracker from kempnerforge.metrics.tracker — per-step metrics + EMA smoothing + WandB / TensorBoard backends.

  • Data pipeline — MemoryMappedDataset or MixtureDataset or an HF streaming / eager dataset (config dispatches), wrapped by DistributedSampler or MixtureSampler and then by StatefulDataLoader from kempnerforge.data.

  • Eval dataloader and torch.profiler (optional).

  • Phase schedule (optional) — for curriculum training, rebalances the mixture at phase.start_step.

Inside the loop: one step

The step body is where the interesting routing happens.

1 · Microbatch fetch

for micro_step in range(grad_accum_steps):
    batch = next(data_iter)
    input_ids = batch["input_ids"].to(device)
    labels    = batch["labels"].to(device)
    doc_ids   = batch.get("doc_ids").to(device) if "doc_ids" in batch else None

When dataloader is None, the loop generates random integer tokens — useful for smoke-testing the parallelism stack without any corpus.

doc_ids is optional: non-None only when the dataset packs multiple documents into one sequence. It triggers the block-diagonal attention mask path (see Model § Attention paths).

2 · maybe_no_sync

with maybe_no_sync(model, micro_step, grad_accum_steps):
    ...

Utility from kempnerforge.training.grad, re-exported from kempnerforge.training. On all microbatches except the last, it disables gradient synchronization so the backward pass accumulates locally and does not fire reduce-scatter N times per optimizer step.

3 · MoE step propagation

if mc.is_moe:
    model.set_moe_step(step, tc.max_steps)

Forwards (step, max_steps) to every SigmoidTopKRouter. Used by the adaptive bias schedule (see docs/moe/). Dense models skip this.

4 · Forward

logits = model(input_ids, doc_ids=doc_ids)
loss   = loss_fn(logits, labels)

The forward pass follows the model page: token embedding → N transformer blocks (RoPE + GQA + SwiGLU or MoE) → final RMSNorm → output head → (batch, seq_len, vocab_size) logits.

Per-dataset metrics, if the dataloader is a mixture, are collected here before backward so the logits are still alive.

5 · MoE auxiliary loss

if mc.is_moe:
    aux_loss = model.get_moe_aux_loss()
    loss = loss + mc.moe_aux_loss_weight * aux_loss

get_moe_aux_loss() sums the MoEMLP.aux_loss attributes of every router: one per MoE layer, or one per modality on a MoT layer. For dense runs it returns 0.0 and the line is a no-op.

6 · Backward with gradient accumulation

scaled_loss = loss / tc.grad_accum_steps
scaled_loss.backward()
total_loss += loss.item()

Scaling by grad_accum_steps keeps the effective learning rate invariant to the accumulation factor.

7 · Clip, NaN check, optimizer step

After the microbatch loop:

grad_norm = clip_grad_norm_(model, tc.grad_clip_norm)

if not nan_detector.check_loss(avg_loss, step, grad_norm=grad_norm):
    optimizer.zero_grad()
    if nan_detector.should_rollback:
        break
    _advance_lr_schedule(session)            # scheduler.step() + phase LR scaling
    step += 1
    _apply_step_transitions(session, step)   # freeze stages, then phases (section 9)
    if _shutdown_if_requested(               # saves only finite weights (section 11)
        session, step, tokens_seen, saved_this_step=False, nan_skipped=True
    ):
        break
    continue                                 # no update, token count, metrics or hooks

optimizer.step()
scheduler.step()
if phase_lr_scale != 1.0:
    for pg in optimizer.param_groups: pg["lr"] *= phase_lr_scale
optimizer.zero_grad()

clip_grad_norm_ wraps PyTorch’s utility so it works with FSDP2 sharded parameters. Under pipeline parallelism pipeline_step also passes the PP group, so the norm covers every stage rather than this rank’s own; under expert parallelism it reads the EP group from the model and covers every rank’s experts. The NaN detector returns False on a NaN/Inf loss or gradient norm, and the loop zeroes the grads. After max_consecutive NaN steps in a row the detector sets should_rollback and the loop stops; rolling back to an earlier checkpoint is manual (see NaN detection).

A skipped step still counts as a step (section 8): it advances the LR schedule, applies the freeze stages and the phases (section 9) due at the new step, and runs the shutdown check (section 11), so all three stay on the step counter. It skips the update, the token count (section 8), metrics and hooks (section 10), and the rest of the periodic work in section 11.

Phase LR scaling applies after the scheduler — it multiplies the base LR that scheduler.step() just computed.

8 · Step accounting

step += 1
tokens_in_step = tc.batch_size * tc.seq_len * tc.grad_accum_steps * dp_size
tokens_seen   += tokens_in_step

tokens_in_step times all ranks in the data-parallel dimension, not the full world_size — TP, PP, and EP don’t multiply unique tokens.

9 · Phase transitions

If a mixture phase is due, sampler.update_weights(...) rebalances the MixtureSampler and data_iter = None forces a refresh next microbatch so the new weights take effect immediately.

10 · Metrics and hooks

step_metrics = tracker.end_step(step=step, loss=avg_loss,
                                grad_norm=grad_norm_val, lr=current_lr,
                                tokens_in_step=tokens_in_step)
hook_runner.on_step_end(StepContext(...))

tracker.end_step dispatches to WandB / TensorBoard at metrics.log_interval. HookRunner (from kempnerforge.training.hooks) runs user-defined TrainingHook callbacks each step.

MoE runs additionally log moe/aux_loss and moe/expert_balance (min/max of per-expert loads, each router’s counts over its own total) at the same cadence, plus moe/expert_balance/<modality> on MoT.

11 · Periodic work

In order:

  • NCCL health (every tc.nccl_health_check_interval steps) — fires a small all-reduce; on failure, break the loop.

  • Eval (every eval_config.interval steps) — runs run_eval on the eval dataloader, logs metrics, fires on_eval_end hooks.

  • Profiler — prof.step() advances the torch.profiler schedule.

  • Checkpoint (checkpoint.should_save(step): dyn_ckpt_window milestones inside the window, every checkpoint.interval steps outside it) — ckpt_mgr.save writes a DCP checkpoint asynchronously and updates the latest symlink.

  • Graceful shutdown — if SIGTERM/SIGUSR1 fired, ckpt_mgr.save(step) as the emergency checkpoint, then exit; a step the checkpoint above already saved is not saved again, and the loop waits for that save before it finishes the shutdown. A NaN-skipped step (section 7) runs this check too; the checkpoint above never saves such a step, so it writes its own emergency checkpoint, but only if an all-reduced check finds every rank’s parameters finite. After a bad update (an optimizer step that passed the NaN check) they are not, and the shutdown logs an ERROR and saves nothing (see NaN detection § A shutdown after a bad update).

Shutdown

After the loop exits:

  • prof.stop() flushes traces.

  • ckpt_mgr.wait() drains the last async save.

  • hook_runner.on_train_end(step, tokens_seen).

  • tracker.close() flushes WandB / TB.

  • destroy_distributed() tears down the process group.