vLLM sleep level 2 silently corrupts Gemma-4 VLM rollouts in all colocated runs

## Summary

Colocated GRPO on `google/gemma-4-E4B-it` (a Gemma-4 VLM with vision + audio
towers) produces garbage rollouts starting on the first refit cycle, with
`gen_kl_error` exploding from step 2. Isolated the cause to **vLLM's sleep
level-2 discard itself** — not the NeMo-RL refit path and not any particular
weight-reload mechanism.

Two completely independent, standalone (no NeMo-RL) reload mechanisms both fail
to restore correct output after a level-2 discard, while a level-1 sleep/wake is
**byte-identical** to baseline. The level-2 discard frees non-weight GPU/engine
state (most likely computed rotary / positional buffers) that re-feeding weights
does not reconstruct.

Since PR #2495, NeMo-RL forces **every** colocated run through level-2 sleep
with no per-model gate:

- `nemo_rl/algorithms/grpo.py` ~L1683:
```python
policy_generation.finish_generation(discard_weights=colocated_inference)
```
- `nemo_rl/models/generation/vllm/vllm_worker.py:1016`:
```python
self.llm.sleep(level=2 if discard_weights else 1)
```

So any model whose state vLLM does not fully restore on a level-2 wake will
silently produce garbage rollouts (manifesting as exploding `gen_kl_error`).
Gemma-4 is confirmed; other multimodal / custom-positional-encoding models are
plausibly affected.

## Environment

- Model: `google/gemma-4-E4B-it` (`Gemma4ForConditionalGeneration`, VLM with
vision + audio towers)
- vLLM 0.20.0
- transformers 5.5.0
- CUDA 13
- DTensor v2 Automodel backend
- Recipe: ` examples/configs/recipes/vlm/vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml` of https://github.com/NVIDIA-NeMo/RL/pull/2224
- Container: `/lustre/fsw/portfolios/coreai/users/shuangy/images/nemorl-gemma4-2026-05-28-cu13-vllm0.20.sqsh`
- Affected training job log: `/lustre/fsw/portfolios/coreai/users/shuangy/src/NeMo-RL/nemo-rl/12286326-logs/ray-driver.log`

## Steps to reproduce

### A. In NeMo-RL (training)

Launch the colocated VLM GRPO recipe and watch `gen_kl_error`:

```bash
uv run examples/run_vlm_grpo.py \
--config examples/configs/recipes/vlm/vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml
```

`gen_kl_error` explodes starting at step 2; sampled rollouts become garbage after
the first refit cycle.

### B. Standalone vLLM repro (no NeMo-RL, no FSDP, no refit)

Two scripts, both run on a single GPU (TP=1) inside the same container. Each
builds the VLM, generates a greedy baseline on a fixed image+text prompt, runs a
**level-1** sleep/wake control (must MATCH), then runs three **level-2**
sleep/wake + reload cycles (each compared to baseline).

`srun` command used for both (TP=1):

```bash
srun -A coreai_dlalgo_nemorl -p batch \
--container-image=$HOME/images/nemorl-gemma4-2026-05-28-cu13-vllm0.20.sqsh \
--container-mounts=/lustre/fs1:/lustre/fs1,/lustre/fsw:/lustre/fsw \
--gres=gpu:1 \
python _vllm_sleep_level2_vlm_test.py --tp 1 \
--model /lustre/fs1/.../shuangy/models/google/gemma-4-E4B-it
```

#### Repro #1 — recovery via vLLM's built-in `reload_weights`

