test(ppo): check that consecutive async checkpoints write fresh optimizer state

[bg51717]

What does this PR do ?

Adds regression checks that consecutive Megatron async checkpoints write the current state when tensors move off the GPU between saves.

With ckpt_assume_constant_structure: true, NVRx caches the persistent writer's CUDA IPC handles across saves. Before NVRx v0.7.0 the cache had no data pointers, so after an offload every checkpoint repeated the first one, and the value worker hung on its second save. #4363 bumped NVRx to v0.7.0, whose NVIDIA/nvidia-resiliency-ext#314 re-sends the handles when the pointers change (filesystem_async.py#L438-L452). This PR first fixed it on the NeMo-RL side; it now only adds tests so a regression shows up in CI.

• tests/check_checkpoint_freshness.py fails when every chunk of an optimizer tensor repeats the previous checkpoint: the Adam moments, plus the master weight (param) when the model weights are bf16/fp16, since TE FusedAdam never updates it with fp32 params. It also fails when every model weight repeats; per-tensor weight checks would trip on bf16 weights that a small lr leaves unchanged.
• ppo_megatron.sh saves policy and value every step with async constant-structure checkpoints and runs the checker on both. The other PPO functional tests train the critic alone first, and PPO rejects policy ckpt_assume_constant_structure with critic warmup (ppo.py#L395-L410), so they are unchanged.
• GPU unit tests save twice around the PPO offload, and around the optimizer offloads of non-colocated distillation (offload_before_refit()), offload_optimizer_for_logprob (prepare_for_lp_inference()) and non-colocated refit with policy.offload_policy_before_refit.

Test

• The GPU tests and ppo_megatron.sh pass on current main.
• The checker passes on two consecutive PPO checkpoints and flags the Adam moments, param and the model weights when a checkpoint is compared with itself.
• On NVRx v0.6.0 the GPU tests catch the stale state: the optimizer-offload cases report repeated exp_avg / exp_avg_sq chunks.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[bg51717]

Summary

This PR fixes two PPO bugs in Megatron async checkpointing when NVRx caches CUDA IPC handles across saves (ckpt_assume_constant_structure: true):

• Stale policy checkpoints with non-colocated generation. PPO offloads the policy every step so the value model can use the GPUs, and the next onload reallocates the tensors whose handles NVRx cached at the first save, but the release in finalize_async_save() only fired for colocated generation. PPO now sets megatron_cfg.offloaded_between_steps in both the legacy and SingleController setup, and _requires_nvrx_cuda_cache_release treats it like colocated generation.
• Value worker hang on the second async save. The worker terminated the NVRx writer without clearing FileSystemWriterAsync._cached_identifiers. It now calls cleanup_tensor_caches() together with the terminate, and terminates only when NVRx caches CUDA handles.

Tests: mock tests for both release paths, setup assertions for the new flag, and a 2-GPU test that saves twice around a PPO offload and compares the checkpoints. Both root causes check out against NVRx v0.6.0, and no defects were found in the PPO changes. Two test suggestions are inline.

PR description

• Root cause → Value worker hang: before this fix the second value save hangs in NVRx's preload_q.join() inside save_checkpoint (the respawned writer raises the cache miss in preload_tensors, filesystem_async.py#L620-L621, before calling task_done()), not in the blocking finalize. Commit 8c7d94a's message says the same.
• When it triggers: async_ckpt_use_cpu_shm: false # default isn't a condition users can set, since _create_checkpoint_config never forwards it to Megatron-Bridge; the line can be dropped.

Reviewed with five agents (RL, upstream NVRx/Bridge/Megatron-LM, bug finder, tests, devil's advocate); findings are read from source, not measured.

Generated by Claude Code

[bg51717]

/ok to test d6086d9

[yuki-97]

This fixes two PPO async-checkpoint bugs that share one root cause in NVRx's CUDA handle cache: stale policy checkpoints without colocated generation, and the value worker hanging on its second save.

Scope: the full diff at d6086d9: the PPO setup flag, the policy and value worker release paths, and the new tests.

Checked:

• Both root causes, against NVRx v0.6.0. The cache key covers only names, shapes and dtypes, with no data pointers, so a save with the same structure reuses the first save's handles. A respawned writer that finds no cache raises in preload_tensors.
• The value worker no longer terminates unconditionally. On the uncached path the handles are locals of a single save, so a writer that stays alive does not pin offloaded value storage.
• PPO sequencing. policy.offload_to_cpu() right after save_checkpoint goes through offload_after_refit → finalize_async_save(), so the cache is released before the storage moves.
• The new 2-GPU test hashes tensor chunks through the .metadata byte ranges instead of diffing files. That keeps it from tripping on the rng_state / common_state shards, which change every step.

Remaining comments:

• Move offloaded_between_steps out of megatron_cfg into a RuntimeConfig field shared by both workers.
• Check whether non-colocated distillation and GRPO save the same stale optimizer state. This one is optional for this PR.

[bg51717]

/ok to test 2e8b351

[yuki-97]

Thanks for resolving the previous comments. The release decision now lives in one place on RuntimeConfig, and non-colocated distillation and offload_optimizer_for_logprob are covered as well.

Scope: the commits since d6086d9, at 2e8b351.

Checked:

• offloaded_between_steps plumbing: the kwarg reaches validate_and_set_config through TQPolicy, the worker builder and both worker subclasses, and the only positional RuntimeConfig construction has the new field in the right slot.
• Storage moves between saves: the PPO policy offload, the non-colocated distillation student's offload_before_refit, prepare_for_lp_inference under offload_optimizer_for_logprob, and colocated refit all turn the release on.
• The new GPU tests: the optimizer state lands in the torch_dist directory that _hash_checkpoint_tensors reads, and the for_logprob case passes only because of the new offload_optimizer_for_logprob term.

Remaining comments are on:

• one storage-move path the release decision still misses: non-colocated Megatron generation with offload_policy_before_refit
• checking consecutive async checkpoints in the PPO functional tests
• a nit on computing the release flag before the RuntimeConfig(...) call

[bg51717]

#4363 (8661b47) bumped nvidia-resiliency-ext from v0.6.0 to v0.7.0, which includes NVIDIA/nvidia-resiliency-ext#314: NVRx now compares the cached tensors' data pointers on every save and re-sends the CUDA IPC handles when they change (filesystem_async.py#L438-L452).

That fixes the stale-checkpoint bug this PR targets, and the value worker no longer hangs on its second async save. On current main with this PR's changes reverted, non-colocated PPO writes fresh policy and value checkpoints every step, and so do the PPO offload, offload_before_refit, offload_optimizer_for_logprob and Megatron-refit paths and the value worker across six saves each.

[bg51717]

/ok to test 406e970