debug w4a4 trtllm rollout ipc by Superjomn · Pull Request #6455 · verl-project/verl

Superjomn · 2026-05-24T12:48:52Z

What does this PR do?

Add concise overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review.

Checklist Before Starting

Search for similar PRs. Paste at least one query link here: ...
Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
- {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
- If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
- {type} is in feat, fix, refactor, chore, test
- If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
- Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

Read the Contribute Guide.
Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
Add / Update the documentation.
Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

…on transformer)

DAPO/PPO with QAT W4A4 (nvfp4) actor and TRT-LLM rollout failed at update_weights because the multiprocessing IPC path cannot move float8 tensors and TRT-LLM's restricted unpickler blocks the rebuild function. Four patches make the path work: 1. trtllm_rollout.py: register a float8 storage map and stub class so reduce_tensor can pickle float8 (fn / fnuz / e5m2). Before reduce_tensor, move CPU-resident params to CUDA (param_offload=True path) and view the float8 tensor as uint8 storage; return a 3-tuple (rebuild_fn, rebuild_args, orig_dtype_str) so the receiver can restore the original dtype. 2. trtllm_worker_extension.py: mirror the float8 storage map and stub class on the receiver. Monkey-patch tensorrt_llm.serialization.loads to extend its restricted- unpickler allowlist with torch.multiprocessing.reductions. rebuild_cuda_tensor. Unpack the 3-tuple and view back to the original float8 dtype. 3. trtllm_async_server.py: route non-VLM TRT-LLM rollouts through verl's WorkerExtension so the receiver patches above actually take effect for non-VLM models. 4. utils/qat/linear.py: minor QAT linear adjustment carried along. Validated on Qwen3-8B-Base + DAPO + FSDP + TRT-LLM rollout + W4A4 (nvfp4), 2-node Lyris gb200, 9 RL training steps, 80 update_weights ENTER events across all WorkerDict ranks, timing_s/update_weights ~12.76s per step-1+ sync, no Tracebacks / OOM / NaN. Known limitations / out of scope for this PR: - The tensorrt_llm.serialization.loads monkey-patch is a workaround; a proper TRT-LLM API change would be cleaner. - The non-VLM -> verl WorkerExtension change in trtllm_async_server.py affects all non-VLM TRT-LLM rollouts in this tree. - W4A4 KL divergence is a separate documented issue.

gemini-code-assist

Code Review

This pull request implements NVFP4 Quantization-Aware Training (QAT) support for the TRT-LLM rollout backend, featuring a workaround for IPC-sharing of float8 tensors by casting them to uint8 and monkey-patching serialization to permit required torch imports. It also introduces a force_dynamic_quantization configuration. Feedback highlights the need to remove extensive debug logging and a performance-degrading out-of-bounds detection block in the model forward pass that causes synchronous GPU-to-CPU stalls.

gemini-code-assist · 2026-05-24T12:51:53Z

+        import sys as _sys
+        try:
+            _iids = input_ids_rmpad
+            if _iids is not None and hasattr(_iids, "numel") and _iids.numel() > 0:
+                _mn = int(_iids.min().item())
+                _mx = int(_iids.max().item())
+                _vs = getattr(model.config, "vocab_size", None) if hasattr(model, "config") else None
+                print(
+                    f"RAY_EXECUTOR_DEBUG: model_forward.input_ids_rmpad: "
+                    f"shape={tuple(_iids.shape)} min={_mn} max={_mx} vocab_size={_vs}",
+                    file=_sys.stderr, flush=True,
+                )
+                if _vs is not None and _mx >= _vs:
+                    print(
+                        f"RAY_EXECUTOR_DEBUG: model_forward.OOB_DETECTED: "
+                        f"max_token_id={_mx} >= vocab_size={_vs}",
+                        file=_sys.stderr, flush=True,
+                    )
+        except Exception as _e:
+            print(f"RAY_EXECUTOR_DEBUG: model_forward.probe_error: {_e}", file=_sys.stderr, flush=True)


This debug block introduces a severe performance regression in the model's forward pass. Calling .min().item() and .max().item() on GPU tensors (lines 288-289) triggers synchronous GPU-to-CPU data transfers, which stalls the computation pipeline. Furthermore, the use of print with flush=True in this hot path will cause significant overhead and log flooding. This entire block should be removed.

gemini-code-assist · 2026-05-24T12:51:53Z

+        print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights ENTER step={self.global_steps}", flush=True)
        self.checkpoint_manager.update_weights(self.global_steps)
+        print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights EXIT step={self.global_steps}", flush=True)


