# kempnerinstitute/kempnerforge architecture

> Architecture diagram of the GitHub repository kempnerinstitute/kempnerforge, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-05). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/kempnerinstitute/kempnerforge
- Repository: https://github.com/kempnerinstitute/kempnerforge

## Overview

#### 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: [`scripts/train.py`](https://github.com/kempnerinstitute/kempnerforge/blob/main/scripts/train.py) (Training runtime), [`kempnerforge/model/transformer.py`](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/transformer.py) (Decoder Transformer), [`kempnerforge/config/loader.py`](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/config/loader.py) (Config loader), [`kempnerforge/data/dataloader.py`](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/data/dataloader.py) (Dataset and loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entry Points"]
  node_train["Training runtime<br/>Python entry point<br/>[train.py]"]
  node_ops_scripts["Operational scripts<br/>Python entry points<br/>[generate.py]"]
  node_launch["Cluster launchers<br/>SLURM wrappers"]
end

subgraph group_assembly["Assembly"]
  node_config_schema["Typed config schema<br/>Python dataclasses<br/>[schema.py]"]
  node_config_loader["Config loader<br/>Python module<br/>[loader.py]"]
  node_train_presets["Training presets<br/>TOML configs"]
end

subgraph group_model["Model &amp; Data"]
  node_data_pipeline["Dataset and loader<br/>checkpointable input pipeline<br/>[dataloader.py]"]
  node_vlm_data["Visual and video data<br/>modality dataset<br/>[vlm_dataset.py]"]
  node_transformer["Decoder Transformer<br/>PyTorch model<br/>[transformer.py]"]
  node_moe["Routed MoE<br/>model component<br/>[moe.py]"]
  node_vlm["VLM wrapper<br/>multimodal model<br/>[vlm.py]"]
end

subgraph group_distributed["Distributed Execution"]
  node_dist_setup["Device mesh lifecycle<br/>distributed runtime<br/>[setup.py]"]
  node_parallel_transforms["Parallel transform chain<br/>model transformation pipeline<br/>[parallel.py]"]
  node_pipeline_parallel["Pipeline parallelism<br/>dense-model parallelism"]
end

subgraph group_operations["Durability &amp; Operations"]
  node_training_components["Training components<br/>composition layer<br/>[optimizer.py]"]
  node_checkpointing[("Distributed checkpointing<br/>DCP persistence<br/>[manager.py]")]
  node_resilience["Resilience controls<br/>fault handling<br/>[health.py]"]
  node_metrics["Metrics and profiling<br/>observability<br/>[tracker.py]"]
end

node_launch -->|"launches"| node_train
node_train_presets -->|"preset inputs"| node_config_loader
node_config_schema -->|"validates against"| node_config_loader
node_config_loader -->|"assembled job config"| node_train
node_train -->|"builds and consumes batches"| node_data_pipeline
node_vlm_data -->|"extends batches"| node_data_pipeline
node_train -->|"constructs"| node_transformer
node_moe -.->|"optional MLP replacement"| node_transformer
node_vlm -->|"wraps text backbone"| node_transformer
node_vlm_data -->|"visual tokens and batches"| node_vlm
node_train -->|"initializes"| node_dist_setup
node_dist_setup -->|"provides meshes and groups"| node_parallel_transforms
node_transformer -->|"transforms after construction"| node_parallel_transforms
node_parallel_transforms -->|"distributed executable model"| node_train
node_pipeline_parallel -.->|"optional dense execution"| node_train
node_train -->|"uses configured behavior"| node_training_components
node_train -->|"save and resume state"| node_checkpointing
node_data_pipeline -->|"sampler and loader state"| node_checkpointing
node_resilience -->|"protects runtime"| node_train
node_train -->|"emits metrics and traces"| node_metrics
node_ops_scripts -.->|"converts checkpoints"| node_checkpointing

