16. Architecture

1. Package and lifecycle overview

NNx is organized around two public entry points (NNModel / Trainer), a training-hook family (train_step_fn, eval_step_fn, and trainer_step_fn), and content-addressed persistence under runs/<id>/. Hook-producing modules inject behavior into the orchestrators; model transforms, exporters, inference helpers, and diagnostics compose around them.

The inference branch, nnx.decisions, sits beside the training path. Typed questions go to a DecisionProvider and come back as validated results (Typed decisions). For a trained classifier, the provider is FixedHeadProvider over NNModel.predict_proba.

The training loop owns callback dispatch, scheduler updates, phase checkpoint cadence, and incremental NNRun persistence. A scheduler updates once per epoch by default. A clock="optimizer_update" scheduler updates once per committed optimizer update.

See Concepts §1 for the full written breakdown.

2. Lifecycle order

For each successfully started training run, NNx calls on_train_begin, then dispatches epoch and batch work. A completed epoch aggregates validation through the built-in path or eval_step_fn, updates each epoch-clock scheduler once, and dispatches on_epoch_end.

An optimizer_update-clock scheduler is stepped instead right after each committed update of its optimizer. It is never stepped per microbatch, masked window or skipped step. One of these reports the committed update:

• the update engine;
• default_train_step;
• finalize_step (the built-in paradigm steps);
• a step function calling ctx.report_update(...).

An objective's on_optimizer_update callbacks see every optimizer's event of a commit before any clock steps.

A run that declares metrics or a named monitor also records the whole-epoch training summary and the monitor's decision before the scheduler update. That one decision drives the plateau scheduler, BEST and any EarlyStopping given the same monitor (Concepts §6.4).

In an objective run the shared update engine accumulates each microbatch's loss terms. At the end of every update window it runs unscale → clip → step. It dispatches on_optimizer_update once per committed update, before the epoch's validation and the durable commit order (Concepts §6.5).

Durable state then commits in order: run history, LAST, phase/BEST, and deferred callback checkpoints. Finalization calls on_train_end in reverse callback order. Both NNModel and Trainer refresh LAST after finalization. Callback mutations and topology-transform metadata are therefore present in the persisted checkpoint. The exception is an NNModel stop at an update boundary (ctx.stop_at_update): LAST then stays at the last committed epoch.

On failure, callbacks whose begin hook completed are still finalized. Every cleanup hook is attempted; cleanup errors do not mask an exception already raised by training. A failed LAST commit rolls history back; failures after LAST retain the durable history/checkpoint pair. On load, history newer than LAST is truncated, while an empty or corrupt LAST is rejected rather than treated as a request to erase history.

Each text file is replaced on its own through an owned temporary file. A failed write leaves that destination's previous bytes and removes its temporary. The history files are not one multi-file transaction, however, and a hard kill can still leave a stale temporary behind.

With an opt-in history journal (history=HistoryJournal(...), Concepts §4.5) the loop keeps a bounded window of records in memory. The "run history" step appends the epoch's records as immutable chunks, indexes them and publishes the journal manifest. This still happens before LAST, which stays the commit marker.

Readers show records up to LAST's epoch only. A crash between the manifest and LAST therefore leaves an uncommitted tail that is ignored. A failed LAST republishes the previous manifest. Callbacks see the window as ctx.idps unless they declare history_access = "full". A continuation writes its own run and chunk files, never its source's.

3. Single-node data parallelism (nnx.distributed)

train(params, distributed=DDP()) runs inside the process group torchrun started — NNx adds no runtime of its own:

torchrun --standalone --nproc_per_node=2 train.py
   ├── rank 0 (writer) ── init_process_group() ── train_loader / validation_loader ── model.train(distributed=DDP()) ── shutdown()
   └── rank 1          ── init_process_group() ── train_loader / validation_loader ── model.train(distributed=DDP()) ── shutdown()
                                  (gloo on CPU, nccl with one GPU per rank)

Teardown. nnx.distributed.shutdown(timeout_seconds=60) ends the recipe. Every rank meets at a teardown barrier, then the group is destroyed. No rank therefore tears down its connections while a peer still uses them.