Leftover debug print statements found. These should be removed to maintain clean logs. Use the logging module if tracking weight updates is necessary.

Suggested change

print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights ENTER step={self.global_steps}", flush=True)

self.checkpoint_manager.update_weights(self.global_steps)

print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights EXIT step={self.global_steps}", flush=True)

self.checkpoint_manager.update_weights(self.global_steps)

gemini-code-assist · 2026-05-24T12:51:53Z

+                        print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights ENTER step={self.global_steps} (critic_warmup)", flush=True)
                        self.checkpoint_manager.update_weights(self.global_steps)
+                        print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights EXIT step={self.global_steps} (critic_warmup)", flush=True)


Leftover debug print statements found. These should be removed.

Suggested change

print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights ENTER step={self.global_steps} (critic_warmup)", flush=True)

self.checkpoint_manager.update_weights(self.global_steps)

print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights EXIT step={self.global_steps} (critic_warmup)", flush=True)

self.checkpoint_manager.update_weights(self.global_steps)

gemini-code-assist · 2026-05-24T12:51:53Z

+                            print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights ENTER step={self.global_steps}", flush=True)
                            self.checkpoint_manager.update_weights(self.global_steps)
+                            print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights EXIT step={self.global_steps}", flush=True)


Leftover debug print statements found. These should be removed.

Suggested change

print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights ENTER step={self.global_steps}", flush=True)

self.checkpoint_manager.update_weights(self.global_steps)

print(f"RAY_EXECUTOR_DEBUG: ray_trainer.update_weights EXIT step={self.global_steps}", flush=True)

self.checkpoint_manager.update_weights(self.global_steps)

gemini-code-assist · 2026-05-24T12:51:53Z

+        print(f"RAY_EXECUTOR_DEBUG: transformer_impl.get_per_tensor_param _qat_enabled={self._qat_enabled} type={type(per_tensor_param).__name__}", flush=True)
        if self._qat_enabled:
            from verl.utils.modelopt import export_qat_weights

            per_tensor_param = export_qat_weights(per_tensor_param, self.module, self._qat_config.mode, self.bridge)
+            print(f"RAY_EXECUTOR_DEBUG: transformer_impl.get_per_tensor_param post_qat type={type(per_tensor_param).__name__}", flush=True)



Debug print statements in the weight export path should be removed to avoid log pollution during training.

Suggested change

print(f"RAY_EXECUTOR_DEBUG: transformer_impl.get_per_tensor_param _qat_enabled={self._qat_enabled} type={type(per_tensor_param).__name__}", flush=True)

if self._qat_enabled:

from verl.utils.modelopt import export_qat_weights

per_tensor_param = export_qat_weights(per_tensor_param, self.module, self._qat_config.mode, self.bridge)

print(f"RAY_EXECUTOR_DEBUG: transformer_impl.get_per_tensor_param post_qat type={type(per_tensor_param).__name__}", flush=True)

if self._qat_enabled:

from verl.utils.modelopt import export_qat_weights

per_tensor_param = export_qat_weights(per_tensor_param, self.module, self._qat_config.mode, self.bridge)

gemini-code-assist · 2026-05-24T12:51:53Z

        peft_config=None, so the rollout receives a standard weight update.
        """

+        print(f"RAY_EXECUTOR_DEBUG: engine_workers.update_weights ENTER global_steps={global_steps}", flush=True)


Remove debug print statement.

gemini-code-assist · 2026-05-24T12:51:53Z

+        print(f"RAY_EXECUTOR_DEBUG: engine_workers.before_rollout_update_weights global_steps={global_steps}", flush=True)
        await self.rollout.update_weights(
            per_tensor_param, peft_config=peft_config, base_sync_done=True, global_steps=global_steps
        )
+        print(f"RAY_EXECUTOR_DEBUG: engine_workers.after_rollout_update_weights global_steps={global_steps}", flush=True)


Remove debug print statements.

await self.rollout.update_weights( per_tensor_param, peft_config=peft_config, base_sync_done=True, global_steps=global_steps )

Pure observational prints across engine_workers, ray_trainer, model_forward, and engine/megatron/transformer_impl. Four functional W4A4 patches in 7c5660b are untouched.

Two status logs added during the debug session (NVFP4 QAT injection and force_dynamic_quantization flag activation). The functional config injections themselves are unchanged.

Factor the duplicated float8 storage-stub registration and NVFP4 quantization-config builder into _w4a4_compat. Inline a stray result-variable refactor and drop the _fdq tempvar. Revert the unrelated recipe submodule bump that slipped in with 7c5660b. No behavior change.

…h live refit (#7971) ### What does this PR do? Add end-to-end native NVFP4 W4A4 training and rollout on NVIDIA Blackwell GPUs, demonstrated with Qwen3-30B-A3B-Base DAPO using Megatron/Transformer Engine and vLLM. This contributes to Sophia's [low-precision roadmap #5972](#5972), specifically **NVFP4 real quant** and **NVFP4 training in Megatron**. In brief: - **Training:** Transformer Engine runs selected routed-expert MLPs in NVFP4 with per-token activation scaling and dequantized backward. Persistent policy weights remain BF16, with full optimizer state. - **Weight synchronization:** following the trainer-side quantization boundary of #8287, the trainer exports the routed experts as NVFP4 with `moe: NVFP4` and `moe_activation: nvfp4_per_token` after each policy update. - **Rollout:** vLLM loads the NVFP4 expert weights and quantizes activations per token, matching training. Refit preserves the storage referenced by CUDA graphs. Attention and the configured boundary-layer MLPs remain BF16. The Qwen3 example keeps the first two and last four MLP layers in BF16 and applies NVFP4 to the remaining 42 routed-expert MLP layers. ### Checklist Before Starting - [x] Search for similar PRs: [NVFP4](https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+NVFP4), [roadmap #5972](https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+5972+in%3Abody). - [x] Format the PR title as `[{modules}] {type}: {description}`. Related work: #6455 targets ModelOpt/TRT-LLM synchronization; #6992 provides FSDP W4A8 simulation; #7895 targets DeepSeek MXFP4 on TE FP8; #7519 adds MXFP8 training and rollout. #8287 adds ModelOpt NVFP4 QAT with static activation scales; this PR adds native TE NVFP4 training with per-token activation scaling, on the same export path. ### Test Unit tests (CPU and GB200) pass; scoped Ruff, formatting and lockfile checks pass. The trainer-side export is bit-identical to vLLM's online `nvfp4_per_token` quantization, and 6-step runs on Qwen3-30B-A3B (8 × 4 GB200) and a small DeepSeek-V3-layout model pass. #### Long run BF16 and W4A4 runs of the example script below (Qwen3-30B-A3B-Base, DAPO, R3, token-level TIS) on 8 nodes × 4 GB200 GPUs, 700 steps with checkpoint resume across jobs. **W4A4 vs BF16:** **Training score (`critic/score/mean`)** <img width="1813" height="957" alt="image" src="https://github.com/user-attachments/assets/d0693c0b-6789-4285-8658-44a885d72a72" /> **Entropy (`actor/entropy`)** <img width="1814" height="952" alt="image" src="https://github.com/user-attachments/assets/e63b552c-fe9b-4b23-9751-62c38976732c" /> **Train–rollout mismatch (`rollout_corr/kl`)** <img width="1811" height="952" alt="image" src="https://github.com/user-attachments/assets/13131fe5-c670-41ea-91e1-68e5fff123ee" /> **Gradient norm (`actor/grad_norm`)** <img width="1813" height="953" alt="image" src="https://github.com/user-attachments/assets/b121c562-b6f7-4432-8a7e-ed5768d78e69" /> **Performance:** at matched response length, NVFP4 rollout generation is about 1.3× faster than BF16 once the mean response exceeds ~2.5k tokens, giving roughly 1.2× shorter end-to-end steps; NVFP4 training itself is still slower than BF16 in Transformer Engine. ### API and Usage Example With a Ray cluster running on eight nodes, provide model and dataset paths shared by all workers: ```bash MODEL_PATH=/shared/models/Qwen3-30B-A3B-Base \ TRAIN_FILE=/shared/data/train.parquet \ TEST_FILE=/shared/data/validation.parquet \ EXP_NAME=qwen3-native-nvfp4 \ PRECISION_MODE=real_nvfp4 \ bash examples/real_nvfp4/run_qwen3_30b_megatron.sh ``` Use `PRECISION_MODE=bf16` for the baseline. NVFP4 training is enabled like FP8 training, through `actor_rollout_ref.actor.megatron.override_transformer_config` (`fp4: e2m1`, `fp4_recipe: nvfp4`, and `first_last_layers_bf16` with `num_layers_at_start_in_bf16` / `num_layers_at_end_in_bf16`), and rollout through `actor_rollout_ref.rollout.quantization.moe=NVFP4` with `moe_activation=nvfp4_per_token`, whose `ignore` list keeps the same layers in BF16; the trainer checks that the two agree. The example uses the standard verl PPO entry point, parquet datasets and built-in reward manager, with the DAPO recipe (clip-higher, token-level loss, dynamic sampling and the overlong penalty), R3 routing replay and token-level TIS. See [the example guide](https://github.com/zhangyimi/verl/blob/feat/native-nvfp4/examples/real_nvfp4/README.md) for configuration and runtime setup. ### Design & Code Changes - **Training precision:** the Megatron per-module recipe (attention BF16, routed-expert MLPs NVFP4, carved-out layers BF16) is generated from the configuration rather than maintained as a file. After model construction, the effective precision of every attention and expert linear layer is checked against it. Packed THD sequences use the TE alignment for FP4 as for FP8. - **Rollout config and refit:** `moe_activation=nvfp4_per_token` selects a compressed-tensors scheme (NVFP4 weights, fully dynamic NVFP4 inputs). The trainer quantizes the experts with its own Triton kernel, sharing one global scale across `gate_proj`/`up_proj`, bit-identical to vLLM's online `nvfp4_per_token` quantization. In vLLM the scheme selects FlashInfer TRT-LLM with per-token activations; refits load each bucket in place into the expert storage captured by CUDA graphs and post-process each MoE layer, so the extra memory is one MoE layer. - **Routing replay:** use vLLM’s native fused-MoE routing capture and align replay data with FP4 training padding. - **Runtime compatibility:** Transformer Engine 2.20.2 from the verl wheelhouse. No vLLM source patch is needed. The validated Megatron + vLLM stack uses: | Component | Version | |---|---| | PyTorch | 2.13 / CUDA 13 | | vLLM | 0.32.0 | | Transformers | 5.16.1 | | Megatron-Core | 0.19.2 | | Megatron-Bridge | 0.6.2 | | Transformer Engine | 2.20.2 | | FlashInfer | 0.7.0.post1 | | CUTLASS DSL / QuACK | 4.8.0 / 0.6.5 | SGLang end-to-end validation remains follow-up work. ### Checklist Before Submitting - [x] Read the [Contribute Guide](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md). - [ ] Run the complete all-files pre-commit suite. Focused lint, formatting, shell and lockfile checks have passed. - [x] Add usage, design and runtime compatibility documentation. - [x] Add focused unit and end-to-end coverage. Project CI integration remains to be completed. - [x] Request project CI and review. - [x] Check recipe coordination: the recipe submodule reference is unchanged; no companion recipe PR is required. AI assistance: OpenAI Codex and Anthropic Claude assisted with implementation, testing and this PR description. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: wuxibin <wuxibin@bytedance.com>

Superjomn added 2 commits May 16, 2026 18:16

wip: 30b-qat-rotary-oob debug session edits (TRT-LLM rollout + megatr…

8ce6b40

…on transformer)

Superjomn requested review from ISEEKYAN, PeterSH6, eric-haibin-lin, tongyx361 and vermouth1992 as code owners May 24, 2026 12:48

Superjomn marked this pull request as draft May 24, 2026 12:49

Superjomn changed the title ~~Chunweiy/w4a4 trtllm rollout ipc~~ debug w4a4 trtllm rollout ipc May 24, 2026

Superjomn closed this May 24, 2026

gemini-code-assist Bot reviewed May 24, 2026

View reviewed changes

Superjomn added 2 commits May 24, 2026 06:24

[chore] remove RAY_EXECUTOR_DEBUG probes from W4A4 debug session

7361044

Pure observational prints across engine_workers, ray_trainer, model_forward, and engine/megatron/transformer_impl. Four functional W4A4 patches in 7c5660b are untouched.

[chore] drop observational logger.info from W4A4 rollout init

eaa7030

Two status logs added during the debug session (NVFP4 QAT injection and force_dynamic_quantization flag activation). The functional config injections themselves are unchanged.

Superjomn reopened this May 24, 2026

Leamonz mentioned this pull request Jun 22, 2026

[rollout, vllm, env] fix: align NVFP4 QAT dependency compatibility #6815

Closed

6 of 8 tasks

This was referenced Sep 16, 2026

[rollout, perf] fix: preserve uint8 expert routes in agent storage #7880

Closed

[megatron, vllm, env] feat: add native NVFP4 training and rollout with live refit #7971

Merged

Navigation Menu

Repository navigation

debug w4a4 trtllm rollout ipc - #6455

debug w4a4 trtllm rollout ipc#6455
Superjomn wants to merge 5 commits into
verl-project:mainfrom
Superjomn:chunweiy/w4a4-trtllm-rollout-ipc

Superjomn commented May 24, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

gemini-code-assist Bot May 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Superjomn commented May 24, 2026

What does this PR do?

Checklist Before Starting

Test

API and Usage Example

Design & Code Changes

Checklist Before Submitting

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 24, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant