kempnerforge.config.training

Training configuration.

Classes

ActivationCheckpointing

TrainConfig

Training hyperparameters.
class kempnerforge.config.training.ActivationCheckpointing[source]

Bases: StrEnum

none = 'none'
full = 'full'
selective = 'selective'
__new__(value)
class kempnerforge.config.training.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