Runtime and Trainer API

Runtime functions

seed_everything(seed: int) -> None

Seeds:

• Python random;
• PYTHONHASHSEED in the current process environment;
• NumPy if it can be imported;
• PyTorch CPU;
• all CUDA devices when CUDA is available.

NumPy failures are ignored.

build_optimizer(model, config, device) -> torch.optim.Optimizer

Builds AdamW from config.optimizer:

lr, betas, eps, weight_decay

When optimizer.fused is null, supports_fused_adamw(device) probes CUDA. A forced fused value is also attempted. TypeError, RuntimeError, or ValueError during construction falls back to ordinary AdamW.

build_scheduler(optimizer, config) -> LambdaLR

Creates warmup and either cosine or constant decay:

warmup: max(1e-8, (step + 1) / warmup_steps)
constant: 1.0
cosine: min_ratio + (1 - min_ratio) * 0.5 * (1 + cos(pi * progress))

The scheduler advances once per global optimizer update, not per microbatch.

build_logger(config) -> MetricLogger

Maps logging.level, file, json_file, and every_steps. It does not attach hooks.

build_runtime(...)

build_runtime(
    config: SpeedtronicConfig | dict[str, Any],
    *,
    resume: bool = False,
    device: str | torch.device | None = None,
    max_steps: int | None = None,
) -> tuple[Trainer, CheckpointManager]

The function validates, seeds, resolves hardware, and constructs every local component. If distributed mode is enabled, it also creates DumbDiLoCoCoordinator.

Effective target precedence:

explicit build_runtime(max_steps=...)
→ data.max_steps
→ run.max_steps

The expression is implemented as max_steps if provided else (data.max_steps or run.max_steps).

Resume lookup happens after the trainer is created. A missing checkpoint is a warning, not an error.

train_from_config(...) -> TrainResult

Builds the runtime and immediately calls Trainer.fit().

Aliases:

run = train_from_config
train = train_from_config

SyncCoordinator

@runtime_checkable
class SyncCoordinator(Protocol):
    def start(self) -> None: ...
    def after_optimizer_step(self, model: Any, step: int) -> bool | None: ...
    def stop(self) -> None: ...
    def state_dict(self) -> dict[str, Any] | None: ...
    def load_state_dict(self, state: dict[str, Any]) -> None: ...

after_optimizer_step() may return true to request inner optimizer state reset. DumbDiLoCo returns true after a successful upload or global installation.

TrainResult

@dataclass
class TrainResult:
    steps: int
    samples: int
    tokens: int
    final_loss: float | None
    elapsed_s: float
    metrics: list[dict[str, Any]]

Field	Meaning
steps	Final absolute global step, including resumed state
samples	Cumulative sample counter
tokens	Cumulative token counter
final_loss	Mean final-update loss for this invocation, or null when no update ran
elapsed_s	Duration of this invocation
metrics	Every metric record generated by this invocation

metrics grows on every step even when the logger suppresses output by cadence.

Trainer

Trainer(
    model,
    optimizer,
    dataloader,
    *,
    device,
    config=None,
    scheduler=None,
    precision=None,
    logger=None,
    checkpoint_manager=None,
    coordinator=None,
    max_steps=None,
    start_step=0,
)

The constructor:

• converts a dictionary config;
• stores the model/optimizer/loader/device;
• creates a default logger;
• resolves a precision plan when omitted;
• establishes cumulative counters;
• configures gradient checkpointing;
• optionally constructs torch.compile(model).

Compatibility aliases:

TrainingEngine = Trainer
SpeedtronicTrainer = Trainer

Batch movement and counting

_move_batch() recursively transfers tensors in dictionaries, tuples, and lists to the selected device with non_blocking=True.

_batch_size_and_tokens() chooses:

• dictionary input_ids, then inputs; or
• the first tuple/list field.

For a multi-dimensional dictionary input with attention_mask, tokens equal the mask sum. Unsupported values report one sample and one token.

