Bitwise Trainer-Sampler Alignment

Bitwise trainer-sampler alignment is a training stability and debugging technique that ensures the computation performed during training (by the trainer) produces bit-exact results to the computation performed during inference (by the sampler). In most neural network training pipelines, small numerical differences accumulate between the forward pass used for training and the actual inference computation — due to differences in operator fusion, precision, randomness, and parallelism strategies. These differences are typically harmless for large-scale training but become problematic when they cause the model to behave differently during training evaluation versus production inference.

The specific sources of trainer-sampler mismatch that bitwise alignment addresses include: different operator fusion patterns between training and inference runtimes, different random number generation sequences (especially in dropout, sampling, and annealing), different accumulation orders in reduction operations (which vary across hardware platforms and parallelism strategies), and different kernel implementations for the same mathematical operation. In large-scale distributed training, even the order of floating-point additions can vary depending on which GPU processes which part of the computation, leading to different rounding patterns.

The technique was found useful for training stability as well as for debugging the various components of the interaction model system. By implementing batch-invariant kernels — kernels whose outputs are independent of the batch dimension — and using deterministic communication primitives (NVLS for all-reduce and reduce-scatter), the training pipeline achieves bit-exact reproducibility across different batch sizes and parallelism configurations. This makes it possible to reproduce any training behavior in inference, and to use inference-time debugging tools on training artifacts.

Interestingly, the batch-invariant kernels were sometimes faster end-to-end than the standard kernels they replaced, because the custom communication kernels were not only batch-invariant but also lower-latency. This illustrates a common pattern in high-performance computing: constraints that improve code quality (determinism, reproducibility) can also improve performance when they force the elimination of unnecessary flexibility in the implementation.