```python
#!/usr/bin/env python
"""Pure-vLLM: does sleep level 2 preserve Gemma-4 E4B VLM inference?
- Level 1: sleep(1) -> wake_up() (weights kept; no reload)
- Level 2: sleep(2) -> wake_up() -> reload_weights -> reset (canonical level-2 recovery)
"""
import argparse, base64, io, sys
from PIL import Image, ImageDraw

DEFAULT_MODEL = "/lustre/fs1/.../shuangy/models/google/gemma-4-E4B-it"

def make_fixed_image():
img = Image.new("RGB", (256, 256), (240, 240, 240))
d = ImageDraw.Draw(img)
d.rectangle([40, 40, 150, 150], fill=(220, 40, 40), outline=(0, 0, 0), width=3)
d.ellipse([120, 120, 220, 220], fill=(40, 80, 220), outline=(0, 0, 0), width=3)
d.line([0, 0, 255, 255], fill=(0, 160, 0), width=4)
d.text((60, 200), "42", fill=(0, 0, 0))
return img

def image_data_url(image):
buf = io.BytesIO(); image.save(buf, format="PNG")
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")

def build_messages(image):
return [{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": image_data_url(image)}},
{"type": "text", "text": "Describe every shape in this image and their colors and positions."},
]}]

def gen(llm, messages, sp):
o = llm.chat(messages, sampling_params=sp)[0].outputs[0]
return o.text, tuple(o.token_ids)

def compare(label, base_text, base_ids, text, ids):
first_div = next((i for i, (a, b) in enumerate(zip(base_ids, ids)) if a != b),
min(len(base_ids), len(ids)) if len(base_ids) != len(ids) else None)
print(f"\n===== {label} =====")
print(f" text_match={text==base_text} ids_match={ids==base_ids} "
f"len(base)={len(base_ids)} len(this)={len(ids)} first_div_tok={first_div}")
print(f" baseline[:240]: {base_text[:240]!r}")
print(f" this [:240]: {text[:240]!r}")
return text == base_text and ids == base_ids

def maybe_clear_mm_cache(llm):
r = getattr(llm, "renderer", None)
if r is not None and hasattr(r, "clear_mm_cache"):
r.clear_mm_cache()

def reload_weights_l2(llm):
for method in ("reload_weights", "load_weights"):
try:
llm.collective_rpc(method); return
except Exception:
pass
raise RuntimeError("level-2 weight reload failed")

def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default=DEFAULT_MODEL)
ap.add_argument("--tp", type=int, default=1)
ap.add_argument("--cycles", type=int, default=3)
ap.add_argument("--max-tokens", type=int, default=256)
args = ap.parse_args()

from vllm import LLM, SamplingParams
llm = LLM(model=args.model, enable_sleep_mode=True, enforce_eager=True,
trust_remote_code=True, dtype="bfloat16", max_model_len=3072,
limit_mm_per_prompt={"image": 1}, tensor_parallel_size=args.tp,
gpu_memory_utilization=0.6, seed=0)

image = make_fixed_image(); messages = build_messages(image)
sp = SamplingParams(temperature=0.0, max_tokens=args.max_tokens, seed=0)
base_text, base_ids = gen(llm, messages, sp)

# Level 1 control
llm.sleep(level=1); llm.wake_up(); maybe_clear_mm_cache(llm)
t, ids = gen(llm, messages, sp)
compare("LEVEL 1 (control, expect MATCH)", base_text, base_ids, t, ids)

# Level 2 cycles
for i in range(args.cycles):
llm.sleep(level=2); llm.wake_up()
reload_weights_l2(llm)
llm.reset_prefix_cache(); maybe_clear_mm_cache(llm)
t, ids = gen(llm, messages, sp)
compare(f"LEVEL 2 cycle {i+1} (vs baseline)", base_text, base_ids, t, ids)

if __name__ == "__main__":
sys.exit(main())
```

#### Repro #2 — recovery via `model.load_weights(disk)` + `process_weights_after_loading`

This mirrors NeMo-RL's two-step training refit
(`vllm_backend.py`: `model.load_weights(weights=...)` followed by
`process_weights_after_loading`), except the `(name, tensor)` pairs are read
from the on-disk HF safetensors instead of the IPC stream. The only difference
from repro #1 is the level-2 recovery body:

```python
import os
os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") # collective_rpc callable

def _reload_from_disk(worker, model_path):
import glob
from safetensors import safe_open
model_runner = getattr(worker, "model_runner", None)
model = model_runner.model
shards = sorted(glob.glob(model_path + "/*.safetensors"))

def weight_iter():
for shard in shards:
with safe_open(shard, framework="pt", device="cpu") as f:
for k in f.keys():
yield k, f.get_tensor(k)

loaded = model.load_weights(weights=weight_iter()) # vLLM's HF->vLLM name mapping

from vllm.config import set_current_vllm_config
from vllm.model_executor.model_loader.utils import process_weights_after_loading
vllm_config = model_runner.vllm_config
model_config = getattr(worker, "model_config", model_runner.model_config)
device = next(model.parameters()).device
with set_current_vllm_config(vllm_config):
process_weights_after_loading(model, model_config, device)
return {"shards": len(shards), "loaded": len(loaded), "pwal": "ok"}

# level-2 recovery body, per cycle:
# llm.sleep(level=2); llm.wake_up()
# llm.collective_rpc(_reload_from_disk, args=(model_path,))
# llm.reset_prefix_cache(); maybe_clear_mm_cache(llm)
```

## Observed behavior

Both standalone runs are identical: **level 1 byte-perfect, all three level-2
cycles diverge from token 0.**

### Repro #1 SUMMARY (`reload_weights`)

```
================ SUMMARY ================
level1 : MATCH
level2_cycle1 : DIVERGED
level2_cycle2 : DIVERGED
level2_cycle3 : DIVERGED

[verdict]
REPRODUCED: level-1 clean but level-2 (with ground-truth disk reload) DIVERGES
-> vLLM level-2 path corrupts Gemma-4 E4B VLM, independent of NeMo-RL.
```

Level-1 control matched exactly:

```
===== LEVEL 1 (control, expect MATCH) =====
text_match=True ids_match=True len(base)=204 len(this)=204 first_div_tok=None
```

