Index
training
¶
BaseTrainer(spec: ExecutionSpec, base: Optional[Node] = None)
¶
Bases: RestoreableJob[C], CometLoggingJob[C], Generic[C, M]
Generic pretrainer for GPT-style models.
DATASET must be declared by concrete trainers (use [] for custom batching). ANALYSIS declares analysis classes run sequentially before validation and checkpointing, including the terminal step. Their CONFIG schemas join the trainer's config; conflicting defaults require explicit configuration. training/analyze disables borrowed analysis creation and execution.
Build a basic trainer
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
spec
|
ExecutionSpec
|
execution specification |
required |
Raises:
| Type | Description |
|---|---|
AssertionError
|
if topology is not provided in spec |
schedule: optax.Schedule
cached
property
¶
Build and cache the declared schedule, or a constant learning rate.
optimizer: optax.GradientTransformation
cached
property
¶
Build and cache the optimizer using the trainer's cached schedule.
initialize() -> None
¶
Initialize a new training state from scratch.
surgery(partial: PyTree[bool]) -> PyTree[Any]
¶
Initialize only state leaves missing from a restored checkpoint.
apply(state: PyTree[Any], metadata: Dict[str, Any]) -> None
¶
Install the restored state, including its completed optimizer step.
evaluator() -> Optional[Evaluator[M]]
¶
define what evaluator to use
trace(state: train_state.TrainState, batch: PyTree[jax.Array], key: jax.Array, *, sharding: ParameterShardingContext) -> Any
classmethod
¶
Pure inspection execution; state, batch and RNG are explicit inputs.
find(module_type: Type[Module]) -> List[Tuple[str, ...]]
¶
Discover paths using trace(); the default reuses the node's cached batch.
debug(path: str | Tuple[str, ...]) -> Debugger
¶
Open the first invocation at an exact path; close after exploration.
batch(slice: str = 'train') -> PyTree[np.ndarray]
¶
get the next batch from the dataset strategy
train_step(state: train_state.TrainState, batch: PyTree[jax.Array], key: jax.Array, accumulate_steps: int, *, sharding: ParameterShardingContext, dtype: str = 'float32', fsdp: bool = False, activation_checkpointing: bool = False) -> Tuple[train_state.TrainState, jax.Array, Any, jax.Array]
classmethod
¶
Compute gradients over S micro-batches and apply one optimizer step.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
state
|
TrainState
|
Current training state |
required |
batch
|
PyTree[Array]
|
(x, y, padding_mask) each with shape (S, B, T) S = accumulation steps, B = batch size, T = sequence length |
required |
key
|
Array
|
PRNG key for dropout |
required |
accumulate_steps
|
int
|
Number of micro-batches (S) |
required |
Returns:
| Type | Description |
|---|---|
Tuple[TrainState, Array, Any, Array]
|
(updated_state, loss, meta, grad_norm); meta is from the last micro-batch |
val_step(state: train_state.TrainState, batch: PyTree[jax.Array], *, sharding: ParameterShardingContext) -> Tuple[jax.Array, jax.Array, Any]
classmethod
¶
Compute validation loss over S micro-batches.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
state
|
TrainState
|
Current training state |
required |
batch
|
PyTree[Array]
|
(x, y, padding_mask) each with shape (S, B, T) S = accumulation size, B = batch size, T = sequence length |
required |
Returns:
| Type | Description |
|---|---|
Tuple[Array, Array, Any]
|
(loss_sum, token_count, last_meta) |
mfu() -> None
¶
Prepare theoretical compute seconds per update; None means unavailable.
checkpoint() -> None
¶
Save the current training state, including its optimizer step.
run() -> None
¶
main entry point to run training, called on all nodes
KLDivergenceTrainer(spec: ExecutionSpec, base: Optional[Node] = None)
¶
Bases: BaseTrainer[C, M], Generic[C, M]
Two-stage trainer with KL-divergence penalty.
- Stage 1 – standard language-model pretraining (cross-entropy only).
- Stage 2 – pretraining loss plus
beta * KL(policy || reference)where the reference policy is a frozen snapshot taken at the stage boundary.
The KL penalty is approximated as the difference in per-token NLL
between the current model and the reference model on the same batch:
kl_penalty = policy_loss - sg(reference_loss).
KLDivergenceTrainerConfig(batch_size: int = field('training/batch_size', default=512), per_device_batch_size: int = field('training/per_device_batch_size', default=(-1)), total_tokens: List[int] = field('training/tokens', default_factory=(lambda: [1000000000, 100000000])), lr: float = field('optimization/lr', default=0.0003), warmup_pct: float = field('training/warmup_pct', default=0.01), decay_pct: float = field('training/decay_pct', default=0.1), validate: bool = field('training/validation', default=True), evaluate: bool = field('training/evaluate', default=True), analyze: bool = field('training/analyze', default=True), block_size: int = field('architecture/block_size', default=512), param_dtype: str = field('architecture/dtype/param', default='float32'), activation_dtype: str = field('architecture/dtype/activation', default='bfloat16'), report_interval: int = field('logging/report_interval', default=32), checkpoint_interval: int = field('logging/checkpoint_interval', default=1024), validation_interval: int = field('logging/validation_interval', default=512), validation_steps: int = field('training/validation_steps', default=2048))
dataclass
¶
Bases: BaseTrainerConfig
Config for two-stage KL-divergence trainer.
total_tokens is a two-element list: [stage1_tokens, stage2_tokens].
Both stages use the trainer's static DATASET declaration.