NVIDIA-NeMo / NVIDIA-NeMo/RL

Speculative decoding silently broken in colocated GRPO: vLLM sleep level 2 discards drafter weights, never restored on refit

Open
#2,646 0 comments 0 reactions 0 assignees View on GitHub
accuracy bug
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

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.