Speculative decoding silently broken in colocated GRPO: vLLM sleep level 2 discards drafter weights, never restored on refit
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
Speculative decoding is silently broken in **colocated** GRPO runs. The vLLM draft model (eagle3 or a standalone draft model) is discarded by `sleep(level=2)` at the end of each generation step and is **never restored** during refit, because refit only re-pushes *policy* weights. After the first generation step the drafter contains uninitialized/dummy weights, so the acceptance rate collapses to exactly `0` and stays there for the rest of training.
This is the speculative-decoding manifestation of the same root cause as NVIDIA-NeMo/RL#2631 (sleep level 2 silently corrupting Gemma-4 VLM rollouts) and NVIDIA-NeMo/RL#2582 (DAPO crash from discarding generation weights). All three trace back to NVIDIA-NeMo/RL#2495, which introduced `discard_weights=colocated_inference`.
## Symptom
`train/vllm/spec_acceptance_rate-pos-*` spikes on the **first** generation step of every run/resume (\~0.45–0.55), then drops to exactly `0` and stays pinned for the rest of the run. Reproduced across all eagle3 and standalone-draft-model ablations (Qwen3-8B, 8 nodes x 4 GPUs, colocated).
vLLM spec-decode metrics confirm it:
```
SpecDecoding metrics: Mean acceptance length: 1.00, Accepted: 0, Drafted: 114264,
Per-position acceptance rate: 0.000, 0.000, 0.000, Avg Draft acceptance rate: 0.0%
```
## Root cause
1. **End of generation discards everything (including the drafter).**
```python
# nemo_rl/algorithms/grpo.py
policy_generation.finish_generation(discard_weights=colocated_inference)
```
```python
# nemo_rl/models/generation/vllm/vllm_worker_async.py
await self.llm.sleep(level=2 if discard_weights else 1) # level 2 = discard, not CPU offload
```
2. **Refit only re-pushes policy weights.** The trainer only exports `draft.`-prefixed weights in the eagle3 *online-training* flow (`policy.draft.enabled: true`). With a frozen/static drafter (`policy.draft.enabled: false`), no draft weights are sent, so `_load_draft_weights()` is a no-op:
```python
# nemo_rl/models/generation/vllm/vllm_backend.py
def _load_draft_weights(self, draft_weights):
if not draft_weights:
return # <- frozen drafter: nothing is ever reloaded
...
draft_model.load_weights(weights=draft_weights)
```
3. **Net effect:** after the first `finish_generation` → `sleep(level=2)` → next-step refit, the policy weights are restored but the drafter is left as freshly-allocated dummy memory. Acceptance → 0 permanently.
The first step (and the first step after every job resume) only works because the container/code version-mismatch guard forces `load_format="auto"`, so vLLM loads the real draft checkpoint once at init. That single good step is exactly the spike seen in the chart; it dies on the very next step.
## Affected code
* `nemo_rl/algorithms/grpo.py` — `finish_generation(discard_weights=colocated_inference)`
* `nemo_rl/models/generation/vllm/vllm_worker_async.py` — `sleep(level=2 if discard_weights else 1)` (and the sync `vllm_worker.py` equivalent)
* `nemo_rl/models/generation/vllm/vllm_backend.py` — `_load_draft_weights` / `_split_policy_and_draft_weights`
## Reproduction
* Model: `Qwen/Qwen3-8B`, colocated GRPO, 8n4g
* `policy.draft.enabled: false`
* `policy.generation.vllm_kwargs.speculative_config: {method: eagle3, model: Qwen3-8B-eagle3_dapo, num_speculative_tokens: 3}`
* Watch `train/vllm/spec_acceptance_rate-pos-*`: nonzero only on step 1 / first step after each resume, then 0.
## Proposed fix
Either of:
1. **Don't discard a frozen drafter.** When spec decoding is enabled but the drafter is not being trained, use `sleep(level=1)` (CPU-backed offload) so the drafter survives the sleep/wake cycle.
2. **Reload the static draft checkpoint after wake.** After `prepare_for_generation(tags=["weights"])` in `refit_policy_generation`, reload the `speculative_config.model` checkpoint into `model_runner.drafter.model`, since those weights are not part of the IPC refit stream and never change.
Option 2 is the more surgical fix and keeps level-2 discard for the policy.
## Related
* Introduced by NVIDIA-NeMo/RL#2495
* Same root cause: NVIDIA-NeMo/RL#2631, NVIDIA-NeMo/RL#2582
Contributor guide
Assessment
This issue has not been assessed yet.