Tools

AI-generated text

Scaling Bitwise-Deterministic Pretraining with NVIDIA Megatron Core

Bitwise determinism ensures independent and checkpoint-resumed training runs follow identical numerical trajectories, which simplifies debugging, validation, and reproducible recovery for very large models.

Scaling Bitwise-Deterministic Pretraining with NVIDIA Megatron Core

Bitwise determinism means that independent runs and checkpoint-resumed runs follow the exact same numerical trajectory when data order, architecture, recipe, parallelism, software, runtime settings, and hardware are fixed. NVIDIA Megatron Core targets two guarantees:

  • Independent reproducibility: two runs launched from the same initial state remain bitwise identical at every training step.
  • Bitwise checkpoint resume: a run that saves and restores one or more checkpoints remains identical to an uninterrupted run.

For very large models trained on thousands of GPUs, bitwise determinism delivers concrete benefits:

  • Replay failures: reproduce a loss spike at the same step to isolate the cause by changing one component at a time.
  • Preserve checkpoint continuity: resume interrupted training without changing the numerical trajectory.
  • Detect unexpected differences: mismatches can reveal checkpoint corruption, software bugs, or hardware issues.
  • Compare changes: separate intended numerical effects from run-to-run variation.

How to verify determinism

Printed loss values are not sensitive enough: two runs can print the same rounded loss while differing in lower-order bits of gradients or parameters. Stronger validation fingerprints selected numerical metrics at every step (training loss, language-modeling loss, load-balancing loss, multi-token prediction loss, gradient norm). Run the same configuration twice and compare these fingerprints step by step; the first mismatching step narrows the investigation.

To validate bitwise checkpoint resume, compare an uninterrupted reference run with checkpoint-resumed runs. After each restore, verify bitwise identity for parameters, optimizer state, RNG state, data position, and subsequent outputs. The Nemotron trillion-parameter checkpoint path preserves a full-precision source of truth and reconstructs the runtime representation using the same conversion path used during training.

For GPU-resident fingerprints the recommended tool is torch.hash_tensor; record each tensor’s shape, dtype, and element count alongside its digest. For MXFP8 or NVFP4 tensors fingerprint both the encoded values and their scale buffers. Whole-tensor XOR fingerprints cannot detect permutations, so for order-sensitive data such as routing maps or MoE dispatch outputs fingerprint by row or chunk. A fingerprint is an efficient check, but byte-level comparisons are required to confirm bitwise equality.

How to localize and fix nondeterminism

At scale, nondeterminism may be intermittent or topology-dependent; visible loss divergence can appear long after the first differing bit. The Megatron-LM PR #7262 workflow records ordered, per-rank tensor fingerprints for offline comparison. First, confirm identical seeds, data order, batch size, parallelism, container, and software stack.

Localize divergence with granular tracing:

  • End-to-end metrics detect the failure and identify the first iteration with a differing serialized metric.
  • Broad semantic tracing covers collectives, pipeline communication, recomputation, optimizer operations, and gradient reductions to identify the divergent training phase.
  • Module and layer tracing narrows the divergence to a particular model component (transformer layer, MoE block, optimizer stage).
  • Operation-level tracing fingerprints ATen operations and targeted extension calls to find the first operation with matching inputs and different outputs.
  • Kernel-level tools identify the responsible kernel, algorithm configuration, and nondeterministic mechanism.

Do not trace only the iteration where loss visibly separates; the first differing bit may have occurred earlier. Use broad tracing to locate the earliest divergent iteration and ranks, then enable detailed operation and kernel tracing only within that narrowed scope.

Compare traces offline: each selected rank writes its trace to a file without adding collectives or cross-rank ordering to avoid masking the race. Align events by a run-independent identity (operation name, occurrence count, module scope). Look for matching input fingerprints and different output fingerprints:

hash(input_run_A) == hash(input_run_B) hash(output_run_A) != hash(output_run_B)

Interpret each rank’s first mismatch separately; sequence numbers cannot establish ordering across ranks. Classify the first mismatch on a rank as candidate origin, downstream receiver, or outside the capture, and extend the capture window or rank set accordingly. If many ranks identify the same originating operation, the operation itself is likely nondeterministic; if only a subset does, investigate topology, rank placement, input distribution, and reduction ordering.

Rule out false alarms: TorchDispatchMode observes ATen operations routed through the PyTorch dispatcher, but custom kernels can escape this scope. If the first mismatch appears at a simple view, slice, or addition, probe the custom kernel that produced the input. Also exclude issues from uninitialized tensors, incomplete asynchronous collectives, and nonblocking copies which may be read before their contents are valid.

Validate the fix with paired-run checks: the original pair should reproduce the divergence and the patched pair should remain bitwise identical. Determine whether the fix corrects the nondeterministic implementation or routes execution around it, and measure the new performance cost. Note that bitwise determinism is validated within the same hardware and software environment; cross-generation GPUs, network configurations, or library versions are outside this validation scope.

Optimizing deterministic training for a trillion-parameter Nemotron

A deterministic recipe with large overhead is useful for debugging but too costly in production. Deterministic and nondeterministic execution should be co-optimized. The reported work used a trillion-parameter Nemotron model that combines Mamba-style state-space model (SSM) layers and Transformer attention layers. Deterministic training must cover SSM and attention kernels, low-precision computation, distributed communication, and transitions between layer types.

Establish a controlled baseline by comparing deterministic execution to the fastest supported nondeterministic recipe using the same model, hardware, batch sizes, parallelism strategy, precision format, software environment, and measurement window. Measure throughput and step time, peak memory and GPU utilization, exposed communication time, and time spent in major kernel groups. Compute the determinism tax as (deterministic step time / baseline step time – 1) × 100%. Collect data only after training performance stabilizes.

The optimization journey reduced determinism overhead from double-digit baselines to low single digits across three workloads. At 2,432 GPUs, the large-scale Nemotron recipe measured approximately a 2% steady-state determinism tax with bitwise determinism demonstrated over 800 steps.

Kernel-level optimization example: in the grouped-GEMM epilogue multiple N-tiles initially accumulated into the same dprob[token] address, making results dependent on arrival order. Serializing those writers restored determinism but reduced parallelism. The solution was to give each N-tile a private output slot, preserve parallel execution inside the kernel, and combine the slots in a fixed order after all writers finish — preserving determinism with more parallelism.

Validate optimizations at scale: a low determinism tax on a few GPUs does not guarantee the same result at production scale because communication, synchronization, pipeline bubbles, expert routing, and load balance can change the relative overhead. For a hypothetical run with a 100-day nondeterministic baseline on 10,000 GPUs, reducing overhead from 15% to 5% would save 10 days, i.e., 100,000 GPU-days.

Maintain determinism as systems evolve

New kernels, fusions, precision formats, or parallelism configurations can introduce nondeterminism. Megatron-LM protects against regressions with:

  • Recipe validation: --deterministic-mode applies canonical environment settings, enables PyTorch deterministic algorithms, and rejects features without a deterministic path.
  • Kernel testing: repeat operations with identical inputs and RNG state, compare outputs and gradients byte for byte.
  • Module validation: module-level tests repeat models and transformer blocks with restored RNG state across parallelism configurations, including FP8 and FP4. Independent end-to-end runs then verify that full-precision training metrics remain bitwise identical.

The documentation includes a determinism status, an operation catalog, and a kernel testing guide to track current coverage and known gaps.

Common nondeterminism sources and next steps

The article lists common symptoms, likely causes, diagnostic methods and resolutions (examples include runtime autotuning selecting different kernel configs, unfixed reduction order in fused MoE losses, incomplete RNG restoration on resume, and low-precision reconstruction differences when loading checkpoints). Use targeted diagnostics to isolate the first scale-dependent operation and apply deterministic reductions, pin validated kernel configurations, or preserve full-precision sources of truth and exact reconstruction paths as appropriate.

Getting started

Bitwise determinism makes large-scale pretraining easier to debug and resume. The Nemotron case study shows how kernel and recipe optimizations can reduce its performance cost. To begin, review the Megatron Core User Guide for a supported recipe, then compare independent and checkpoint-resumed runs to validate reproducibility and measure overhead in the intended hardware and software environment.