Loss resolution

Mapping output

• {"loss": tensor} uses the loss directly.
• {"logits": 3d_tensor} without loss triggers trainer-side causal cross-entropy.

Tensor output

A bare tensor is a direct loss, even if it is 3-D.

Tuple/list output

• A scalar first tensor is a direct loss.
• A 3-D first tensor is treated as causal logits when labels have compatible shapes.

Inferred causal loss

For logits (B, T, V):

• same-length labels are treated as already next-token aligned and use the full logit row;
• logits are truncated to logits[:, :-1] only for (B, T-1) labels;
• -100 is ignored;
• all-ignored labels produce a differentiable zero tied to the logits.

The bundled datasets and the bundled model share this one-label-per-position contract, so the first target is not skipped.

Accumulation and optimization

For K = accumulation_steps:

1. Run each of K microbatches.
2. Mean any non-scalar loss.
3. Backpropagate loss / K immediately.
4. On the final microbatch, optionally unscale and clip.
5. Perform one optimizer update.
6. Clear gradients with set_to_none=True.

The reported loss is the arithmetic mean of the K unscaled microbatch losses.

Feature setup

Gradient checkpointing

If enabled, the trainer calls:

model.set_gradient_checkpointing(True)

A missing hook, attribute error, or any hook exception logs a warning and continues.

Compilation

torch.compile construction failure leaves the eager model active. Any exception from a compiled forward disables compilation and reruns that batch eagerly. The fallback is broad: it can hide a non-compilation model exception if eager execution succeeds.

Global-step lifecycle

The target is max(self.step, requested_target). fit() always calls coordinator stop() in finally, even on failure.

fit(max_steps=None) -> TrainResult

1. Validate positive target.
2. Move model to the device.
3. Select original or compiled model.
4. Create a new GradScaler and restore deferred scaler state.
5. Zero gradients.
6. Capture invocation start time/counters.
7. Wrap loader in infinite_batches.
8. Start coordinator if needed.
9. Apply deferred coordinator checkpoint state after startup.
10. Emit train_start.
11. Run complete global steps.
12. Emit train_end and return.

Trainer.train is an alias for fit.

Resume

Trainer.resume(state) immediately restores:

• model state;
• optimizer state;
• scheduler state;
• global step/sample/token counters;
• RNG state.

Scaler and coordinator state are staged for fit(); v2 coordinators can apply prepared baseline state during startup before any remote refresh.

Saved precision and config copies are informational; they are not compared with or applied to the current runtime.

Per-step metrics

Key	Meaning
step	New cumulative step
loss	Mean microbatch loss
lr	First parameter-group LR after scheduler advancement
steps_per_sec	Invocation-local step rate
samples_per_sec	Invocation-local sample rate
tokens_per_sec	Invocation-local token rate
samples	Cumulative samples
tokens	Cumulative tokens
micro_batches	Microbatches in this update

The reported LR is one scheduler position ahead of the LR that produced the just-completed update.

v2 integration

build_runtime() resolves precision, validates startup shapes, builds AdamW or hybrid Muon, and passes a resolved PrecisionPlan to the trainer. The trainer starts the opt-in conservative CUDA stage scheduler when ooo_backprop is enabled and removes its hooks when the fit ends.

See v2 optimizers, out-of-order backprop, and shape validation.

Important caveats

• A GradScaler may skip an overflowing optimizer update; the trainer emits an event and keeps the global step unchanged for that attempt.
• The trainer explicitly switches the model to train() mode before fitting.
• fit() creates a new scaler each invocation.
• A stopped v2 coordinator resets its started flag; a later fit() can start fresh background lanes, subject to the bounded shutdown timeout.
• A compile runtime failure reruns the same batch eagerly, which can duplicate side effects.
• Every per-microbatch loss is copied to CPU, synchronizing CUDA on each microbatch.

The generated source inventory lists all private methods and exact AST-derived signatures.