kempnerforge.config

Configuration system for KempnerForge.

class kempnerforge.config.CheckpointConfig[source]

Bases: object

Checkpointing settings.

dir: str = 'checkpoints'
interval: int = 1000
dyn_ckpt_window: DynamicCheckpointWindow | None = None
async_mode: AsyncCheckpointMode = 'disabled'
keep_last_n: int = 3
load_path: str | None = None
export_dtype: Literal['float32', 'bfloat16'] = 'bfloat16'
exclude_from_loading: list[str]
ignore_freeze_mismatch: bool = False
ignore_rank_state_mismatch: bool = False
should_save(step)[source]

Whether to write a checkpoint at step.

Inside dyn_ckpt_window: the registered strategy decides (default "power2" saves at start and each start + 2^k while <= stop). Outside the window: every interval steps. Dynamic milestones are exempt from keep_last_n (see CheckpointManager._cleanup).

Parameters:

step (int)

Return type:

bool

is_dynamic_milestone(step)[source]

True if step is a milestone of the configured dyn_ckpt_window.

CheckpointManager._cleanup excludes these from keep_last_n so the dense early-window checkpoints survive a finite retention.

Parameters:

step (int)

Return type:

bool

__init__(dir='checkpoints', interval=1000, dyn_ckpt_window=None, async_mode=AsyncCheckpointMode.disabled, keep_last_n=3, load_path=None, export_dtype='bfloat16', exclude_from_loading=<factory>, ignore_freeze_mismatch=False, ignore_rank_state_mismatch=False)
Parameters:
Return type:

None

class kempnerforge.config.DataConfig[source]

Bases: object

Data pipeline settings.

token_dtype is the integer type of the raw .bin token shards, "uint16" or "uint32", for dataset_path and every datasets source that does not set its own. Left empty, the metadata.yaml next to the shards decides, as MemoryMappedDataset describes. .npy shards carry their dtype in their header and ignore it.

dataset_path: str = ''
file_pattern: str = '*.npy'
token_dtype: str = ''
tokenizer_path: str = ''
num_workers: int = 4
pin_memory: bool = True
prefetch_factor: int = 2
hf_dataset_name: str | None = None
hf_dataset_config: str | None = None
hf_dataset_split: str = 'train'
hf_dataset_text_field: str = 'text'
hf_dataset_image_field: str = 'image'
hf_dataset_prompt_field: str = ''
hf_image_size: int = 224
hf_streaming: bool = False
pack_sequences: bool = False
datasets: list[DatasetSource]
mix_temperature: float = 1.0
phases: list[TrainingPhase]
anneal_start_step: int = 0
anneal_weights: dict[str, float]
__init__(dataset_path='', file_pattern='*.npy', token_dtype='', tokenizer_path='', num_workers=4, pin_memory=True, prefetch_factor=2, hf_dataset_name=None, hf_dataset_config=None, hf_dataset_split='train', hf_dataset_text_field='text', hf_dataset_image_field='image', hf_dataset_prompt_field='', hf_image_size=224, hf_streaming=False, pack_sequences=False, datasets=<factory>, mix_temperature=1.0, phases=<factory>, anneal_start_step=0, anneal_weights=<factory>)
Parameters:
Return type:

None

class kempnerforge.config.DistributedConfig[source]

Bases: object

Parallelism dimensions and distributed settings.

dp_shard: int = -1
dp_replicate: int = 1
tp: int = 1
pp: int = 1
pp_schedule: PipelineSchedule = '1f1b'
cp: int = 1
ep: int = 1
nccl_timeout_sec: int = 1800
backend: str = 'cpu:gloo,cuda:nccl'
validate_world_size(world_size)[source]

Validate that parallelism dimensions match world size.

Parameters:

world_size (int)

Return type:

None

resolve(world_size)[source]

Return a copy with dp_shard resolved to a concrete value.

Parameters:

world_size (int)

Return type:

DistributedConfig

__init__(dp_shard=-1, dp_replicate=1, tp=1, pp=1, pp_schedule=PipelineSchedule.schedule_1f1b, cp=1, ep=1, nccl_timeout_sec=1800, backend='cpu:gloo,cuda:nccl')
Parameters:
Return type:

