FT local checkpointing: corrupt replica-retrieved local ckpt (deterministic NaN at load+3 iters) has no integrity validation / global-ckpt fallback; rerun_state_machine loops forever under ft_launcher in-job restart

[cxiong-ship-it]

Summary

After an in-job node-loss recovery that loads a replica-retrieved local (non-persistent) checkpoint (nvidia-resiliency-ext local checkpointing with --replication), the restored training state deterministically produces NaN in the forward loss 3 iterations after the load — reproduced in two independent runs, same iteration both times. Because nothing validates the local checkpoint's integrity/quality, and the rerun state machine's only remedy is "reschedule on a different GPU" (impossible under ft_launcher in-job restart, which returns the same GPUs by construction), the job enters an infinite restart loop (we observed 44 identical ~90 s cycles before cancelling).

Two concrete asks:

1. Validate local/replica-assembled checkpoints and fall back to the persistent global checkpoint when they fail (integrity check at load, and/or treat a reproducible NaN immediately after a local-tier load as "local checkpoint bad → reload global tier"). Today the corrupt local tier is trusted unconditionally and there is no fallback path.
2. rerun_state_machine needs an escape hatch when GPU diversity is impossible: under --enable-ft-package + ft_launcher worker-group restart, "Got rescheduled on the same GPU. Need to resume again from the same checkpoint" repeats forever. A bounded retry → abort (or → global-checkpoint fallback) would convert an infinite loop into a recoverable failure.

A related crash-safety hazard we hit on the way out: each loop cycle re-saves the diagnostic checkpoint in place at the same iteration; interrupting one of those saves (our scancel) left latest_checkpointed_iteration.txt pointing at a partial directory (5.7 GB of ~344 GB) — an atomic temp-dir + rename (tracker updated only after completion) would remove that class of failure.

Environment

• Megatron-LM (mcore dev image 648b916, 2026-05), PyTorch 2.12.0a0+...nv26.04, Python 3.12
• nvidia-resiliency-ext 0.6.0 (ft_launcher, c10d rendezvous; FT section timeouts configured via --ft-cfg-path)
• 5 nodes × 8 GPUs (B300, EFA), SLURM + pyxis
• ~38B MoE (150B-A12B family cut to 12 layers), --bf16 + MXFP8 (--fp8-format=e4m3 --fp8-recipe=mxfp8), EP=8, mock data
• Relevant flags: --enable-ft-package, --ckpt-format=torch_dist, --exit-signal-handler,
  --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
• Hot-spare topology: 5 ft_launcher agents, --nnodes 4 (1 parked spare)

Sequence (run A = job 250184; run B = job 249886 — identical signature)

1. Training healthy; local checkpoint saved at iteration 120 (/dev/shm, replicated), global /fsx checkpoint at iteration 100.
2. Fault injection: one active node is killed entirely (agent + workers + its /dev/shm shards).
3. ft_launcher restarts the worker group; the spare node joins; the group loads the local iteration-120 checkpoint — the dead node's shards are retrieved from replicas. Load reports success:

   successfully loaded checkpoint from <ckpt_dir> [ t 1/1, p 1/1 ] at iteration 120
   

