| name | explain-architecture |
| description | Walk through KempnerForge's subsystems in the order a forward pass encounters them. Starting point for anyone new to the codebase. |
When to use
- New contributor asking "how does training work" or "where does X live".
- Before a non-trivial change, to confirm which subsystem owns the behavior in question.
- When reviewing a PR that touches multiple subsystems and the user wants to understand the interaction.
This skill is read-only. It does not run anything.
Preflight
Run:
uv run python scripts/check_env.py
Baseline only (uv, repo layout). If the exit code is non-zero, fix the baseline first, otherwise file paths below may not resolve.
Context (auto-generated, do not edit)
Top-level package: kempnerforge/
Subsystems: config, model, distributed, data, training, checkpoint, resilience, metrics, profiling
Entry point: scripts/train.py (does NOT use torchrun internally; caller launches it)
Config system: kempnerforge/config/{schema,job,data,eval,optimizer}.py with registry.py for pluggable components
Model stack: embedding -> transformer.py (blocks of attention + mlp/moe) -> norm -> output projection
Parallelism layers: FSDP2 (distributed/parallel.py), TP (distributed/tensor_parallel.py), PP (distributed/pipeline_parallel.py), EP (distributed/expert_parallel.py)
Data path: dataset.py (MemoryMapped | HF | Mixture) -> sampler.py (DistributedSampler) -> dataloader.py (StatefulDataLoader)
Checkpoint: checkpoint/manager.py orchestrates DCP save/load, state.py owns state dict layout, async_save.py handles non-blocking writes
Resilience: resilience/ (signal handlers, NaN detection, NCCL health)
Metrics: metrics/tracker.py dispatches to backends (wandb, tensorboard), mfu.py computes MFU, memory.py tracks peak memory
Procedure
Walk the user through the path a single forward+backward pass takes. Use the outline below and follow up only on the subsystems they ask about.
1. Config is loaded
scripts/train.py calls load_config(toml_path, cli_args) from kempnerforge/config/.- Layered resolution: dataclass defaults, then TOML overrides, then
--section.key=value CLI overrides.
JobConfig is the top-level container. Inspect the sections: ModelConfig, TrainConfig, OptimizerConfig, SchedulerConfig, DistributedConfig, DataConfig, CheckpointConfig, MetricsConfig.- Pluggable components (optimizers, schedulers, routers, losses) come from
kempnerforge/config/registry.py via @register_* decorators. Read the registry to see what names are valid for each component.
2. Distributed process group is initialized
kempnerforge/distributed/setup.py::init_distributed() reads MASTER_ADDR / MASTER_PORT (or SLURM env vars) and calls torch.distributed.init_process_group.DeviceMesh from distributed/parallel.py gives the training loop a multi-dimensional view (DP x TP x PP) that each subsystem queries.- On multi-node: the shell script (
scripts/slurm/multinode.sh) sets env vars; this function only consumes them.
3. Model is built and sharded
kempnerforge/model/transformer.py::Transformer composes: embedding -> N blocks (TransformerBlock) -> final norm -> output projection.- Each
TransformerBlock has an Attention (model/attention.py, SDPA-based GQA with RoPE from position.py) and an MLP (model/mlp.py SwiGLU, or model/moe.py for mixture of experts with a router from router.py).
- After construction,
distributed/parallel.py applies fully_shard() (FSDP2) per block. If TP or PP is enabled, they compose via parallelize_module and DeviceMesh.
4. Data pipeline is constructed
kempnerforge/data/dataset.py holds three dataset types: MemoryMappedDataset (pre-tokenized .npy, fastest), HuggingFaceDataset / StreamingHuggingFaceDataset (HF Hub), MixtureDataset (weighted combination of sub-datasets with phase scheduling).data/sampler.py::DistributedSampler partitions indices across data-parallel ranks. set_skip(n) supports exact resume from mid-epoch.data/dataloader.py::StatefulDataLoader wraps the torch DataLoader and tracks enough state (epoch, batch index, sampler skip) to restore the exact batch after checkpoint load.
5. Optimizer and scheduler
kempnerforge/training/optimizer.py builds the optimizer (AdamW fused, Lion, Muon, schedule-free) from [optimizer] config via the registry.training/scheduler.py builds the LR scheduler. Registered names: cosine, linear, wsd, constant, rex, none.training/grad.py owns gradient utilities: clipping, accumulation (maybe_no_sync context manager for FSDP), NaN detection hooks.
6. Training loop
scripts/train.py runs the loop directly (no framework wrapper). Each step: forward, loss (training/loss.py), backward, optional gradient accumulation, optimizer step, scheduler step, metrics log, checkpoint if interval hit.- MoE runs apply an aux loss term inside the model (see
model/transformer.py::get_moe_aux_loss), which train.py adds to the main loss before backward().
7. Checkpointing
kempnerforge/checkpoint/manager.py::CheckpointManager orchestrates save and load. Uses torch.distributed.checkpoint (DCP), which automatically handles FSDP sharded state and reshards on load.checkpoint/state.py defines what goes in the state dict: model, optimizer, scheduler, dataloader position, RNG state.checkpoint/async_save.py wraps dcp.async_save so training continues while writes are in flight.- Resume follows the
latest symlink in checkpoint.dir, falling back to the highest step_N directory.
8. Metrics and resilience
kempnerforge/metrics/tracker.py::MetricsTracker collects per-step scalars with optional EMA smoothing, then dispatches to enabled backends (wandb, tensorboard).metrics/mfu.py computes model FLOPs utilization from the model config and measured step time.metrics/memory.py tracks peak GPU memory, with an optional snapshot export for pytorch.org/memory_viz.resilience/ owns signal handlers (SIGTERM/SIGUSR1 for SLURM preemption), NaN action policies, and GPU/NCCL liveness probes that run every N steps.
Verification
None. This skill is informational. Ask the user which subsystem they want to dig into, then point them at the specific files and classes listed above.
Gotchas
- The training loop is in
scripts/train.py, not inside a Trainer class. Deliberately flat, easy to read top to bottom. Do not try to find a class Trainer.
- FSDP2 uses
fully_shard() (composable), NOT FullyShardedDataParallel (the v1 class). Old tutorials do not apply.
- Context parallelism (CP) is stubbed but not wired up yet (PyTorch 2.11 has an experimental ring-attention API). If the user asks "how do I enable CP", explain it is pending, not a docs/user error.
scripts/train.py takes the TOML path as a positional arg, not via --config. CLI overrides are --section.key=value, double dash.
Related skills
/kempnerforge:component-gaps — ask-it-what-is-built-vs-planned for each subsystem./kempnerforge:add-optimizer — concrete walkthrough of the training/optimizer.py + registry pattern.
