The Stale Batch

I learned this one from a run that did exactly what we asked and still stopped being the same experiment.

We resumed from a checkpoint. Loss kept falling. Nothing crashed, no dashboard turned red, and the continuation looked ordinary enough to trust. It was not. An asynchronous loader worker had already fetched one batch beyond the cursor recorded in the checkpoint. Restarting destroyed the in-memory queue, so the resumed process consumed a different next batch. From that step onward we had a new experiment wearing the old run's name.

I would like to blame the thread. The less flattering answer is that I had been treating “the data position” as one number because one number is pleasant to serialize. We had already found config drift, pipeline changes that never reached the checkpoint format, and supposedly comparable runs built with different tokenizers. The stale batch finally made the general failure impossible to ignore: a research trainer is not trustworthy because it can save weights. It is trustworthy only if it can say what happens next.

The batch that existed only in RAM

Prefetch creates two honest positions in the data stream. There is the last batch the model consumed, and there is the point the producer has reached while selecting data, advancing cursors, materializing tensors, and filling its queue. The old checkpoint knew the second position but not the tensors promised between the two.

That is how the bug worked. The cursor advanced, the promised batch lived only in memory, and the checkpoint saved the cursor. On restart the cursor returned but the batch did not. The loader continued after data the model had never seen (a very efficient way to achieve exact restoration of the wrong state).

There are three direct contracts I know how to trust: checkpoint the queue and producer state; stop the producer and drain it at a known boundary before saving; or checkpoint the consumer cursor and deterministically regenerate any prefetched work after restart. The prototypes did none of them. The current nmoe loader instead sets prefetch_depth=0 on its supported training paths. We give up overlap because the queue is not yet part of the checkpoint contract. If prefetch returns, it has to bring one of those contracts with it.

A checkpoint is a claim about the next step

The stale batch changed how I draw the boundary. At training step t, the relevant state is closer to

where θ is model state, o is every optimizer state, r is the random-number-generator state, c is the consumed and producer-side data position, q is already-produced data, τ is the training clock, and p is the resolved execution plan.

Exact resume means restoring enough of S_t that the next transition receives the same inputs and makes the same stochastic choices as an uninterrupted run, inside whatever determinism envelope we actually promise. In the current loader, q is intentionally empty. This keeps the implemented state aligned with the state we claim to restore, at the cost of overlap.

nmoe/checkpoint.py divides the serialized pieces along the distributed model boundary. Rank 0 writes the replicated dense and router state, token count, run metadata, configuration fingerprint, and resolved data-plan bundle. Every rank writes its local expert shard, optimizer state, loader state, and Torch and CUDA RNG state. Separately stepped optimizers such as Muon are explicit checkpoint entries; if one is active and its state is missing, resume fails.

The loader contributes mixture accumulators, per-source cursors, emitted-sequence counts, and the global sequence index. The checkpoint's run_info also carries a Git SHA field (which can fall back to unknown) and the world size. Resume records those facts but does not enforce either one. The fuller physical topology, kernel behavior, shard inventory, and dataset bytes remain outside the checkpoint contract. A seed pins none of them.

The same mistake keeps changing clothes

The stale batch was one instance of a larger error: allowing the process to continue after the meaning of the run has changed. Checkpoints are the obvious boundary, but configuration, distributed evaluation, and telemetry can all fail the same way.

On resume, nmoe checks a fingerprint of the resolved public configuration fields and, for mixture-backed flows, the hash of the resolved mixture plan. Either mismatch stops the run before it loads a different experiment into familiar weights. Those hashes do not enforce equivalent code, environment, world size, shard inventory, or dataset contents; the direct data_path route currently uses a sentinel plan hash. They are specific guards, not a notarized universe.

Expert-parallel evaluation belongs to the distributed program too. A forward pass reaches experts owned across the group, so ranks must execute compatible collectives, shapes, and ordering. The live CORE path takes the conservative route: rank 0 broadcasts the example index, every rank evaluates the same prompt, and only rank 0 accumulates correctness. That avoids variable-length examples producing incompatible traces.

There is still a correctness hole. A rank-local exception becomes a wrong answer without distributed error consensus. An exception inside a collective can strand peers; one after the collective can leave ranks disagreeing about whether the example succeeded. A serious score needs synchronized failure reporting and fail-fast abort semantics. The current implementation does not finish that contract.

Metrics can leave the same kind of counterfeit continuation. nmoe records loss, throughput, router health, and gradient statistics as one parquet file per logged step through a DuckDB-backed writer. Today writer creation, insertion, flush, and shutdown can fail while training continues. That may be useful during bring-up; it is not acceptable for a run whose evidence is required. Telemetry has to be declared optional before launch, or a write failure has to invalidate the run.

The spine should be boring

After a bad run I need four answers without an expedition: which path loaded the data, which state entered the checkpoint, which distributed evaluation ran, and whether the metrics writer survived. Much of nmoe is shaped by keeping those answers short. Precision is explicit. Data paths and mixture plans are explicit. The supported expert transport has one path instead of a menu of distributed modes that can quietly stop being comparable.

The trainer still needs room for weird experiments. The loader, checkpoint, evaluation, and evidence path should be boring enough that the weirdness belongs to the experiment, rather than to whether the same batch existed after restart.