NVIDIA-NeMo / NVIDIA-NeMo/RL

bug: SGLang refit ignores recompute_kv_cache_after_weight_updates and rejects in_place pause

Open
#3,747 0 comments 0 reactions 1 assignee Claimed by @youngeunkwon0405 View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

`grpo.async_grpo.recompute_kv_cache_after_weight_updates` selects between two documented strategies, per the comment in `nemo_rl/algorithms/async_utils/trajectory_collector.py`:

> Invalidate&recompute vLLM caches after the weight updates (in-flight or not) if `recompute_kv_cache_after_weight_updates` is True (AREAL-style implementation). Otherwise, keep using the stale KV caches (Magistral-style implementation).

The default is `False` — Magistral-style, keep the caches.

The SGLang refit path added in #3612 is hardwired to AREAL-style in two places, neither of which consults that setting.

**1. Unconditional invalidation on every refit:**

```python
# nemo_rl/weight_sync/sglang_weight_synchronizer.py:161
if not self._generation.invalidate_kv_cache():
raise RuntimeError("SGLang KV cache invalidation failed before refit.")
```

`grep -rn "recompute_kv_cache_after_weight_updates" nemo_rl/weight_sync/ nemo_rl/models/generation/sglang/` returns nothing, so the setting cannot reach this code.

**2. The one pause mode that would preserve KV is rejected at construction:**

```python
# nemo_rl/weight_sync/sglang_weight_synchronizer.py:80
if self._generation.pause_generation_mode == "in_place":
raise ValueError(
"pause_generation_mode='in_place' is unsafe for weight refit because "
"it preserves KV cache entries created by the previous weights."
)
```

Calling that "unsafe" reads as a correctness claim, but preserving KV entries across a weight update is exactly what the repo's default (`False`) does elsewhere, and the code calls it Magistral-style rather than a bug. The rationale and the default disagree.

## Where the setting is and is not honored

`grep -rn "invalidate_kv_cache" nemo_rl/ --include=*.py` gives four call sites:

| call site | gated on the setting? |
|---|---|
| `nemo_rl/algorithms/single_controller.py:1015` | ✅ yes, `:1014` |
| `nemo_rl/algorithms/async_utils/trajectory_collector.py:603` | ✅ yes, `:598` |
| `nemo_rl/weight_sync/sglang_weight_synchronizer.py:161` | ❌ **no** |
| `nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py:133` | ❌ no — see #3746 |

Both orchestrator-level call sites honor the setting; both synchronizer-level call sites ignore it.

For contrast, Megatron generation honors it (`nemo_rl/algorithms/grpo.py:1080` maps it onto `kv_cache_management_mode = "recompute"`) and TRT-LLM reads it directly (`nemo_rl/models/generation/trtllm/trtllm_generation.py:471`).

## Scope caveat

In **synchronous** GRPO the rollouts finish before the refit, so there are no in-flight requests whose KV must be preserved. What is lost by always flushing is cross-step prefix-cache reuse — a throughput cost, not a correctness one. The setting matters most on the **async** path, where generation continues across a weight update.

I could not find a guard either permitting or rejecting async GRPO together with the SGLang backend, so whether the ignored setting changes behaviour today depends on whether that combination is reachable. Someone who has run it would know.

## Suggested resolution

1. Gate the `invalidate_kv_cache()` call on `recompute_kv_cache_after_weight_updates`, matching `trajectory_collector.py:598`.
2. If SGLang genuinely cannot support the Magistral-style mode, say so in the comment and **raise at setup** when the setting is `False`, rather than silently overriding it.
3. Revisit the `in_place` rejection message so its rationale matches whichever of the above is true.

## Related

- #3612 — adds the SGLang refit path
- #3746 — same defect on the vLLM remote-sparse synchronizer
- Setting defined at `nemo_rl/algorithms/grpo.py:196`, mirrored at `nemo_rl/algorithms/single_controller_utils/config.py:188`

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.