MXFP8 rollout support for Qwen3.5 (grouped-MoE expert refit)
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
# **Summary**
#2744 adds FP8 rollout + training for Qwen3.5 (dense and MoE), but only for the **blockwise** recipe (e4m3, `[128, 128]` weight blocks) . MXFP8 currently does not work for Qwen3.5: the grouped-MoE expert refit path only implements blockwise quantization, so an MXFP8 refit of a Qwen3.5 MoE model fails at the first refit with:
```
NotImplementedError: MXFP8 refit does not support grouped MoE expert weights.
```
(guard + unit test added in #2744). Dense Qwen3.5 (9B) does not hit the guard but is unvalidated under MXFP8.
# **Proposal**
End-to-end MXFP8 rollout support for Qwen3.5, which needs:
1. An MXFP8 counterpart of `_expand_grouped_moe_expert_to_fp8` in `nemo_rl/models/generation/vllm/quantization/fp8.py`: per-expert e4m3 data + E8M0 scales (32-element blocks along the reduction dim), emitted under the `*_scale_from_checkpoint` naming the dense MXFP8 path already uses (including `pad_flashinfer_scale_k` handling).
2. Verification that per-expert MXFP8 scale entries route correctly through vLLM's `RoutedExperts` expert mapping with the ModelOpt MXFP8 MoE method (the blockwise path needed per-expert re-emission because the fused loader has no scale path — the same question applies here).
3. The same exclusion behavior as blockwise: `quantization_ignored_layer_kws` (GDN projections + vision tower) and bf16 pass-through for ignored experts.
4. Tests and a recipe: extend the grouped-expert unit tests in `tests/unit/models/generation/test_vllm_fp8_quantization.py` (currently a `pytest.raises` guard for MXFP8) with value-level checks, and add/validate an MXFP8 recipe variant once the path works.
Contributor guide
Research direction
Start in nemo_rl/models/generation/vllm/quantization/fp8.py, alongside _expand_grouped_moe_expert_to_fp8, and run the grouped-expert tests in tests/unit/models/generation/test_vllm_fp8_quantization.py. Trace how per-expert MXFP8 scales map through vLLM's RoutedExperts and review the existing blockwise exclusion behavior. Done means value-level tests pass and an MXFP8 Qwen3.5 recipe is added or validated.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, testing
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100