4. Training resumes (iterations 121, 122 fine), then:

   ERROR:megatron.core.rerun_state_machine:Unexpected result nan on rank 9 at iteration #123 invocation #1 (message='found NaN in local forward loss calculation')
   WARNING:megatron.core.rerun_state_machine:Need to rerun step to check reproducibility of initial result
   ERROR:megatron.core.rerun_state_machine:... First rerun: unexpected result is reproducible within the tolerance (nan = nan). Need to rerun on a different GPU to verify ...
   WARNING:megatron.core.rerun_state_machine:Saving a checkpoint and exiting now. Please resume the job from the checkpoint to rerun the last iteration and establish a diagnostic
   

   (Run B: NaN also at iteration fix typo in mappings.py #123, also reproducible — deterministic, not a transient.)
5. The workers exit per the rerun protocol; ft_launcher restarts the group on the same nodes/GPUs (that is what in-job restart does):

   ERROR:megatron.core.rerun_state_machine:Got rescheduled on the same GPU. Need to resume again from the same checkpoint (node: ..., gpu: 0)
   

   → exit → restart → same GPU → exit → ... We cancelled after 44 identical cycles (~90 s each); --max-restarts 100 would have burned ~2.5 h of 40 idle GPUs.

Why we believe the local/replica checkpoint is the trigger

• An identical drill without --replication (same overlay otherwise, same node-kill, spare swap, /fsx fallback load) recovered cleanly and trained to completion (no NaN, no rerun events).
• Multiple earlier kill-tests loading non-replicated local checkpoints at 38B/75B/150B resumed cleanly past the load point.
• Both replication runs NaN at the same iteration (fix typo in mappings.py #123) shortly after a replica-retrieved load, reproducibly under rerun.
• We have separately verified (different campaign, no replication) that local and global checkpoints are bitwise-identical — so the local mechanism is sound; the replica retrieval/assembly path is the suspect. We are happy to file/redirect that part against nvidia-resiliency-ext if that's the right home (related rendezvous issue already filed: ft_launcher 0.6.0 hot-spare: re-rendezvous with live agents > max_nodes permanently closes the rendezvous ('Maximum number of nodes exceeded') and kills the whole job nvidia-resiliency-ext#348).

Impact

A single node loss with replication enabled converts into either total job death (run B, via the rendezvous race in nvrx#348) or an unbounded restart loop (run A). The layered design (local tier → global tier → requeue) cannot help because nothing ever decides the local tier is bad: ask (1) above is the missing piece — check local checkpoint quality; if bad, use the global checkpoint.

Full (lightly redacted) worker/agent logs for both runs available on request.

[guihong-nv]

@cxiong-ship-it Thanks for reporting this issue. May you please provide a minimal reproducer?

[cxiong-ship-it]

Thanks for the quick response! Evidence bundle here: https://gist.github.com/cxiong-ship-it/bc5a300c9449649834cd7466ac88cfa6 — contains (1) the exact ft_launcher invocation + full Megatron argv + the FT config YAML + versions, (2) curated agent/worker timelines for both failing runs (node-kill variant and hang variant — both NaN at iteration 123, i.e. load+3), and (3) the control run that is byte-identical minus the three --replication* flags and recovers cleanly through the same node-kill + spare swap. Full unsanitized logs available privately on request.

On a minimal reproducer: the bug's preconditions put a floor on size — it needs the replica-retrieval path, so >= 3 nodes (2 active + 1 spare, replication_jump = GPUs/node, world %% (jump x factor) == 0), a local save with --replication, then loss of one active node's shards before the next persistent save. We're working on two things now and will report back here:

1. A downscaled reproducer (small model, 2+1 nodes, mock data, scripted node-kill). One open question is whether the corruption is MoE/MXFP8-state-specific (our run carries EP=8 expert shards, expert bias, FP8 scale state — prime suspects for misassembly with replication_jump=8); if a dense-bf16 downscale doesn't reproduce, the minimal config will need --num-experts + the fp8 flags and we'll note that.
2. A forensic bitwise diff of the replica-assembled local checkpoint against a same-iteration persistent checkpoint, to name exactly which tensors differ. We have tooling that previously verified local==global bitwise (37,593/37,593 tensors / 757 GB) without replication, so any diff that appears with replication enabled is the corruption. First step is comparing a replication-enabled local save with all shards present (no failure) to separate save-path vs retrieval-path corruption.

[cxiong-ship-it]

Root cause found — better than a minimal reproducer: a deterministic single-run forensic, no fault injection needed.

We ran the bitwise diff promised above: train 110 iters with replication on (no failure at all), force a local+persistent checkpoint pair at the same iteration, snapshot the /dev/shm shard files, and compare every tensor in every shard file — primaries AND the replica files each rank stores for its jump-partner — against the persistent checkpoint.

Verdict: 28,128/32,796 tensors bitwise-equal; all 4,668 mismatches are in REPLICA files (0 primary mismatches, all 32 ranks affected) — and the replica data contains NaNs at save time. The corrupt tensors are exclusively the TE/MXFP8-quantized GEMM weights (expert fc1/fc2: 4,608 of 4,668; plus qkv/proj/shared-expert); norms, biases, router, embeddings and optimizer state replicate correctly.

So the chain in this issue is fully explained: node loss → recovery assembles state from replicas → loads NaN-poisoned quantized weights (load reports success) → forward loss NaNs ~3 iterations later → rerun_state_machine confirms deterministic → same-GPU loop under ft_launcher.

Filed the replica-corruption root cause against nvidia-resiliency-ext with the full forensic: NVIDIA/nvidia-resiliency-ext#349.

One correction to my original report: the worker interpreter in our image carries nvrx 0.5.0 (the 0.6.0 I cited lives only in the separate ft_launcher venv) — so all replication save/load code ran on 0.5.0. Given the v0.6.0 changelog entry "MXFP8/TE quantized tensor handling in IPC cache (#276)", the corruption may be 0.5.0-specific; we can retest with 0.6.0 in the worker env if useful.

The asks in this issue stand independently of the nvrx fix, and this incident is a concrete argument for them: a checkpoint that loads "successfully" carried NaN weights, and nothing in the load path or recovery stack could detect it or fall back to the (perfectly good) global checkpoint.