Training

StableTested and in active use. Breaking changes will be called out. CUDA-graph capture in the nn Trainer is off by default — see below.

Trainer is the supervised training loop: classification and regression on a dataset, distinct from the RL drivers in deep_agents/training.

from noeira.nn.training.trainer import Trainer

The signature

Trainer[
    MODEL: Module,
    NC: Int,                 # number of classes / output width
    IN: Int,                 # input dimension
    BATCH: Int,
    target: StaticString,    # "cpu" | "gpu"
    POLICY: AMPPolicy = NoAMP,
    OPT: Optimizer = Adam,
    USE_TRAIN_CUDA_GRAPH: Bool = False,
]

Construct with make, passing the initializer:

var trainer = Trainer[Net, NC, IN_DIM, BATCH, "gpu"].make[Kaiming](
    Optional(ctx), lr=1e-3
)

The trainer owns the model’s parameters, the optimizer state and the loss.

Whole-dataset GPU training

train_gpu uploads the dataset to the device once and runs the entire training loop on-device, slicing each batch as a zero-copy sub-buffer view:

var result = trainer.train_gpu[MNIST.N_TRAIN, MNIST.N_TEST](
    ds.train_images, train_y,      # training inputs, one-hot targets
    ds.test_images, ds.test_labels,
    Optional(ctx),
    epochs=5,
    shuffle=True,                  # on-device per-epoch shuffle
)

for acc in result.epoch_test_top1:
    print(acc)

No per-batch host-to-device transfer. The shuffle is a device kernel, so it does not force a round trip either.

For finer control there is train_epoch and eval_top1.

Mixed precision (AMP)

Set the POLICY parameter to run in bf16:

var trainer = Trainer[Net, NC, IN_DIM, BATCH, "gpu", BF16Policy].make[Kaiming](
    Optional(ctx), lr=1e-3
)

The design is bf16-only true autocast — no loss scaling, because bf16’s exponent range makes it unnecessary. Activations are stored at the flow dtype, which is where the memory saving comes from: a module declares comptime ACT_DT = bfloat16 and combinators derive theirs from their children, typing inter-module buffers accordingly.

Casts happen only at boundaries. The fp32 path (ACT_DT == DT) is the existing code with no casts and no extra buffers — bit-identical to a non-AMP build.

pixi run -e nvidia mojo run -I . examples/nn/mlp/mlp_mnist_bf16_amp_gpu.mojo
pixi run -e nvidia mojo run -I . examples/nn/resnet/resnet20_cifar10_bf16_amp_gpu.mojo

CUDA-graph capture

USE_TRAIN_CUDA_GRAPH=True captures the per-batch device compute — batch gather, zero-grad, forward, loss accumulate, vjp, optimizer step — as a graph and replays it once per batch. The gather reads its offset from a device counter, so one capture serves every batch of the epoch. Per-epoch shuffling and augmentation run eagerly, in place, before the replays.

An eager run and a captured run of the same seed give equal losses and top-1 every epoch and byte-identical checkpoints (tests/nn/test_trainer_capture_shuffle_aug.mojo, with real replays on an RTX 5090).

Off by default, and for a specific reason

Stream capture forbids any device allocation inside the captured region, and MAX allocates in places the trainer does not control. The best-known one is split-K: linalg.matmul allocates its reduction workspace on every call that partitions K. The six weight-gradient (dW) GEMMs in Linear, LinearAct, NoisyLinear, Embedding, Conv2D and Conv2DTranspose now own that workspace (noeira/nn/core/splitk_gemm.mojo), but only where MAX itself would take the multistage path — sm_80, sm_89 and sm_120. H100 and datacenter Blackwell fall back to linalg.matmul, which still allocates, and other per-call allocations inside MAX’s memory manager also abort capture. Enabling capture can therefore still abort in a shape- and hardware-dependent way.

Owning the workspace also pays outside capture: on an RTX 5090, ACT’s dW GEMMs went from 3.05 to 0.35 ms per step once the partition count was tuned.