None

class kempnerforge.config.EvalConfig[source]

Bases: object

Evaluation pipeline settings (disabled by default).

token_dtype is the integer type of the raw .bin eval shards, as DataConfig.token_dtype is for the training shards.

enabled: bool = False
interval: int = 1000
steps: int = 50
dataset_path: str = ''
file_pattern: str = '*.npy'
token_dtype: str = ''
hf_dataset_name: str | None = None
hf_dataset_config: str | None = None
hf_dataset_split: str = 'validation'
__init__(enabled=False, interval=1000, steps=50, dataset_path='', file_pattern='*.npy', token_dtype='', hf_dataset_name=None, hf_dataset_config=None, hf_dataset_split='validation')
Parameters:

  • enabled (bool)

  • interval (int)

  • steps (int)

  • dataset_path (str)

  • file_pattern (str)

  • token_dtype (str)

  • hf_dataset_name (str | None)

  • hf_dataset_config (str | None)

  • hf_dataset_split (str)

Return type:

None

class kempnerforge.config.JobConfig[source]

Bases: object

Top-level configuration aggregating all sub-configs.

The VLM stack lives as three sibling top-level sections, [vision_encoder] / [adapter] / [vlm]. A pure text run leaves all three at None and writes only [model] (plus the other operational sections). A VLM run sets [vlm] and is required to provide [vision_encoder]; [adapter] defaults to the registered 2-layer MLP and may be omitted.

model: ModelConfig
train: TrainConfig
optimizer: OptimizerConfig
scheduler: SchedulerConfig
data: DataConfig
eval: EvalConfig
distributed: DistributedConfig
checkpoint: CheckpointConfig
metrics: MetricsConfig
profiling: ProfilingConfig
vision_encoder: VisionEncoderConfig | None = None
adapter: AdapterConfig | None = None
vlm: VLMConfig | None = None
video: VideoConfig | None = None
plugins: list[str]
property is_vlm: bool

Whether this job builds a VLMWrapper around the text backbone.

property is_video: bool

Whether this job trains on a video dataset (a VLM sub-mode).

property views_per_sample: int

a clip’s frames, a tiled image’s views, or one.

Type:

Images the encoder sees per sample

validate(world_size=1)[source]

Run cross-config validations that depend on the world size.

Parameters:

world_size (int)

Return type:

None

__init__(model=<factory>, train=<factory>, optimizer=<factory>, scheduler=<factory>, data=<factory>, eval=<factory>, distributed=<factory>, checkpoint=<factory>, metrics=<factory>, profiling=<factory>, vision_encoder=None, adapter=None, vlm=None, video=None, plugins=<factory>)
Parameters:
Return type:

None

class kempnerforge.config.MetricsConfig[source]

Bases: object

Logging and metrics settings.

log_interval: int = 10
enable_wandb: bool = False
enable_tensorboard: bool = False
enable_mlflow: bool = False
wandb_project: str = 'kempnerforge'
wandb_run_name: str | None = None
wandb_run_id: str = ''
tensorboard_dir: str = 'tb_logs'
mlflow_tracking_uri: str | None = None
mlflow_experiment: str | None = None
mlflow_run_name: str | None = None
mlflow_description: str | None = None
mlflow_tags: dict[str, str]
mlflow_run_id: str = ''
mlflow_log_system_metrics: bool = True
mlflow_system_metrics_interval: float = 10.0
mlflow_register_checkpoints: bool = False
resolve_mlflow_tracking_uri()[source]

The tracking URI in force: the field, else $MLFLOW_TRACKING_URI, else Databricks.

Return type:

str

