How kempnerinstitute/kempnerforge works
Purpose and entry points
KempnerForge is a Python 3.12+/PyTorch 2.4+ framework for fault-tolerant, distributed training of decoder LLMs, MoE models, and vision-language/video models on CUDA clusters. The primary runtime entry point is scripts/train.py; distributed launches use torchrun or SLURM wrappers under scripts/slurm/. Supporting operational entry points include scripts/generate.py, scripts/eval.py, scripts/prepare_data.py, and scripts/convert_checkpoint.py.
Configuration as the assembly boundary
Training is assembled from TOML presets plus CLI overrides. Typed configuration dataclasses define model, data, distributed, optimizer, checkpoint, observability, and VLM choices; the loader validates unsupported combinations before allocating a job.
• kempnerforge/config/schema.py
• kempnerforge/config/loader.py
• kempnerforge/config/registry.py
• configs/train/
Training orchestration and state
scripts/train.py owns job startup and coordinates model/data construction, forward-backward-update iterations, evaluation, hooks, checkpointing, and metrics. Training behavior is composition-driven: optimizers, schedulers, losses, gradient processing, freezing, and hooks are selected from config. Stateful training/data state enables resume rather than treating checkpoints as model-only artifacts.
• scripts/train.py
• kempnerforge/training/optimizer.py
• kempnerforge/training/scheduler.py
• kempnerforge/training/hooks.py
Model and modality layer
The core model is a decoder Transformer with RoPE, grouped-query attention, SwiGLU, RMSNorm, generation, activation hooks, and optional torch.compile. Dense MLP blocks can be replaced with routed MoE. VLM wraps the text backbone with HF vision encoders, connectors, and architecture-specific modality strategies: joint decoder, cross-attention, Mixture-of-Transformers, or MoMa. Video data becomes frame-derived visual tokens through the same VLM boundary.
• kempnerforge/model/transformer.py
• kempnerforge/model/moe.py
• kempnerforge/model/vlm.py
• kempnerforge/model/vision.py
Distributed transformation pipeline
After construction, model execution is transformed in an enforced order: tensor parallelism, expert parallelism, FP8 handling, activation checkpointing, then FSDP2; pipeline parallelism is separately supported for dense models. This ordering is a correctness boundary, not merely an optimization. Device-mesh/process-group setup and distributed lifecycle are centralized rather than embedded in model code.
• kempnerforge/distributed/setup.py
• kempnerforge/distributed/parallel.py
• kempnerforge/distributed/tensor_parallel.py
• kempnerforge/distributed/expert_parallel.py
• kempnerforge/distributed/pipeline_parallel.py
Data-to-batch flow
Inputs originate from memory-mapped/tokenized corpora or HuggingFace datasets, optionally mixed and annealed. Dataset/sampler/dataloader state is checkpointable, preserving ordering/progress across restarts. VLM data extends this path with image/video decoding and modality-aware batches; video decoding uses optional PyAV, while vision encoders are loaded from Hugging Face.
• kempnerforge/data/dataset.py
• kempnerforge/data/sampler.py
• kempnerforge/data/dataloader.py
• kempnerforge/data/vlm_dataset.py
• kempnerforge/data/video_dataset.py
Durability, resilience, and observability
Distributed Checkpoint (DCP) persistence supports synchronous/asynchronous saves, resumable training state, and checkpoint conversion. Resilience handles SIGTERM/preemption, NaN detection, GPU health, NCCL liveness, and elastic behavior. Metrics record training/MFU/memory data through WandB, TensorBoard, or MLflow; PyTorch profiler and CUDA timing provide performance traces.
• kempnerforge/checkpoint/manager.py
• kempnerforge/checkpoint/async_save.py
• kempnerforge/resilience/signal_handler.py
• kempnerforge/resilience/health.py
• kempnerforge/metrics/tracker.py
Main entry points: train.py Training runtime · transformer.py Decoder Transformer · loader.py Config loader · dataloader.py Dataset and loader
Components 18
Entry Points
- Training runtime · Coordinates construction, iterations, evaluation, hooks, checkpoints, and metrics.
scripts/train.py - Operational scripts · Generation, evaluation, data preparation, and checkpoint conversion.
scripts/generate.py - Cluster launchers · Launch distributed jobs through torchrun or SLURM.
Assembly
- Typed config schema · Defines model, data, distributed, optimizer, checkpoint, observability, and VLM choices.
kempnerforge/config/schema.py - Config loader · Applies TOML presets and CLI overrides, then validates combinations.
kempnerforge/config/loader.py - Training presets
Model & Data
- Dataset and loader · Builds mixed or annealed corpora, samplers, and resumable dataloaders.
kempnerforge/data/dataloader.py - Visual and video data · Decodes image/video inputs into modality-aware batches.
kempnerforge/data/vlm_dataset.py - Decoder Transformer · Text backbone with RoPE, GQA, SwiGLU, RMSNorm, generation, and optional compilation.
kempnerforge/model/transformer.py - Routed MoE · Replaces dense MLP blocks with routed experts.
kempnerforge/model/moe.py - VLM wrapper · Adds vision encoders, connectors, and modality strategies around the text backbone.
kempnerforge/model/vlm.py
Distributed Execution
Model transformations are applied in a required order.
- Device mesh lifecycle · Centralizes process groups, device meshes, and distributed lifecycle.
kempnerforge/distributed/setup.py - Parallel transform chain · Applies tensor parallelism, expert parallelism, FP8, activation checkpointing, then FSDP2.
kempnerforge/distributed/parallel.py - Pipeline parallelism · Separate pipeline execution support for dense models.
kempnerforge/distributed/pipeline_parallel.py
Durability & Operations
- Training components · Selects optimizers, schedulers, hooks, losses, gradient processing, and freezing.
kempnerforge/training/optimizer.py - Distributed checkpointing · Persists model plus resumable training and data state, synchronously or asynchronously.
kempnerforge/checkpoint/manager.py - Resilience controls · Handles preemption, NaNs, GPU health, NCCL liveness, and elastic behavior.
kempnerforge/resilience/health.py - Metrics and profiling · Reports training, MFU, memory, and performance traces.
kempnerforge/metrics/tracker.py
Connections 21
- Cluster launchers → Training runtime: launches
- Training presets → Config loader: preset inputs
- Typed config schema → Config loader: validates against
- Config loader → Training runtime: assembled job config
- Training runtime → Dataset and loader: builds and consumes batches
- Visual and video data → Dataset and loader: extends batches
- Training runtime → Decoder Transformer: constructs
- Routed MoE → Decoder Transformer: optional MLP replacement
- VLM wrapper → Decoder Transformer: wraps text backbone
- Visual and video data → VLM wrapper: visual tokens and batches
- Training runtime → Device mesh lifecycle: initializes
- Device mesh lifecycle → Parallel transform chain: provides meshes and groups
- Decoder Transformer → Parallel transform chain: transforms after construction
- Parallel transform chain → Training runtime: distributed executable model
- Pipeline parallelism → Training runtime: optional dense execution
- Training runtime → Training components: uses configured behavior
- Training runtime → Distributed checkpointing: save and resume state
- Dataset and loader → Distributed checkpointing: sampler and loader state
- Resilience controls → Training runtime: protects runtime
- Training runtime → Metrics and profiling: emits metrics and traces
- Operational scripts → Distributed checkpointing: converts checkpoints