Level-2 cycle 1 (identical for cycles 2 and 3) — diverges at token 0 with garbage:

```
===== LEVEL 2 cycle 1 (vs baseline) =====
text_match=False ids_match=False len(base)=204 len(this)=256 first_div_tok=0
baseline[:240]: 'This image contains three distinct shapes: a square, a circle, and a line.\n\nHere is a detailed description of each:\n\n1. **Square:**\n * **Color:** Red (filled).\n * **Position:** Located in the upper-left to central area of the ima'
this [:240]: '...L/captionty:T-data \\_-data \\_-data \\_-data \\_-data \\_-data \\_-d'
```

During each `reload_weights` cycle, vLLM emits these warnings (it cannot reload
the rotary / positional / head modules):

```
WARNING [layerwise.py:225] Gemma4VisionRotaryEmbedding: Failed to load weights
WARNING [layerwise.py:225] Gemma4AudioRelPositionalEncoding: Failed to load weights
WARNING [layerwise.py:225] Gemma4Model: Failed to load weights
WARNING [layerwise.py:225] RotaryEmbedding: Failed to load weights
WARNING [layerwise.py:225] Gemma4RotaryEmbedding: Failed to load weights
WARNING [layerwise.py:225] ParallelLMHead: Failed to load weights
```

### Repro #2 SUMMARY (`model.load_weights(disk)` + `process_weights_after_loading`)

```
================ SUMMARY ================
level1 : MATCH
level2_cycle1 : DIVERGED
level2_cycle2 : DIVERGED
level2_cycle3 : DIVERGED

[verdict]
Level-2 via model.load_weights(disk) STILL DIVERGES (like reload_weights).
=> the bug is in the level-2 DISCARD itself, independent of the reload mechanism.
```

Here **all 2214 tensors were accepted and there were NO `layerwise.py` warnings**:

```
[worker 0] shards=1 load_weights_returned=2214 pwal=ok
```

…yet the output is the same garbage from token 0:

```
===== LEVEL 2 cycle 1 (load_weights, vs baseline) =====
text_match=False ids_match=False len(base)=204 len(this)=256 first_div_tok=0
this [:240]: '...L/captionty:T-data \\_-data \\_-data \\_-data ...'
```

## Expected behavior

After a level-2 sleep + a complete weight reload, greedy generation on the same
prompt should be byte-identical to baseline (as it is for level 1). It is not —
the engine produces garbage regardless of how weights are reloaded.

## Root-cause analysis

- **Level 1 is byte-perfect.** Level-1 sleep keeps weights resident (only KV
cache / activations freed) and restores output exactly.
- **Two independent, complete reload mechanisms both fail after level 2:**
1. vLLM's own `collective_rpc("reload_weights")` — diverges, and reports
`layerwise.py:225 "Failed to load weights"` for the rotary / positional /
`Gemma4Model` / `ParallelLMHead` modules.
2. `model.load_weights(disk)` + `process_weights_after_loading` (the exact
two-step refit NeMo-RL uses, fed from disk) — diverges too, even though all
2214 tensors load cleanly with **no** warnings.
- Since the reload mechanism is irrelevant to the outcome, the defect is in the
**level-2 discard itself**: it frees non-weight GPU/engine state that
re-feeding weights does not reconstruct.

## Impact

PR #2495 made `discard_weights=colocated_inference` unconditional, so **every
colocated run** now sleeps at level 2 with **no per-model gate**. Any model whose
engine state vLLM does not fully restore on a level-2 wake will silently emit
garbage rollouts — there is no crash, just exploding `gen_kl` and meaningless
training signal. Gemma-4 VLM is confirmed broken; other multimodal models and
models with custom positional encodings are plausibly affected and currently
unguarded.

## Suggested fix / mitigation

- **NeMo-RL (short term):** default the colocated sleep to **level 1**, or gate
level-2 behind a per-model validated allowlist (e.g. only models confirmed to
survive a level-2 sleep/wake round-trip). Concretely, stop hard-wiring
`level=2 if discard_weights else 1` in
`nemo_rl/models/generation/vllm/vllm_worker.py:1016` without a model-level
guard.
- **Upstream vLLM:** this reproduces with zero NeMo-RL involvement, so a vLLM-side
fix is also warranted — level-2 discard should either preserve or correctly
rebuild the computed rotary/positional/engine buffers (or `reload_weights`
should stop silently failing on those modules) for Gemma-4 and similar models.

## Artifacts
Under `/lustre/fsw/portfolios/coreai/users/shuangy/src/NeMo-RL/nemo-rl/`
- `_vllm_sleep_level2_vlm_test.py` / `_vllm_sleep_test_tp1.log` (repro 1)
- `_vllm_sleep_level2_vlm_test_loadweights.py` / `_vllm_sleep_test_loadweights_tp1.log` (repro 2)