NVIDIA-NeMo / NVIDIA-NeMo/RL

GRPO broken for MoE models if vLLM's moe_backend is FLASHINFER_CUTLASS

Open
#2,618 1 comment 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

**Describe the bug**

Currently when an MoE model choses FLASHINFER_CUTLASS as the moe backend, GRPO doesn't work, on h100 it gives abnormal logprob errors and on GB200 it crashes during refit.

We had to pin moe_backend to Triton, the backend that was default in vLLM 0.18.0, to fix the regression (PR: https://github.com/NVIDIA-NeMo/RL/pull/2586), but we should fix FLASHINFER_CUTLASS codepath and practically any moe_backend.

H100 behavior:
https://github.com/NVIDIA-NeMo/RL/blob/fbe2e429169368aa46cd3212486924bd3a7550e5/ISSUE_vllm_020_regressions.md#tldr row 3

GB200 behavior:
```bash
[36m(VllmGenerationWorker pid=2616579, ip=10.66.4.71)[0m Error: Worker failed to update weights. Result: False
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m Error in VllmInternalWorkerExtension.update_weights_via_ipc_zmq: shard_dim=0 is not a valid data dimension for a 3D tensor (expected 1 or 2).
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m Traceback (most recent call last):
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/nemo-rl/nemo_rl/models/generation/vllm/vllm_backend.py", line 306, in update_weights_via_ipc_zmq
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m self._load_weights(weights)
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ~~~~~~~~~~~~~~~~~~^^^^^^^^^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/nemo-rl/nemo_rl/models/generation/vllm/vllm_backend.py", line 237, in _load_weights
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m self.model_runner.model.load_weights(weights=policy_weights)
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.13/site-packages/vllm/model_executor/models/deepseek_v2.py", line 1636, in load_weights
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m success = weight_loader(
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m param,
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ...<4 lines>...
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m return_success=True,
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m )
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.13/site-packages/vllm/model_executor/layers/fused_moe/layer.py", line 1374, in weight_loader
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m self._load_model_weight_or_group_weight_scale(
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m shard_id=shard_id,
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ^^^^^^^^^^^^^^^^^^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ...<3 lines>...
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m tp_rank=self.tp_rank,
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ^^^^^^^^^^^^^^^^^^^^^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m )
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.13/site-packages/vllm/model_executor/layers/fused_moe/layer.py", line 851, in _load_model_weight_or_group_weight_scale
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m self._load_w13(
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ~~~~~~~~~~~~~~^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m shard_id=shard_id,
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ^^^^^^^^^^^^^^^^^^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ...<3 lines>...
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m tp_rank=tp_rank,
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ^^^^^^^^^^^^^^^^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m )
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ^
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.13/site-packages/vllm/model_executor/layers/fused_moe/layer.py", line 980, in _load_w13
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.13/site-packages/vllm/model_executor/layers/fused_moe/layer.py", line 904, in _get_hidden_dim
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m raise ValueError(
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ...<2 lines>...
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m )
[36m(RayWorkerWrapper pid=2617116, ip=10.66.4.71)[0m ValueError: shard_dim=0 is not a valid data dimension for a 3D tensor (expected 1 or 2)
```

**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

**Expected behavior**

A clear and concise description of what you expected to happen.

**Additional context**

Add any other context about the problem here.

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.