click node_train "https://github.com/kempnerinstitute/kempnerforge/blob/main/scripts/train.py"
click node_ops_scripts "https://github.com/kempnerinstitute/kempnerforge/blob/main/scripts/generate.py"
click node_config_schema "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/config/schema.py"
click node_config_loader "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/config/loader.py"
click node_data_pipeline "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/data/dataloader.py"
click node_vlm_data "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/data/vlm_dataset.py"
click node_transformer "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/transformer.py"
click node_moe "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/moe.py"
click node_vlm "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/vlm.py"
click node_dist_setup "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/distributed/setup.py"
click node_parallel_transforms "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/distributed/parallel.py"
click node_pipeline_parallel "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/distributed/pipeline_parallel.py"
click node_training_components "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/training/optimizer.py"
click node_checkpointing "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/checkpoint/manager.py"
click node_resilience "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/resilience/health.py"
click node_metrics "https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/metrics/tracker.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train,node_ops_scripts,node_launch toneBlue
class node_config_schema,node_config_loader,node_train_presets toneAmber
class node_data_pipeline,node_vlm_data,node_transformer,node_moe,node_vlm toneMint
class node_dist_setup,node_parallel_transforms,node_pipeline_parallel toneRose
class node_training_components,node_checkpointing,node_resilience,node_metrics toneIndigo
```

## Components

### Entry Points

- [Training runtime](https://github.com/kempnerinstitute/kempnerforge/blob/main/scripts/train.py): Python entry point · Coordinates construction, iterations, evaluation, hooks, checkpoints, and metrics. · `scripts/train.py`
- [Operational scripts](https://github.com/kempnerinstitute/kempnerforge/blob/main/scripts/generate.py): Python entry points · Generation, evaluation, data preparation, and checkpoint conversion. · `scripts/generate.py`
- **Cluster launchers**: SLURM wrappers · Launch distributed jobs through torchrun or SLURM.

### Assembly

- [Typed config schema](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/config/schema.py): Python dataclasses · Defines model, data, distributed, optimizer, checkpoint, observability, and VLM choices. · `kempnerforge/config/schema.py`
- [Config loader](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/config/loader.py): Python module · Applies TOML presets and CLI overrides, then validates combinations. · `kempnerforge/config/loader.py`
- **Training presets**: TOML configs

### Model & Data

- [Dataset and loader](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/data/dataloader.py): checkpointable input pipeline · Builds mixed or annealed corpora, samplers, and resumable dataloaders. · `kempnerforge/data/dataloader.py`
- [Visual and video data](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/data/vlm_dataset.py): modality dataset · Decodes image/video inputs into modality-aware batches. · `kempnerforge/data/vlm_dataset.py`
- [Decoder Transformer](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/transformer.py): PyTorch model · Text backbone with RoPE, GQA, SwiGLU, RMSNorm, generation, and optional compilation. · `kempnerforge/model/transformer.py`
- [Routed MoE](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/moe.py): model component · Replaces dense MLP blocks with routed experts. · `kempnerforge/model/moe.py`
- [VLM wrapper](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/model/vlm.py): multimodal model · 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](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/distributed/setup.py): distributed runtime · Centralizes process groups, device meshes, and distributed lifecycle. · `kempnerforge/distributed/setup.py`
- [Parallel transform chain](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/distributed/parallel.py): model transformation pipeline · Applies tensor parallelism, expert parallelism, FP8, activation checkpointing, then FSDP2. · `kempnerforge/distributed/parallel.py`
- [Pipeline parallelism](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/distributed/pipeline_parallel.py): dense-model parallelism · Separate pipeline execution support for dense models. · `kempnerforge/distributed/pipeline_parallel.py`

### Durability & Operations

- [Training components](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/training/optimizer.py): composition layer · Selects optimizers, schedulers, hooks, losses, gradient processing, and freezing. · `kempnerforge/training/optimizer.py`
- [Distributed checkpointing](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/checkpoint/manager.py): DCP persistence · Persists model plus resumable training and data state, synchronously or asynchronously. · `kempnerforge/checkpoint/manager.py`
- [Resilience controls](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/resilience/health.py): fault handling · Handles preemption, NaNs, GPU health, NCCL liveness, and elastic behavior. · `kempnerforge/resilience/health.py`
- [Metrics and profiling](https://github.com/kempnerinstitute/kempnerforge/blob/main/kempnerforge/metrics/tracker.py): observability · Reports training, MFU, memory, and performance traces. · `kempnerforge/metrics/tracker.py`

## Connections

- 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