base
base
¶
a very basic trainer
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