NVIDIA-NeMo / NVIDIA-NeMo/RL

Megatron async checkpointing silently saves stale weights under colocate mode

Open
#3,312 2 comments 0 reactions 1 assignee Assigned to @terrykong View on GitHub
accuracy bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

With the current exemplar defaults (`policy.megatron_cfg.checkpoint.async_save: true`
and `ckpt_assume_constant_structure: true`, added in #2828), a colocated Megatron GRPO
run writes checkpoints that look fine but contain the wrong tensors. The first
checkpoint a job saves is correct. Every checkpoint after that silently re-writes the
first save's model weights and optimizer state, no matter how far training has
progressed. Saves report success, loads report success, and training metrics look
healthy the whole time. The problem only becomes visible when you resume: the run
restarts from a much older policy and the reward collapses.

An example error symptom:
Image

- The job saved checkpoints at steps 10, 20, ..., 120 (save_period=10, keep_top_k=3).
- After the job ended cleanly and a new job resumed from step_120, train reward fell abnormally.
- Comparing checkpoint tensors directly: step_90, step_110, and step_120 are
**bit-identical** to each other.

Root cause seems to be (traced through the stack):

1. `ckpt_assume_constant_structure` is forwarded to Megatron-Core, which passes
`use_cached_data_structure=True` to nvidia-resiliency-ext's async writer
(`megatron/core/dist_checkpointing/strategies/torch.py`, nvrx is the default
async strategy).
2. NVRx 0.6.0's `FileSystemWriterAsync.prepare_write_data` sends the GPU tensors to
the persistent checkpoint worker via CUDA IPC **only on the first save** and caches
the handles in the worker (`_worker_data_cache`). On every later save it sends
nothing ("signal to reuse cached data") and the worker re-reads GPU memory through
the first save's IPC handles.
3. A CUDA IPC handle names a physical allocation, not a tensor. Colocated GRPO
offloads/reloads the model and optimizer around every generation phase, so the
parameters move to new allocations every step. The worker's cached handles keep
pointing at the old allocations which PyTorch's IPC refcounting pins alive and
unmodified (that's also why the stale data is a coherent old snapshot rather than
garbage, and why nothing crashes).
4. A new job spawns a fresh worker with an empty cache, so each job's first save is
correct and then it starts poisoning its own checkpoints again.

NVRx already has an open fix for exactly this: NVIDIA/nvidia-resiliency-ext#314
("Refresh async checkpoint IPC cache on pointer change", open since 2026-04-28,
unmerged). NeMo-RL pins NVRx v0.6.0, which does not contain it.

**Steps/Code to reproduce bug**

Any Megatron-backend recipe with colocated generation, run long enough to save some
checkpoints in one job, e.g.:

```bash
uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml \
policy.megatron_cfg.enabled=true \
checkpointing.enabled=true checkpointing.save_period=10 \
grpo.max_num_steps=30
# (async_save: true and ckpt_assume_constant_structure: true are the defaults)
```

Then compare tensor content across the saved checkpoints, e.g.

- will see: step_20 and step_30 are bit-identical to step_10.

Equivalent black-box repro: kill the run after step_30, resume, and watch train
reward / entropy jump back to their early-training values.

**Expected behavior**

A clear and concise description of what you expected to happen.

**Additional context**

We may need to:

1. Flip the exemplar defaults back to `async_save: false` (or
`ckpt_assume_constant_structure: false`) until NVRx #314 merges and the pin is
bumped.
2. Guard in code: refuse or auto-disable `ckpt_assume_constant_structure` for the
async path when colocated generation / optimizer offload is enabled.

Besides, I think it's better to have the tests cover:
1. Saving more than one checkpoints, compare to check they are different.
2. Resume from a checkpoint and watch the key metrics such as reward.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.