It is fp32-only: under a bf16 POLICY the flag is ignored. On non-NVIDIA targets it is a no-op and the run is bit-identical.

The RL drivers in deep_agents/training are a different code path and do capture by default; see CPU and GPU paths.

Other trainers

Trainer	Purpose
AutoregressiveTrainer	sequence models — GPT on TinyShakespeare
DistributedAutoregressiveTrainer	the same over N ranks, DDP or ZeRO-1/2/3 — see Distributed training
Augmenter	data augmentation for image training

AutoregressiveTrainer builds its batches on the device (DEVICE_BATCH=True, the default). The corpus is resident as int32. One kernel draws each row’s window start from a hash of (seed, step, row) and writes the input and target one-hots, so under USE_TRAIN_CUDA_GRAPH the batch build is inside the graph. DEVICE_BATCH=False keeps the host path as the baseline.

Against PyTorch

Each nn example has a PyTorch twin in tools/nn/torch_nn_reference.py: same recipe, batch size and data files, with the data resident on the GPU as ours is. It runs eager, under torch.compile, or captured in CUDA graphs (reduce-overhead), with TF32 on and fused Adam. Time per training step on one RTX 5090:

Model	noeira (device batch + CUDA graph)	PyTorch 2.12, compiled + CUDA graphs
TinyShakespeare GPT, fp32	19.40 ms	21.35 ms
TinyShakespeare GPT, bf16 AMP	10.74 ms	10.81 ms
ResNet-20, CIFAR-10	2.84 ms	2.84 ms
CIFAR CNN	1.35 ms	1.33 ms

The GPT benchmark is examples/nn/benchmarks/gpt_cuda_graph_benchmark_gpu.mojo (-D GPT_BENCH_DTYPE=bf16 for the AMP column). Final validation losses over three seeds agree within seed noise: fp32 GPT 1.824 / 1.847 / 1.824 with the TF32 attention, against 1.836 / 1.837 / 1.828 with the fp32 kernels.

Datasets

from noeira.nn.datasets import MNIST, CIFAR10, TinyShakespeare

Dataset	Notes
MNIST	60k train / 10k test, 28×28 grayscale
CIFAR10	50k / 10k, 32×32 colour — Hugging Face parquet by default
TinyShakespeare	character-level text
lewm_pusht	PushT demonstrations, for the world-model work

Datasets download on first use and cache locally.

These are supervised data — images plus labels, or a token stream — and they stay here rather than moving to the trajectory store on purpose: they have no episode axis, so they are not trajectories. Episodic data (offline RL datasets, demonstrations, replay dumps) belongs in the data layer instead.

pyarrow crashes the embedded interpreter

The CIFAR-10 loader reads Hugging Face parquet, and pyarrow cannot run inside Mojo’s embedded Python — it must be subprocess-isolated. This is handled in the loader; it matters if you write your own parquet-backed dataset.

More generally, Python temporaries created inside a Mojo function are not released until the function returns. A long-running loop holding PythonObjects leaks in a way that RSS sampling will not reveal. Keep hot Python loops short, or move them into an exec’d block.

Worked examples

# MLP
examples/nn/mlp/mlp_mnist_training_storage_gpu.mojo
examples/nn/mlp/mlp_mnist_training_storage_cpu.mojo

# convolutional
examples/nn/conv2d/conv2d_mnist_training_storage_gpu.mojo
examples/nn/resnet/resnet20_cifar10_training_storage_gpu.mojo

# transformer
examples/nn/transformer/gpt_tinyshakespeare_training_gpu.mojo
examples/nn/vit/vit_cifar_training_storage_gpu.mojo

Benchmarking discipline

Benchmark on NVIDIA; treat Apple as a parity target

Kernel performance rankings invert between Metal and CUDA. An optimization that wins on one can lose on the other, and Apple’s small-N kernel launch overhead distorts microbenchmarks further.

Use Apple to prove correctness and NVIDIA to prove speed.

Next

• Optimizers and losses — what the Trainer drives
• Checkpoints — saving the result
• Modules and Params — the storage being trained