feat(single-controller): pause generation during refit

[youngeunkwon0405]

Summary

Implementing the same feature as #3839 for the single controller path.

• apply the common generation pause/resume lifecycle around SingleController weight refits
• pass recompute_kv_cache_after_weight_updates into native pause and retain cache invalidation as the unsupported-backend fallback
• keep rollout dispatch closed when pause, refit, or resume fails
• make vLLM pause/resume target the surviving refit leaders after fleet membership recovery

Implementation

SingleControllerActor._sync_weights() continues to close the rollout-dispatch gate and cancel sampler-stale native requests first. After refit membership reconciliation and optional FP8 KV-scale calibration, it asks the generation backend to pause through pause_generation_for_refit(clear_cache=...), runs the existing WeightSynchronizer refit (including its one-shot recovery retry), and calls resume_generation_after_refit() before reopening dispatch.

vLLM preserves active request state with native keep-mode pause. Backends that return unsupported keep their existing in-flight refit behavior, receive the shared warning, and retain the existing post-refit cache-invalidation path when recompute is enabled.

The lifecycle is in the shared SingleController _sync_weights() path, so it applies to GRPO and PPO and to native and NeMo-Gym rollouts.

Validation

• Ruff 0.9.9 check and format verification passed for all modified Python files
• git diff --check passed
• CW-DFW Slurm job 17015948 in nemo_rl.0825.sqsh: 92 focused tests passed
  • full test_single_controller_actor.py
  • full test_refit_recovery.py
  • native vLLM worker and generation pause/resume tests
• after L0 exposed a checkpointing fixture without the new lifecycle hooks, CW-DFW Slurm job 2134 in the same image/account ran the full test_checkpointing.py: 62 passed
• after Megatron L0 exposed a training-pump fixture with gen_handle=None, CW-DFW Slurm job 17510464 in the same image/account ran the full test_train_pump_e2e.py under --hf-gated --mcore-only with a forced worker-environment rebuild: 1 passed, 66 warnings in 123.19s; Slurm COMPLETED / 0:0 in 6:27
• tests cover native and NeMo-Gym selection, recompute on/off, supported and unsupported pause, resume/refit failure gating, and a single pause spanning a recovery retry
• full-size real-GPU SingleController GRPO recipe on CW-DFW at production-code head e16a857503c0843df20d3b58f737be6bcde6bf76 (current head differs only by two test-fixture compatibility fixes found by L0 CI; production source is unchanged)
  • unchanged grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml workload: 2 nodes / 16 H100 GPUs, 64 prompts x 32 generations, global batch size 2048, max sequence length 4096
  • only runtime overrides were 10 steps, recompute mode, W&B destinations, and vLLM generation metrics logging
  • KV-cache preserve: Slurm 17016712, COMPLETED / 0:0, 10 optimizer steps and 11 successful pause/refit/resume cycles; W&B
  • KV-cache recompute: Slurm 17016711, COMPLETED / 0:0, 10 optimizer steps and 11 successful pause/refit/resume cycles; W&B
  • the initial sync had zero in-flight groups; every one of the 10 training refits paused and preserved nonzero in-flight groups in both runs

