NVIDIA-NeMo / NVIDIA-NeMo/RL

bug: VllmRemoteSparseWeightSynchronizer ignores recompute_kv_cache_after_weight_updates

Open
#3,746 2 comments 0 reactions 0 assignees 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.

`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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.