The wait is bounded with Gloo; with NCCL on older torch, it is bounded by the group's own timeout. A rank that does not arrive in time makes the others raise RuntimeError instead of blocking in the call. The group is then left for process exit. The process may still wait at exit until the missing rank arrives or the group's timeout expires. Without a process group, or after a successful call, it does nothing.

Partitions. nnx.distributed.train_loader(dataset, batch_size, policy=, seed=) deals each epoch's permutation (seeded by seed + epoch) to the ranks. It first pads ("pad" repeats leading rows) or drops ("drop") the tail, so every rank takes the same number of steps. A mismatch fails on every rank before the epoch. validation_loader(dataset, batch_size) shards rows rank::world unpadded: every row id is scored exactly once, a rank may have none, and the validation loop runs no collective.

One update. Per microbatch:

forward + backward through DDP  (no_sync until the window's last microbatch)
   │   loss numerator only; finiteness agreed on every rank before backward
   ▼
window end:  W = all_reduce(sum of the ranks' loss denominators)
   │   W == 0 for a task model → skip the update on every rank
   ▼
grad = DDP average × world_size / W        (an additive loss: × world_size)
   ▼
clip → optimizer.step()  — identical parameters on every rank

With this order, each committed update equals one process training on the union of the ranks' batches, unequal valid counts included.

Records and decisions. Each batch's outputs, targets and loss terms are gathered and scored the same way on every rank. Validation gathers each rank's scored shard once. Records are therefore global and identical on every rank. Monitors, EarlyStopping, BEST and the plateau scheduler decide alike. Every rank returns the same NNRun (id, records and the writer's provenance).

Each step gathers the batch's outputs and targets, and validation gathers every scored output once. That traffic suits supervised models, not models with huge outputs such as large-vocabulary logits.

One writer. Only the writer rank (DDP(writer_rank=0)) holds the run lease. It alone writes history, LAST / BEST / phase checkpoints, deferred callback checkpoints and the provenance attempt. Each epoch's commit is agreed, so a failed write stops every rank.

Checkpoints hold the canonical module's keys. DDP wraps model.net per fit and is never assigned to it. Checkpoints also hold the world size, the partitions and every rank's RNG streams. A resume with the same world size and partition restores them all. Any other resume fails before anything is restored. So does a single-process stateful resume of a distributed checkpoint; resume_mode="weights_only" starts fresh from its weights.

The writer rank alone verifies the resume point (its files' SHA-256). It shares the manifest or the refusal, so the other ranks never hash it. A torn, mixed or corrupted point is refused on every rank with the same reason. Resume, component restore and epoch-end callback failures are agreed too, and so are each default step's loss finiteness and target checks. Anything else failing on one rank only ends through the process group's timeout and torchrun stopping the other ranks.

            writer rank                        other ranks
lease       runs/.leases/<id>.lock             —
history     runs/<id>/idps.csv, run.yaml       in memory (same records)
checkpoint  LAST / BEST / phase / deferred     —
callbacks   "all" + "writer" + writer_only()   "all"

Callbacks declare their rank behaviour (Callback.distributed). "all" callbacks (EarlyStopping, LRMonitor) run everywhere and must not write. "writer" callbacks (ModelCheckpoint, plain function callbacks) run on the writer only and may not carry checkpointed component state. TensorBoardCallback and WandbCallback write when constructed, so a borrowed instance is refused; nnx.distributed.writer_only(lambda: TensorBoardCallback(...)) builds one on the writer alone. A callback that declares nothing is refused before training, on every rank.

Exactness holds for deterministic forwards. Batch normalization is refused, because its running statistics would differ per rank. Dropout draws each rank's own stream, matching one process only in distribution.

Out of scope: multi-node or elastic runs, FSDP / DeepSpeed, mixed precision, compile=, custom train or validation steps and objectives, history journals, graph neighbour sampling. ExperimentPlan.fit, nnx.search.search, Trainer.train and NNModel.train without distributed=DDP() train in one process and refuse a group of more than one rank. See examples/ddp_supervised.py.