bug: VllmRemoteSparseWeightSynchronizer ignores recompute_kv_cache_after_weight_updates
- 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.
`VllmRemoteSparseWeightSynchronizer.sync_weights` invalidates the KV cache **unconditionally**, without consulting that setting:
```python
# nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py:133
if not self._generation.invalidate_kv_cache():
raise RuntimeError(
f"vLLM KV cache invalidation failed before {self._transport} "
"weight update."
)
```
So a run using the remote-sparse transport gets AREAL-style behaviour regardless of what the user configured, with no warning.
## 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/vllm_remote_sparse_weight_synchronizer.py:133` | ❌ **no** |
| `nemo_rl/weight_sync/sglang_weight_synchronizer.py:161` | ❌ no — see the companion issue |
The pattern is clean: both orchestrator-level call sites honor the setting; both synchronizer-level call sites ignore it.
For contrast, the Megatron generation backend does honor it — `nemo_rl/algorithms/grpo.py:1080` translates it into `mcore_generation_config.kv_cache_management_mode = "recompute"`, and TRT-LLM reads it directly at `nemo_rl/models/generation/trtllm/trtllm_generation.py:471`.
## Question rather than assertion
It is possible this is deliberate and the sparse transport genuinely requires invalidation — a partial or delta weight update could leave KV entries inconsistent in a way a full refit does not. If so, that is worth stating.
But it should then **reject** the non-default value rather than silently overriding it, so a user who sets `recompute_kv_cache_after_weight_updates: false` learns their setting does not apply on this transport. Right now the config says one thing and the code does another.
### Suggested resolution, either way
1. If the setting should apply: gate the call, matching `trajectory_collector.py:598`.
2. If invalidation is required by the transport: keep it, add a comment saying why, and raise at setup when `recompute_kv_cache_after_weight_updates` is `False` together with this transport.
## Related
- Companion issue for the same defect on the SGLang refit path (filed alongside this one)
- The setting is defined at `nemo_rl/algorithms/grpo.py:196` and mirrored at `nemo_rl/algorithms/single_controller_utils/config.py:188`
Contributor guide
Research direction
Start with nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py:133 and compare it with the gated call at nemo_rl/algorithms/async_utils/trajectory_collector.py:598. Check how recompute_kv_cache_after_weight_updates is defined in nemo_rl/algorithms/grpo.py:196 and mirrored in single_controller_utils/config.py:188. Done means the remote-sparse transport either honors the setting or explicitly rejects false at setup, with the transport-specific behavior documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100