📺 Info Siaran: verl project/verl/pull/

PeterYang12 · 2026-08-19T08:52:57Z

What does this PR do?

Fix FP8 weight refits for the ROCm AITER MoE backend.

AITER shuffles expert weights into an MFMA runtime layout without changing their shape or dtype. The existing staging check therefore treated those parameters as still being in checkpoint layout and skipped post-load reprocessing. Detect is_shuffled parameters so refits stage the layer and reapply the AITER shuffle before inference.

This does not duplicate an existing PR. Searches for AITER FP8 refit, is_shuffled FP8, and ROCm FP8 staging found no open PR addressing this fix.

Checklist Before Starting

Searched for similar PRs using the links above.
Formatted the PR title as [vllm] fix: ....

Test

uvx --with hydra-core pre-commit run --all-files --show-diff-on-failure --color=never — all hooks passed.
python3 -m py_compile verl/utils/vllm/vllm_fp8_utils.py — passed.
Inline CPU smoke test for _layer_needs_fp8_staging — passed, covering unchanged parameters, shape mismatch, non-parameter values, and the new is_shuffled=True path.

A full ROCm AITER MoE refit was not run in this environment.

API and Usage Example

No API or configuration changes. FP8 refits automatically restage AITER-shuffled MoE weights.

Design & Code Changes

Keep the existing shape and dtype staging checks.
Also stage a layer when a live FP8 parameter is marked is_shuffled, since AITER's MFMA permutation preserves shape and dtype.

Checklist Before Submitting

Read the Contribute Guide.
Applied the full pre-commit checks.
Documentation was not updated because this is an internal correctness fix with no user-facing API or configuration change.
No CI test was added because the end-to-end failure requires the ROCm AITER MoE backend. The new predicate path was covered by an isolated CPU smoke test.
CI will be requested in the verl community channel when the PR is ready for maintainer review.
Recipe submodule update is not applicable.

AI assistance

OpenAI Codex was used to review the human-authored change, run checks, prepare the commit and PR text, and submit the PR. The human submitter authored and reviewed every changed line and owns the change end-to-end.

Signed-off-by: PeterYang12 yuhanya@amd.com

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: PeterYang12 <yuhanya@amd.com>