__init__(log_interval=10, enable_wandb=False, enable_tensorboard=False, enable_mlflow=False, wandb_project='kempnerforge', wandb_run_name=None, wandb_run_id='', tensorboard_dir='tb_logs', mlflow_tracking_uri=None, mlflow_experiment=None, mlflow_run_name=None, mlflow_description=None, mlflow_tags=<factory>, mlflow_run_id='', mlflow_log_system_metrics=True, mlflow_system_metrics_interval=10.0, mlflow_register_checkpoints=False)
Parameters:

  • log_interval (int)

  • enable_wandb (bool)

  • enable_tensorboard (bool)

  • enable_mlflow (bool)

  • wandb_project (str)

  • wandb_run_name (str | None)

  • wandb_run_id (str)

  • tensorboard_dir (str)

  • mlflow_tracking_uri (str | None)

  • mlflow_experiment (str | None)

  • mlflow_run_name (str | None)

  • mlflow_description (str | None)

  • mlflow_tags (dict[str, str])

  • mlflow_run_id (str)

  • mlflow_log_system_metrics (bool)

  • mlflow_system_metrics_interval (float)

  • mlflow_register_checkpoints (bool)

Return type:

None

class kempnerforge.config.ModelConfig[source]

Bases: object

Architecture hyperparameters for a transformer model.

dim: int = 4096
n_layers: int = 32
n_heads: int = 32
n_kv_heads: int | None = None
head_dim_override: int = 0
vocab_size: int = 32000
ffn_dim_multiplier: float = 1.0
ffn_hidden_dim: int | None = None
norm_type: NormType = 'rmsnorm'
norm_eps: float = 1e-05
activation: Activation = 'silu'
max_seq_len: int = 2048
rope_theta: float = 10000.0
tie_embeddings: bool = False
qk_norm: bool = False
init_std: float = 0.02
model_type: str = 'transformer'
sdpa_backend: str = 'auto'
attention_backend: str = 'sdpa'
num_experts: int = 0
moe_top_k: int = 2
moe_frequency: int = 1
moe_router: str = 'softmax_topk'
moe_shared_experts: int = 0
moe_aux_loss_weight: float = 0.01
moe_capacity_factor: float = 0.0
moe_sequence_aux_loss_weight: float = 0.0
moe_gradient_scale: bool = False
moe_bias_schedule: str = 'constant'
moe_packed_experts: bool = False
moe_expert_ffn_multiplier: float = 1.0
moe_router_z_loss_weight: float = 0.0
property is_moe: bool

Whether this config uses Mixture-of-Experts.

property tp_sequence_parallel: bool

Whether tensor parallelism runs this model with SequenceParallel.

apply_tensor_parallel shards activations along the sequence between blocks unless the embeddings are tied, since sharding the output head’s columns would also shard the shared embedding weight, or the model is MoE, since boolean indexing in expert dispatch breaks sequence-sharded DTensors. Pipeline stages also run without it; that depends on the stage, not on this config.

property head_dim: int

Per-head attention width.

Derived rather than stored so a later dim / n_heads change (a TOML or CLI overlay onto a default config) cannot leave a stale value behind. head_dim_override decouples it from dim.

property computed_ffn_hidden_dim: int

FFN hidden dimension, rounded to nearest multiple of 256 for hardware efficiency.

property computed_expert_ffn_hidden_dim: int

Per-expert FFN hidden dim = computed_ffn_hidden_dim * moe_expert_ffn_multiplier.

Rounded to a multiple of 16 for tensor-core alignment. With the default multiplier 1.0 this equals computed_ffn_hidden_dim (zero behavior change); set 0.5 for fine-grained experts so top-2 routing matches the dense FFN’s activated FLOPs (2 * F/2 = F). Routed experts are this wide (build_moe and MoMa’s ExpertChoiceMoE); build_moe’s shared experts are one MLP moe_shared_experts times this wide.

property num_params_estimate: int

Rough total parameter count estimate (excluding embedding if tied).

For MoE models, counts all expert parameters (total, not active).

__init__(dim=4096, n_layers=32, n_heads=32, n_kv_heads=None, head_dim_override=0, vocab_size=32000, ffn_dim_multiplier=1.0, ffn_hidden_dim=None, norm_type=NormType.rmsnorm, norm_eps=1e-05, activation=Activation.silu, max_seq_len=2048, rope_theta=10000.0, tie_embeddings=False, qk_norm=False, init_std=0.02, model_type='transformer', sdpa_backend='auto', attention_backend='sdpa', num_experts=0, moe_top_k=2, moe_frequency=1, moe_router='softmax_topk', moe_shared_experts=0, moe_aux_loss_weight=0.01, moe_capacity_factor=0.0, moe_sequence_aux_loss_weight=0.0, moe_gradient_scale=False, moe_bias_schedule='constant', moe_packed_experts=False, moe_expert_ffn_multiplier=1.0, moe_router_z_loss_weight=0.0)
Parameters:

  • dim (int)

  • n_layers (int)

  • n_heads (int)

  • n_kv_heads (int | None)

  • head_dim_override (int)

  • vocab_size (int)

  • ffn_dim_multiplier (float)

  • ffn_hidden_dim (int | None)

  • norm_type (NormType)

  • norm_eps (float)

  • activation (Activation)

  • max_seq_len (int)

  • rope_theta (float)

  • tie_embeddings (bool)

  • qk_norm (bool)

  • init_std (float)

  • model_type (str)

  • sdpa_backend (str)

  • attention_backend (str)

  • num_experts (int)

  • moe_top_k (int)

  • moe_frequency (int)

  • moe_router (str)

  • moe_shared_experts (int)

  • moe_aux_loss_weight (float)

  • moe_capacity_factor (float)

  • moe_sequence_aux_loss_weight (float)

  • moe_gradient_scale (bool)

  • moe_bias_schedule (str)

  • moe_packed_experts (bool)

  • moe_expert_ffn_multiplier (float)

  • moe_router_z_loss_weight (float)

Return type:

None

class kempnerforge.config.OptimizerConfig[source]

Bases: object

Optimizer settings.

name: str = 'adamw'
lr: float = 0.0003
weight_decay: float = 0.1
betas: tuple[float, float] = (0.9, 0.95)
eps: float = 1e-08
fused: bool = True
muon_momentum: float = 0.95
muon_ns_steps: int = 5
muon_adam_lr: float | None = None
muon_nesterov: bool = False
muon_adjust_lr: str = 'original'
muon_orthogonalize: str = 'local'
muon_router: bool = False
schedule_free_warmup_steps: int = 0
weight_decay_scope: str = 'exclude_norms_and_biases'
lr_multipliers: dict[str, float]
__init__(name='adamw', lr=0.0003, weight_decay=0.1, betas=(0.9, 0.95), eps=1e-08, fused=True, muon_momentum=0.95, muon_ns_steps=5, muon_adam_lr=None, muon_nesterov=False, muon_adjust_lr='original', muon_orthogonalize='local', muon_router=False, schedule_free_warmup_steps=0, weight_decay_scope='exclude_norms_and_biases', lr_multipliers=<factory>)
Parameters:
Return type:

None

class kempnerforge.config.ProfilingConfig[source]

Bases: object

Performance profiling settings.

enable: bool = False
start_step: int = 5
end_step: int = 8
trace_dir: str = 'profiler_traces'
__init__(enable=False, start_step=5, end_step=8, trace_dir='profiler_traces')
Parameters:
Return type:

None

class kempnerforge.config.SchedulerConfig[source]

Bases: object

Learning rate schedule settings.

name: SchedulerType = 'cosine'
warmup_steps: int = 2000
warmup_start_factor: float = 0.0
decay_steps: int | None = None
min_lr_ratio: float = 0.1
stable_steps: int | None = None
wsd_decay_type: str = 'cosine'
rex_alpha: float = 1.0
__init__(name=SchedulerType.cosine, warmup_steps=2000, warmup_start_factor=0.0, decay_steps=None, min_lr_ratio=0.1, stable_steps=None, wsd_decay_type='cosine', rex_alpha=1.0)
Parameters:
Return type:

None

class kempnerforge.config.TrainConfig[source]

Bases: object

Training hyperparameters.

"fp16" stays in the mixed_precision Literal only so the config loader passes it on to __post_init__, which rejects it with the reason; without it the loader would fail first with a generic Literal error.

