Local checkpointing replication writes NaN-corrupted replica shards for TE/MXFP8 quantized weights (0.5.0) — primaries bitwise-clean; root cause of Megatron-LM#5281

[cxiong-ship-it]

Summary

Local checkpointing with replication (LocalCheckpointManager + --replication --replication-jump 8 --replication-factor 2, driven through Megatron-LM --non-persistent-ckpt-type local) writes corrupt replica shard files for TE/MXFP8-quantized GEMM weight tensors — the replica copies contain NaN values at save time, with no failure or retrieval involved. Primary (own-rank) shard files are bitwise-perfect.

Forensic method: train ~38B MoE for 110 iters on 4 nodes × 8 GPUs with a local save (replication on) and a persistent torch_dist save landing on the same iteration (100); snapshot all /dev/shm shard files; compare every tensor in every shard file (primaries AND replicas) bitwise against the persistent checkpoint, one rank per shard, gloo CPU.

Verdict: 28,128/32,796 tensors bitwise-equal, 4,668 mismatched, 765.9 GB compared — every single mismatched tensor is in a REPLICA file; 0 primary mismatches; all 32 ranks hold corrupt replicas; the replica data contains NaNs (max_abs_diff = nan against a clean reference). Mismatch breakdown by key:

key	mismatched shards
decoder.layers.mlp.experts.experts.linear_fc1.weight	3072
decoder.layers.mlp.experts.experts.linear_fc2.weight	1536
decoder.layers.mlp.shared_experts.linear_fc1.weight	24
decoder.layers.self_attention.linear_qkv.weight	12
decoder.layers.self_attention.linear_proj.weight	12
decoder.layers.mlp.shared_experts.linear_fc2.weight	12

Note the pattern: only the TE/MXFP8-quantized GEMM weights (run uses --fp8-format=e4m3 --fp8-recipe=mxfp8 --fp8-param-gather --reuse-grad-buf-for-mxfp8-param-ag). LayerNorms, biases, router weights, embeddings and the entire optimizer state replicate correctly. A plausible mechanism: the replica copy is taken from a source buffer that has already been reused/recycled by the time the replica exchange runs (e.g. interplay with --reuse-grad-buf-for-mxfp8-param-ag / fp8 param-gather buffers), while the primary shard is serialized earlier from live data.

Impact (this is the root cause of NVIDIA/Megatron-LM#5281)

Any recovery that actually reads a replica (node loss; or a post-restart world that admits a previously-parked spare) assembles a state with NaN weights. The load reports success, training resumes, and the forward loss NaNs within ~3 iterations — then rerun_state_machine confirms it as deterministic and, under ft_launcher in-job restart, the job loops forever (or dies via the rendezvous issue #348). Reproduced in two independent recovery runs; a control run identical minus the --replication* flags recovers cleanly through the same fault.

Environment

• nvidia-resiliency-ext: 0.5.0 in the worker interpreter (/usr/local python — this is what executes all local-ckpt save/load/replication code via Megatron); 0.6.0 in the separate ft_launcher venv. We note v0.6.0's changelog mentions "MXFP8/TE quantized tensor handling in IPC cache (Handle quantized tensors to be dequantized in the trainer without being their IPC handles cached #276)" — same suspect class, different path; we have not yet tested a worker env with 0.6.0.
• Megatron-LM mcore dev 648b916 (2026-05), PyTorch 2.12.0a0+...nv26.04, Python 3.12
• 4 nodes × 8 GPUs (B300, EFA), SLURM + pyxis
• Model: ~38B MoE (e128 top-8 + shared expert, EP=8, seq 8192, bf16 + MXFP8), mock data
• Relevant flags: --non-persistent-ckpt-type local --non-persistent-local-ckpt-dir /dev/shm/<run> --non-persistent-save-interval 20 --non-persistent-local-ckpt-algo fully_parallel --replication --replication-jump 8 --replication-factor 2 --ckpt-format=torch_dist --enable-ft-package + the fp8 flags above

Reproducer (deterministic, single run, no fault injection needed)

1. Train any TE/MXFP8 model (we used a 12-layer cut of a 150B-A12B MoE; suspect fp8 is required — see tensor pattern) on ≥2 nodes with the local-ckpt + replication flags above until one local save completes.
2. Snapshot /dev/shm/<run> — each rank dir holds iter_N_<rank>_local.pt (primary) + iter_N_<rank±jump>_local.pt (replica it stores).
3. torch.load each file and compare each ShardedTensor's data against a same-iteration persistent checkpoint (we used a collective dist_checkpointing.load with the local shards' metadata as templates).
4. Replica files contain NaN-poisoned data for the quantized GEMM weights; primaries are bitwise-correct.

We can share the comparison harness (~120-line script), per-rank JSON reports, and full logs on request. Baseline sanity: the identical harness on a non-replicated run previously verified local==global bitwise (37,593/37,593 tensors, 757 GB).

[tusharg1993]

The local checkpoint replication strategy treats the quantized tensor as ordinary tensors. It transfers them without first calling  dequantize(), so the receiving rank lacks the internal quantizer metadata needed to reconstruct valid values. The invalid replica is then passed to  torch.save() .

I reproduced the underlying failure on two H100 GPUs using an AI-generated FP8 test script. The unfixed path failed while constructing the replica. After dequantizing before replication, all primary and replica tensors were bitwise-equal.

I can check Blackwell hardware too but I think this should fix it.

@hexinw-nvidia would you be open if I send a patch to fix it? If yes, please assign it to me. The fix should be a small one in the strategies.py