… the checkpoint shape (#7986) ### What does this PR do? The FP8 refit staging cycle in `verl/utils/vllm/vllm_fp8_utils.py` misses MoE layers whose kernel prep **rewrites the expert weights without changing their shape or dtype**. This PR makes the patched `replace_parameter` record such rewrites so the layer is staged and re-processed on every refit. **Mechanism.** `_layer_needs_fp8_staging` compares each FP8 parameter's live shape/dtype with the checkpoint layout recorded at load, and honours ROCm AITER's `is_shuffled` mark (#7470). Both signals are blind to a prep that returns a same-shaped, same-typed *copy*: - FlashInfer **CUTLASS** block-FP8 MoE prep (`prepare_fp8_moe_layer_for_fi`): `swap_w13_to_w31` on `w13_weight` and `w13_weight_scale_inv` (gate/up halves swapped, shape kept). - FlashInfer **TRT-LLM MXFP8** prep: W13→W31 + row interleave + tile shuffle of weights and scales (shape kept) — reached through #7519, not on `main` yet. For such a layer `stage_fp8_params_for_loading` returned nothing, so on a refit the sync wrote canonical (checkpoint-layout) expert weights straight into the repacked live buffer and `process_weights_after_loading` was **not** re-run: the kernel read the new weights as if they were already in its layout. **Who hits it on `main`.** `rollout.quantization=fp8` (block `[128,128]`) with the FlashInfer CUTLASS MoE backend — which vLLM 0.24 and 0.29 auto-select on **Hopper when `expert_parallel_size > 1`** (`fused_moe/oracle/fp8.py::_get_priority_backends`, source-verified on both tags). On Blackwell the same config selects TRT-LLM, whose block-FP8 prep changes the rank (4-D BlockMajorK) and is therefore caught by the existing shape check. Hopper TP-only picks Triton (no repack). **Fix.** `replace_parameter_preserve_subclass` is the one place the rewrite is visible (a same-shaped, same-typed tensor whose storage is not the old parameter's): it records the param name in `layer._verl_fp8_repacked`, and `_layer_needs_fp8_staging` treats a recorded name like `is_shuffled`. Identity replacements (Triton, vLLM-CUTLASS pass the same tensor back) record nothing and stay unstaged; shape-changing repacks (DeepGEMM, Marlin, TRT-LLM block-FP8) are detected as before. Related: #7470 (AITER `is_shuffled` restage — the same class of blind spot, fixed for one backend by trusting its mark), #7519 (where this was found; carries an additional patch of ModelOpt's own `replace_parameter` namespace for the MXFP8 method, which does not exist on `main`). ### Checklist Before Starting - [x] Similar PRs: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+restage (0), https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+is_shuffled (0), https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+fp8+staging+refit (only #7519, the feature branch this was found on). Not a duplicate: #7519 adds the MXFP8 feature; this PR fixes the pre-existing staging predicate that `main`'s block-FP8 path already depends on. - [x] Title format `[rollout, vllm] fix: ...` ### Test ```bash python -m pytest tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py tests/utils/test_vllm_quant_utils_moe_on_cpu.py -q # 7 passed uvx ruff@0.12.2 check / format --check on the two changed files: clean ``` New `tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py` (loads `vllm_fp8_utils` against a stub `vllm...quantization.fp8`, no vLLM needed): - `test_same_shape_repack_is_restaged_and_reprocessed_into_the_live_storage` — CUTLASS-like swap on `w13_weight` / `w13_weight_scale_inv`: recorded at load, staged on refit, prep re-run, result folded into the original storage (`data_ptr` unchanged), kernel references captured after the replace calls point at the live parameters. - `test_layout_keeping_backend_is_not_staged` — identity replacement records nothing, layer stays unstaged. - `test_shape_changing_repack_is_still_staged` — the pre-existing detection path is unchanged. #### Hardware evidence (1×B200, 2026-09-21) Setup: RunPod 1×B200 (SM100, 183 GB), image `verlai/verl:vllm024.dev2` — vLLM 0.24.0, FlashInfer 0.6.12, Megatron-Core 0.18, Megatron-Bridge 0.5.2. Model: random-init 2-layer / 8-expert Qwen3-MoE (hidden 512, `moe_intermediate_size` 256, top-k 2, per-expert HF layout). Megatron MXFP8 training + vLLM MXFP8 rollout via the #7519 stack (`ModelOptMxFp8FusedMoE` → FlashInfer TRT-LLM MoE backend), TP1 EP1, GRPO on gsm8k, 4 steps, same seed and config across the three runs below. The MoE refit probe from #7519 routes every probe row to local expert 0 and compares the kernel's output with a dequantized reference (default tolerance 0.25). **Before (`79caec3c`, first weight sync — no training step has happened yet).** The dense probe passes, the MoE probe fails and the run dies: ```text INFO:2026-09-21 09:43:32,540:FP8 weights loaded (async), loaded_params: 29, loaded_buffers: 0 INFO:2026-09-21 09:43:33,246:mxfp8 refit self-check passed on model.layers.0.self_attn.o_proj (vllm): rel err 0.0264 <= tol 0.25 ERROR 09-21 09:43:33 [multiproc_executor.py:1000] RuntimeError: MXFP8 refit self-check failed on vllm MoE layer 'model.layers.0.mlp.experts.routed_experts', expert 0: the layer routed to this expert disagrees with the dequantized reference by 1.739 relative error (tolerance 0.25). This is the signature of expert scales (w13_/w2_weight_scale) whose kernel layout was not re-derived after a weight sync, ... Exception: Call to collective_rpc method failed: Worker failed with error 'MXFP8 refit self-check failed on vllm MoE layer ...' ``` (`trainer.init → checkpoint_engine/base.py:514 update_weights → update_weights_from_ipc → process_quanted_weights_after_loading → self_check_mxfp8_moe`.) **Same code, probe tolerance opened (`VERL_MXFP8_REFIT_CHECK_MOE_TOL=1000`) so training proceeds.** The error is deterministic (1.7394 again) and the engine's expert weights are wrong — `rollout_corr/kl` is 25× the bf16 baseline — while the layer-set audit is clean, because the audit checks *which* layers are fp8, not their layout: ```text mxfp8 MoE refit self-check passed on model.layers.0.mlp.experts.routed_experts expert 0 (vllm): rel err 1.7394 <= tol 1000.00 quantized-layer audit: the rollout engine's live parameters agree with the training side (69 parameters re-checked against the engine itself) rollout_corr/kl step 1-4: 0.04655 0.04668 0.04753 0.04719 rollout_probs_diff_mean step 1-4: 1.90e-06 1.92e-06 1.92e-06 1.89e-06 ``` **After (`0161d15f`, the record in place; same weights and inputs).** ```text mxfp8 refit self-check passed on model.layers.0.self_attn.o_proj (vllm): rel err 0.0264 <= tol 0.25 mxfp8 MoE refit self-check passed on model.layers.0.mlp.experts.routed_experts expert 0 (vllm): rel err 0.0520 <= tol 0.25 quantized-layer audit: training and the engine blacklist as configured (quantization_config.ignored_layers) quantize the same layers (2 of 2 decoder layers on this rank in fp8, 69 parameters checked) quantized-layer audit: the rollout engine's live parameters agree with the training side (69 parameters re-checked against the engine itself) rollout_corr/kl step 1-4: 0.00225 0.00214 0.00201 0.00229 rollout_probs_diff_mean step 1-4: 1.52e-07 1.52e-07 1.55e-07 1.51e-07 TRAIN_EXIT=0 ``` All four syncs passed the probe at the default tolerance. Forcing the tolerance to 1e-6 on the fixed code makes the same layer report `disagrees with the dequantized reference by 0.052 relative error (tolerance 1e-06)` — 0.052 is the genuine kernel-vs-reference error of a two-GEMM MXFP8 gated MLP, not a leftover layout mismatch. **bf16 rollout baseline** (same model and config; `moe_backend=triton` because of #7978): kl 0.00190 / 0.00187 / 0.00185 / 0.00186, `rollout_probs_diff_mean` ≈ 4.7e-08, `TRAIN_EXIT=0`. | run | MoE probe rel err | `rollout_corr/kl` (steps 1–4) | `rollout_probs_diff_mean` | |---|---|---|---| | bf16 baseline | — | 0.0019 | 4.7e-08 | | mxfp8, before | 1.739 | 0.047 (25× bf16) | 1.9e-06 | | mxfp8, after | **0.052** | **0.0022** (1.2× bf16) | 1.5e-07 | Caveat, as above: the measured instance is MXFP8 / TRT-LLM under the #7519 stack (which additionally installs the same `replace_parameter` stand-in in ModelOpt's namespace, not needed on `main`). The block-FP8 / FlashInfer CUTLASS instance reachable from `main` (Hopper, EP>1) shares the predicate and the blind spot but is source-verified and CPU-tested only; I do not have a Hopper EP>1 setup at hand. ### API and Usage Example No API or config change. ### Design & Code Changes - `verl/utils/vllm/vllm_fp8_utils.py` - `_FP8_REPACKED_ATTR = "_verl_fp8_repacked"`: set of param names a prep handed back as a same-shape rewrite. - `replace_parameter_preserve_subclass`: on the non-refit path, record `param_name` when `old_param` exists, shapes and dtypes match, and `data_ptr()` differs. - `_layer_needs_fp8_staging`: `is_shuffled or name in repacked`. - `tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py`: the three tests above. ### Checklist Before Submitting - [x] Read the Contribute Guide - [x] ruff 0.12.2 check/format on the changed files (full `pre-commit run --all-files` not run locally) - [x] Documentation: no user-facing behaviour change beyond the fix; the mechanism is documented in the code comments - [x] Unit tests added (CPU, run under the `cpu_unit_tests` workflow's `tests/utils` path) - [ ] ci-request AI assistance was used for the diagnosis (source reading of vLLM 0.24/0.29 and verl) and for drafting the code, tests and this description (Claude, with Codex on the initial source analysis). The author reviewed every changed line and ran the tests above. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

… QAT, FP8 or BF16 training (#8146) ### What does this PR do? Makes DeepSeek-V4-Flash GRPO run end to end on AMD MI355X with an FP8 rollout, and gives it one recipe with three training modes. Follow-up to #7050 (ROCm enablement) and #7455 / #7470 (Megatron-Bridge path). Two commits: 1. **`[docker]` image.** `docker/rocm/Dockerfile.rocm.deepseek-v4` builds the versions the runs were validated with: Transformer Engine 2.18 (ROCm build: Triton grouped GEMM path that avoids a bf16 grouped-GEMM backward hang on gfx950, float-scale blockwise FP8 GEMM), a libhsa-runtime that returns memory under `expandable_segments` on ROCm 7.2.x, Megatron-LM `fd1121b8f` plus the `deepseek-v4-rocm-support` commits (config-driven FP8 fake quantization, FP8-training alignment options, and the `index_select` gather of the unfused CSA sparse attention that upstream introduced with the reworked SBHD path in NVIDIA/Megatron-LM#6400, which this older base predates), Megatron-Bridge 0.5.0 plus a fix that keeps the FP8 weight export on the checkpoint's power-of-two scale grid, and four vLLM 0.26.0 patches applied at build time (NaN guard on the FP8 indexer logits; CuMem wake cache release, vllm#49857; indexer K-cache block-flat layout on gfx950/gfx942, vllm#50455; SwiGLU clamp on the FP8 Triton fused MoE path). 2. **`[recipe]` one script, three modes.** `examples/grpo_trainer/run_deepseek_v4_flash_megatron_rocm.sh`, `TRAIN_MODE=qat|fp8|bf16`. `qat` is BF16 training with Megatron's FP8 fake quantization (weights on 128x128 power-of-two scales, inputs on 1x128 float scales, straight-through estimator) and an FP8 W8A8 vLLM rollout; `fp8` is Transformer Engine blockwise FP8 training with float activation scales; `bf16` is BF16 training with a BF16 rollout. All modes round the KV latent and the indexer q/k the way the rollout's FP8 KV cache and FP8 indexer do. Everything is a transformer-config override; the only environment variables are the two Transformer Engine kernel selectors. Why the alignment matters: on ROCm, vLLM quantizes FP8 activations with float scales (amax / 448) while Transformer Engine's blockwise recipe uses powers of two, so the same training code that aligns on NVIDIA (DeepGEMM) does not align here. With the rollout's rounding reproduced in training, `training/rollout_probs_diff_mean` at step 1 drops to the BF16-forward level and the train/rollout KL stays flat over 300 steps. The Megatron-LM and Megatron-Bridge changes live in the author's forks while they are proposed upstream; the Transformer Engine wheel and the libhsa-runtime are published as release assets because no official ROCm package carries those fixes yet. ### Checklist Before Starting - [x] Search for similar PRs. https://github.com/verl-project/verl/pulls?q=is%3Apr+deepseek-v4+rocm - [x] Format the PR title as `[{modules}] {type}: {description}` ### Test DeepSeek-V4-Flash (FP8 checkpoint), 4 nodes x 8 MI355X, ROCm 7.2.3, GRPO/GSPO on DAPO-Math, validated on AIME-2024 every 10 steps: | Mode | step-1 `rollout_probs_diff_mean` | step-1 `rollout_corr/kl` | 300 steps | AIME-2024 acc | | --- | --- | --- | --- | --- | | qat | 0.0096 | 0.0021 | KL 0.0022 -> 0.0029, no NaN/hang | 0.60 -> 0.76 | | fp8 | 0.0097 | 0.0021 | KL 0.0022 -> 0.0039 | 0.60 -> 0.76 | | bf16 | 0.0060 | 0.0008 | KL 0.0008 -> 0.0011 | 0.59 -> 0.79 peak | On the image built from this branch, a 7-layer cut of the model (single node) gives diff / KL of 0.0129 / 0.0068 in `qat` and 0.0130 / 0.0069 in `fp8` over 3 steps, matching the earlier runs. Per step on 32 GPUs: ~490 s (`qat`), ~420 s (`fp8`), ~330 s (`bf16`), of which ~85 s is the weight sync. ### API and Usage Example ```bash DOCKER_BUILDKIT=1 docker build -f docker/rocm/Dockerfile.rocm.deepseek-v4 -t rocm-verl-dsv4 . # Ray on every node (commands in the script header), then on the head node: WK=<wandb key> TRAIN_MODE=qat bash examples/grpo_trainer/run_deepseek_v4_flash_megatron_rocm.sh TRAIN_MODE=fp8 TOTAL_TRAINING_STEPS=5 bash examples/grpo_trainer/run_deepseek_v4_flash_megatron_rocm.sh trainer.val_before_train=False ``` ### Design & Code Changes - `docker/rocm/Dockerfile.rocm.deepseek-v4`: version pins and build-time patches as described; `docker/rocm/patches/vllm-*.patch` carry the upstream reference in their headers. - `examples/grpo_trainer/run_deepseek_v4_flash_megatron_rocm.sh`: the recipe; the header documents the modes, the Ray commands and the reference values. The mode only changes a set of `override_transformer_config` keys (`fp8_fake_quant`, `fp8_fake_quant_kv_cache`, `fp8_fake_quant_indexer_qk`, and for `fp8`: `fp8=e4m3 fp8_recipe=blockwise fp8_blockwise_act_pow2_scale=False`), applied to both the actor and the reference policy. ### Checklist Before Submitting - [x] Read the [Contribute Guide](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md). - [x] Apply [pre-commit checks](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting) - [ ] Add / Update the documentation: the recipe header is the documentation; happy to add a docs page if preferred. - [ ] Add unit or end-to-end test(s): needs 32 MI355X and the DeepSeek-V4 checkpoint; validated by the runs above. - [ ] Once your PR is ready for CI, send a message in the `ci-request` channel. --------- Signed-off-by: PeterYang12 <yuhanya@amd.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

[vllm] fix: restage shuffled AITER FP8 weights

c95175d

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: PeterYang12 <yuhanya@amd.com>

PeterYang12 requested review from HollowMan6 and wuxibin89 as code owners August 19, 2026 08:52

wuxibin89 approved these changes Aug 19, 2026

View reviewed changes

wuxibin89 merged commit 45cbfcf into verl-project:main Aug 19, 2026

PeterYang12 deleted the dsv4-fp8-rfix branch August 20, 2026 08:46

wengeezhang mentioned this pull request Sep 22, 2026

[rollout, vllm] fix: restage FP8 MoE layers whose kernel repack keeps the checkpoint shape #7986

Merged

6 of 7 tasks

PeterYang12 mentioned this pull request Oct 8, 2026

[hardware, recipe] feat: DeepSeek-V4-Flash GRPO on ROCm (MI355X) with QAT, FP8 or BF16 training #8146

Merged

4 of 7 tasks

Navigation Menu

Repository navigation

[vllm] fix: restage shuffled AITER FP8 weights - #7470

[vllm] fix: restage shuffled AITER FP8 weights#7470
wuxibin89 merged 1 commit into
verl-project:mainfrom
PeterYang12:dsv4-fp8-rfix

PeterYang12 commented Aug 19, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

PeterYang12 commented Aug 19, 2026

What does this PR do?

Checklist Before Starting

Test

API and Usage Example

Design & Code Changes

Checklist Before Submitting

AI assistance

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants