feat: pause SingleController refits and enable KV recomputation by default

[youngeunkwon0405]

Summary

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

• apply the common generation pause/resume lifecycle around SingleController weight refits
• make recompute_kv_cache_after_weight_updates effective for active vLLM requests by clearing their KV state during native pause; the previous post-refit prefix-cache reset could leave running-request KV untouched while reporting success
• keep rollout dispatch closed on unrecovered lifecycle failures; tolerate confirmed resume-time actor deaths when enough healthy engines survive
• make vLLM pause/resume target the surviving refit leaders after fleet membership recovery
• enable KV recomputation by default, with explicit backend-aware opt-outs and compatibility fixes

Behavior change

This PR also enables KV-cache recomputation by default, beyond the SingleController pause/resume lifecycle:

• grpo.async_grpo.recompute_kv_cache_after_weight_updates (legacy async GRPO), ppo.async_ppo.recompute_kv_cache_after_weight_updates (legacy async PPO), and async_rl.recompute_kv_cache_after_weight_updates (SingleController) now default to true. For vLLM's native async path this clears paused requests' KV, repeating prefill after refit; it can reduce throughput.
• Opt out with the corresponding algorithm-specific flag set explicitly to false. For legacy async GRPO with Megatron generation and for SingleController, pair true with policy.generation.mcore_generation_config.kv_cache_management_mode: recompute, or false with persist/offload. Setup rejects inconsistent pairs.
• Synchronous legacy GRPO does not consume the async recompute flag and is exempt from that consistency check. However, synchronous Megatron recipes/nightlies that inherit the shared exemplar now inherit kv_cache_management_mode: recompute instead of persist; explicitly selecting persist or offload remains valid.
• The TRT-LLM exemplar's generation recompute setting interpolates the legacy GRPO flag, so its inherited default changes too; explicitly setting the GRPO flag to false propagates the opt-out.
• SGLang's existing weight synchronizer always pauses and flushes KV before transferring weights, even with the flag off. Its paired lifecycle acknowledgements prevent a second cache flush after generation resumes; they do not add another pause/resume RPC or timeout.
• The legacy async Omni video recipe explicitly keeps recomputation off to preserve its existing persist-mode/CUDA-graph and long-video-prefill behavior. The SingleController CLEVR recipe remains true/recompute.

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.

SingleController passes the existing refit timeout-plus-grace budget to pause and resume, with a 360-second fallback when refit_timeout_s=None. The collective's existing disabled-watchdog behavior and _sync_weights_within helper/docstring remain unchanged. Pause and refit failures enter the existing one-shot fleet recovery, but survivors are marked as holding partial weights only if a transfer actually started. After rebuilding membership, the controller pauses every retry participant, including restarted engines, before transferring weights again.

Resume waits for all leader RPCs under one shared timeout. With fleet health enabled, confirmed ActorDiedError failures remove only the dead engines; dispatch reopens if enough survivors successfully resumed. A late error from an old actor does not condemn its replacement. Timeouts, temporary actor unavailability, false acknowledgements, and exhausted fleets still stop the run with dispatch closed. Resume never retries the weight transfer.

vLLM preserves active request state with native keep-mode pause. With clear_cache=True, vLLM clears the KV state owned by those paused requests so they recompute it after refit. This fixes a pre-existing correctness gap: the former async reset_prefix_cache path did not clear KV blocks held by running requests, and its None worker results could still be reduced to success. This PR does not change the None-reduction in invalidate_kv_cache; the native async pause path instead no longer calls that function. Backends without native pause keep their existing in-flight refit behavior and receive the shared warning; when recompute is requested, they must confirm equivalent engine-side invalidation or the controller raises NotImplementedError instead of silently claiming recomputation.

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

Historical performance snapshot

For the previously linked Llama 3.1 8B SingleController runs at e16a857503c0843df20d3b58f737be6bcde6bf76 (not the current head), the median timing/train/total_step_time over optimizer steps 1–10 was:

KV mode	Median step time
Preserve	58.43 s
Recompute	60.28 s

Recompute was 3.16% slower by this median. This includes step 1 and is one historical pair, not a steady-state estimate, statistical result, or current-head performance claim. These runs did not record generation-tokens/s, so that requested metric cannot be supplied from their history. The logged valid-token throughput includes training behavior and must not be relabeled as generation throughput.

Validation

Current review fixes (d001be04e610f3a1b2d5c5900606cb8d7fde1114):