batch_size: int = 8
seq_len: int = 2048
max_steps: int = 100000
grad_accum_steps: int = 1
grad_clip_norm: float = 1.0
seed: int = 42
data_seed: int | None = None
compile_model: bool = True
mixed_precision: Literal['bf16', 'fp16', 'fp32', 'fp8'] = 'bf16'
fp32_master_weights: bool = False
tf32: bool = False
activation_checkpointing: ActivationCheckpointing = 'none'
loss_fn: str = 'cross_entropy'
z_loss_weight: float = 0.0
ce_chunk_size: int = 0
fp32_cross_entropy: bool = False
shutdown_timeout_sec: float = 600.0
nccl_health_check_interval: int = 0
property param_dtype: torch.dtype

Resolve mixed_precision to the compute dtype, also the master weight dtype by default.

FP8 resolves to bf16 – FP8 is a compute mode, not a storage dtype. The storage dtype is master_dtype, which is fp32 with fp32_master_weights. A config switched to "fp16" after construction skips __post_init__, so "fp16" raises the same ValueError here.

property master_dtype: torch.dtype

Dtype the parameters, their gradients and the optimizer state are kept in.

property is_fp8: bool

Whether FP8 mixed precision is enabled.

property effective_data_seed: int

Seed for data shuffling / batch composition.

Falls back to seed when data_seed is unset, so existing configs reproduce their current trajectory. Kept independent from seed (parameter init) so stability studies can vary batch order while holding initialization fixed. Must stay identical across data-parallel ranks so the global shuffle is consistent before rank partitioning.

__init__(batch_size=8, seq_len=2048, max_steps=100000, grad_accum_steps=1, grad_clip_norm=1.0, seed=42, data_seed=None, compile_model=True, mixed_precision='bf16', fp32_master_weights=False, tf32=False, activation_checkpointing=ActivationCheckpointing.none, loss_fn='cross_entropy', z_loss_weight=0.0, ce_chunk_size=0, fp32_cross_entropy=False, shutdown_timeout_sec=600.0, nccl_health_check_interval=0)
Parameters:

  • batch_size (int)

  • seq_len (int)

  • max_steps (int)

  • grad_accum_steps (int)

  • grad_clip_norm (float)

  • seed (int)

  • data_seed (int | None)

  • compile_model (bool)

  • mixed_precision (Literal['bf16', 'fp16', 'fp32', 'fp8'])

  • fp32_master_weights (bool)

  • tf32 (bool)

  • activation_checkpointing (ActivationCheckpointing)

  • loss_fn (str)

  • z_loss_weight (float)

  • ce_chunk_size (int)

  • fp32_cross_entropy (bool)

  • shutdown_timeout_sec (float)

  • nccl_health_check_interval (int)

Return type:

None

kempnerforge.config.load_config(config_path=None, cli_args=None)[source]

Load a JobConfig from optional TOML file + CLI overrides.

The returned config has all sub-config __post_init__ validations applied. Cross-config validation (e.g., parallelism vs world_size) requires calling config.validate(world_size=…) separately at distributed setup time.

Modules named by the top-level plugins key are imported before the overlay, so components they register are visible to those validations.

Parameters:

  • config_path (str | Path | None) – Path to a TOML config file (or None for defaults).

  • cli_args (list[str] | None) – CLI arguments to parse (defaults to sys.argv[1:]).

Returns:

A JobConfig with layered defaults → TOML → CLI overrides.

Return type:

JobConfig

Modules

adapter

Adapter (connector) configuration.

checkpoint

Checkpoint configuration.

data

Data pipeline configuration.

distributed

Distributed parallelism configuration.

eval

Evaluation configuration.

job

Top-level job configuration aggregating all sub-configs.

loader

Config loading: TOML files → dataclass configs with CLI overrides.

metrics

Metrics configuration.

model

Model architecture configuration.

optimizer

Optimizer configuration.

profiling

Profiling configuration.

registry

Central registry for named components.

scheduler

LR scheduler configuration.

schema

Backward-compatible re-exports.

training

Training configuration.

video

Video input configuration.

vision

Vision-encoder configuration.

vlm

VLM (vision-language model) configuration.