[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.

[youngeunkwon0405]

/ok to test e16a857

[youngeunkwon0405]

Real-GPU validation completed at exact PR head e16a857503c0843df20d3b58f737be6bcde6bf76.

I ran the full-size native SingleController vLLM GRPO recipe (grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml) on CW-DFW with the requested coreai_dlalgo_genai account and nemo_rl.0825.sqsh image. The workload dimensions were unchanged: 2 nodes / 16 H100 GPUs, 64 prompts per step x 32 generations, global batch size 2048, and max sequence length 4096. The only runtime overrides were 10 steps, recompute mode, logging destinations, and enable_vllm_metrics_logger=true.

• recompute_kv_cache_after_weight_updates=false: Slurm 17016712, COMPLETED / 0:0 in 17:55; 10 optimizer steps; 11 pause/refit/resume cycles (initial sync + 10 training refits); all 10 training refits paused nonzero in-flight groups. W&B run
• recompute_kv_cache_after_weight_updates=true: Slurm 17016711, COMPLETED / 0:0 in 18:14; 10 optimizer steps; 11 pause/refit/resume cycles (initial sync + 10 training refits); all 10 training refits paused nonzero in-flight groups. W&B run

The only traceback in each log was a post-success W&B service BrokenPipeError during Ray process teardown; both Slurm jobs exited successfully after W&B had emitted the final run summary.

[youngeunkwon0405]

/ok to test daa045d

[youngeunkwon0405]

Addressed the L0_Unit_Tests_Other failure at new head daa045de89a5a0071a1b09cc4b50adcf262a1534.

The failure was PR-caused but test-only: TestCounterRestore.test_resumed_pump_continues_to_max_steps used gen_handle=object() and now reaches the real _sync_weights() lifecycle, so the fixture lacked pause_generation_for_refit() / resume_generation_after_refit(). I replaced it with an unsupported-backend fake implementing the common lifecycle contract; no production code changed from the GPU-validated head.

Validation:

• pinned Ruff 0.9.9 check and format check passed
• CW-DFW Slurm 2134, nemo_rl.0825.sqsh, coreai_dlalgo_genai: COMPLETED / 0:0; full tests/unit/single_controller/test_checkpointing.py: 62 passed, 17 warnings in 20.93s
• exact-head GitHub CI was reauthorized with /ok to test daa045de89a5a0071a1b09cc4b50adcf262a1534

[youngeunkwon0405]

Addressed the remaining L0_Unit_Tests_Megatron failure at new head 15074e4ccfe626a71df6972d247e6bf394f77eef.

The failure was PR-caused but test-only: test_train_pump_drives_mcore_training_step passed gen_handle=None and now reaches the real _sync_weights() lifecycle, so the fixture lacked the common pause/resume hooks. I replaced it with an unsupported-backend fake implementing pause_generation_for_refit() and resume_generation_after_refit(). No production code changed from the two 10-step GPU-validated runs.

Validation:

• pinned Ruff 0.9.9 check and format verification passed; git diff --check passed
• CW-DFW Slurm 17510464, nemo_rl.0825.sqsh, coreai_dlalgo_genai: COMPLETED / 0:0 in 6:27
• full tests/unit/single_controller/test_train_pump_e2e.py under --hf-gated --mcore-only: 1 passed, 66 warnings in 123.19s

[youngeunkwon0405]

/ok to test 15074e4

[youngeunkwon0405]

Team review: 4 agents (rl-expert, test-agent, bug-finder, devil's advocate) analyzed this PR; every finding below was independently re-verified against source, including the pinned vLLM 0.25.1 wheel.

Verified clean (no action): clear_cache-at-pause genuinely subsumes post-refit invalidation (vLLM keep-mode pause_scheduler resets caches at pause); _refit_leader_workers() targets exactly the set the old run_rank_0_only_axes dispatch selected (incl. pp>1) and matches the existing refit dispatches; interface default hooks keep unsupported backends safe; fail-stop resume semantics mirror trajectory_collector; removing the drain-gate TODO is consistent (docs + issue #2625 still track it). Ruff 0.9.9 check+format clean. The A/B W&B evidence (preserve vs recompute, 11 pause/refit/resume cycles each) is solid.

Findings: 1 hardening suggestion, 1 docs note, 2 optional test suggestions (both verified passing locally on CPU) — all inline.

Generated by Claude Code

[youngeunkwon0405]

/ok to test 1f68285

[youngeunkwon0405]

Post-review validation for exact head 1f68285:

• Local: Ruff 0.9.9 check and format check, compileall, and git diff --check passed.
• Focused CW-DFW GPU tests: Slurm 17531796, 1 node / 8 H100, COMPLETED 0:0. SingleController recovery/lifecycle suites: 92 passed. Focused vLLM native-pause/cache-result cases: 8 passed.
• Real GRPO validation: Slurm 17531931, 2 nodes / 16 H100, account coreai_dlalgo_genai, full grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2 recipe, COMPLETED 0:0 in 19m30s.
  • 10/10 steps completed; 64 prompts x 32 generations, global batch 2048, max generation/total sequence length 4096. No batch-size or sequence-length reduction.
  • async_rl.recompute_kv_cache_after_weight_updates=true.
  • vLLM generation metrics logging enabled.
  • Observed 11 successful pause/refit/resume cycles (initial sync plus 10 optimizer steps), 11 successful refits, and prefix-cache reset logs on every cycle. The 10 training-cycle pauses preserved 58-64 in-flight rollout groups.
  • No RefitAborted or NotImplementedError; final receipt: SC run complete with train_steps=10 and trainer_version=10.
• W&B: https://wandb.ai/nvidia/async-grpo-pr-validation/runs/gabgq6rm

Claude review: all four threads were addressed in 1f68285, replied to, and resolved: bounded lifecycle calls, timeout-budget documentation, recovery ordering coverage, and pause-failure coverage. Exact-head CI was retriggered; branch-up-to-date, lint, docs, and script-coverage checks are passing, with container builds still running.

[youngeunkwon0405]

/ok to test 463f247

[youngeunkwon0405]

Simplified the post-review implementation at exact head 463f247737b65a7257a5f3dfaec94b1b3881cdbf.

• Removed the generalized _run_refit_call_within abstraction and restored the existing specialized _sync_weights_within implementation and its docstring unchanged.
• Pause, resume, and fallback cache invalidation again use direct asyncio.to_thread calls in _sync_weights().
• Removed the unrelated vLLM prefix-reset result plumbing and its tests.
• Kept the correctness guard: when native pause is unavailable and recomputation is requested, the backend must confirm equivalent active-request KV handling or the run raises NotImplementedError instead of silently preserving stale KV state.
• The simplification follow-up removed 139 lines and added 43.

Validation at this exact head:

• Ruff 0.9.9 check and format check, Python compile, and git diff --check: passed.
• Focused CW-DFW GPU tests: Slurm 17534872, 1 node / 8 H100, coreai_dlalgo_genai, COMPLETED 0:0 in 2m38s. SingleController suites: 91 passed; focused vLLM native-pause lifecycle cases: 3 passed.
• Real GRPO validation: Slurm 17535168, 2 nodes / 16 H100, coreai_dlalgo_genai, full grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2 recipe, COMPLETED 0:0 in 19m10s.
  • 10/10 optimizer steps completed; final receipt: SC run complete: {'train_steps': 10, 'trainer_version': 10}.
  • No batch-size or sequence-length reduction: 64 prompts x 32 generations, global batch 2048, max generation/total sequence length 4096.
  • async_rl.recompute_kv_cache_after_weight_updates=true and vLLM generation metrics logging were enabled.
  • Observed 11 successful pause/refit/resume cycles (initial sync plus 10 training refits). Every training refit paused nonzero in-flight work (42-64 rollout groups).
  • No RefitAborted or NotImplementedError.
  • W&B run

The only post-success traceback was the same W&B service BrokenPipeError during Ray teardown seen in prior successful validations; it occurred after the final W&B summary and SC run complete, and the Slurm job exited successfully.

[youngeunkwon0405]

/ok to test 463f247

[youngeunkwon0405]

/ok to test df75b8c

[youngeunkwon0405]

/ok to test c6a082b

[youngeunkwon0405]

/ok to test 40ca90a

[youngeunkwon0405]

/ok to test ae0ca27

[youngeunkwon0405]

/ok to test ae0ca27