• Addressed the checkpointing fake-backend failure, the equivalent training-pump fixture failure in L0, the legacy async video recipe mismatch, SGLang's duplicate post-resume flush, sync-GRPO overvalidation, the stale MCore default comment, and launcher coverage. vLLM resume now uses one snapshot of shard/worker/leader triples; this removes duplicated mapping logic rather than claiming a demonstrated current mapping race.
• 363 selected full-module CPU tests passed locally in two runs (309 + 54) on macOS/Python 3.13.14 with Torch 2.13.0 and Ray 2.58.0. This includes the full checkpointing, controller, refit-recovery, and refit-resume modules; the 17 Megatron launcher combinations; 15 setup-level flag/mode cases; and the affected native vLLM lifecycle tests. Tests used an isolated dependency environment and --noconftest, not the repository GPU test environment. Existing dependency/marker/timeout-thread warnings and mock-object shutdown messages remain; no GPU engine or distributed training runtime is validated by these CPU checks.
• Ruff 0.9.9 lint/format, recipe minimization, and git diff --check passed.
• Exact-head CI has been requested for this commit; results are pending. Do not treat historical green runs below as validation of this head.
• The requested H100 Megatron nightly performance check and full-mode SingleController/router lane are being prepared but have not run on this head. Historical W&B runs do not establish current-head performance.

Historical validation

The following evidence is tied to its stated revisions and does not establish current-head CI, performance, or GPU recovery readiness.

Review follow-up (70f3b3967c8026863b4ffe4bd108a5803f7a1d45):

• Docs and tests only; no production behavior changes. Restored non-colocated timeout-sizing guidance, including separate lifecycle budgets, recovery/retry time, and the unbounded collective wait when refit_timeout_s=null.
• Added a pause-timeout/no-absent-shard regression and covered confirmed actor death during the retry pause. 64 focused source-extracted control-flow tests pass locally on macOS/Python 3.12; this is not full-module or GPU validation. Existing daemon-thread/event-loop-close warnings remain. An in-memory mutation restoring the old pause classification makes the new healthy-survivor assertion fail.
• Ruff 0.9.9 lint/format, Python 3.13.14 syntax, and diff checks pass. CI has been requested for the full follow-up SHA; results are pending.
• CI on preceding head cbf5d3be97d3e64f4d2a03585c61a0137d9f67ad passed: 63 successful jobs, 9 skipped.
• Separate reviewer-reported fault-injection results are pinned to cbf5d3be97d3e64f4d2a03585c61a0137d9f67ad against main 9c8778b1404e84c64bc73f671427234ba304e41c: nccl_reshard step-boundary kill passed 24/24 on both; nccl step-boundary recovery failed with the same rebuild traceback on both. PR-only kill-during-refit and restart variants also failed with that rebuild error. These are the reviewer's reported results, not independently rerun here. They do not replace the unresolved historical manual GPU evidence or constitute a complete nine-variant suite on the follow-up SHA.

Review fixes (bad0cb1691e9e20f0130219ddaee22b76a7a07b4):

• Addressed initial-pause failure classification, confirmed actor death during resume, bounded lifecycle calls when the collective watchdog is disabled, and synchronous-vLLM unsupported-hook behavior.
• Ruff 0.9.9 lint/format and git diff --check pass.
• GitHub CI for this head has been triggered and is queued.
• 60 focused tests pass with source-extracted controller/backend methods and the real fleet-health state machine on macOS/Python 3.11, using Ray 2.56.1 (the lockfile version at that historical revision; the current lock pins Ray 2.58.0 and vLLM 0.29.0). This bypasses the full import stack and is not full-runtime or GPU validation. Two existing daemon-thread/event-loop-close warnings remain in deadline tests.
• A supplemental CPU-only check with real Ray 2.56.1 actors passed: an intentionally dying leader was recorded dead, a delayed survivor acknowledged resume before success, and a separate pending resume raised RefitAborted within its deadline. The controller/backend methods were source-extracted; no vLLM engine or GPU was exercised. The first attempt stopped during automatic uv workspace packaging; disabling that local-only packaging allowed the test to run.
• The requested nine-variant GPU fleet-health fault-injection suite has not run on this head. The historical GPU runs below do not validate these changes.

Previous merge (d8e713dbbf6ba3d08f481c53d25e0a3e9fff41f1):

• Merged main at feff70f5b6239d90dc703ae79b973ea809550732, including feat(sc): restart dead generation shards and re-admit them at the next refit #3592. Preserved participant-aware refit completion and shard re-admission. Recovery retries now re-pause the rebuilt fleet, covering both returning shards and replacement engines that reuse a shard index.
• Ruff 0.9.9 lint/format, Python syntax, and diff checks pass. The existing bounded refit helper and docstring are unchanged.
• 46 isolated recovery control-flow tests passed locally using source-extracted controller methods and Ray 2.55.1 on macOS/Python 3.11. This bypasses the full GPU import stack and is not full-module or GPU validation. The run reported two daemon-thread event-loop-close warnings in existing deadline tests and one temporary pytest-cache warning.
• Both new restart/retry regression cases fail when the retry pause is removed in memory, and pass with the fix.
• GitHub CI for the preceding merged head passed. Its fast functional lane does not cover the multi-shard fleet-health fault-injection variants.
• Previous-head CW-DFW job 18608218 failed during environment setup (/opt/nemo_rl_venv/bin/python was absent from the selected image), before pytest started. It provides no runtime validation.

Earlier 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 keeping generation paused until a recovery retry completes
• earlier full-size real-GPU SingleController GRPO validation on CW-DFW at commit e16a857503c0843df20d3b58f737be6bcde6bf76. These runs predate the subsequent merges from main and do not validate the current PR head.
  • 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 cbf5d3b

[mehraakash]

Review (review-pr-team): PR #3929 at cbf5d3be9

2 action items, neither blocking, posted inline: a docs regression and two test gaps (with verified tests). Six agents (RL expert, bug finder, test reviewer, design reviewer, comment reviewer, devil's advocate) re-verified every finding against head and against the pinned Ray 2.58.0 / vLLM 0.29.0 sources.

Both earlier blocking items are resolved at head. A pause-origin RefitAborted no longer marks survivors partial (refit_started gates it, pinned by the ledger snapshot in test_failed_pause_recovers_and_pauses_survivors_before_refit), and a confirmed ActorDiedError on one leader at resume is recorded and the run continues (pinned at both the vLLM and controller level). All other "Addressed in" threads verify fixed. The reconcile_communicator missing-argument crash from the September GPU runs is fixed on main by #4233 and is in this head.

Fault-injection results on this head, 4×GB300, 2 train + 2 generation shards, matched against main at the merge-base 9c8778b14:

Variant	main @ merge-base	PR head
step-boundary kill, nccl_reshard	PASS 24/24 steps	PASS 24/24 steps; 25 pause / 25 resume requests across the rebuild
step-boundary kill, nccl	FAIL: NCCLError: unhandled system error in MegatronPolicyWorker.init_collective during the rebuild	FAIL, identical traceback
kill during refit, nccl	not run	FAIL, same rebuild error
restart dead shards, nccl	not run	FAIL, same rebuild error

The nccl collective-transport recovery failure reproduces on main with the same traceback, so it is pre-existing on this platform and not introduced here. It is worth tracking separately. The reshard transport recovers with the pause/resume lifecycle in place.

Unit tests at head on the same node: 288 passed in the SingleController bucket; 11/12 in the vLLM bucket (the one failure is test_vllm_policy_generation_with_fastokens_enabled, an nixl_ep ABI import error in the test environment, unrelated to this PR); the two suggested tests pass.

PR description nits (no code change): it cites Ray 2.56.1 as the lockfile version, but the lock at head pins Ray 2.58.0 and vLLM 0.29.0; and the sentence saying the None-reduction in invalidate_kv_cache was fixed does not match the diff, which leaves that function untouched (the async path simply no longer calls it).

Still open from earlier threads, not re-raised: the resume-timeout design discussion with @asolergi-nv, a KILL_DURING_PAUSE functional variant, and setup-time validation of backend_timeout_s against the pause budget.

Generated by Claude Code

[youngeunkwon0405]

/ok to test 70f3b39

[mehraakash]

Could we gate pause/resume behind an async_rl.pause_generation_during_refit flag? Defaulting to true would preserve this PR’s behavior while allowing users to retain the previous behavior and compare performance. recompute_kv_cache_after_weight_updates currently controls cache clearing, not whether generation pauses.

[youngeunkwon0405]

Could we gate pause/resume behind an async_rl.pause_generation_during_refit flag? Defaulting to true would preserve this PR’s behavior while allowing users to retain the previous behavior and compare performance. recompute_kv_cache_after_weight_updates currently controls cache clearing, not whether generation pauses.

I would like to double check about this with you. I personally think not pausing during refit is a bug and this PR is sort of a bug fix. Do you prefer to make it configurable?

[mehraakash]

1 action item.

TL;DR — now that recompute_kv_cache_after_weight_updates=true actually clears in-flight request KV on the native vLLM pause path, please make true the default.

Before this PR the flag was only a prefix-cache reset that could leave running-request KV untouched, so a stale-KV default was the safe choice. With native keep-mode pause plus clear_cache=True, every preserved request recomputes against the new weights after resume, which is the behavior most users expect from an on-policy-ish async loop. Keeping it off by default means a fresh config silently continues old-weight prefixes into new requests.

Flipping it touches three defaults that should move together so the two loops agree:

• AsyncRLConfig.recompute_kv_cache_after_weight_updates (

  RL/nemo_rl/algorithms/single_controller_utils/config.py

  Line 497 in 70f3b39

  recompute_kv_cache_after_weight_updates: bool = False

  )
• grpo.async_grpo (

  RL/nemo_rl/algorithms/grpo.py

  Line 289 in 70f3b39

  recompute_kv_cache_after_weight_updates: bool = False

  )
• PPO (

  RL/nemo_rl/algorithms/ppo.py

  Line 147 in 70f3b39

  recompute_kv_cache_after_weight_updates: bool = False

  )

Two things to handle alongside it:

• Megatron generation raises at setup when the flag disagrees with kv_cache_management_mode (

  RL/nemo_rl/models/generation/megatron/megatron_generation.py

  Line 231 in 70f3b39

  if recompute_kv_cache_after_weight_updates != (kv_cache_mode == "recompute"):

  ). The Megatron-generation recipes that omit the flag would need kv_cache_management_mode: recompute, or an explicit false.
• Recompute redoes the full prefill of every in-flight request on each refit, so a before/after throughput number on one of the async recipes would be good to have in the PR description.

If you would rather keep this PR to the pause mechanics and do the default flip in a follow-up, say so here and I will open an issue to track it.

[youngeunkwon0405]

/ok to test 0f623c1

[youngeunkwon0405]

/ok to test 2505b20

[youngeunkwon0405]

/ok to test 7de9e6d

[youngeunkwon0405]

/ok to test a1677bc

[youngeunkwon0405]

Hi @mehraakash, I updated the PR based on your comment. Mostly just changing default setting of recompute_kv_cache to true.
I also did a benchmark run to investigate the perf overhead of this change.
QWEN3-30B GRPO batchsize 512, 8K seq lenght run(wandb link)

• Average step time is 150s-ish in this run; recomputing kv-cache introduced a 3s-ish (2%) increase in the step time.

[youngeunkwon0405]

Hi @terrykong, @mehraakash asked me to ask for the final approval from you. Thanks!

[youngeunkwon0405]

/ok to test a5052e4

[terrykong]

Thanks for this. The pause/clear fix is right and needed: before it, recompute_kv_cache_after_weight_updates: true on the single controller silently did nothing, because vLLM refuses a prefix-cache reset while requests hold blocks. Two blockers come from the default flip in 7de9e6d, not from the pause code: 36 failing unit tests in test_checkpointing.py, and the GB200 L1 video script that now fails at setup (comment on grpo.py#L299). Please fix both and re-run CI. The green CI in the description is from before the flip.

All three of my earlier comments are fixed at a5052e4: the pause is bounded and inside the recovery try, the Gym test axis has a real assertion, and the drain-gate TODO is back. Reusing the shared pause/resume hooks from your vLLM pause PR (#3839) and the refit_started flag are clean.

Behavior change. The title and description never mention the default flip. Please add a "Behavior change" section that says:

• the default is now true for legacy async GRPO/PPO and for the single controller;
• Megatron generation needs kv_cache_management_mode: recompute, or an explicit false with persist/offload;
• how to opt out;
• the sync Megatron-generation nightlies move to recompute;
• TRT-LLM recipes pass recompute through interpolation.

Consider widening the title scope to match.

Numbers. Please put median step time and generation tokens/s from the linked preserve vs recompute runs in the description. Also run the H100 nightly grpo-llama3.2-1b-instruct-1n8g-megatron_generation on this head; it gates median step time < 12 s (L41).

Router. Pause → refit → resume works with async_rl.generation_router.enabled, and the new docs paragraph (single-controller.md#L427) already covers sizing backend_timeout_s for the pause. Please run the full-mode L1 SingleController lane that turns the router on (L1_Functional_Tests_SingleController_2.sh#L52-L54) on this head before merge: Lfast skips it, and recompute is now on by default.

Not for this PR. Two router issues (a restarted engine never rejoins the router; router timeouts during a hung refit can mark a healthy shard as suspect) will be filed separately. Keeping running requests' KV while stopping new requests from reusing old prefix blocks (a cache_salt per weight version) is also a separate follow-up. No separate pause flag is needed: pausing during refit is the fix, and recompute_kv_cache_after_weight_updates stays the one setting, now default true, as agreed.

Overlapping PRs. #4018 and #4647 touch the same flag and reset path; they will be asked to rebase on this one.

Generated by Claude Code

[youngeunkwon0405]

/ok to test d001be0

[youngeunkwon0405]

/ok to test 9c15e56

[youngeunkwon0405]

/ok to test 064a49e