diff --git a/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml b/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml index 4d9c2f33fca..2da83e9c72b 100644 --- a/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml +++ b/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml @@ -274,34 +274,6 @@ jobs: - name: clean up run: | rm -rf checkpoints - - name: Prepare GSM8K dataset (data_preprocess) - run: | - python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_save_dir ${PWD}/data/gsm8k - - name: Running GRPO E2E with FP8 TRT-LLM rollout (Qwen3-0.6B, gsm8k) - run: | - unset TORCH_CUDA_ARCH_LIST # Let TRT-LLM CUTLASS DSL auto-detect runtime GPU arch - export TRTLLM_ENABLE_PDL=0 # Disable Programmatic Dependent Launch (uses SM 9.0 griddepcontrol) - ray stop --force - export INFER_TP=2 ACTOR_TP=2 ACTOR_PP=2 ACTOR_VPP=2 ACTOR_EP=1 ACTOR_CP=2 REF_TP=2 REF_PP=2 REF_VPP=2 REF_EP=1 REF_CP=2 GEN_MOE_TP=null GEN_MOE_EP=null - export NNODES=1 GPUS_PER_NODE=8 TRTLLM_MOE_BACKEND=TRITON - export DATA_DIR=${PWD} DAPO_MATH_TRAIN=${PWD}/data/gsm8k/train.parquet AIME_VAL=${PWD}/data/gsm8k/test.parquet MODEL_PATH=${HOME}/models/Qwen/Qwen3-0.6B - INFER_BACKEND=trtllm ROLLOUT_QUANTIZATION=fp8 bash examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh \ - reward_model.reward_kwargs.overlong_buffer_cfg.len=258 \ - reward_model.reward_kwargs.max_resp_len=512 \ - data.max_prompt_length=128 \ - data.max_response_length=64 \ - data.train_batch_size=32 \ - actor_rollout_ref.rollout.n=1 \ - actor_rollout_ref.rollout.max_num_seqs=16 \ - actor_rollout_ref.rollout.max_num_batched_tokens=1024 \ - actor_rollout_ref.rollout.max_model_len=1024 \ - actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=False \ - actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=False \ - trainer.total_training_steps=1 \ - trainer.logger='["console"]' - - name: clean up - run: | - rm -rf checkpoints cleanup: runs-on: ubuntu-latest diff --git a/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml b/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml index 5874604fa4f..30929b926a3 100644 --- a/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml +++ b/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml @@ -203,8 +203,7 @@ jobs: - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout run: | $UV_RUN ray stop --force - export VLLM_USE_V1=1 - ROLLOUT_QUANTIZATION=fp8 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + ROLLOUT_QUANTIZATION=FP8_BLOCK TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh - name: clean up run: | rm -rf checkpoints diff --git a/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml b/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml index 3b6ab8fd79a..56058886773 100644 --- a/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml +++ b/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml @@ -209,7 +209,7 @@ jobs: run: | $UV_RUN ray stop --force export VLLM_USE_V1=1 - ROLLOUT_QUANTIZATION=fp8 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + ROLLOUT_QUANTIZATION=FP8_BLOCK TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh - name: Running GSM8K E2E PPO training tests on 8 L20 GPUs with Megatron-FSDP (Qwen3) run: | $UV_RUN ray stop --force diff --git a/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml b/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml index 31bd1b4ad02..188522fc2af 100644 --- a/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml +++ b/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml @@ -183,7 +183,7 @@ jobs: MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 \ MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True \ COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=4 COMMON_ETP=1 INFER_TP=2 \ - USE_DIST_CKPT=False ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 ROLLOUT_QUANTIZATION=fp8 bash tests/special_e2e/run_ppo_trainer_megatron.sh + USE_DIST_CKPT=False ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 ROLLOUT_QUANTIZATION=FP8_BLOCK bash tests/special_e2e/run_ppo_trainer_megatron.sh - name: clean up run: | rm -rf checkpoints diff --git a/.github/workflows/e2e_ppo_trainer_megatron_vllm_rocm.yml b/.github/workflows/e2e_ppo_trainer_megatron_vllm_rocm.yml index 053e808d18f..f729d6851b3 100644 --- a/.github/workflows/e2e_ppo_trainer_megatron_vllm_rocm.yml +++ b/.github/workflows/e2e_ppo_trainer_megatron_vllm_rocm.yml @@ -253,7 +253,7 @@ jobs: MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 \ MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True \ COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=4 COMMON_ETP=1 INFER_TP=2 \ - USE_DIST_CKPT=False ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 ROLLOUT_QUANTIZATION=fp8 \ + USE_DIST_CKPT=False ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 ROLLOUT_QUANTIZATION=FP8_BLOCK \ bash tests/special_e2e/run_ppo_trainer_megatron.sh \ global_profiler.tool=null actor_rollout_ref.actor.profiler.tool=null actor_rollout_ref.ref.profiler.tool=null actor_rollout_ref.rollout.profiler.tool=null critic.profiler.tool=null "${VLLM_ROCM_OVERRIDE}" - name: clean up diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f72c7d1c28c..81e06364a43 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,3 +1,5 @@ +exclude: ^\.venv[^/]*/ + repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: "v0.12.2" @@ -58,7 +60,7 @@ repos: - id: check-naming-conventions name: Check naming conventions - entry: sh -c 'fail=0; if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=.venv --exclude-dir=__pycache__ --exclude=.pre-commit-config.yaml "veRL" .; then echo "Please use verl instead of veRL"; fail=1; fi; if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=.venv --exclude-dir=__pycache__ --exclude=ascend_sglang_best_practices.rst --exclude=.pre-commit-config.yaml -E "Sglang|sgLang|sglAng|sglaNg|sglanG" .; then echo "Please use SGLang or sglang"; fail=1; fi; exit $fail' + entry: sh -c 'fail=0; if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=.venv* --exclude-dir=__pycache__ --exclude=.pre-commit-config.yaml "veRL" .; then echo "Please use verl instead of veRL"; fail=1; fi; if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=.venv* --exclude-dir=__pycache__ --exclude=ascend_sglang_best_practices.rst --exclude=.pre-commit-config.yaml -E "Sglang|sgLang|sglAng|sglaNg|sglanG" .; then echo "Please use SGLang or sglang"; fail=1; fi; exit $fail' language: system pass_filenames: false @@ -76,6 +78,6 @@ repos: - id: compileall name: Compile all python files - entry: sh -c 'PYTHONWARNINGS=error python3 -m compileall -q . -x "(^|[\\/])(\.venv|venv|\.git|verl[\\/]experimental[\\/]vla)([\\/]|$)"' + entry: sh -c 'PYTHONWARNINGS=error python3 -m compileall -q . -x "(^|[\\/])(\.venv[^\\/]*|venv|\.git|verl[\\/]experimental[\\/]vla)([\\/]|$)"' language: python pass_filenames: false diff --git a/docs/ascend_tutorial/zh/dev_guide/model_dev/parameter_and_metrics.md b/docs/ascend_tutorial/zh/dev_guide/model_dev/parameter_and_metrics.md index dea09d53a1a..ced8e188c91 100644 --- a/docs/ascend_tutorial/zh/dev_guide/model_dev/parameter_and_metrics.md +++ b/docs/ascend_tutorial/zh/dev_guide/model_dev/parameter_and_metrics.md @@ -120,8 +120,8 @@ verl 通过层级化的 YAML 配置文件管理所有参数，涉及到的所有 | `actor_rollout_ref.rollout.calculate_log_probs` | `false` | 是否在 rollout 阶段计算 log probs | | `actor_rollout_ref.rollout.skip_tokenizer_init` | `true` | 是否跳过分词器初始化 | | `actor_rollout_ref.rollout.enable_rollout_routing_replay` | `false` | 是否启用 rollout 路由重放 | -| `actor_rollout_ref.rollout.quantization` | `null` | 量化方式 | -| `actor_rollout_ref.rollout.quantization_config_file` | `null` | 量化配置文件路径 | +| `actor_rollout_ref.rollout.quantization.linear` / `.moe` | `null` | compressed-tensors 量化 preset（如 `FP8_BLOCK`） | +| `actor_rollout_ref.rollout.quantization.config_file` | `null` | compressed-tensors 量化配置文件路径 | | `actor_rollout_ref.rollout.layered_summon` | `false` | 是否启用分层召唤（仅 FSDP 方案） | #### 1.1.5 Rollout 验证采样配置 diff --git a/docs/ascend_tutorial/zh/feature_support/ascend_backend_features.md b/docs/ascend_tutorial/zh/feature_support/ascend_backend_features.md index 79f79a769e9..8b6830a0098 100644 --- a/docs/ascend_tutorial/zh/feature_support/ascend_backend_features.md +++ b/docs/ascend_tutorial/zh/feature_support/ascend_backend_features.md @@ -37,7 +37,7 @@ Last updated: 03/03/2026. | `max_num_batched_tokens`| `actor_rollout_ref.rollout.max_num_batched_tokens` |在一次批处理（batch）中可以处理的最大总Token数| | `skip_tokenizer_init`| `actor_rollout_ref.rollout.skip_tokenizer_init` |跳过初始化分词器并将 input_ids 传递到推理请求中| | `enable_prefix_caching` | `actor_rollout_ref.rollout.enable_prefix_caching` |用于启用自动前缀缓存 | -| `quantization`| `actor_rollout_ref.rollout.quantization`，默认为None|`量化方法`| +| `quantization`| `actor_rollout_ref.rollout.engine_kwargs.vllm.quantization`，默认不设置|`量化方法`| ### 2. sglang: @@ -74,7 +74,7 @@ verl中通过rollout config管理推理后端参数使能，包含通用参数 | log_level| 无，默认设置为error |日志记录器的日志级别| | skip_tokenizer_init| actor_rollout_ref.rollout.skip_tokenizer_init |跳过初始化分词器并将 input_ids 传递到推理请求中| | skip_server_warmup| 无，默认设置为True |跳过预热| -| quantization| actor_rollout_ref.rollout.quantization，默认为None|量化方法| +| quantization| actor_rollout_ref.rollout.engine_kwargs.sglang.quantization，默认不设置|量化方法| | attention_backend|actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend|attention内核,NPU应该设置为ascend| --- diff --git a/docs/ascend_tutorial/zh/feature_support/npu_advance_features.md b/docs/ascend_tutorial/zh/feature_support/npu_advance_features.md index c61f7157215..ca0027f49ab 100644 --- a/docs/ascend_tutorial/zh/feature_support/npu_advance_features.md +++ b/docs/ascend_tutorial/zh/feature_support/npu_advance_features.md @@ -61,7 +61,7 @@ Last updated: 05/13/2026. | SGLang 参数 | verl 对应通用参数 | 功能说明 | |:---|:---|:---| | `attention_backend` | `actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend` | **注意力后端选择** — NPU 上应设置为 `ascend` 以调用昇腾优化内核 | -| `quantization` | `actor_rollout_ref.rollout.quantization` | **量化支持** — 支持模型量化加载与推理 | +| `quantization` | `actor_rollout_ref.rollout.engine_kwargs.sglang.quantization` | **量化支持** — 支持模型量化加载与推理 | > 更多 SGLang NPU 特性参数请参考 [sglang 社区 NPU 特性支持文档](https://docs.sglang.io/docs/hardware-platforms/ascend-npus/ascend_npu_support_features) @@ -229,7 +229,7 @@ repatch | 模型路径 | `model_path` | `model_path` | `actor_rollout_ref.model.path` | | 显存控制 | `gpu_memory_utilization` | `mem_fraction_static` | `actor_rollout_ref.rollout.gpu_memory_utilization` | | 图模式 | `enforce_eager` | `disable_cuda_graph` | `actor_rollout_ref.rollout.enforce_eager` | -| 量化 | `quantization` | `quantization` | `actor_rollout_ref.rollout.quantization` | +| 量化 | `quantization` | `quantization` | `actor_rollout_ref.rollout.engine_kwargs.{vllm,sglang}.quantization` | | 最大序列长度 | `max_model_len` | — | `actor_rollout_ref.rollout.max_model_len` | | 最大并发数 | `max_num_seqs` | `max_running_requests` | `actor_rollout_ref.rollout.max_num_seqs` | | 分词器 | `skip_tokenizer_init` | `skip_tokenizer_init` | `actor_rollout_ref.rollout.skip_tokenizer_init` | diff --git a/docs/index.rst b/docs/index.rst index 584d919cd43..4bf2a03a776 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -139,6 +139,7 @@ verl is fast with: low_precision/fp8.md low_precision/nvfp4_qat.md + low_precision/qat.md .. toctree:: :maxdepth: 1 diff --git a/docs/low_precision/fp8.md b/docs/low_precision/fp8.md index d91a0020447..a9dda1021f6 100644 --- a/docs/low_precision/fp8.md +++ b/docs/low_precision/fp8.md @@ -20,19 +20,15 @@ FP8 rollout-only mode keeps training in BF16 and quantizes rollout inference to ### Implementation -We monkey patch several vLLM functions to enable FP8 rollout for reinforcement learning: - -1. **Quantize weights**: Quantize model weights on-the-fly from higher-precision formats to FP8. -2. **Process weights after loading**: For vLLM, we replace the `vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.process_weights_after_loading` function to handle weight processing after quantization. For SGLang, this patch is not needed as it natively supports loading quantized weights. +The rollout engine starts from a compressed-tensors `quantization_config` with dummy weights, and the trainer quantizes every weight sync to that config with the compressed-tensors compressors (`verl/utils/qat/compressed_tensors.py`). On the vLLM side we monkey patch `process_weights_after_loading` of the compressed-tensors FP8 linear/MoE methods so refits are written back into the live parameters (`verl/utils/vllm/vllm_fp8_utils.py`). ### Support Matrix -- FP8 blockwise quantization for rollout +- FP8 blockwise quantization for rollout (compressed-tensors preset `FP8_BLOCK`) - Used in Deepseek, which is 1x128 quantization for activations and 128x128 quantization for model weights - Dense models and MoE models - Async rollout interfaces -- vLLM 0.10.x & vLLM 0.11 & vLLM 0.12 & SGLang 0.5.5 -- FSDP and Megatron training backends +- vLLM or SGLang rollout (SGLang validated on Hopper only) with the Megatron training backend; TRT-LLM and FSDP are not supported ### Usage @@ -40,52 +36,20 @@ Enable in config file: ```yaml rollout: - quantization: "fp8" + quantization: + linear: FP8_BLOCK # dense linear layers (attention, dense MLP, shared experts) + moe: FP8_BLOCK # routed experts + ignore: [] # extra "re:" patterns kept in BF16 ``` Or via command line: ```bash -actor_rollout_ref.rollout.quantization=fp8 -``` - -#### Skipping layers in SGLang FP8 rollout - -When using SGLang FP8 rollout, you can skip FP8 weight quantization for -selected modules. Skipped modules stay in the rollout model dtype instead -of being converted to FP8. This is useful for layers that are not -compatible with block-wise FP8 weight quantization, or for modules that -you prefer to keep in higher precision. - -Set `SGLANG_FP8_IGNORED_LAYERS` before starting training: - -```bash -SGLANG_FP8_IGNORED_LAYERS=linear_attn \ -python3 -m verl.trainer.main_ppo \ - actor_rollout_ref.rollout.name=sglang \ - actor_rollout_ref.rollout.quantization=fp8 \ - ... -``` - -Multiple entries can be separated by commas: - -```bash -SGLANG_FP8_IGNORED_LAYERS=linear_attn,visual -``` - -You can also use the model `quantization_config`: - -```json -{ - "quantization_config": { - "ignored_layers": ["re:.*linear_attn.*"] - } -} +actor_rollout_ref.rollout.quantization.linear=FP8_BLOCK \ +actor_rollout_ref.rollout.quantization.moe=FP8_BLOCK ``` -Plain module names, full module paths, and `re:` regex patterns are -supported. verl applies the same ignored-layer rules when launching -SGLang and when syncing updated actor weights into the rollout engine. +Routers, `lm_head`, embeddings, vision towers and MTP layers stay in BF16 by default, as do layers whose weight shape is not a multiple of the 128x128 block. To provide a complete compressed-tensors `quantization_config` instead, set `actor_rollout_ref.rollout.quantization.config_file=/path/to/quantization_config.json` (exclusive with `linear` / `moe` / `ignore`). ### Experiments and Outcomes @@ -183,8 +147,9 @@ actor_rollout_ref.actor.optim.override_optimizer_config: fp8_recipe: "blockwise" # FP8 rollout inference (vLLM) -actor_rollout_ref.rollout: - quantization: fp8 +actor_rollout_ref.rollout.quantization: + linear: FP8_BLOCK + moe: FP8_BLOCK ``` ### Support Matrix diff --git a/docs/low_precision/nvfp4_qat.md b/docs/low_precision/nvfp4_qat.md index f7c2f19a23e..43d99001b65 100644 --- a/docs/low_precision/nvfp4_qat.md +++ b/docs/low_precision/nvfp4_qat.md @@ -1,51 +1,17 @@ -# NVFP4 QAT (Quantization-Aware Training) in verl +# NVFP4 QAT -Last updated: 04/02/2026 +Last updated: 10/19/2026 +verl supports NVFP4 as **W4A4 routed experts** on Blackwell: -verl supports NVFP4 Quantization-Aware Training (QAT), which applies fake quantization during training so the model learns to tolerate NVFP4 quantization error. At rollout time, weights are packed into real NVFP4 format for vLLM inference. This closes the precision gap between training and inference, preventing KL divergence explosion. +- trainer: Megatron + ModelOpt QAT +- rollout: vLLM compressed-tensors +- kernel: FlashInfer CUTLASS +- dense linear layers: BF16 or MXFP8; NVFP4 linear is not supported -| Training Backend | Training Precision | Rollout Precision | vLLM Quant Method | -|---|---|---|---| -| **FSDP** | BF16 + fake quantization | NVFP4 W4A16 | `compressed-tensors` | -| **Megatron** | BF16 + fake quantization | NVFP4 W4A16 | `modelopt` | +W4A16, FSDP NVFP4 and non-CUTLASS NVFP4 backends are not supported by this +path. -> [!TIP] -> For ready-to-run scripts, environment setup, and experimental results, see the [QAT recipe](https://github.com/verl-project/verl-recipe/tree/main/qat). - ---- - -## Key Configuration - -### FSDP Backend - -Configured under `actor_rollout_ref.actor.fsdp_config.qat`: - -```yaml -actor_rollout_ref: - actor: - fsdp_config: - qat: - enable: true - mode: "w4a16" - group_size: 16 - ignore_patterns: - - "lm_head" - - "embed_tokens" - - "re:.*mlp.gate$" - quantization_config_path: "recipe/qat/config/nvfp4_w4a16.json" -``` - -| Parameter | Description | Default | -|-----------|-------------|---------| -| `fsdp_config.qat.enable` | Enable QAT | `False` | -| `fsdp_config.qat.mode` | Quantization mode | `"w4a16"` | -| `fsdp_config.qat.group_size` | Quantization group size | `16` | -| `fsdp_config.qat.ignore_patterns` | Layers to skip. Supports `re:` prefix for regex, otherwise substring match | `["lm_head", "embed_tokens", "re:.*mlp.gate$"]` | -| `fsdp_config.qat.quantization_config_path` | vLLM quantization config JSON path | Required | - -### Megatron Backend - -Configured under `actor_rollout_ref.actor.megatron.qat`: +## Configuration ```yaml actor_rollout_ref: @@ -53,35 +19,40 @@ actor_rollout_ref: megatron: qat: enable: true - mode: "w4a16" - group_size: 16 - ignore_patterns: - - "lm_head" - - "*mlp.gate" - quantization_config_path: "recipe/qat/config/nvfp4_w4a16_megatron.json" + mode: nvfp4 + rollout: + quantization: + linear: null + moe: NVFP4 ``` -| Parameter | Description | Default | -|-----------|-------------|---------| -| `megatron.qat.enable` | Enable QAT | `False` | -| `megatron.qat.mode` | Quantization mode | `"w4a16"` | -| `megatron.qat.group_size` | Quantization group size | `16` | -| `megatron.qat.ignore_patterns` | Layers to skip. Uses `fnmatch` glob syntax | `["lm_head", "*mlp.gate"]` | -| `megatron.qat.quantization_config_path` | vLLM quantization config JSON path | Required | +The weight uses E2M1 values in groups of 16, FP8 E4M3 local scales and one +tensor-level global divisor. Activations use the same NVFP4 hierarchy with +dynamic group scales. Before ModelOpt observes an activation, verl exports the +neutral tensor divisor `1.0`; subsequent syncs use the running QAT value. ---- +ModelOpt fake quantization and compressed-tensors export share the same QDQ +implementation. The exporter also aligns gate/up global scales, as required by +vLLM's fused MoE loader. -## Support Matrix +## Validation -- NVFP4 W4A16 (weight-only FP4 quantization) -- Dense models and MoE models -- FSDP and Megatron training backends -- Full quantization and FFN-only quantization strategies -- Verified on Qwen3-8B-Base and Qwen3-30B-A3B-Base +Use CUDA Graph mode; eager-mode results are not representative: ---- - -## Notes +```bash +.venv-vllm/bin/python scripts/rollout_logprob_diff.py \ + --model-path /mnt/efs/model/Qwen/Qwen3.5-35B-A3B --ep 8 \ + --quant-moe NVFP4 --qat-mode nvfp4 \ + '+actor_rollout_ref.rollout.engine_kwargs.vllm.enable_flashinfer_autotune=false' +``` -- FSDP backend has scalability limitations for very large models. For large-scale training, use the Megatron backend. -- FSDP uses `re:` prefix regex for `ignore_patterns`, while Megatron uses `fnmatch` glob syntax. The two are not interchangeable. +On Qwen3.5-35B-A3B with EP=8, CUDA Graph and no R3, the final validation +measured `k3_kl=0.25778` without QAT and `k3_kl=0.32112` with NVFP4 QAT +(Pearson correlation `0.991763` and `0.981248`, respectively). FlashInfer's +fused finalize is nondeterministic; repeated results ranged from `0.219–0.258` +and `0.263–0.380`. A synthetic single-expert CUTLASS comparison confirmed that +the exported weights and activation QDQ use the expected NVFP4 format, but the +fused kernel still differed from dequantize-then-BF16 matmul by mean absolute +`0.0084` per MoE layer (correlation `0.99855`). This validates the refit +lifecycle and kernel execution, but not numerical acceptance: the accumulated +model-level mismatch remains too large for training. diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md new file mode 100644 index 00000000000..211448c07b6 --- /dev/null +++ b/docs/low_precision/qat.md @@ -0,0 +1,223 @@ +# Low-Precision Rollout and Quantization-Aware Training + +Last updated: 10/16/2026 + +In RL training, rollout usually dominates the step time, and for large MoE models it is bound by memory capacity and bandwidth. Running rollout in low precision: + +- shrinks weight memory, leaving more room for KV cache and higher concurrency; +- reduces memory traffic during decode, improving throughput; +- uses FP8/FP4 Tensor Cores on Hopper/Blackwell to speed up prefill. + +The cost is a numerical gap between the high-precision trainer and the low-precision rollout (training-inference mismatch). It makes importance sampling and off-policy correction less accurate and, in bad cases, destabilizes training. Quantization-Aware Training (QAT) narrows this gap by training the model against the same quantization the rollout uses. + +## Background: model weight formats + +HuggingFace checkpoints come in two flavors: + +1. **High precision**: BF16 weights, no `quantization_config` in `config.json`. +2. **Low precision**: `config.json` carries a `quantization_config` describing how each layer is quantized. Models often use different formats for linear layers (attention, dense MLP, shared experts) and MoE (routed experts). + +| Model | Architecture | `quant_method` | Linear | MoE (routed experts) | Not quantized | +|---|---|---|---|---|---| +| [Qwen/Qwen3.8-2.4T-A95B](https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B/blob/main/config.json) | 92 layers, 512 experts, top-10 | none | BF16 | BF16 | - | +| [zai-org/GLM-5.2](https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json) | 78 layers, 256 experts, top-8, first 3 dense | none | BF16 | BF16 | - | +| [zai-org/GLM-5.3](https://huggingface.co/zai-org/GLM-5.3/blob/main/config.json) | same as GLM-5.2 | `fp8` | FP8 E4M3, 128x128 block, FP32 scale (`weight_scale_inv`), dynamic activation | same as linear | router `mlp.gate`, `indexers_proj`, `embed_tokens`, `lm_head`, MTP `eh_proj`, norms (`modules_to_not_convert`) | +| [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3/blob/main/config.json) | 93 layers, 896 experts, first layer dense | `compressed-tensors` | BF16 | MXFP4 W4A16: FP4 E2M1, 1x32 group, UE8M0 scale (`weight_packed` + `weight_scale`) | `self_attn`, `shared_experts`, dense `mlp`, `lm_head`, `vision_tower`, `mm_projector` (`ignore` regex) | +| [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/main/config.json) | 40 layers, 384 experts, top-6 | `fp8` | FP8 E4M3, 32x32 block (`wo_a` may be 128x128), UE8M0 scale (`scale_fmt: ue8m0`) | FP4 E2M1, 1x32 group, UE8M0 scale (`expert_dtype: fp4`), same layout as MXFP4 | router `ffn.gate`, indexer `wk`/`weights_proj`, `compressor` | + +- GLM-5.3 uses the standard HF FineGrainedFP8 format (same as DeepSeek-V3), which both vLLM and SGLang parse. +- Kimi-K3 uses compressed-tensors: a `Linear`-targeted config group plus an `ignore` list that excludes non-expert layers yields "linear BF16 + MoE MXFP4". +- DeepSeek-V4.1-Flash's `fp8` config has private fields (`scale_fmt`, `expert_dtype`) and needs dedicated engine support. + +The trainer always trains in high precision; low-precision checkpoints are dequantized before training. The rollout uses the model's own `quantization_config` if it has one; a BF16 model can run rollout in BF16 or, for higher throughput, in low precision through `rollout.quantization`. + +## How it works + +- **Trainer**: always trains in high precision (BF16). With QAT enabled, the forward pass uses fake-quantized weights (quantize then dequantize) and gradients flow through a straight-through estimator. +- **Rollout**: the inference engine starts with dummy weights and either the model's native + `quantization_config` or a `rollout.quantization` override. It allocates parameters as if + it were loading that low-precision checkpoint (for example FP8 `weight` + + `weight_scale_inv`). +- **Weight refit export**: the training backend first exports weights under canonical + HuggingFace names. A trainer-side `WeightQuantizer`, selected from the effective rollout + `quantization_config`, converts each sync bucket into exactly the checkpoint tensors + expected by the rollout engine. Standard native FP8 is quantized to the `Fp8Config` + layout; native and overridden compressed-tensors configs use the same + compressed-tensors exporter. DeepSeek-V4 is already quantized by Megatron-Bridge, so its + exporter only normalizes the packed byte dtype. The rollout worker therefore loads + pre-quantized tensors and is not responsible for quantization. + +![weight_quant](https://github.com/verl-project/verl-data/blob/main/images/weight_quantizer_refit.png?raw=true) + +This exporter/loader boundary is intended to decouple the two sides. The target architecture +is the Cartesian product of training backends `[Megatron, VeOmni, TorchTitan]` and rollout +engines `[vLLM, SGLang]`: each trainer exports canonical named tensors, each +`WeightQuantizer` defines a checkpoint wire format, and each rollout engine only implements +loading/refit for that format. The support matrix below records the combinations currently +validated; it does not imply that the full target matrix is implemented yet. + +For an override-based rollout, the HuggingFace checkpoint stays in BF16; no offline +quantization or calibration step is needed. The override format is described using the open +[compressed-tensors](https://github.com/neuralmagic/compressed-tensors) schema, which can +express different schemes for linear and routed MoE layers. + +## Support matrix + +The trainer must be Megatron (FSDP is not yet supported for this feature). The rollout engine can be vLLM or SGLang; TRT-LLM raises an error if `rollout.quantization` is set. The engine columns list the compressed-tensors scheme each engine implements. + +| Precision | Preset | QAT mode | GPU | vLLM Linear | vLLM MoE | SGLang Linear | SGLang MoE | vLLM status in verl | SGLang status in verl | +|---|---|---|---|---|---|---|---|---|---| +| FP8 block (E4M3, 128x128 weight blocks, dynamic 1x128 activations) | `FP8_BLOCK` | `fp8` | Hopper / Blackwell | `W8A8Fp8` | `W8A8Fp8MoE` | `W8A8Fp8` | `W8A8Fp8MoE` | ✅ Hopper | ✅ Hopper | +| MXFP8 | `MXFP8` | `mxfp8` | Blackwell | `W8A8Mxfp8` | `W8A8Mxfp8MoE` | ❌ (only whole-model `quant_method: mxfp8`) | ❌ | ✅ Blackwell | Blocked by engine support | +| MXFP4 (MoE only) | `MXFP4` | `mxfp4` / `mxfp4_mxfp8` | Blackwell | ❌ | MXFP4×MXFP8
`(moe_backend=flashinfer_cutlass)` | ❌ | ❌ | ✅ Blackwell (MoE only) | Blocked by engine support | +| NVFP4 (MoE only) | `NVFP4` | `nvfp4` | Blackwell | ❌ | W4A4
`(moe_backend=flashinfer_cutlass)` | ❌ | ❌ | ✅ Blackwell (MoE only) | Blocked by engine support | + +MXFP8, MXFP4 and NVFP4 need Blackwell (SM100). On Hopper, vLLM falls back to Marlin weight-only kernels for these formats, whose numerics differ from Blackwell W8A8 / W4A4; verl does not support that path. + +MXFP4 and NVFP4 support is scoped to routed experts. The supported combinations are BF16 linear + FP4 MoE or MXFP8 linear + FP4 MoE; FP4 linear QAT, export and refit are out of scope. verl defines its `MXFP4` preset as MXFP4 expert weights with dynamic MXFP8 expert activations (group size 32 and E8M0 scales for both). This is intentionally different from compressed-tensors' `MXFP4` preset, which is W4A4, and from vLLM's native `Mxfp4Config`, whose checkpoint format is nominally W4A16. + +## Configuration + +### Rollout quantization + +```yaml +actor_rollout_ref: + rollout: + quantization: + linear: FP8_BLOCK # attention, dense MLP, shared experts + moe: FP8_BLOCK # routed experts + ignore: [] # extra layers kept in BF16: exact names or "re:" + config_file: null # path to a compressed-tensors quantization_config JSON +``` + +- When `linear`, `moe` and `config_file` are all `null` (the default), the rollout runs in the model's native precision. +- `linear` and `moe` can be set independently. Setting one of them to `null` keeps those layers in BF16, e.g. `linear: null, moe: FP8_BLOCK` quantizes only the routed experts. +- `MXFP4` and `NVFP4` are valid only for `moe`. Use `linear: null` for BF16 dense layers or `linear: MXFP8` for MXFP8 dense layers. +- The following layers are always kept in BF16: `lm_head`, embeddings, MoE routers (`mlp.gate`, `shared_expert_gate`), vision towers and MTP layers. +- `ignore` adds more layers on top of these, using HuggingFace layer names. Layers incompatible with the selected scheme are not detected automatically: list them explicitly. For example, MXFP8 on Qwen3.5 requires `["re:.*linear_attn\\.in_proj_[ab]$"]` because these projections are too small for the vLLM kernel. + +From the command line: + +```bash +actor_rollout_ref.rollout.quantization.linear=FP8_BLOCK \ +actor_rollout_ref.rollout.quantization.moe=FP8_BLOCK +``` + +For Qwen3.5 MXFP8: + +```bash +actor_rollout_ref.rollout.quantization.linear=MXFP8 \ +actor_rollout_ref.rollout.quantization.moe=MXFP8 \ +'actor_rollout_ref.rollout.quantization.ignore=["re:.*linear_attn\\.in_proj_[ab]$"]' +``` + +### Custom `quantization_config` + +For full control, point `config_file` at a compressed-tensors `quantization_config`. It can be either the `quantization_config` object itself or a HuggingFace `config.json` that contains one. `config_file` cannot be combined with `linear`, `moe` or `ignore`, and it is used as-is (the default ignore list is not added). Every config group must numerically match a supported preset. + +Example equivalent to `linear: FP8_BLOCK, moe: FP8_BLOCK` for Qwen3.5: + +```json +{ + "quant_method": "compressed-tensors", + "format": "float-quantized", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "float", "symmetric": true, + "strategy": "block", "block_structure": [128, 128], "dynamic": false}, + "input_activations": {"num_bits": 8, "type": "float", "symmetric": true, + "strategy": "group", "group_size": 128, "dynamic": true}, + "format": "float-quantized" + } + }, + "ignore": [ + "re:.*lm_head$", "re:.*embed_tokens$", "re:.*visual\\..*", "re:^mtp\\..*", + "re:.*mlp\\.gate$", "re:.*shared_expert_gate$", + "re:.*linear_attn\\.in_proj_a$", "re:.*linear_attn\\.in_proj_b$" + ] +} +``` + +### QAT (Megatron) + +QAT is optional. Enable it on the actor and pick the mode that matches the rollout preset: + +```yaml +actor_rollout_ref: + actor: + megatron: + qat: + enable: true + mode: fp8 # FP8_BLOCK -> fp8, MXFP8 -> mxfp8, MXFP4 -> mxfp4, NVFP4 -> nvfp4 +``` + +QAT requires `rollout.quantization` to be set, and verl checks at startup that the QAT mode matches the rollout preset. Megatron QAT uses [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer) for fake quantization. MXFP8 uses ModelOpt's native `MXFP8_DEFAULT_CFG`; FP8_BLOCK uses a verl backend because ModelOpt's `FP8_DEFAULT_CFG` is per-tensor and its 2D-block preset is weight-only. MXFP4 and NVFP4 use compressed-tensors-compatible QDQ backends so fake quantization is bit-exact with the exported checkpoint format. Both FP4 modes apply only to routed experts. Use `mxfp4_mxfp8` when dense linear layers also use MXFP8. +Its `ignore_patterns` use Megatron module names and `fnmatch` syntax. Keep them aligned with the rollout ignore list; for Qwen3.5 MXFP8, add `*linear_attn.in_proj_a` and `*linear_attn.in_proj_b`. + +### Engine notes + +- **vLLM**: `VLLM_USE_DEEP_GEMM_E8M0` defaults to `0` for `FP8_BLOCK`, so its block scales stay in FP32. MX formats retain the E8M0 path. +- **vLLM MXFP4**: vLLM 0.29's compressed-tensors path otherwise ignores the MXFP8 input scheme and selects W4A4 CUTLASS. verl routes this explicit W4A8(MX) scheme through vLLM's MXFP4 backend oracle with `kMxfp8Dynamic` and defaults to FlashInfer CUTLASS. The TRTLLM MXFP4×MXFP8 backend is not used because it produced a large numerical mismatch on Qwen3.5-35B-A3B. +- **vLLM NVFP4**: verl supports only W4A4 routed experts and forces FlashInfer CUTLASS. W4A16, dense NVFP4 and other MoE backends are rejected. +- **SGLang**: the first launch of a quantized MoE model can take several minutes while Triton compiles FP8 MoE kernels during CUDA graph capture. Later launches use the cache. + +## Example + +Qwen3.5-35B-A3B on 8 GPUs with FP8 block rollout and FP8 QAT: + +```bash +python3 -m verl.trainer.main_ppo \ + model_engine=megatron \ + actor_rollout_ref.model.path=Qwen/Qwen3.5-35B-A3B \ + actor_rollout_ref.actor.megatron.expert_model_parallel_size=8 \ + actor_rollout_ref.actor.megatron.qat.enable=True \ + actor_rollout_ref.actor.megatron.qat.mode=fp8 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.quantization.linear=FP8_BLOCK \ + actor_rollout_ref.rollout.quantization.moe=FP8_BLOCK \ + ... +``` + +Use `actor_rollout_ref.rollout.name=sglang` for SGLang. For complete scripts, see the [low-precision recipe directory](https://github.com/verl-project/verl-recipe/tree/main/low_precision). + +## Measuring the mismatch + +`scripts/rollout_logprob_diff.py` runs one rollout, recomputes the log-probs with the Megatron actor, and reports the gap (`kl`, `k3_kl`, absolute log-prob difference, etc.). Use it to check a configuration before a long run: + +```bash +python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --ep 8 \ + --quant-linear FP8_BLOCK --quant-moe FP8_BLOCK \ + --quant-ignore 're:.*linear_attn\.in_proj_[ab]$' \ + --qat-mode fp8 \ + --qat-ignore '*linear_attn.in_proj_a' --qat-ignore '*linear_attn.in_proj_b' \ + --r3-router-replay +``` + +Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arguments. Reference `k3_kl` on Qwen3.5-35B-A3B (EP=8, rollout TP=2): + +| GPU | Engine | R3 | Rollout | MoE backend | k3_kl | +|---|---|---|---|---|---| +| Hopper | vLLM | ✅ | BF16 | - | 0.00031 | +| Hopper | SGLang | ✅ | BF16 | - | 0.00031 | +| Hopper | vLLM | ✅ | FP8_BLOCK | - | 0.00225 | +| Hopper | SGLang | ✅ | FP8_BLOCK | - | 0.00231 | +| Hopper | vLLM | ✅ | FP8_BLOCK + FP8 QAT | - | 0.00187 | +| Hopper | SGLang | ✅ | FP8_BLOCK + FP8 QAT | - | 0.00188 | +| Blackwell | vLLM | ❌ | MXFP8 linear + MoE | auto | 0.00148 | +| Blackwell | vLLM | ❌ | MXFP8 linear + MoE + MXFP8 QAT | auto | 0.00315 | +| Blackwell | vLLM | ❌ | MXFP4 MoE | FlashInfer CUTLASS | 0.00858 | +| Blackwell | vLLM | ❌ | MXFP4 MoE + MXFP4 QAT | FlashInfer CUTLASS | 0.00210 | +| Blackwell | vLLM | ✅ | NVFP4 MoE | FlashInfer CUTLASS | 0.00813 | +| Blackwell | vLLM | ✅ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS | 0.00574 | + +## Limitations + +- `FP8_BLOCK` is supported on Hopper. `MXFP8` with vLLM is supported on Blackwell; SGLang has no per-layer MXFP8 support. +- Native vLLM `Fp8Config` rollout is supported without a rollout override. Standard + 128x128 block-FP8 models are quantized by the trainer-side exporter; DeepSeek-V4 + keeps the private mixed FP8/FP4 layout already exported by Megatron-Bridge. +- Native compressed-tensors configs are also supported without an override. The + same trainer-side dispatcher handles native FP8, native compressed-tensors and + `rollout.quantization` compressed-tensors overrides. +- Megatron QAT and the rollout use separate ignore lists and naming conventions. Users must keep them aligned. diff --git a/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh b/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh index 984af8edbab..fef7f239d4e 100644 --- a/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh +++ b/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh @@ -88,7 +88,7 @@ FP8=( # rollout actor_rollout_ref.actor.megatron.dtype=${dtype} actor_rollout_ref.rollout.dtype=${dtype} - # +actor_rollout_ref.rollout.quantization="fp8" + # actor_rollout_ref.rollout.quantization.linear=FP8_BLOCK actor_rollout_ref.rollout.quantization.moe=FP8_BLOCK ) DATA=( diff --git a/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh index 8f7648c2b42..2ac76893c74 100644 --- a/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -4,7 +4,7 @@ # # Knobs: # INFER_BACKEND rollout backend: vllm | sglang | trtllm (default: vllm) -# ROLLOUT_QUANTIZATION fp8 to enable TRT-LLM FP8 rollout (default: unset) +# ROLLOUT_QUANTIZATION compressed-tensors preset for a vLLM low-precision rollout, e.g. FP8_BLOCK (default: unset) # # Ascend NPU users: see examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_megatron.sh. @@ -81,8 +81,8 @@ experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_${INFER_BACKEND}_megatron${ROLL ########################### end user-adjustable ########################### ########################### derived defaults ########################### -if [ -n "${ROLLOUT_QUANTIZATION}" ] && [ "${INFER_BACKEND}" != trtllm ]; then - echo "ROLLOUT_QUANTIZATION is only supported with INFER_BACKEND=trtllm, got: ${INFER_BACKEND}" >&2 +if [ -n "${ROLLOUT_QUANTIZATION}" ] && [ "${INFER_BACKEND}" != vllm ]; then + echo "ROLLOUT_QUANTIZATION is only supported with INFER_BACKEND=vllm, got: ${INFER_BACKEND}" >&2 exit 1 fi @@ -234,13 +234,11 @@ if [ "${INFER_BACKEND}" = trtllm ]; then ) fi -if [ "${ROLLOUT_QUANTIZATION}" = fp8 ]; then - # Avoid stray launcher env leaking into Ray workers under FP8 TRT-LLM. - for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do - unset "$v" - done - export RAY_DEDUP_LOGS=0 - EXTRA+=(+actor_rollout_ref.rollout.quantization=fp8) +if [ -n "${ROLLOUT_QUANTIZATION}" ]; then + EXTRA+=( + actor_rollout_ref.rollout.quantization.linear=${ROLLOUT_QUANTIZATION} + actor_rollout_ref.rollout.quantization.moe=${ROLLOUT_QUANTIZATION} + ) fi if [ -n "$MCORE_MODEL_PATH" ]; then diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py new file mode 100644 index 00000000000..103dcdde4dd --- /dev/null +++ b/scripts/rollout_logprob_diff.py @@ -0,0 +1,307 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Rollout with vLLM / SGLang, recompute old_log_probs with Megatron, and report the rollout-vs-actor log-prob gap. + +Exits non-zero if k3_kl exceeds --max-k3-kl. Trailing arguments are passed to Hydra as config overrides. + +Examples (8 GPUs): + # bf16 rollout + .venv/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-0.8B + # fp8 blockwise rollout on a MoE model + .venv/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --ep 8 \\ + --quant-linear FP8_BLOCK --quant-moe FP8_BLOCK + # MXFP8 rollout on Blackwell (linear only for a dense model, linear + MoE for a MoE model) + .venv-vllm/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-0.8B \\ + --quant-linear MXFP8 --quant-ignore 're:.*linear_attn\\.in_proj_[ab]$' + .venv-vllm/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --ep 8 \\ + --quant-linear MXFP8 --quant-moe MXFP8 --qat-mode mxfp8 \\ + --quant-ignore 're:.*linear_attn\\.in_proj_[ab]$' \\ + --qat-ignore '*linear_attn.in_proj_a' --qat-ignore '*linear_attn.in_proj_b' \\ + --r3-router-replay + # MXFP4 expert weights with dynamic MXFP8 expert activations + .venv-vllm/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --ep 8 \\ + --quant-moe MXFP4 --qat-mode mxfp4 --r3-router-replay + # NVFP4 W4A4 experts with FlashInfer CUTLASS + .venv-vllm/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --ep 8 \\ + --quant-moe NVFP4 --qat-mode nvfp4 + # same, with Megatron ModelOpt QAT fake-quantizing to the rollout scheme + .venv/bin/python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --ep 8 \\ + --quant-linear FP8_BLOCK --quant-moe FP8_BLOCK --qat-mode fp8 + # SGLang rollout + .venv-sglang/bin/python scripts/rollout_logprob_diff.py --engine sglang --model-path ~/models/Qwen3.5-35B-A3B \\ + --ep 8 --quant-linear FP8_BLOCK --quant-moe FP8_BLOCK +""" + +import argparse +import json +import os +import sys + +import numpy as np +import ray +import torch +from hydra import compose, initialize_config_dir +from omegaconf import DictConfig + +import verl +from verl.checkpoint_engine import CheckpointEngineManager +from verl.experimental.agent_loop import AgentLoopManager +from verl.experimental.reward_loop import RewardLoopManager +from verl.protocol import DataProto +from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup +from verl.single_controller.ray.base import create_colocated_worker_cls +from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role +from verl.trainer.ppo.rollout_corr_helper import compute_offpolicy_metrics +from verl.utils import omega_conf_to_dataclass +from verl.utils import tensordict_utils as tu +from verl.utils.debug.metrics import calculate_debug_metrics +from verl.utils.device import get_device_name +from verl.workers.engine_workers import ActorRolloutRefWorker +from verl.workers.rollout.llm_server import LLMServerManager +from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding + +QUESTIONS = [ + "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. " + "How many clips did Natalia sell altogether in April and May?", + "Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?", + "Betty is saving money for a new wallet which costs $100. Betty has only half of the money she needs. " + "Her parents decided to give her $15 for that purpose, and her grandparents twice as much as her parents. " + "How much more money does Betty need to buy the wallet?", + "James writes a 3-page letter to 2 different friends twice a week. How many pages does he write a year?", + "Write a short poem about the ocean.", + "Explain the difference between TCP and UDP in a few sentences.", + "Let's play a role playing game. Your name is Alice, your favorite color is blue. Introduce yourself.", + "What is the derivative of x^3 * sin(x)?", +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-path", default="~/models/Qwen3.5-35B-A3B") + parser.add_argument("--engine", default="vllm", choices=["vllm", "sglang"]) + parser.add_argument("--n-gpus", type=int, default=8) + parser.add_argument("--tp", type=int, default=2, help="Megatron tensor parallel size") + parser.add_argument("--pp", type=int, default=1, help="Megatron pipeline parallel size") + parser.add_argument("--cp", type=int, default=1, help="Megatron context parallel size") + parser.add_argument("--ep", type=int, default=1, help="Megatron expert parallel size") + parser.add_argument("--etp", type=int, default=1, help="Megatron expert tensor parallel size") + parser.add_argument("--gen-tp", type=int, default=2, help="rollout tensor parallel size") + parser.add_argument("--quant-linear", default=None, help="rollout.quantization.linear preset: FP8_BLOCK or MXFP8") + parser.add_argument( + "--quant-moe", + default=None, + help="rollout.quantization.moe preset: FP8_BLOCK, MXFP8, MXFP4, or NVFP4", + ) + parser.add_argument( + "--quant-ignore", + action="append", + default=[], + help='extra rollout ignore pattern; repeat as needed, e.g. "re:.*linear_attn\\\\.in_proj_[ab]$"', + ) + parser.add_argument("--qat-mode", default=None, help="Megatron ModelOpt QAT mode matching the rollout") + parser.add_argument( + "--qat-ignore", + action="append", + default=[], + help="extra Megatron ModelOpt fnmatch ignore pattern; repeat as needed", + ) + parser.add_argument( + "--r3-router-replay", + action="store_true", + help="replay rollout expert choices in the Megatron log-prob recomputation", + ) + parser.add_argument( + "--refit-rounds", + type=int, + default=1, + help="generate and measure after this many consecutive live-engine weight refits", + ) + parser.add_argument("--max-k3-kl", type=float, default=1e-2) + parser.add_argument("overrides", nargs="*", help="extra Hydra overrides, e.g. actor_rollout_ref.rollout.n=8") + return parser.parse_args() + + +def build_config(args: argparse.Namespace) -> DictConfig: + config_dir = os.path.join(os.path.dirname(verl.__file__), "trainer", "config") + with initialize_config_dir(config_dir=config_dir, version_base=None): + config = compose( + config_name="ppo_trainer", + overrides=[ + "model_engine=megatron", + f"trainer.n_gpus_per_node={args.n_gpus}", + "trainer.nnodes=1", + "actor_rollout_ref.actor.optim.total_training_steps=1", + "actor_rollout_ref.actor.use_dynamic_bsz=True", + "actor_rollout_ref.actor.megatron.use_remove_padding=True", + f"actor_rollout_ref.actor.megatron.tensor_model_parallel_size={args.tp}", + f"actor_rollout_ref.actor.megatron.pipeline_model_parallel_size={args.pp}", + f"actor_rollout_ref.actor.megatron.context_parallel_size={args.cp}", + f"actor_rollout_ref.actor.megatron.expert_model_parallel_size={args.ep}", + f"actor_rollout_ref.actor.megatron.expert_tensor_parallel_size={args.etp}", + "actor_rollout_ref.actor.megatron.param_offload=True", + "actor_rollout_ref.actor.megatron.optimizer_offload=True", + "++actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend=auto", + "+actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True", + "+actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True", + f"actor_rollout_ref.rollout.name={args.engine}", + "actor_rollout_ref.rollout.mode=async", + f"actor_rollout_ref.rollout.tensor_model_parallel_size={args.gen_tp}", + "actor_rollout_ref.rollout.gpu_memory_utilization=0.6", + "actor_rollout_ref.rollout.prompt_length=1024", + "actor_rollout_ref.rollout.response_length=2048", + "actor_rollout_ref.rollout.n=4", + "actor_rollout_ref.rollout.agent.num_workers=2", + "actor_rollout_ref.rollout.skip_tokenizer_init=True", + *args.overrides, + ], + ) + + config.actor_rollout_ref.model.path = os.path.expanduser(args.model_path) + config.actor_rollout_ref.rollout.quantization.linear = args.quant_linear + config.actor_rollout_ref.rollout.quantization.moe = args.quant_moe + config.actor_rollout_ref.rollout.quantization.ignore = args.quant_ignore + if args.qat_mode: + config.actor_rollout_ref.actor.megatron.qat.enable = True + config.actor_rollout_ref.actor.megatron.qat.mode = args.qat_mode + config.actor_rollout_ref.actor.megatron.qat.ignore_patterns = [ + *config.actor_rollout_ref.actor.megatron.qat.ignore_patterns, + *args.qat_ignore, + ] + if args.r3_router_replay: + config.actor_rollout_ref.actor.megatron.router_replay.mode = "R3" + config.actor_rollout_ref.rollout.enable_rollout_routing_replay = True + return config + + +def init_workers(config: DictConfig) -> tuple[AgentLoopManager, RayWorkerGroup, CheckpointEngineManager]: + """Hybrid actor/rollout workers + agent loop, without reward model.""" + pool_id = "global_pool" + resource_pool_manager = ResourcePoolManager( + resource_pool_spec={pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes}, + mapping={Role.ActorRollout: pool_id}, + ) + resource_pool_manager.create_resource_pool() + resource_pool = resource_pool_manager.get_resource_pool(Role.ActorRollout) + + class_dict = { + "actor_rollout": RayClassWithInitArgs( + cls=ray.remote(ActorRolloutRefWorker), config=config.actor_rollout_ref, role="actor_rollout" + ) + } + wg_dict = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=create_colocated_worker_cls(class_dict=class_dict), + device_name=get_device_name(), + ) + actor_rollout_wg = wg_dict.spawn(prefix_set=class_dict.keys())["actor_rollout"] + actor_rollout_wg.init_model() + + reward_loop_manager = RewardLoopManager(config=config, rm_resource_pool=None) + llm_server_manager = LLMServerManager.create(config=config, worker_group=actor_rollout_wg) + agent_loop_manager = AgentLoopManager.create( + config=config, + llm_client=llm_server_manager.get_client(), + reward_loop_worker_handles=reward_loop_manager.reward_loop_workers, + ) + checkpoint_manager = CheckpointEngineManager( + config=omega_conf_to_dataclass(config.actor_rollout_ref.rollout.checkpoint_engine), + actor_wg=actor_rollout_wg, + replicas=llm_server_manager.get_replicas(), + ) + checkpoint_manager.sleep_replicas() + checkpoint_manager.update_weights() + return agent_loop_manager, actor_rollout_wg, checkpoint_manager + + +def compute_old_log_prob(actor_rollout_wg: RayWorkerGroup, batch: DataProto) -> DataProto: + """Same as RayPPOTrainer._compute_old_log_prob.""" + batch_td = left_right_2_no_padding(batch.to_tensordict()) + tu.assign_non_tensor(batch_td, calculate_entropy=True, compute_loss=False) + output = actor_rollout_wg.compute_log_prob(batch_td) + log_probs = no_padding_2_padding(tu.get(output, "log_probs"), batch_td) + return DataProto.from_tensordict(tu.get_tensordict({"old_log_probs": log_probs.float()})) + + +def main() -> int: + args = parse_args() + if args.refit_rounds < 1: + raise ValueError("--refit-rounds must be at least 1") + config = build_config(args) + + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + "VLLM_ALLREDUCE_USE_SYMM_MEM": "0", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + } + } + ) + agent_loop_manager, actor_rollout_wg, checkpoint_manager = init_workers(config) + + raw_prompts = [[{"role": "user", "content": q}] for q in QUESTIONS] + batch = DataProto( + non_tensor_batch={ + "raw_prompt": np.array(raw_prompts), + "agent_name": np.array(["single_turn_agent"] * len(raw_prompts)), + "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)), + "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)), + }, + ).repeat(config.actor_rollout_ref.rollout.n) + + refit_runs = [] + for round_index in range(args.refit_rounds): + if round_index: + checkpoint_manager.update_weights() + result = agent_loop_manager.generate_sequences(prompts=batch) + checkpoint_manager.sleep_replicas() + result.meta_info.pop("timing", None) + + result.meta_info["temperature"] = config.actor_rollout_ref.rollout.temperature + result.meta_info["global_token_num"] = torch.sum(result.batch["attention_mask"], dim=-1).tolist() + result = result.union(compute_old_log_prob(actor_rollout_wg, result)) + + old_log_probs = result.batch["old_log_probs"] + rollout_log_probs = result.batch["rollout_log_probs"] + response_mask = result.batch["response_mask"].bool() + + metrics = compute_offpolicy_metrics( + old_log_prob=old_log_probs, rollout_log_prob=rollout_log_probs, response_mask=response_mask + ) + metrics.update(calculate_debug_metrics(result)) + log_prob_abs_diff = (old_log_probs - rollout_log_probs).abs()[response_mask] + metrics["log_prob_abs_diff_mean"] = log_prob_abs_diff.mean().item() + metrics["log_prob_abs_diff_max"] = log_prob_abs_diff.max().item() + metrics["num_response_tokens"] = int(response_mask.sum().item()) + refit_runs.append(metrics) + + metrics = dict(refit_runs[-1]) + if len(refit_runs) > 1: + metrics["refit_runs"] = refit_runs + print(json.dumps(metrics, indent=2)) + ray.shutdown() + + k3_kl = metrics["k3_kl"] + if not (np.isfinite(metrics["kl"]) and np.isfinite(k3_kl)) or k3_kl >= args.max_k3_kl: + print(f"FAILED: k3_kl {k3_kl} (threshold {args.max_k3_kl})") + return 1 + print("PASSED") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/experimental/agent_loop/agent_utils.py b/tests/experimental/agent_loop/agent_utils.py index 452b9a50452..b20578e8228 100644 --- a/tests/experimental/agent_loop/agent_utils.py +++ b/tests/experimental/agent_loop/agent_utils.py @@ -28,6 +28,13 @@ def init_agent_loop_manager(config: DictConfig) -> AgentLoopManager | RayWorkerGroup: + agent_loop_manager, _, _ = init_agent_loop_manager_and_workers(config) + return agent_loop_manager + + +def init_agent_loop_manager_and_workers( + config: DictConfig, +) -> tuple[AgentLoopManager, RayWorkerGroup, CheckpointEngineManager]: # =========================== 1. Create hybrid ActorRollout workers =========================== # The unified model-engine ActorRolloutRefWorker supports both sync and async rollout modes. actor_rollout_cls = ActorRolloutRefWorker @@ -103,4 +110,4 @@ def init_agent_loop_manager(config: DictConfig) -> AgentLoopManager | RayWorkerG checkpoint_manager.sleep_replicas() checkpoint_manager.update_weights() - return agent_loop_manager + return agent_loop_manager, actor_rollout_wg, checkpoint_manager diff --git a/tests/special_e2e/run_ppo_trainer_megatron.sh b/tests/special_e2e/run_ppo_trainer_megatron.sh index a02502e341b..5d7c1fa10c7 100644 --- a/tests/special_e2e/run_ppo_trainer_megatron.sh +++ b/tests/special_e2e/run_ppo_trainer_megatron.sh @@ -224,7 +224,8 @@ common_params=( actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP actor_rollout_ref.rollout.gpu_memory_utilization=0.6 actor_rollout_ref.rollout.n=${n_resp_per_prompt} - ++actor_rollout_ref.rollout.quantization=${ROLLOUT_QUANTIZATION} + actor_rollout_ref.rollout.quantization.linear=${ROLLOUT_QUANTIZATION} + actor_rollout_ref.rollout.quantization.moe=${ROLLOUT_QUANTIZATION} actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} actor_rollout_ref.rollout.enable_rollout_routing_replay=${ENABLE_ROLLOUT_ROUTING_REPLAY} actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} diff --git a/tests/special_sanity/check_device_api_usage.py b/tests/special_sanity/check_device_api_usage.py index 59c804ccbbe..7db1a7a8878 100644 --- a/tests/special_sanity/check_device_api_usage.py +++ b/tests/special_sanity/check_device_api_usage.py @@ -56,7 +56,6 @@ "verl/workers/rollout/trtllm_rollout/trtllm_async_server.py", # appear in config.cudagraph_capture_sizes "verl/workers/rollout/replica.py", # appear in default device_name "verl/checkpoint_engine", # checkpoint engine backend are device specific - "verl/utils/modelopt/megatron_qat_patch.py", # appear in torch.cuda.empty_cache() "verl/models/mcore/patch.py", # checkpoint patch only on cuda ] diff --git a/tests/utils/test_modelopt_mxfp8_on_cpu.py b/tests/utils/test_modelopt_mxfp8_on_cpu.py new file mode 100644 index 00000000000..d0e63a05679 --- /dev/null +++ b/tests/utils/test_modelopt_mxfp8_on_cpu.py @@ -0,0 +1,455 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("compressed_tensors") +pytest.importorskip("modelopt") + +# Hide an unusable optional FlashAttention-4 package before ModelOpt imports Megatron. +# isort: off +import verl.models.mcore # noqa: F401 +import modelopt.torch.quantization as mtq +from verl.utils.qat.modelopt import ( + _Mxfp4FakeQuant, + _Nvfp4FakeQuant, + _nvfp4_global_scale, + _nvfp4_input_backend, + _nvfp4_weight_backend, + build_quantize_config, +) +from verl.utils.qat.compressed_tensors import ( + CompressedTensorsWeightQuantizer, + Fp8ConfigWeightQuantizer, + _compress_weight, + _scheme_for_preset, + build_compressed_tensors_config, + dispatch_rollout_weight_quantizer, + qat_modes_for_config, +) +# isort: on + + +def _entry(config: dict, name: str): + return next(entry for entry in config["quant_cfg"] if entry["quantizer_name"] == name) + + +def test_mxfp8_qat_uses_modelopt_native_preset(): + config = build_quantize_config("mxfp8") + expected = mtq.normalize_quant_cfg_list(mtq.MXFP8_DEFAULT_CFG["quant_cfg"]) + + assert config["quant_cfg"] == expected + assert config["algorithm"] is None + weight_cfg = _entry(config, "*weight_quantizer")["cfg"] + assert weight_cfg["num_bits"] == (4, 3) + assert weight_cfg["block_sizes"] == {-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + assert weight_cfg.get("backend") is None + + +def test_mxfp8_qat_appends_user_ignores_without_mutating_modelopt_preset(): + original = mtq.normalize_quant_cfg_list(mtq.MXFP8_DEFAULT_CFG["quant_cfg"]) + + config = build_quantize_config("mxfp8", ["*linear_attn.in_proj_a"]) + + assert config["quant_cfg"][:-1] == original + assert config["quant_cfg"][-1] == { + "quantizer_name": "*linear_attn.in_proj_a", + "enable": False, + } + assert mtq.normalize_quant_cfg_list(mtq.MXFP8_DEFAULT_CFG["quant_cfg"]) == original + + +def test_mxfp8_rollout_requires_explicit_shape_ignore(): + pattern = r"re:.*linear_attn\.in_proj_[ab]$" + + config = build_compressed_tensors_config( + linear="MXFP8", + moe="MXFP8", + ignore=[pattern], + config_file=None, + # Config construction must not inspect the checkpoint. + model_path="/path/does/not/exist", + ) + + assert pattern in config["ignore"] + + +def test_mxfp4_rollout_is_moe_only_w4a8_mx(): + config = build_compressed_tensors_config( + linear=None, + moe="MXFP4", + ignore=[], + config_file=None, + ) + + scheme = config["config_groups"]["group_moe"] + assert scheme["format"] == "mxfp4-pack-quantized" + assert scheme["weights"]["num_bits"] == 4 + assert scheme["weights"]["group_size"] == 32 + assert scheme["weights"]["scale_dtype"] == "torch.uint8" + assert scheme["input_activations"]["num_bits"] == 8 + assert scheme["input_activations"]["group_size"] == 32 + assert scheme["input_activations"]["dynamic"] is True + assert qat_modes_for_config(config) == {"mxfp4"} + + with pytest.raises(ValueError, match="only for routed MoE experts"): + build_compressed_tensors_config( + linear="MXFP4", + moe=None, + ignore=[], + config_file=None, + ) + + +def test_mxfp4_rollout_quantizes_packed_experts(): + config = build_compressed_tensors_config( + linear=None, + moe="MXFP4", + ignore=[], + config_file=None, + ) + quantizer = CompressedTensorsWeightQuantizer(config) + weights = torch.linspace(-1, 1, 2 * 64 * 32, dtype=torch.bfloat16).reshape(2, 64, 32) + + result = dict(quantizer([("model.layers.0.mlp.experts.gate_up_proj", weights)])) + + assert set(result) == { + "model.layers.0.mlp.experts.0.gate_proj.weight_packed", + "model.layers.0.mlp.experts.0.gate_proj.weight_scale", + "model.layers.0.mlp.experts.0.up_proj.weight_packed", + "model.layers.0.mlp.experts.0.up_proj.weight_scale", + "model.layers.0.mlp.experts.1.gate_proj.weight_packed", + "model.layers.0.mlp.experts.1.gate_proj.weight_scale", + "model.layers.0.mlp.experts.1.up_proj.weight_packed", + "model.layers.0.mlp.experts.1.up_proj.weight_scale", + } + assert result["model.layers.0.mlp.experts.0.gate_proj.weight_packed"].dtype == torch.uint8 + assert result["model.layers.0.mlp.experts.0.gate_proj.weight_packed"].shape == (32, 16) + assert result["model.layers.0.mlp.experts.0.gate_proj.weight_scale"].dtype == torch.uint8 + assert result["model.layers.0.mlp.experts.0.gate_proj.weight_scale"].shape == (32, 1) + + +def test_native_fp8_config_export_uses_vllm_scale_names(): + quantizer = Fp8ConfigWeightQuantizer( + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + "modules_to_not_convert": ["model.layers.0.mlp.gate"], + } + ) + weight = torch.linspace(-1, 1, 256 * 256, dtype=torch.bfloat16).reshape(256, 256) + gate = torch.randn(16, 256, dtype=torch.bfloat16) + + result = dict( + quantizer( + [ + ("model.layers.0.self_attn.q_proj.weight", weight), + ("model.layers.0.mlp.gate.weight", gate), + ] + ) + ) + + assert set(result) == { + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.q_proj.weight_scale_inv", + "model.layers.0.mlp.gate.weight", + } + assert result["model.layers.0.self_attn.q_proj.weight"].dtype == torch.float8_e4m3fn + assert result["model.layers.0.self_attn.q_proj.weight_scale_inv"].shape == (2, 2) + assert result["model.layers.0.mlp.gate.weight"] is gate + + +def test_native_fp8_config_export_keeps_default_ignored_layers_in_high_precision(): + quantizer = Fp8ConfigWeightQuantizer( + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + ) + embedding = torch.randn(256, 256, dtype=torch.bfloat16) + lm_head = torch.randn(256, 256, dtype=torch.bfloat16) + router = torch.randn(16, 256, dtype=torch.bfloat16) + + result = dict( + quantizer( + [ + ("model.embed_tokens.weight", embedding), + ("lm_head.weight", lm_head), + ("model.layers.0.mlp.gate.weight", router), + ] + ) + ) + + expected = { + "model.embed_tokens.weight": embedding, + "lm_head.weight": lm_head, + "model.layers.0.mlp.gate.weight": router, + } + assert result.keys() == expected.keys() + assert all(result[name] is tensor for name, tensor in expected.items()) + + +def test_native_fp8_config_export_unpacks_fused_experts(): + quantizer = Fp8ConfigWeightQuantizer( + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + ) + weights = torch.randn(2, 256, 128, dtype=torch.bfloat16) + + result = dict(quantizer([("model.layers.0.mlp.experts.down_proj", weights)])) + + assert set(result) == { + "model.layers.0.mlp.experts.0.down_proj.weight", + "model.layers.0.mlp.experts.0.down_proj.weight_scale_inv", + "model.layers.0.mlp.experts.1.down_proj.weight", + "model.layers.0.mlp.experts.1.down_proj.weight_scale_inv", + } + assert result["model.layers.0.mlp.experts.0.down_proj.weight"].dtype == torch.float8_e4m3fn + + +def _rollout_config(**quantization): + return SimpleNamespace( + quantization=SimpleNamespace( + linear=quantization.get("linear"), + moe=quantization.get("moe"), + ignore=quantization.get("ignore", []), + config_file=quantization.get("config_file"), + ) + ) + + +def test_weight_quantizer_dispatches_native_fp8_and_deepseek_passthrough(): + fp8_config = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + standard_model = SimpleNamespace( + local_path="/model", + hf_config=SimpleNamespace(model_type="glm_moe_dsa", text_config=None, quantization_config=fp8_config), + ) + deepseek_model = SimpleNamespace( + local_path="/model", + hf_config=SimpleNamespace( + model_type="deepseek_v4", text_config=None, quantization_config={"quant_method": "fp8"} + ), + ) + + standard = dispatch_rollout_weight_quantizer(model_config=standard_model, rollout_config=_rollout_config()) + deepseek = dispatch_rollout_weight_quantizer(model_config=deepseek_model, rollout_config=_rollout_config()) + + assert isinstance(standard, Fp8ConfigWeightQuantizer) + assert standard.passthrough is False + assert isinstance(deepseek, Fp8ConfigWeightQuantizer) + assert deepseek.passthrough is True + weight = torch.tensor([-1, 2], dtype=torch.int8) + scale = torch.tensor([127, -1], dtype=torch.float8_e8m0fnu) + dense = torch.ones(2, dtype=torch.bfloat16) + result = dict( + deepseek( + [ + ("model.layers.0.mlp.experts.0.weight", weight), + ("model.layers.0.mlp.experts.0.scale", scale), + ("model.layers.0.self_attn.q_proj.weight", dense), + ] + ) + ) + assert result["model.layers.0.mlp.experts.0.weight"].dtype == torch.uint8 + assert result["model.layers.0.mlp.experts.0.weight"].data_ptr() == weight.data_ptr() + assert result["model.layers.0.mlp.experts.0.scale"].dtype == torch.uint8 + assert result["model.layers.0.mlp.experts.0.scale"].data_ptr() == scale.data_ptr() + assert result["model.layers.0.self_attn.q_proj.weight"] is dense + + +def test_weight_quantizer_dispatches_native_and_override_compressed_tensors(): + native_config = build_compressed_tensors_config(linear="FP8_BLOCK", moe="FP8_BLOCK", ignore=[], config_file=None) + native_model = SimpleNamespace( + local_path="/model", + hf_config=SimpleNamespace(model_type="qwen3_5_moe", text_config=None, quantization_config=native_config), + ) + bf16_model = SimpleNamespace( + local_path="/model", + hf_config=SimpleNamespace(model_type="qwen3_5_moe", text_config=None, quantization_config=None), + ) + + native = dispatch_rollout_weight_quantizer(model_config=native_model, rollout_config=_rollout_config()) + override = dispatch_rollout_weight_quantizer( + model_config=bf16_model, + rollout_config=_rollout_config(linear="FP8_BLOCK", moe="FP8_BLOCK"), + ) + + assert isinstance(native, CompressedTensorsWeightQuantizer) + assert isinstance(override, CompressedTensorsWeightQuantizer) + + +def test_mixed_mxfp8_linear_mxfp4_moe_uses_composite_qat_mode(): + config = build_compressed_tensors_config( + linear="MXFP8", + moe="MXFP4", + ignore=[], + config_file=None, + ) + + assert config["format"] == "mixed-precision" + assert qat_modes_for_config(config) == {"mxfp4_mxfp8"} + + +@pytest.mark.parametrize( + ("qat_mode", "dense_mxfp8"), + [("mxfp4", False), ("mxfp4_mxfp8", True)], +) +def test_mxfp4_qat_combines_mxfp4_weights_with_mxfp8_inputs(qat_mode, dense_mxfp8): + config = build_quantize_config(qat_mode) + weight_cfg = _entry(config, "*mlp.experts*weight_quantizer")["cfg"] + input_cfg = _entry(config, "*mlp.experts*input_quantizer")["cfg"] + + assert weight_cfg["num_bits"] == (2, 1) + assert weight_cfg["block_sizes"] == {-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + assert weight_cfg["backend"] == "verl_mxfp4" + assert input_cfg["num_bits"] == (4, 3) + assert input_cfg["block_sizes"] == {-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + assert config["algorithm"] is None + + generic_weight_entries = [entry for entry in config["quant_cfg"] if entry["quantizer_name"] == "*weight_quantizer"] + assert bool(generic_weight_entries) is dense_mxfp8 + + +def test_mxfp4_qat_weight_qdq_exactly_matches_export(): + from compressed_tensors.compressors import BaseCompressor + from compressed_tensors.compressors.format import infer_module_format + + weight = torch.randn(64, 64, dtype=torch.bfloat16) + fake_quantized = _Mxfp4FakeQuant.apply(weight) + scheme = _scheme_for_preset("MXFP4", ["Linear"]) + compressed = _compress_weight(weight, scheme) + compressor = BaseCompressor.get_value_from_registry(infer_module_format(torch.nn.Linear, scheme).value) + exported_qdq = compressor.decompress(compressed, scheme)["weight"] + + torch.testing.assert_close(fake_quantized, exported_qdq, rtol=0, atol=0) + + +def test_nvfp4_rollout_is_moe_only_w4a4(): + config = build_compressed_tensors_config( + linear=None, + moe="NVFP4", + ignore=[], + config_file=None, + ) + + scheme = config["config_groups"]["group_moe"] + assert config["format"] == "nvfp4-pack-quantized" + assert scheme["weights"]["strategy"] == "tensor_group" + assert scheme["weights"]["group_size"] == 16 + assert scheme["weights"]["scale_dtype"] == "torch.float8_e4m3fn" + assert scheme["input_activations"]["strategy"] == "tensor_group" + assert scheme["input_activations"]["dynamic"] == "local" + assert qat_modes_for_config(config) == {"nvfp4"} + + with pytest.raises(ValueError, match="only for routed MoE experts"): + build_compressed_tensors_config( + linear="NVFP4", + moe=None, + ignore=[], + config_file=None, + ) + + +def test_nvfp4_qat_uses_experts_only_export_compatible_backends(): + config = build_quantize_config("nvfp4") + weight_cfg = _entry(config, "*mlp.experts*weight_quantizer")["cfg"] + input_cfg = _entry(config, "*mlp.experts*input_quantizer")["cfg"] + + assert weight_cfg["backend"] == "verl_nvfp4_weight" + assert input_cfg["backend"] == "verl_nvfp4_input" + assert weight_cfg["block_sizes"] == {-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + assert input_cfg["block_sizes"] == {-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + assert config["algorithm"] is None + assert not any( + entry["quantizer_name"] == "*weight_quantizer" and entry.get("enable", True) for entry in config["quant_cfg"] + ) + + +def test_nvfp4_qat_weight_qdq_exactly_matches_export(): + from compressed_tensors.compressors import BaseCompressor + from compressed_tensors.compressors.format import infer_module_format + + weight = torch.randn(64, 64, dtype=torch.bfloat16) + global_scale = _nvfp4_global_scale(weight) + fake_quantized = _Nvfp4FakeQuant.apply(weight, global_scale) + scheme = _scheme_for_preset("NVFP4", ["Linear"]) + compressed = _compress_weight(weight, scheme, global_scale=global_scale) + compressor = BaseCompressor.get_value_from_registry(infer_module_format(torch.nn.Linear, scheme).value) + exported_qdq = compressor.decompress(compressed, scheme)["weight"] + + torch.testing.assert_close(fake_quantized, exported_qdq, rtol=0, atol=0) + + +def test_nvfp4_grouped_expert_qdq_uses_one_global_scale_per_expert(): + weight = torch.randn(3, 64, 32, dtype=torch.bfloat16) + actual = _nvfp4_weight_backend(weight, None) + expected = torch.stack([_Nvfp4FakeQuant.apply(expert, _nvfp4_global_scale(expert)) for expert in weight]) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_nvfp4_input_backend_initializes_from_first_observation_then_tracks_minimum(): + quantizer = SimpleNamespace() + first = torch.full((16, 16), 0.5, dtype=torch.bfloat16) + second = first * 2 + + _nvfp4_input_backend(first, quantizer) + first_scale = _nvfp4_global_scale(first) + torch.testing.assert_close(quantizer._verl_input_global_scale, first_scale) + + _nvfp4_input_backend(second, quantizer) + torch.testing.assert_close( + quantizer._verl_input_global_scale, + torch.minimum(first_scale, _nvfp4_global_scale(second)), + ) + + +def test_nvfp4_export_uses_neutral_then_observed_input_scales(): + config = build_compressed_tensors_config( + linear=None, + moe="NVFP4", + ignore=[], + config_file=None, + ) + weights = torch.randn(2, 64, 32, dtype=torch.bfloat16) + quantizer = CompressedTensorsWeightQuantizer(config) + neutral = dict(quantizer([("model.layers.0.mlp.experts.gate_up_proj", weights)])) + assert neutral["model.layers.0.mlp.experts.0.gate_proj.input_global_scale"].item() == 1.0 + + quantizer.set_input_global_scales( + { + "layers.0.mlp.experts.gate_up": torch.tensor([0.25]), + "layers.0.mlp.experts.down": torch.tensor([0.5]), + } + ) + result = dict(quantizer([("model.layers.0.mlp.experts.gate_up_proj", weights)])) + gate_prefix = "model.layers.0.mlp.experts.0.gate_proj." + up_prefix = "model.layers.0.mlp.experts.0.up_proj." + assert result[gate_prefix + "input_global_scale"].item() == 0.25 + assert result[up_prefix + "input_global_scale"].item() == 0.25 + torch.testing.assert_close( + result[gate_prefix + "weight_global_scale"], + result[up_prefix + "weight_global_scale"], + ) diff --git a/tests/utils/test_sglang_fp8_utils.py b/tests/utils/test_sglang_fp8_utils.py deleted file mode 100644 index 8be69a3181a..00000000000 --- a/tests/utils/test_sglang_fp8_utils.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from types import SimpleNamespace - -from verl.utils.sglang.sglang_fp8_utils import SGLangFP8QuantizerHelper, build_sglang_fp8_quant_config - - -class MappingLikeConfig: - def __init__(self, values): - self.values = values - - def get(self, key, default=None): - return self.values.get(key, default) - - -def test_build_sglang_fp8_quant_config_preserves_defaults(monkeypatch): - monkeypatch.delenv("SGLANG_FP8_IGNORED_LAYERS", raising=False) - - quant_config = build_sglang_fp8_quant_config() - - assert quant_config == { - "activation_scheme": "dynamic", - "fmt": "e4m3", - "quant_method": "fp8", - "weight_block_size": [128, 128], - } - - -def test_sglang_fp8_quant_config_merges_hf_ignored_layers(monkeypatch): - monkeypatch.delenv("SGLANG_FP8_IGNORED_LAYERS", raising=False) - hf_config = SimpleNamespace( - quantization_config={ - "ignored_layers": ["model.layers.0.self_attn.q_proj"], - "modules_to_not_convert": ["model.layers.1.mlp.down_proj"], - } - ) - - quant_config = build_sglang_fp8_quant_config(hf_config) - helper = SGLangFP8QuantizerHelper(quant_config) - - assert quant_config["ignored_layers"] == [ - "model.layers.0.self_attn.q_proj", - "model.layers.1.mlp.down_proj", - ] - assert not helper.should_quantize_param("model.layers.0.self_attn.q_proj.weight") - assert not helper.should_quantize_param("model.layers.1.mlp.down_proj.weight") - assert helper.should_quantize_param("model.layers.2.mlp.down_proj.weight") - - -def test_sglang_fp8_quant_config_accepts_mapping_like_config(monkeypatch): - monkeypatch.delenv("SGLANG_FP8_IGNORED_LAYERS", raising=False) - hf_config = MappingLikeConfig( - { - "quantization_config": MappingLikeConfig( - { - "ignored_layers": ["model.layers.0.linear_attn"], - } - ) - } - ) - - quant_config = build_sglang_fp8_quant_config(hf_config) - helper = SGLangFP8QuantizerHelper(quant_config) - - assert quant_config["ignored_layers"] == ["model.layers.0.linear_attn"] - assert not helper.should_quantize_param("model.layers.0.linear_attn.in_proj_ba.weight") - - -def test_sglang_fp8_quantizer_matches_regex_ignored_layers(monkeypatch): - monkeypatch.delenv("SGLANG_FP8_IGNORED_LAYERS", raising=False) - hf_config = SimpleNamespace( - quantization_config={ - "ignored_layers": ["re:.*linear_attn.*"], - } - ) - - quant_config = build_sglang_fp8_quant_config(hf_config) - helper = SGLangFP8QuantizerHelper(quant_config) - - assert quant_config["ignored_layers"] == ["re:.*linear_attn.*"] - assert not helper.should_quantize_param("model.layers.0.linear_attn.in_proj_ba.weight") - assert not helper.should_quantize_param("model.layers.0.linear_attn.g_proj.weight") - assert helper.should_quantize_param("model.layers.0.mlp.experts.0.up_proj.weight") - - -def test_sglang_fp8_quantizer_reads_sglang_env_ignored_layers(monkeypatch): - monkeypatch.setenv("SGLANG_FP8_IGNORED_LAYERS", "linear_attn") - - quant_config = build_sglang_fp8_quant_config() - helper = SGLangFP8QuantizerHelper(quant_config) - - assert quant_config["ignored_layers"] == ["linear_attn"] - assert not helper.should_quantize_param("model.layers.0.linear_attn.in_proj_ba.weight") - assert not helper.should_quantize_param("model.layers.0.linear_attn.g_proj.weight") - assert helper.should_quantize_param("model.layers.0.mlp.experts.0.up_proj.weight") diff --git a/tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py b/tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py index 64dd52009b3..7421ce73b70 100644 --- a/tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py +++ b/tests/utils/test_vllm_fp8_restage_same_shape_on_cpu.py @@ -76,6 +76,11 @@ def _repacking_process(self, layer): _fp8().replace_parameter(layer, "w13_weight_scale_inv", layer.w13_weight_scale_inv.data.reshape(E, -1).clone()) +def _direct_repacking_process(self, layer): + """MXFP8-linear-like prep that assigns a same-layout replacement without replace_parameter.""" + layer.w13_weight = torch.nn.Parameter(_swap_halves(layer.w13_weight.data), requires_grad=False) + + class _StubVllmFp8: NAMES = ( "vllm", @@ -202,3 +207,18 @@ def test_shape_changing_repack_is_still_staged(): finally: for p in patchers: p.stop() + + +def test_direct_same_shape_parameter_replacement_is_restaged(): + module = _load_fp8_utils() + with _StubVllmFp8(_direct_repacking_process) as stub: + patchers = module.build_fp8_method_patchers(version.parse("0.24.0")) + for p in patchers: + p.start() + try: + model, layer = _loaded_layer(module, stub) + assert layer._verl_fp8_repacked == {"w13_weight"} + assert module.stage_fp8_params_for_loading(model) == [layer] + finally: + for p in patchers: + p.stop() diff --git a/tests/utils/test_vllm_mxfp4_restage_on_cpu.py b/tests/utils/test_vllm_mxfp4_restage_on_cpu.py new file mode 100644 index 00000000000..54ab034b610 --- /dev/null +++ b/tests/utils/test_vllm_mxfp4_restage_on_cpu.py @@ -0,0 +1,168 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exercise FP4 refit helpers without importing the optional vLLM package.""" + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _load_fp4_utils(): + path = _REPO_ROOT / "verl/utils/vllm/vllm_fp4_utils.py" + spec = importlib.util.spec_from_file_location("vllm_fp4_utils_under_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_FP4_UTILS = _load_fp4_utils() +_MXFP4_LIVE_ATTR = _FP4_UTILS._MXFP4_LIVE_ATTR +_NVFP4_LIVE_ATTR = _FP4_UTILS._NVFP4_LIVE_ATTR +_restore_nvfp4_input_scales_in_place = _FP4_UTILS._restore_nvfp4_input_scales_in_place +_restore_nvfp4_kernel_config_in_place = _FP4_UTILS._restore_nvfp4_kernel_config_in_place +_stage_mxfp4_moe_params = _FP4_UTILS._stage_mxfp4_moe_params +_stage_nvfp4_moe_params = _FP4_UTILS._stage_nvfp4_moe_params + + +class _PackedMxfp4Method: + _verl_uses_packed_checkpoint_names = True + num_experts = 2 + intermediate_size = 32 + hidden_size = 32 + + +def test_compressed_tensors_mxfp4_restage_exposes_packed_weight_names(): + module = torch.nn.Module() + module.quant_method = _PackedMxfp4Method() + module.weight_loader = lambda *args, **kwargs: None + module.w13_weight = torch.nn.Parameter(torch.empty(2, 64, 16, dtype=torch.uint8), requires_grad=False) + module.w2_weight = torch.nn.Parameter(torch.empty(2, 32, 16, dtype=torch.uint8), requires_grad=False) + module.w13_weight_scale = torch.nn.Parameter(torch.empty(2, 64, 1, dtype=torch.uint8), requires_grad=False) + module.w2_weight_scale = torch.nn.Parameter(torch.empty(2, 32, 1, dtype=torch.uint8), requires_grad=False) + + live_w13 = module.w13_weight + live_w2 = module.w2_weight + _stage_mxfp4_moe_params(module) + + assert module.w13_weight is live_w13 + assert module.w2_weight is live_w2 + assert module.w13_weight_packed.data_ptr() == live_w13.data_ptr() + assert module.w2_weight_packed.data_ptr() == live_w2.data_ptr() + assert module.w13_weight_packed.weight_loader == module.weight_loader + assert getattr(module, _MXFP4_LIVE_ATTR)["w13_weight"] is live_w13 + assert getattr(module, _MXFP4_LIVE_ATTR)["w2_weight"] is live_w2 + assert module.w13_weight_scale.quant_method == "block" + assert module.w2_weight_scale.quant_method == "block" + + +class _Nvfp4Moe: + is_act_and_mul = True + + +class _Nvfp4Method: + moe = _Nvfp4Moe() + + +def test_compressed_tensors_nvfp4_restage_exposes_checkpoint_layout(): + module = torch.nn.Module() + module.quant_method = _Nvfp4Method() + module.weight_loader = lambda *args, **kwargs: None + module.num_experts = 2 + module.intermediate_size_per_partition = 32 + module.hidden_size = 32 + module.w13_weight = torch.nn.Parameter(torch.empty(2, 64, 16, dtype=torch.uint8), requires_grad=False) + module.w2_weight = torch.nn.Parameter(torch.empty(2, 32, 16, dtype=torch.uint8), requires_grad=False) + # CUTLASS swizzles/pads scales; use checkpoint shape here to exercise the + # same-storage path without requiring CUDA. + module.w13_weight_scale = torch.nn.Parameter(torch.empty(2, 64, 2, dtype=torch.float8_e4m3fn), requires_grad=False) + module.w2_weight_scale = torch.nn.Parameter(torch.empty(2, 32, 2, dtype=torch.float8_e4m3fn), requires_grad=False) + module.w13_weight_scale_2 = torch.nn.Parameter(torch.empty(2), requires_grad=False) + module.w2_weight_scale_2 = torch.nn.Parameter(torch.empty(2), requires_grad=False) + module.w13_weight_global_scale = torch.nn.Parameter(torch.ones(2, 2), requires_grad=False) + module.w2_weight_global_scale = torch.nn.Parameter(torch.ones(2), requires_grad=False) + module.w13_input_global_scale = torch.nn.Parameter(torch.ones(2, 2), requires_grad=False) + module.w2_input_global_scale = torch.nn.Parameter(torch.ones(2), requires_grad=False) + module.w13_input_scale = torch.ones(2) + module.w2_input_scale = torch.ones(2) + + live_w13 = module.w13_weight + live_global = module.w13_weight_global_scale + live_input_scale = module.w13_input_scale + _stage_nvfp4_moe_params(module) + + assert module.w13_weight_packed.data_ptr() == live_w13.data_ptr() + assert module.w13_weight_global_scale is live_global + assert module.w13_weight_scale.quant_method == "group" + assert getattr(module, _NVFP4_LIVE_ATTR)["w13_weight"] is live_w13 + assert "w13_weight_scale_2" in getattr(module, _NVFP4_LIVE_ATTR) + assert getattr(module, _NVFP4_LIVE_ATTR)["w13_input_scale"] is live_input_scale + + +def test_nvfp4_restage_updates_activation_scales_without_changing_addresses(): + module = torch.nn.Module() + live_w13 = torch.zeros(2) + live_w2 = torch.zeros(2) + setattr( + module, + _NVFP4_LIVE_ATTR, + {"w13_input_scale": live_w13, "w2_input_scale": live_w2}, + ) + converted = [None] * 8 + converted[3] = torch.tensor([2.0, 3.0]) + converted[7] = torch.tensor([4.0, 5.0]) + + result = _restore_nvfp4_input_scales_in_place(module, converted) + + assert result[3] is live_w13 + assert result[7] is live_w2 + torch.testing.assert_close(live_w13, torch.tensor([2.0, 3.0])) + torch.testing.assert_close(live_w2, torch.tensor([4.0, 5.0])) + assert getattr(module, _NVFP4_LIVE_ATTR) == {} + + +def test_nvfp4_restage_retains_kernel_config_and_updates_derived_scales(): + module = torch.nn.Module() + module.w13_weight_scale = torch.nn.Parameter(torch.ones(2)) + module.w2_weight_scale = torch.nn.Parameter(torch.ones(2)) + module.w13_weight_scale_2 = torch.nn.Parameter(torch.ones(2)) + module.w2_weight_scale_2 = torch.nn.Parameter(torch.ones(2)) + module.w13_input_scale = torch.tensor([0.5, 0.25]) + module.w2_input_scale = torch.tensor([0.2, 0.1]) + retained_config = SimpleNamespace( + g1_alphas=module.w13_weight_scale_2, + g2_alphas=module.w2_weight_scale_2, + w1_scale=module.w13_weight_scale, + w2_scale=module.w2_weight_scale, + a1_gscale=torch.zeros(2), + a2_gscale=torch.zeros(2), + ) + retained_kernel = SimpleNamespace(fused_experts=SimpleNamespace(quant_config=retained_config)) + processing_config = SimpleNamespace( + a1_gscale=1.0 / module.w13_input_scale, + a2_gscale=1.0 / module.w2_input_scale, + ) + module.quant_method = SimpleNamespace(moe_kernel=SimpleNamespace(), moe_quant_config=processing_config) + + _restore_nvfp4_kernel_config_in_place(module, retained_kernel, retained_config) + + assert module.quant_method.moe_kernel is retained_kernel + assert module.quant_method.moe_quant_config is retained_config + torch.testing.assert_close(retained_config.a1_gscale, torch.tensor([2.0, 4.0])) + torch.testing.assert_close(retained_config.a2_gscale, torch.tensor([5.0, 10.0])) diff --git a/tests/utils/test_vllm_quant_utils_moe_on_cpu.py b/tests/utils/test_vllm_quant_utils_moe_on_cpu.py deleted file mode 100644 index 71e5f6d1a48..00000000000 --- a/tests/utils/test_vllm_quant_utils_moe_on_cpu.py +++ /dev/null @@ -1,240 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CPU tests for fused-MoE handling in ``verl/utils/vllm/vllm_quant_utils.py``. - -vLLM 0.24.0 (MoE refactor, vllm-project/vllm#41184) turned ``FusedMoE`` from an -``nn.Module`` class into a factory *function* that returns a ``MoERunner``, and -moved the fused expert weights onto a ``RoutedExperts`` submodule -(``experts`` -> ``experts.routed_experts``). The old code called -``isinstance(module, FusedMoE)`` which raised -``TypeError: isinstance() arg 2 must be a type`` once ``FusedMoE`` became a -function. These tests exercise both the pre-0.24 and post-0.24 module layouts -with lightweight fakes (no real vLLM required). -""" - -import importlib.util -import sys -import types -from pathlib import Path - -import torch - -_REPO_ROOT = Path(__file__).resolve().parents[2] -_VLLM_UTILS_DIR = _REPO_ROOT / "verl/utils/vllm" -_MODULE_PATH = _VLLM_UTILS_DIR / "vllm_quant_utils.py" - - -def _make_module(name: str) -> types.ModuleType: - return types.ModuleType(name) - - -def _load_quant_utils(fused_moe_is_function: bool): - """Load ``vllm_quant_utils`` with vLLM/verl heavyweight deps stubbed. - - Args: - fused_moe_is_function: when True, emulate vLLM >= 0.24 where ``FusedMoE`` - is a factory function and the expert weights live on ``RoutedExperts`` - owned by a ``MoERunner``. When False, emulate vLLM < 0.24 where - ``FusedMoE`` is the ``nn.Module`` class holding the weights. - - Returns: - (module, namespace) where ``namespace`` exposes the fake classes used to - build test models so callers can assert against the exact class objects - the loaded module resolved. - """ - - class _FakeLinearBase(torch.nn.Module): - def __init__(self, dtype=torch.float8_e4m3fn): - super().__init__() - self.weight = torch.empty(2, 2, dtype=dtype) - - ns: dict = {"LinearBase": _FakeLinearBase} - - fused_moe_pkg = _make_module("vllm.model_executor.layers.fused_moe") - fused_moe_layer = _make_module("vllm.model_executor.layers.fused_moe.layer") - linear_mod = _make_module("vllm.model_executor.layers.linear") - linear_mod.LinearBase = _FakeLinearBase - - if fused_moe_is_function: - - class _FakeRoutedExperts(torch.nn.Module): - def __init__(self, dtype=torch.float8_e4m3fn): - super().__init__() - self.w13_weight = torch.empty(2, 2, dtype=dtype) - self.w2_weight = torch.empty(2, 2, dtype=dtype) - - class _FakeMoERunner(torch.nn.Module): - def __init__(self, dtype=torch.float8_e4m3fn): - super().__init__() - self.routed_experts = _FakeRoutedExperts(dtype) - - def _fused_moe_factory(*args, **kwargs): # noqa: N802 - mirrors vLLM name - return _FakeMoERunner() - - fused_moe_layer.FusedMoE = _fused_moe_factory - fused_moe_pkg.FusedMoE = _fused_moe_factory - fused_moe_pkg.RoutedExperts = _FakeRoutedExperts - fused_moe_pkg.MoERunner = _FakeMoERunner - ns.update(RoutedExperts=_FakeRoutedExperts, MoERunner=_FakeMoERunner) - else: - - class _FakeFusedMoE(torch.nn.Module): - def __init__(self, dtype=torch.float8_e4m3fn): - super().__init__() - self.w13_weight = torch.empty(2, 2, dtype=dtype) - self.w2_weight = torch.empty(2, 2, dtype=dtype) - - fused_moe_layer.FusedMoE = _FakeFusedMoE - fused_moe_pkg.FusedMoE = _FakeFusedMoE - ns.update(FusedMoE=_FakeFusedMoE) - - # verl leaf modules that pull in torch/triton or vLLM at import time. - fake_kernel = _make_module("verl.utils.kernel.fp8_kernel") - fake_kernel.scaled_fp8_blockwise = lambda *a, **k: (None, None) - - # The vllm_fp8_utils / vllm_fp4_utils siblings only touch vLLM inside - # function bodies, so they can be imported for real. Stub the package - # ``__init__`` (which does pull heavyweight deps) but point its search path - # at the real directory so the submodule imports still resolve. - fake_vllm_pkg = _make_module("verl.utils.vllm") - fake_vllm_pkg.__path__ = [str(_VLLM_UTILS_DIR)] - - fakes = { - "vllm": _make_module("vllm"), - "vllm.model_executor": _make_module("vllm.model_executor"), - "vllm.model_executor.layers": _make_module("vllm.model_executor.layers"), - "vllm.model_executor.layers.fused_moe": fused_moe_pkg, - "vllm.model_executor.layers.fused_moe.layer": fused_moe_layer, - "vllm.model_executor.layers.linear": linear_mod, - "verl.utils.kernel": _make_module("verl.utils.kernel"), - "verl.utils.kernel.fp8_kernel": fake_kernel, - "verl.utils.vllm": fake_vllm_pkg, - } - - saved = {name: sys.modules.get(name) for name in fakes} - preexisting = set(sys.modules) - try: - sys.modules.update(fakes) - spec = importlib.util.spec_from_file_location("verl_vllm_quant_utils_under_test", _MODULE_PATH) - module = importlib.util.module_from_spec(spec) - assert spec is not None and spec.loader is not None - spec.loader.exec_module(module) - finally: - # The sibling modules imported above are bound to the stubbed package, - # so drop them rather than leave them for the next importer. - for name in set(sys.modules) - preexisting: - if name.startswith("verl.utils.vllm."): - sys.modules.pop(name, None) - for name, prev in saved.items(): - if prev is None: - sys.modules.pop(name, None) - else: - sys.modules[name] = prev - return module, ns - - -PACKED_MODULES_MAPPING = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"], -} - - -def _build_model(ns: dict, moe_dtype=torch.float8_e4m3fn): - """Build ``model.layers.0`` with an attention linear and a fused-MoE block.""" - linear_cls = ns["LinearBase"] - moe_cls = ns.get("MoERunner") or ns.get("FusedMoE") - - attn = torch.nn.Module() - attn.qkv_proj = linear_cls(torch.float8_e4m3fn) - - mlp = torch.nn.Module() - mlp.experts = moe_cls(moe_dtype) - mlp.gate = linear_cls(torch.bfloat16) # router: must NOT be treated as fp8 - - layer0 = torch.nn.Module() - layer0.self_attn = attn - layer0.mlp = mlp - - inner = torch.nn.Module() - inner.layers = torch.nn.ModuleList([layer0]) - - model = torch.nn.Module() - model.model = inner - model.packed_modules_mapping = PACKED_MODULES_MAPPING - return model - - -def test_new_vllm_resolves_routed_experts_and_does_not_crash(): - """vLLM >= 0.24: per-expert names must resolve to RoutedExperts, no TypeError.""" - mod, ns = _load_quant_utils(fused_moe_is_function=True) - model = _build_model(ns) - - routed_experts = model.model.layers[0].mlp.experts.routed_experts - assert isinstance(routed_experts, ns["RoutedExperts"]) - - # Class resolution landed on the concrete post-refactor classes. - assert ns["RoutedExperts"] in mod._EXPERT_WEIGHT_CLASSES - assert ns["MoERunner"] in mod._MOE_STOP_CLASSES - - # The exact name that crashed before the fix: a per-expert HF weight that - # walks into the fused-MoE block (index + proj remapped to gate_up_proj). - name = "model.layers.0.mlp.experts.0.gate_proj.weight" - resolved = mod.get_module_from_param_name(model, name) - assert resolved is routed_experts - - # Already-fused expert weight name resolves to the same holder. - resolved_fused = mod.get_module_from_param_name(model, "model.layers.0.mlp.experts.w13_weight") - assert resolved_fused is routed_experts - - -def test_new_vllm_is_fp8_weight_detects_moe_and_linear(): - mod, ns = _load_quant_utils(fused_moe_is_function=True) - model = _build_model(ns, moe_dtype=torch.float8_e4m3fn) - - mod.fp8_state.seen_params.clear() - mod.fp8_state.fp8_param_names.clear() - - # fp8 expert weight (per-expert HF name) -> detected as fp8 - assert mod.is_fp8_weight("model.layers.0.mlp.experts.0.gate_proj.weight", model) is True - # fp8 attention linear -> detected as fp8 - assert mod.is_fp8_weight("model.layers.0.self_attn.q_proj.weight", model) is True - # bf16 router gate -> NOT fp8 - assert mod.is_fp8_weight("model.layers.0.mlp.gate.weight", model) is False - - -def test_new_vllm_bf16_experts_not_flagged_fp8(): - mod, ns = _load_quant_utils(fused_moe_is_function=True) - model = _build_model(ns, moe_dtype=torch.bfloat16) - - mod.fp8_state.seen_params.clear() - mod.fp8_state.fp8_param_names.clear() - - assert mod.is_fp8_weight("model.layers.0.mlp.experts.0.gate_proj.weight", model) is False - - -def test_old_vllm_fusedmoe_class_still_supported(): - """vLLM < 0.24: FusedMoE is an nn.Module class holding the weights.""" - mod, ns = _load_quant_utils(fused_moe_is_function=False) - model = _build_model(ns, moe_dtype=torch.float8_e4m3fn) - - assert ns["FusedMoE"] in mod._EXPERT_WEIGHT_CLASSES - - experts = model.model.layers[0].mlp.experts - resolved = mod.get_module_from_param_name(model, "model.layers.0.mlp.experts.0.gate_proj.weight") - assert resolved is experts - - mod.fp8_state.seen_params.clear() - mod.fp8_state.fp8_param_names.clear() - assert mod.is_fp8_weight("model.layers.0.mlp.experts.0.gate_proj.weight", model) is True diff --git a/tests/utils/test_vllm_weight_name_normalization_on_cpu.py b/tests/utils/test_vllm_weight_name_normalization_on_cpu.py index 947add577cf..fb5d42fb3e9 100644 --- a/tests/utils/test_vllm_weight_name_normalization_on_cpu.py +++ b/tests/utils/test_vllm_weight_name_normalization_on_cpu.py @@ -140,6 +140,7 @@ def hijack(): fake_vllm_quant = types.ModuleType("verl.utils.vllm.vllm_quant_utils") fake_vllm_quant.apply_vllm_quant_patches = lambda: None + fake_vllm_quant.configure_fp8_ignored_layers_match_mode = lambda config: None fake_vllm_quant.is_quantized_model = lambda config: False fake_vllm_quant.load_quanted_weights = lambda *a, **k: [] fake_vllm_quant.prepare_quanted_weights_for_loading = lambda model: None @@ -296,6 +297,7 @@ def __init__(self, speculative_config=None, quant_config=None): def _make_worker(model): worker = object.__new__(vLLMColocateWorkerExtension) worker.model_runner = types.SimpleNamespace(model=model, vllm_config=_FakeVllmConfig(), drafter=None) + worker._is_nvfp4_packed = False return worker @@ -731,8 +733,7 @@ def _add_lora(request): worker.remove_lora = lambda lora_id: None worker.device = torch.device("cpu") worker.local_rank = 0 - worker._is_qat_model = False - worker._is_modelopt_qat = False + worker._is_nvfp4_packed = False worker._get_zmq_handle = lambda: "ipc:///tmp/test-bucketed-lora.sock" worker.update_weights_from_ipc(peft_config={"r": 1}, base_sync_done=True) @@ -765,8 +766,7 @@ def test_update_weights_from_ipc_standard_loads_per_bucket(monkeypatch): worker = _make_worker(model) worker.device = torch.device("cpu") worker.local_rank = 0 - worker._is_qat_model = False - worker._is_modelopt_qat = False + worker._is_nvfp4_packed = False worker._get_zmq_handle = lambda: "ipc:///tmp/test-bucketed-std.sock" # Bypass the post-load process_weights_after_loading (needs real vLLM). @@ -909,8 +909,7 @@ def _staging_sync_worker(model): worker = _make_worker(model) worker.device = torch.device("cpu") worker.local_rank = 0 - worker._is_qat_model = False - worker._is_modelopt_qat = False + worker._is_nvfp4_packed = False worker._get_zmq_handle = lambda: "ipc:///tmp/test-bucketed-moe-staging.sock" return worker diff --git a/tests/workers/rollout/test_vllm_weight_update_utils_on_cpu.py b/tests/workers/rollout/test_vllm_weight_update_utils_on_cpu.py index df2c2a05f07..5b0e7691062 100644 --- a/tests/workers/rollout/test_vllm_weight_update_utils_on_cpu.py +++ b/tests/workers/rollout/test_vllm_weight_update_utils_on_cpu.py @@ -115,6 +115,7 @@ def hijack(): fake_vllm_quant = types.ModuleType("verl.utils.vllm.vllm_quant_utils") fake_vllm_quant.apply_vllm_quant_patches = lambda: None + fake_vllm_quant.configure_fp8_ignored_layers_match_mode = lambda config: None fake_vllm_quant.is_quantized_model = lambda config: False fake_vllm_quant.load_quanted_weights = lambda weights, runner, is_drafter=False: weights fake_vllm_quant.prepare_quanted_weights_for_loading = lambda model: None diff --git a/tests/workers/test_megatron_liger_fused_kernels_on_cpu.py b/tests/workers/test_megatron_liger_fused_kernels_on_cpu.py index 8ffc93af11f..d8df894898c 100644 --- a/tests/workers/test_megatron_liger_fused_kernels_on_cpu.py +++ b/tests/workers/test_megatron_liger_fused_kernels_on_cpu.py @@ -169,7 +169,12 @@ def test_worker_reserves_both_consumers_before_ref_reset(monkeypatch): config.ref.log_prob_micro_batch_size_per_gpu = 1 model = SimpleNamespace(get=lambda key, default=None: config.model.get(key, default)) ref = SimpleNamespace(engine=McoreEngineConfig(tensor_model_parallel_size=4), optim=None, checkpoint=None) - actor = SimpleNamespace(engine=McoreEngineConfig(tensor_model_parallel_size=2), optim=None, checkpoint=None) + actor = SimpleNamespace( + engine=McoreEngineConfig(tensor_model_parallel_size=2), + optim=None, + checkpoint=None, + policy_loss={}, + ) def convert(value): if value is config.model: diff --git a/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_8_8_mis_trtllm.sh b/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_8_8_mis_trtllm.sh index ad2aba4f28c..9b2241b71ab 100755 --- a/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_8_8_mis_trtllm.sh +++ b/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_8_8_mis_trtllm.sh @@ -116,13 +116,6 @@ rollout_is=null rollout_rs=seq_mean_k1 rollout_rs_threshold="0.999_1.001" -# To use fp8 rollout, add the following flags -# actor_rollout_ref.rollout.expert_parallel_size=2 \ -# +actor_rollout_ref.rollout.moe_tensor_parallel_size=2 \ -# +actor_rollout_ref.rollout.quantization="fp8" \ -# +actor_rollout_ref.rollout.engine_kwargs.trtllm.moe_config.backend=TRTLLM \ - - python -m verl.experimental.fully_async_policy.fully_async_main \ --config-path=config \ --config-name='fully_async_ppo_megatron_trainer.yaml'\ diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index 3f001c45fa5..c4d6ef0b80f 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -89,7 +89,7 @@ actor_rollout_ref: qat: _target_: verl.workers.config.QATEngineConfig enable: false - mode: w4a16 + mode: fp8 group_size: 16 ignore_patterns: - lm_head @@ -290,7 +290,7 @@ actor_rollout_ref: qat: _target_: verl.workers.config.QATEngineConfig enable: false - mode: w4a16 + mode: fp8 group_size: 16 ignore_patterns: - lm_head @@ -434,10 +434,13 @@ actor_rollout_ref: port: 9090 file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml served_model_name: ${oc.select:actor_rollout_ref.model.path,null} - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: ${oc.select:actor_rollout_ref.model.mtp, null} - qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,${oc.select:actor_rollout_ref.actor.megatron.qat,null}} disaggregation: enabled: false prefill_replicas: 1 @@ -631,7 +634,7 @@ critic: qat: _target_: verl.workers.config.QATEngineConfig enable: false - mode: w4a16 + mode: fp8 group_size: 16 ignore_patterns: - lm_head @@ -864,8 +867,12 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: null disaggregation: enabled: false diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index 30cd4186978..ef5dbc68306 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -385,10 +385,13 @@ actor_rollout_ref: port: 9090 file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml served_model_name: ${oc.select:actor_rollout_ref.model.path,null} - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: ${oc.select:actor_rollout_ref.model.mtp, null} - qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,${oc.select:actor_rollout_ref.actor.megatron.qat,null}} disaggregation: enabled: false prefill_replicas: 1 @@ -770,8 +773,12 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: null disaggregation: enabled: false diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index 13729550d5f..bf41fa143a6 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -422,10 +422,13 @@ actor_rollout_ref: port: 9090 file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml served_model_name: ${oc.select:actor_rollout_ref.model.path,null} - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: ${oc.select:actor_rollout_ref.model.mtp, null} - qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,${oc.select:actor_rollout_ref.actor.megatron.qat,null}} disaggregation: enabled: false prefill_replicas: 1 @@ -833,8 +836,12 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: null disaggregation: enabled: false diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index e9cbc4978d3..481c245f1ff 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -403,10 +403,13 @@ actor_rollout_ref: port: 9090 file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml served_model_name: ${oc.select:actor_rollout_ref.model.path,null} - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: ${oc.select:actor_rollout_ref.model.mtp, null} - qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,${oc.select:actor_rollout_ref.actor.megatron.qat,null}} disaggregation: enabled: false prefill_replicas: 1 @@ -804,8 +807,12 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: null disaggregation: enabled: false diff --git a/verl/trainer/config/engine/megatron.yaml b/verl/trainer/config/engine/megatron.yaml index 0ed6528e80b..0ee54390a31 100644 --- a/verl/trainer/config/engine/megatron.yaml +++ b/verl/trainer/config/engine/megatron.yaml @@ -140,8 +140,9 @@ qat: _target_: verl.workers.config.QATEngineConfig # Whether to enable QAT enable: false - # Quantization mode: "w4a16" (weight-only FP4) or "w4a4" (weight + activation FP4) - mode: "w4a16" + # Quantization mode: "fp8", "mxfp8", "mxfp4", "mxfp4_mxfp8", or "nvfp4". + # NVFP4 is W4A4 on routed experts only. + mode: "fp8" # Group size for blockwise quantization (NVFP4 requires 16) group_size: 16 # Module name patterns to exclude from quantization (fnmatch style for Megatron) diff --git a/verl/trainer/config/reward/reward.yaml b/verl/trainer/config/reward/reward.yaml index 6553e56737e..afb4bd0ed16 100644 --- a/verl/trainer/config/reward/reward.yaml +++ b/verl/trainer/config/reward/reward.yaml @@ -72,8 +72,12 @@ reward_model: prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + _target_: verl.workers.config.RolloutQuantizationConfig + linear: null + moe: null + ignore: [] + config_file: null mtp: null disaggregation: enabled: false diff --git a/verl/trainer/config/rollout/rollout.yaml b/verl/trainer/config/rollout/rollout.yaml index dbf59f10ded..d503665c4ba 100644 --- a/verl/trainer/config/rollout/rollout.yaml +++ b/verl/trainer/config/rollout/rollout.yaml @@ -481,18 +481,28 @@ prometheus: # Specify served_model_name to avoid displaying overly long model paths in Grafana served_model_name: ${oc.select:actor_rollout_ref.model.path,null} -# type of quantization in vllm, currently support fp8 and torchao -quantization: null +# Low-precision rollout as a compressed-tensors quantization_config: the engine starts from it with dummy +# weights and the trainer quantizes every weight sync (vLLM / SGLang + Megatron only for now). Off when all are null. +quantization: -# extra quantization information serialized in a config file, e.g. torchao_config.json -quantization_config_file: null + # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs + _target_: verl.workers.config.RolloutQuantizationConfig + + # compressed-tensors preset for dense linear layers (attention, dense MLP, shared experts), e.g. FP8_BLOCK + linear: null + + # compressed-tensors preset for routed experts, e.g. FP8_BLOCK + moe: null + + # extra "re:" patterns of layers kept in high precision + ignore: [] + + # compressed-tensors quantization_config JSON used as is, exclusive with linear/moe/ignore + config_file: null # MTP configuration, reuse model configuration mtp: ${oc.select:actor_rollout_ref.model.mtp, null} -# QAT configuration (inherited from actor's engine config) -qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,${oc.select:actor_rollout_ref.actor.megatron.qat,null}} - # Prefill-Decode disaggregation (SGLang only; keep rollout.name=sglang and toggle enabled=True). disaggregation: diff --git a/verl/utils/modelopt/__init__.py b/verl/utils/modelopt/__init__.py deleted file mode 100644 index cc151f98f20..00000000000 --- a/verl/utils/modelopt/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""ModelOpt integration for NVFP4 quantization with Megatron QAT training and vLLM inference.""" - -from verl.utils.modelopt.qat_utils import apply_qat_to_modules, export_qat_weights -from verl.utils.modelopt.qat_weight_exporter import QATWeightExporter -from verl.utils.modelopt.quantize import apply_qat, build_quantize_config -from verl.utils.modelopt.vllm_modelopt_patch import ( - apply_modelopt_nvfp4_patches, - modelopt_process_weights_after_loading, - prepare_modelopt_for_weight_reload, -) - -__all__ = [ - "build_quantize_config", - "apply_qat", - "QATWeightExporter", - "apply_modelopt_nvfp4_patches", - "prepare_modelopt_for_weight_reload", - "modelopt_process_weights_after_loading", - "apply_qat_to_modules", - "export_qat_weights", -] diff --git a/verl/utils/modelopt/qat_utils.py b/verl/utils/modelopt/qat_utils.py deleted file mode 100644 index ca002c5132c..00000000000 --- a/verl/utils/modelopt/qat_utils.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""High-level QAT workflow helpers for Megatron backend.""" - - -def _get_qat_field(qat_config, key, default=None): - """Extract a field from qat_config, supporting both dict and object-style access.""" - if isinstance(qat_config, dict): - return qat_config.get(key, default) - return getattr(qat_config, key, default) - - -def apply_qat_to_modules(modules, qat_config): - """Apply ModelOpt fake quantization to a list of Megatron module chunks.""" - from verl.utils.modelopt.quantize import apply_qat - - qat_mode = _get_qat_field(qat_config, "mode", "w4a16") - ignore_patterns = _get_qat_field(qat_config, "ignore_patterns", None) - if ignore_patterns is not None: - ignore_patterns = list(ignore_patterns) - - for i in range(len(modules)): - modules[i] = apply_qat(modules[i], qat_mode, ignore_patterns=ignore_patterns) - return modules - - -def export_qat_weights(per_tensor_param, modules, qat_mode, bridge): - """Process exported weights through QATWeightExporter for quantized weight sync.""" - from verl.utils.modelopt.qat_weight_exporter import QATWeightExporter - - qat_weight_exporter = QATWeightExporter(modules, bridge, qat_mode) - return qat_weight_exporter.process_weights_iterator(per_tensor_param) diff --git a/verl/utils/modelopt/qat_weight_exporter.py b/verl/utils/modelopt/qat_weight_exporter.py deleted file mode 100644 index d578e306390..00000000000 --- a/verl/utils/modelopt/qat_weight_exporter.py +++ /dev/null @@ -1,275 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""QAT weight exporter for Megatron-to-vLLM NVFP4 quantized weight sync.""" - -import re -from dataclasses import dataclass -from typing import Any, Iterator, Optional - -import torch -from modelopt.torch.export.quant_utils import ( - QUANTIZATION_NONE, - QUANTIZATION_NVFP4, - get_quantization_format, - get_weight_block_size, - to_quantized_weight, -) -from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor - -from verl.utils.megatron_utils import unwrap_model - -# NVFP4 two-level scaling denominator: FP4_MAX (6.0) * FP8_MAX (448.0). -_NVFP4_AMAX_DENOMINATOR = 6.0 * 448.0 - - -@dataclass -class _QuantMeta: - """Quantization metadata for a single parameter.""" - - qformat: str - block_size: int - weight_amax: Optional[torch.Tensor] - input_amax: Optional[torch.Tensor] = None - input_quantizer: Any = None - - -class QATWeightExporter: - """Export QAT-trained bf16 weights as quantized weights (e.g. NVFP4).""" - - def __init__( - self, - actor_module: list, - bridge: Any, - qat_mode: str = "w4a16", - ): - self.qat_mode = qat_mode - self._actor_module = actor_module - - self._registry = self._get_mapping_registry(bridge) - - from megatron.core import parallel_state as mpu - - self._pp_size = mpu.get_pipeline_model_parallel_world_size() - self._pp_rank = mpu.get_pipeline_model_parallel_rank() - self._pp_group = mpu.get_pipeline_model_parallel_group() if self._pp_size > 1 else None - - self._ep_size = mpu.get_expert_model_parallel_world_size() - self._ep_rank = mpu.get_expert_model_parallel_rank() if self._ep_size > 1 else 0 - self._ep_group = mpu.get_expert_model_parallel_group() if self._ep_size > 1 else None - - self._config = self._get_model_config(actor_module) - self._num_local_experts = self._count_local_experts(actor_module) - - self._metadata: dict[str, _QuantMeta] = {} - self._collect_metadata(actor_module) - - if self._pp_size > 1 and self._pp_group is not None: - self._sync_metadata(self._pp_group) - if self._ep_size > 1 and self._ep_group is not None: - self._sync_metadata(self._ep_group) - - def process_weights_iterator( - self, - per_tensor_param: Iterator[tuple[str, torch.Tensor]], - ) -> Iterator[tuple[str, torch.Tensor]]: - """Wrap a weight iterator to apply quantization. - - For each ``(hf_name, bf16_weight)`` from the iterator, yields the - quantized weight plus its scaling factors when the parameter is - quantized, or the original tensor unchanged otherwise. - """ - for hf_name, weight in per_tensor_param: - if "_quantizer." in hf_name: - continue - meta = self._resolve_quant_metadata(hf_name) - if meta is None: - yield (hf_name, weight) - else: - assert meta.qformat == QUANTIZATION_NVFP4, f"Unsupported qformat: {meta.qformat}" - yield from self._quantize_nvfp4(hf_name, weight, meta) - - @staticmethod - def _get_mapping_registry(bridge): - return bridge._model_bridge.mapping_registry() - - @staticmethod - def _get_model_config(actor_module): - model = unwrap_model(actor_module[0]) - return getattr(model, "config", None) - - @staticmethod - def _count_local_experts(actor_module) -> int: - indices: set[int] = set() - for module in actor_module: - model = unwrap_model(module) - for name, _ in model.named_modules(): - m = re.search(r"local_experts\.(\d+)", name) - if m: - indices.add(int(m.group(1))) - return max(indices) + 1 if indices else 0 - - def _collect_metadata(self, actor_module: list) -> None: - for vpp_idx, module in enumerate(actor_module): - model = unwrap_model(module) - for name, submodule in model.named_modules(): - qformat = get_quantization_format(submodule) - if qformat == QUANTIZATION_NONE: - continue - block_size = get_weight_block_size(submodule) - if block_size == 0: - continue - - w_q = getattr(submodule, "weight_quantizer", None) - i_q = getattr(submodule, "input_quantizer", None) - w_amax = w_q._amax.clone().cpu() if w_q and getattr(w_q, "_amax", None) is not None else None - i_amax = i_q._amax.clone().cpu() if i_q and getattr(i_q, "_amax", None) is not None else None - - meta = _QuantMeta( - qformat=qformat, - block_size=block_size, - weight_amax=w_amax, - input_amax=i_amax, - input_quantizer=i_q, - ) - - for pname, _ in submodule.named_parameters(recurse=False): - full_name = f"{name}.{pname}" if name else pname - global_name = self._local_to_global_param_name(full_name, vpp_idx) - self._metadata[global_name] = meta - - def _local_to_global_param_name(self, name: str, vpp_idx: int) -> str: - if self._config is None: - return name - - from megatron.bridge.models.conversion.model_bridge import _megatron_local_name_to_global - - return _megatron_local_name_to_global(self._actor_module, self._config, name, vpp_idx) - - def _sync_metadata(self, group) -> None: - world_size = torch.distributed.get_world_size(group=group) - - local_info = { - name: { - "qformat": m.qformat, - "block_size": m.block_size, - "weight_amax": m.weight_amax, - "input_amax": m.input_amax, - } - for name, m in self._metadata.items() - } - - gathered: list[dict | None] = [None] * world_size - torch.distributed.all_gather_object(gathered, local_info, group=group) - - for rank_info in gathered: - if rank_info is None: - continue - for name, info in rank_info.items(): - if name in self._metadata: - continue - self._metadata[name] = _QuantMeta( - qformat=info["qformat"], - block_size=info["block_size"], - weight_amax=info["weight_amax"], - input_amax=info["input_amax"], - input_quantizer=None, - ) - - def _resolve_quant_metadata(self, hf_name: str) -> Optional[_QuantMeta]: - if not hf_name.endswith(".weight") or "norm" in hf_name: - return None - - for resolved in _iter_hf_to_megatron_matches(self._registry, hf_name): - meta = self._metadata.get(resolved.megatron_param) - if meta is not None: - return meta - - return None - - def _quantize_nvfp4( - self, - name: str, - weight: torch.Tensor, - meta: _QuantMeta, - ) -> Iterator[tuple[str, torch.Tensor]]: - """NVFP4 two-level quantization. - - Produces up to four tensors: - ``(name, packed_uint8_weight)`` - ``(weight_scale, per_block_fp8_scale)`` - ``(weight_scale_2, global_scale_from_amax)`` - ``(input_scale, activation_scale)`` -- only when available - """ - w_amax = meta.weight_amax.to(weight.device) - w_scale_2 = w_amax.float() / _NVFP4_AMAX_DENOMINATOR - - w_scale = NVFP4QTensor.get_weights_scaling_factor( - weight, - meta.block_size, - weights_scaling_factor_2=w_scale_2.to(weight.device), - )[0] - - quantized = to_quantized_weight(weight, w_scale, meta.qformat, w_scale_2, meta.block_size) - - yield (name, quantized) - yield (_derive_scale_name(name, "weight_scale"), w_scale) - yield (_derive_scale_name(name, "weight_scale_2"), w_scale_2) - - input_scale = _compute_input_scale(meta) - if input_scale is not None: - yield (_derive_scale_name(name, "input_scale"), input_scale) - - -def _iter_hf_to_megatron_matches(registry, hf_name: str): - """Yield all resolved mappings whose HF pattern matches *hf_name*.""" - for pattern_info, mapping in registry._reverse_patterns: - if isinstance(mapping.hf_param, str): - pattern = pattern_info - if pattern is None: - if mapping.hf_param == hf_name: - yield mapping - else: - match = pattern.match(hf_name) - if match: - yield mapping.resolve(match.groups()) - else: - patterns_dict = pattern_info - for key, pattern in patterns_dict.items(): - if pattern is None: - if mapping.hf_param[key] == hf_name: - yield mapping.resolve(()) - else: - match = pattern.match(hf_name) - if match: - yield mapping.resolve(match.groups()) - - -def _derive_scale_name(weight_name: str, suffix: str) -> str: - result = weight_name.replace(".weight", f".{suffix}") - return result if result != weight_name else f"{weight_name}_{suffix}" - - -def _compute_input_scale(meta: _QuantMeta) -> Optional[torch.Tensor]: - if meta.input_quantizer is not None: - if hasattr(NVFP4QTensor, "get_activation_scaling_factor"): - return NVFP4QTensor.get_activation_scaling_factor(meta.input_quantizer) - if hasattr(meta.input_quantizer, "_amax") and meta.input_quantizer._amax is not None: - return meta.input_quantizer._amax.float() / _NVFP4_AMAX_DENOMINATOR - - if meta.input_amax is not None: - return meta.input_amax.float() / _NVFP4_AMAX_DENOMINATOR - - return None diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py deleted file mode 100644 index dee7274099c..00000000000 --- a/verl/utils/modelopt/quantize.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""ModelOpt NVFP4 quantization config and application for Megatron QAT.""" - -import copy - -import modelopt.torch.quantization as mtq -import torch.nn as nn -from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg - -_NVFP4_W4A16_QUANTIZER_CFG = { - "*weight_quantizer": { - "num_bits": (2, 1), - "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, - "axis": None, - "enable": True, - }, - "*input_quantizer": {"enable": False}, -} - - -def _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: - cfg = [] - mapping = { - "lm_head": "*output_layer*", - "*mlp.gate": "*router*", - "*self_attn*": "*self_attention*", - } - for pattern in ignore_patterns: - key = pattern - if key in mapping: - key = mapping[key] - cfg.append({"quantizer_name": key, "enable": False}) - return cfg - - -def build_quantize_config( - qat_mode: str, - ignore_patterns: list[str] | None = None, -) -> dict: - """Build a complete ModelOpt quantization config for ``mtq.quantize``.""" - if qat_mode != "w4a16": - raise ValueError(f"Only 'w4a16' is supported, got: {qat_mode}") - - if ignore_patterns is None: - ignore_patterns = [] - - ignore_cfg = _ignore_patterns_to_quant_cfg(ignore_patterns) - - quant_cfg = mtq.normalize_quant_cfg_list(_NVFP4_W4A16_QUANTIZER_CFG) - disabled_cfg = copy.deepcopy(_default_disabled_quantizer_cfg) - if isinstance(disabled_cfg, dict): - disabled_cfg = mtq.normalize_quant_cfg_list(disabled_cfg) - quant_cfg.extend(disabled_cfg) - quant_cfg.extend(ignore_cfg) - return {"quant_cfg": quant_cfg, "algorithm": "max"} - - -def apply_qat( - model: nn.Module, - qat_mode: str, - ignore_patterns: list[str] | None = None, -) -> nn.Module: - """Apply Quantization-Aware Training to a Megatron model.""" - config = build_quantize_config(qat_mode, ignore_patterns) - mtq.quantize(model, config) - return model diff --git a/verl/utils/modelopt/vllm_modelopt_patch.py b/verl/utils/modelopt/vllm_modelopt_patch.py deleted file mode 100644 index 9cdb69b5ae6..00000000000 --- a/verl/utils/modelopt/vllm_modelopt_patch.py +++ /dev/null @@ -1,619 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""vLLM ModelOpt NVFP4 patches for dynamic weight updates (Marlin backend).""" - -from typing import Optional - -import torch -from torch.nn import Parameter - -from verl.utils.device import get_device_name - - -def _save_param_meta(layer: torch.nn.Module, param_name: str): - if not hasattr(layer, "_hf_param_meta"): - layer._hf_param_meta = {} - - param = getattr(layer, param_name, None) - if param is None: - return - - meta = { - "shape": tuple(param.shape), - "dtype": param.dtype, - "device": str(param.device), - "param_class": type(param), - } - - if hasattr(param, "_input_dim"): - meta["input_dim"] = param._input_dim - if hasattr(param, "_output_dim"): - meta["output_dim"] = param._output_dim - - layer._hf_param_meta[param_name] = meta - - -def _create_param_from_meta( - module: torch.nn.Module, - param_name: str, - meta: dict, - device: Optional[torch.device] = None, -) -> Parameter: - shape = meta["shape"] - dtype = meta["dtype"] - dev = device or meta.get("device", get_device_name()) - param_class = meta.get("param_class", Parameter) - - weight_loaders = getattr(module, "_weight_loaders", {}) - weight_loader = weight_loaders.get(param_name) - - data = torch.empty(shape, dtype=dtype, device=dev) - - if param_class is not Parameter and weight_loader is not None: - kwargs = {"data": data, "weight_loader": weight_loader} - if "input_dim" in meta: - kwargs["input_dim"] = meta["input_dim"] - if "output_dim" in meta: - kwargs["output_dim"] = meta["output_dim"] - new_param = param_class(**kwargs) - else: - new_param = Parameter(data, requires_grad=False) - if weight_loader is not None: - new_param.weight_loader = weight_loader - - return new_param - - -def _check_first_call(layer: torch.nn.Module) -> bool: - count = getattr(layer, "_process_weights_call_count", 0) - layer._process_weights_call_count = count + 1 - return count == 0 - - -def _save_weight_loaders(layer: torch.nn.Module, param_names: list[str]): - if not hasattr(layer, "_weight_loaders"): - layer._weight_loaders = {} - for pname in param_names: - param = getattr(layer, pname, None) - if param is not None and hasattr(param, "weight_loader"): - layer._weight_loaders[pname] = param.weight_loader - - -def _update_ref_or_create(layer, ref_name, new_data): - refs = getattr(layer, "_marlin_tensor_refs", {}) - ref = refs.get(ref_name) - if ref is not None: - ref.copy_(new_data) - setattr(layer, ref_name, Parameter(ref, requires_grad=False)) - else: - t = new_data.clone() if isinstance(new_data, torch.Tensor) else torch.tensor(new_data) - setattr(layer, ref_name, Parameter(t, requires_grad=False)) - - -def _unwrap_marlin_scale(value): - return value[0] if isinstance(value, tuple) else value - - -def _split_marlin_scale(value): - return value if isinstance(value, tuple) else (value, 1.0) - - -class ModelOptParamMetaDict(dict): - """Dict-like parameter store with metadata-based rebuild and tensor swap.""" - - def __init__(self, model: torch.nn.Module, device: Optional[torch.device] = None): - super().__init__() - self.device = device - - actual_model = model - if hasattr(model, "model"): - actual_model = model.model - self._model = actual_model - - self._layer_meta_cache: dict[str, dict] = {} - self._tensor_swap_layers: dict[str, dict] = {} - - self._build_mappings() - - for name, param in actual_model.named_parameters(): - self[name] = param - - def _build_mappings(self): - for layer_name, module in self._model.named_modules(): - if not hasattr(module, "_hf_param_meta"): - continue - - self._layer_meta_cache[layer_name] = { - "module": module, - "meta": module._hf_param_meta, - } - - marlin_refs = getattr(module, "_marlin_tensor_refs", {}) - for param_name, meta in module._hf_param_meta.items(): - if param_name in marlin_refs: - key = f"{layer_name}.{param_name}" if layer_name else param_name - self._tensor_swap_layers[key] = { - "module": module, - "param_name": param_name, - "marlin_ref": marlin_refs[param_name], - "hf_meta": meta, - } - - def _try_rebuild(self, key: str) -> Optional[Parameter]: - parts = key.rsplit(".", 1) - if len(parts) != 2: - return None - layer_name, param_name = parts - if layer_name not in self._layer_meta_cache: - return None - cache_entry = self._layer_meta_cache[layer_name] - module = cache_entry["module"] - meta = cache_entry["meta"] - if param_name not in meta: - return None - if hasattr(module, param_name): - param = getattr(module, param_name) - if param is not None: - return param - new_param = _create_param_from_meta(module, param_name, meta[param_name], self.device) - module.register_parameter(param_name, new_param) - return new_param - - def prepare_for_reload(self) -> None: - """Replace kernel-format tensors with HF-shape tensors for reload.""" - for _key, swap_info in self._tensor_swap_layers.items(): - module = swap_info["module"] - param_name = swap_info["param_name"] - hf_meta = swap_info["hf_meta"] - if hasattr(module, param_name): - new_param = _create_param_from_meta(module, param_name, hf_meta, self.device) - setattr(module, param_name, new_param) - - def __getitem__(self, key: str) -> Parameter: - if key in dict.keys(self): - return super().__getitem__(key) - param = self._try_rebuild(key) - if param is not None: - self[key] = param - return param - raise KeyError(f"Parameter not found: {key}") - - def __contains__(self, key: str) -> bool: - if super().__contains__(key): - return True - parts = key.rsplit(".", 1) - if len(parts) == 2: - layer_name, param_name = parts - if layer_name in self._layer_meta_cache: - if param_name in self._layer_meta_cache[layer_name]["meta"]: - return True - return False - - def get(self, key: str, default=None): - try: - return self[key] - except KeyError: - return default - - -_DENSE_HF_PARAMS = ["weight", "weight_scale", "input_scale", "weight_scale_2"] - - -def _modelopt_dense_process_weights(self, layer: torch.nn.Module) -> None: - """ - Replacement for ModelOptNvFp4LinearMethod.process_weights_after_loading. - - First call: save metadata + weight_loaders, convert HF→Marlin format, - save _marlin_tensor_refs for CUDA Graph stability. - Subsequent: read reloaded HF data, convert, copy_ into saved refs. - """ - import vllm._custom_ops as ops - from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - marlin_make_workspace_new, - marlin_permute_scales, - ) - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - nvfp4_marlin_process_global_scale, - nvfp4_marlin_process_scales, - ) - - is_first_call = _check_first_call(layer) - - if is_first_call: - for pname in _DENSE_HF_PARAMS: - _save_param_meta(layer, pname) - _save_weight_loaders(layer, _DENSE_HF_PARAMS) - - weight_data = layer.weight.data - weight_scale_data = layer.weight_scale.data - weight_scale_2_data = layer.weight_scale_2.data - - assert weight_scale_data.dtype == torch.float8_e4m3fn - - device = weight_data.device - part_size_n = layer.output_size_per_partition - part_size_k = layer.input_size_per_partition - param_dtype = layer.params_dtype - group_size = 16 - weight_scale_2_max = weight_scale_2_data.max().to(torch.float32) - - if is_first_call: - layer.workspace = marlin_make_workspace_new(device) - - perm = torch.empty(0, dtype=torch.int, device=device) - qweight = weight_data.view(torch.int32).T.contiguous() - marlin_weight = ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=part_size_k, - size_n=part_size_n, - num_bits=4, - ) - - weight_scale = weight_scale_data.T.contiguous().to(param_dtype) - weight_scale = marlin_permute_scales( - s=weight_scale, - size_k=part_size_k, - size_n=part_size_n, - group_size=group_size, - ) - marlin_weight_scale, scale_factor = _split_marlin_scale( - nvfp4_marlin_process_scales(weight_scale, a_dtype=param_dtype) - ) - marlin_weight_global_scale = ( - nvfp4_marlin_process_global_scale(weight_scale_2_max.to(torch.float32), param_dtype) / scale_factor - ) - - if is_first_call: - layer.weight = Parameter(marlin_weight, requires_grad=False) - layer.weight_scale = Parameter(marlin_weight_scale, requires_grad=False) - layer.weight_global_scale = Parameter(marlin_weight_global_scale, requires_grad=False) - layer.weight_scale_2 = Parameter(layer.weight_global_scale.data, requires_grad=False) - layer._marlin_tensor_refs = { - "weight": layer.weight.data, - "weight_scale": layer.weight_scale.data, - "weight_scale_2": layer.weight_global_scale.data, - } - else: - _update_ref_or_create(layer, "weight", marlin_weight) - _update_ref_or_create(layer, "weight_scale", marlin_weight_scale) - _update_ref_or_create(layer, "weight_scale_2", marlin_weight_global_scale) - layer.weight_global_scale = Parameter(layer._marlin_tensor_refs["weight_scale_2"], requires_grad=False) - - if hasattr(layer, "input_scale"): - input_global_scale = layer.input_scale.data.max().to(torch.float32) - layer.input_global_scale = Parameter(input_global_scale, requires_grad=False) - layer.alpha = Parameter(input_global_scale * weight_scale_2_max, requires_grad=False) - layer.input_global_scale_inv = Parameter((1.0 / input_global_scale).to(torch.float32), requires_grad=False) - - for attr in ["input_scale"]: - if hasattr(layer, attr): - delattr(layer, attr) - - -def _marlin_repack_experts(packed, perm, size_k, size_n, num_experts): - import vllm._custom_ops as ops - - result = [] - for i in range(num_experts): - qweight = packed[i].view(torch.int32).T.contiguous() - result.append( - ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=size_k, - size_n=size_n, - num_bits=4, - ) - ) - return torch.stack(result) - - -def _marlin_process_scales_experts(scale_hf, param_dtype, size_k, size_n, group_size, num_experts): - from vllm.model_executor.layers.quantization.utils.marlin_utils import marlin_permute_scales - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import nvfp4_marlin_process_scales - - result = [] - scales = scale_hf.to(param_dtype) - for i in range(num_experts): - s = marlin_permute_scales(s=scales[i].T, size_k=size_k, size_n=size_n, group_size=group_size) - result.append(_unwrap_marlin_scale(nvfp4_marlin_process_scales(s))) - return torch.stack(result) - - -_MOE_HF_PARAMS = [ - "w13_weight", - "w2_weight", - "w13_weight_scale", - "w2_weight_scale", - "w13_weight_scale_2", - "w2_weight_scale_2", - "w13_input_scale", - "w2_input_scale", -] - - -def _modelopt_moe_marlin_convert(self, layer: torch.nn.Module, is_first_call: bool) -> None: - from vllm.model_executor.layers.quantization import modelopt as vllm_modelopt - - prepare_for_marlin = vllm_modelopt.convert_to_nvfp4_moe_kernel_format.__globals__[ - "prepare_nvfp4_moe_layer_for_marlin" - ] - ( - w13_weight_marlin, - w13_weight_scale_marlin, - w13_scale_2_processed, - w2_weight_marlin, - w2_weight_scale_marlin, - w2_scale_2_processed, - ) = prepare_for_marlin( - layer=layer, - w13=layer.w13_weight.data, - w13_scale=layer.w13_weight_scale.data, - w13_scale_2=layer.w13_weight_scale_2.data, - w2=layer.w2_weight.data, - w2_scale=layer.w2_weight_scale.data, - w2_scale_2=layer.w2_weight_scale_2.data, - is_act_and_mul=self.moe.is_act_and_mul, - ) - - if is_first_call: - layer.w13_weight = Parameter(w13_weight_marlin, requires_grad=False) - layer.w2_weight = Parameter(w2_weight_marlin, requires_grad=False) - layer.w13_weight_scale = Parameter(w13_weight_scale_marlin, requires_grad=False) - layer.w2_weight_scale = Parameter(w2_weight_scale_marlin, requires_grad=False) - layer.w13_weight_scale_2 = Parameter(w13_scale_2_processed, requires_grad=False) - layer.w2_weight_scale_2 = Parameter(w2_scale_2_processed, requires_grad=False) - if not hasattr(layer, "_marlin_tensor_refs"): - layer._marlin_tensor_refs = {} - for rn in [ - "w13_weight", - "w2_weight", - "w13_weight_scale", - "w2_weight_scale", - "w13_weight_scale_2", - "w2_weight_scale_2", - ]: - layer._marlin_tensor_refs[rn] = getattr(layer, rn).data - else: - for rn, nd in [ - ("w13_weight", w13_weight_marlin), - ("w2_weight", w2_weight_marlin), - ("w13_weight_scale", w13_weight_scale_marlin), - ("w2_weight_scale", w2_weight_scale_marlin), - ("w13_weight_scale_2", w13_scale_2_processed), - ("w2_weight_scale_2", w2_scale_2_processed), - ]: - _update_ref_or_create(layer, rn, nd) - - layer.w13_input_scale = None - layer.w2_input_scale = None - - -def _modelopt_moe_process_weights(self, layer: torch.nn.Module) -> None: - """ - Replacement for ModelOptNvFp4FusedMoE.process_weights_after_loading (Marlin). - - First call: save metadata + weight_loaders, convert HF→Marlin format, - save _marlin_tensor_refs for CUDA Graph stability. - Subsequent: read reloaded HF data, convert, copy_ into saved refs. - """ - is_first_call = _check_first_call(layer) - - if is_first_call: - for pname in _MOE_HF_PARAMS: - _save_param_meta(layer, pname) - _save_weight_loaders(layer, _MOE_HF_PARAMS) - - w13_weight_scale_2 = layer.w13_weight_scale_2.data - if w13_weight_scale_2.dim() == 2: - w13_weight_scale_2 = w13_weight_scale_2[:, 0] - layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2, requires_grad=False) - - _modelopt_moe_marlin_convert(self, layer, is_first_call) - - from vllm.model_executor.layers.quantization.modelopt import make_nvfp4_moe_kernel - - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - assert self.experts_cls is not None - routing_tables = ( - layer._maybe_init_expert_routing_tables() if hasattr(layer, "_maybe_init_expert_routing_tables") else None - ) - self.moe_kernel = make_nvfp4_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - shared_experts=getattr(layer, "shared_experts", None), - routing_tables=routing_tables, - ) - self.moe_kernel.fused_experts.process_weights_after_loading(layer) - - -def _modelopt_kv_process_weights(self, layer) -> None: - """ - Replacement for BaseKVCacheMethod.process_weights_after_loading. - Doesn't delete k_scale, v_scale, q_scale, prob_scale to allow - for dynamic updates during refit. - """ - from vllm.platforms import current_platform - - if layer.kv_cache_dtype != "auto" and not layer.calculate_kv_scales: - if layer.k_scale > 0.0 and layer.v_scale > 0.0: - k_scale = layer.k_scale.to("cpu").tolist() - v_scale = layer.v_scale.to("cpu").tolist() - if current_platform.is_fp8_fnuz(): - k_scale *= 2 - v_scale *= 2 - elif layer.k_scale < 0.0 and layer.v_scale < 0.0: - k_scale = 1.0 - v_scale = 1.0 - else: - assert layer.k_scale > 0.0 - scale_to_duplicate = max(layer.k_scale, layer.v_scale) - k_scale = scale_to_duplicate.to("cpu").tolist() - v_scale = scale_to_duplicate.to("cpu").tolist() - if current_platform.is_fp8_fnuz(): - k_scale *= 2 - v_scale *= 2 - - if not isinstance(k_scale, float) or not isinstance(v_scale, float): - raise ValueError("Only support per-tensor scaling factor for fp8 KV cache") - - if layer.q_scale < 0.0: - layer._q_scale.copy_(k_scale) - layer._q_scale_float = k_scale - - layer._k_scale.copy_(k_scale) - layer._v_scale.copy_(v_scale) - layer._k_scale_float = k_scale - layer._v_scale_float = v_scale - - if layer.q_scale > 0.0: - q_scale = layer.q_scale - if current_platform.is_fp8_fnuz(): - q_scale *= 2 - layer.calculate_kv_scales = False - else: - q_scale = 1.0 - if layer.prob_scale > 0.0: - prob_scale = layer.prob_scale - if current_platform.is_fp8_fnuz(): - prob_scale *= 2 - else: - prob_scale = 1.0 - - is_singleton_float = ( - lambda x: isinstance(x, float) or isinstance(x, torch.Tensor) and x.numel() == 1 and x.is_floating_point() - ) - if not is_singleton_float(q_scale) or not is_singleton_float(prob_scale): - raise ValueError("Only support per-tensor scaling factor for fp8-quantized Q/prob") - - layer._q_scale.copy_(q_scale) - layer._q_scale_float = q_scale.item() if isinstance(q_scale, torch.Tensor) else q_scale - layer._prob_scale.copy_(prob_scale) - - -_patched = False -_original_dense_init = None -_original_moe_init = None - - -def _require_fp4_marlin_supported() -> None: - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import is_fp4_marlin_supported - - if not is_fp4_marlin_supported(): - raise RuntimeError("ModelOpt NVFP4 QAT weight reload requires vLLM FP4 Marlin support.") - - -def _modelopt_dense_init_marlin(self, quant_config) -> None: - _original_dense_init(self, quant_config) - _require_fp4_marlin_supported() - self.backend = "marlin" - - -def _modelopt_moe_init_marlin(self, quant_config, moe_config) -> None: - _original_moe_init(self, quant_config, moe_config) - _require_fp4_marlin_supported() - from vllm.model_executor.layers.fused_moe.fused_marlin_moe import MarlinExperts - from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend - - self.allow_flashinfer = False - self.cutlass_nvfp4_supported = False - self.flashinfer_moe_backend = None - self.use_marlin = True - self.backend = "marlin" - self.nvfp4_backend = NvFp4MoeBackend.MARLIN - self.experts_cls = MarlinExperts - self.use_global_sf = False - - -def prepare_modelopt_for_weight_reload(model, device=None): - """Prepare ModelOpt model for weight reloading. Call ONCE before each reload cycle.""" - inner_model = model - if hasattr(model, "model"): - inner_model = model.model - - param_meta = ModelOptParamMetaDict(inner_model, device=device) - - param_meta.prepare_for_reload() - - for layer_name, cache_entry in param_meta._layer_meta_cache.items(): - module = cache_entry["module"] - for param_name, pm in cache_entry["meta"].items(): - existing = getattr(module, param_name, None) - if existing is not None: - hf_shape = tuple(pm["shape"]) - hf_dtype = pm["dtype"] - if ( - tuple(existing.shape) == hf_shape - and existing.dtype == hf_dtype - and hasattr(existing, "weight_loader") - ): - continue - new_param = _create_param_from_meta(module, param_name, pm, device) - module.register_parameter(param_name, new_param) - - inner_model._param_meta_for_restore = param_meta - return param_meta - - -def modelopt_process_weights_after_loading(model): - """Trigger weight post-processing for all quantized layers after load_weights.""" - dense_count = 0 - moe_count = 0 - - actual_model = model - if hasattr(model, "model"): - actual_model = model.model - - for module in actual_model.modules(): - if hasattr(module, "scheme"): - module.scheme.process_weights_after_loading(module) - dense_count += 1 - - quant_method = getattr(module, "quant_method", None) - if quant_method is not None and not hasattr(module, "scheme"): - if hasattr(quant_method, "process_weights_after_loading"): - if "KVCache" in quant_method.__class__.__name__: - continue - quant_method.process_weights_after_loading(module) - moe_count += 1 - - return dense_count + moe_count - - -def apply_modelopt_nvfp4_patches(): - """Apply ModelOpt NVFP4 patches to support dynamic weight updates. Call before model loading.""" - global _patched - global _original_dense_init - global _original_moe_init - - if _patched: - return - - from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod - from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4FusedMoE, ModelOptNvFp4LinearMethod - - _original_dense_init = ModelOptNvFp4LinearMethod.__init__ - _original_moe_init = ModelOptNvFp4FusedMoE.__init__ - - ModelOptNvFp4LinearMethod.__init__ = _modelopt_dense_init_marlin - ModelOptNvFp4FusedMoE.__init__ = _modelopt_moe_init_marlin - ModelOptNvFp4LinearMethod.process_weights_after_loading = _modelopt_dense_process_weights - ModelOptNvFp4FusedMoE.process_weights_after_loading = _modelopt_moe_process_weights - BaseKVCacheMethod.process_weights_after_loading = _modelopt_kv_process_weights - - _patched = True diff --git a/verl/utils/qat/__init__.py b/verl/utils/qat/__init__.py index 6f2a85c814d..fc62bfac3e7 100644 --- a/verl/utils/qat/__init__.py +++ b/verl/utils/qat/__init__.py @@ -15,21 +15,29 @@ """ QAT (Quantization-Aware Training) module for verl. -Supports NVFP4 (W4A4 and W4A16) quantization modes for FSDP training. +Provides compressed-tensors rollout presets and trainer-side export helpers. Module Structure: - core.py: QATConfig, apply_qat, enable_qat_fuse (training setup) - linear.py: QATLinear layer with Triton kernels for fake quantization - quantizer.py: QATQuantizer for true quantization + scale computation utilities -- vllm_patch.py: Patches for vLLM dynamic weight loading +- compressed_tensors.py: compressed-tensors rollout config and trainer-side weight quantizer +- modelopt.py: ModelOpt fake quantization and Megatron QAT workflow helpers Usage: from verl.utils.qat import apply_qat, QATConfig - config = QATConfig(enable=True, mode="w4a16") + config = QATConfig(enable=True, mode="w4a4") model = apply_qat(model, config) # Before FSDP wrapping """ +from verl.utils.qat.compressed_tensors import ( + CompressedTensorsWeightQuantizer, + Fp8ConfigWeightQuantizer, + build_compressed_tensors_config, + dispatch_rollout_weight_quantizer, + qat_modes_for_config, +) from verl.utils.qat.core import ( QATConfig, apply_qat, @@ -37,21 +45,16 @@ invalidate_all_scales, load_quantization_config, ) -from verl.utils.qat.vllm_patch import ( - apply_qat_patches, - manual_process_weights_after_loading, - prepare_qat_for_load_weights, -) __all__ = [ - # Core "QATConfig", "apply_qat", "load_quantization_config", "enable_qat_fuse", "invalidate_all_scales", - # vLLM Patch - "apply_qat_patches", - "manual_process_weights_after_loading", - "prepare_qat_for_load_weights", + "CompressedTensorsWeightQuantizer", + "Fp8ConfigWeightQuantizer", + "build_compressed_tensors_config", + "dispatch_rollout_weight_quantizer", + "qat_modes_for_config", ] diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py new file mode 100644 index 00000000000..3408132742c --- /dev/null +++ b/verl/utils/qat/compressed_tensors.py @@ -0,0 +1,583 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Low-precision rollout expressed as a compressed-tensors ``quantization_config``. + +The rollout config either selects a compressed-tensors preset scheme per layer kind:: + + quantization: + linear: FP8_BLOCK # dense linear layers (attention, dense MLP, shared experts) + moe: FP8_BLOCK # routed experts + ignore: [] # extra "re:" patterns kept in high precision + +or provides a complete ``quantization_config`` JSON via ``quantization.config_file``. + +``RolloutQuantizationConfig.build_quantization_config`` (via +``build_compressed_tensors_config``) turns that into the HF ``quantization_config`` +the inference engine is started with (``load_format=dummy``), and +``CompressedTensorsWeightQuantizer`` quantizes the trainer's exported HF weights +with exactly the same targets and ignore list, using the compressed-tensors +compressors, so every refit ships tensors in the layout the engine loads. +""" + +import json +import re +from collections.abc import Iterable, Iterator, Mapping +from typing import TYPE_CHECKING, Any, Optional + +import torch + +# compressed-tensors ships with the inference engines rather than verl, so it is +# imported lazily to keep this module importable without them. +if TYPE_CHECKING: + from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme + +__all__ = [ + "MOE_TARGET", + "SUPPORTED_SCHEMES", + "CompressedTensorsWeightQuantizer", + "Fp8ConfigWeightQuantizer", + "dispatch_rollout_weight_quantizer", + "build_compressed_tensors_config", + "qat_modes_for_config", +] + +# Matched against vLLM's unfused expert names (``...mlp.experts.0.gate_proj``), +# which is how compressed-tensors resolves fused MoE layers. +MOE_TARGET = r"re:.*mlp\.experts\..*" + +# Presets whose rollout refit path (weight reload + process_weights_after_loading) +# has been validated, mapped to the ModelOpt QAT mode that fake-quantizes identically. +SUPPORTED_SCHEMES: dict[str, Optional[str]] = { + "FP8_BLOCK": "fp8", + "MXFP8": "mxfp8", + # There is no compressed-tensors preset for this combination. verl's + # MXFP4 mode deliberately means MXFP4 expert weights with dynamically + # quantized MXFP8 expert inputs, rather than the upstream MXFP4 W4A4 + # preset. + "MXFP4": "mxfp4", + # NVFP4 is deliberately W4A4 and routed-expert-only. Unlike MX formats, + # its tensor-level activation scale is static and must come from QAT. + "NVFP4": "nvfp4", +} + +# Layers that stay in high precision regardless of model. Patterns must start +# with "re:": vLLM rewrites plain dotted names through the model's +# WeightsMapper, but leaves regexes alone, so these match on both sides. +DEFAULT_IGNORE = [ + r"re:.*lm_head$", + r"re:.*embed_tokens$", + r"re:.*visual\..*", + r"re:^mtp\..*", + r"re:.*mlp\.gate$", + r"re:.*shared_expert_gate$", +] + + +_HIGH_PRECISION_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + + +def _check_scheme(kind: str, scheme: Optional[str]) -> Optional[str]: + if scheme is None: + return None + scheme = scheme.upper() + if scheme not in SUPPORTED_SCHEMES: + raise ValueError(f"Unsupported rollout quantization {kind}={scheme!r}, choose from {sorted(SUPPORTED_SCHEMES)}") + return scheme + + +_NUMERIC_FIELDS = { + "num_bits", + "type", + "symmetric", + "group_size", + "strategy", + "block_structure", + "dynamic", + "scale_dtype", +} + + +def _numerics(args: Optional["QuantizationArgs"]) -> Optional[dict[str, Any]]: + return None if args is None else args.model_dump(include=_NUMERIC_FIELDS) + + +def _preset_of(scheme: "QuantizationScheme") -> Optional[str]: + """The supported preset ``scheme`` quantizes like, ignoring observers and targets.""" + + for preset in SUPPORTED_SCHEMES: + reference = _scheme_for_preset(preset, scheme.targets) + if all( + _numerics(getattr(scheme, kind)) == _numerics(getattr(reference, kind)) + for kind in ("weights", "input_activations", "output_activations") + ): + return preset + return None + + +def _group_schemes(ct_config: Mapping[str, Any]) -> list["QuantizationScheme"]: + from compressed_tensors.quantization import QuantizationScheme + + return [ + QuantizationScheme.model_validate({k: v for k, v in group.items() if k != "format"}) + for group in ct_config["config_groups"].values() + ] + + +def qat_modes_for_config(ct_config: Mapping[str, Any]) -> set[Optional[str]]: + """ModelOpt QAT modes whose fake quant matches the config groups of a rollout ``quantization_config``.""" + presets = {_preset_of(scheme) for scheme in _group_schemes(ct_config)} + if presets == {"MXFP4", "MXFP8"}: + return {"mxfp4_mxfp8"} + return {SUPPORTED_SCHEMES[preset] for preset in presets} + + +def _matches(name: str, targets: Iterable[str]) -> bool: + for target in targets: + if target.startswith("re:"): + if re.match(target[3:], name): + return True + elif name == target: + return True + return False + + +def _scheme_for_preset(preset: str, targets: list[str]) -> "QuantizationScheme": + """Resolve verl presets, including the MXFP4-weight/MXFP8-input combination.""" + from compressed_tensors.quantization import preset_name_to_scheme + + scheme = preset_name_to_scheme(preset, targets) + if preset == "MXFP4": + mxfp8 = preset_name_to_scheme("MXFP8", targets) + scheme = scheme.model_copy(update={"input_activations": mxfp8.input_activations}) + return scheme + + +def _config_group(targets: list[str], preset: str) -> dict[str, Any]: + from compressed_tensors.compressors.format import infer_module_format + + scheme = _scheme_for_preset(preset, targets) + group = scheme.model_dump(mode="json", exclude_none=True) + group["format"] = infer_module_format(torch.nn.Linear, scheme).value + return group + + +def _load_config_file(config_file: str) -> dict[str, Any]: + """Read a compressed-tensors ``quantization_config`` (bare, or inside an HF ``config.json``).""" + with open(config_file) as f: + config = json.load(f) + config = config.get("quantization_config", config) + if config.get("quant_method") != "compressed-tensors" or not config.get("config_groups"): + raise ValueError(f"{config_file} is not a compressed-tensors quantization_config with config_groups") + for name, scheme in zip(config["config_groups"], _group_schemes(config), strict=True): + if _preset_of(scheme) is None: + raise ValueError( + f"{config_file}: config group {name!r} does not match any supported scheme {sorted(SUPPORTED_SCHEMES)}" + ) + return config + + +def build_compressed_tensors_config( + *, + linear: Optional[str], + moe: Optional[str], + ignore: Iterable[str], + config_file: Optional[str], + model_path: Optional[str] = None, +) -> dict[str, Any]: + """Build the HF compressed-tensors ``quantization_config`` for the rollout engine. + + Prefer ``RolloutQuantizationConfig.build_quantization_config``; the arguments mirror its fields. + """ + ignore = list(ignore or []) + if config_file: + if linear or moe or ignore: + raise ValueError("rollout.quantization.config_file is exclusive with linear / moe / ignore") + return _load_config_file(config_file) + + linear = _check_scheme("linear", linear) + moe = _check_scheme("moe", moe) + if linear in ("MXFP4", "NVFP4"): + raise ValueError(f"{linear} is supported only for routed MoE experts; use BF16 or MXFP8 dense linear layers") + if linear is None and moe is None: + raise ValueError("rollout quantization needs at least one of `linear` / `moe`") + + ignore = DEFAULT_IGNORE + ignore + + config_groups: dict[str, Any] = {} + if linear == moe: + config_groups["group_0"] = _config_group(["Linear"], linear) + else: + # A regex target wins over the class name, so experts always get the MoE scheme. + if moe is not None: + config_groups["group_moe"] = _config_group([MOE_TARGET], moe) + else: + ignore.append(MOE_TARGET) + if linear is not None: + config_groups["group_linear"] = _config_group(["Linear"], linear) + + formats = {g["format"] for g in config_groups.values()} + return { + "quant_method": "compressed-tensors", + "format": formats.pop() if len(formats) == 1 else "mixed-precision", + "quantization_status": "compressed", + "config_groups": config_groups, + "ignore": ignore, + "kv_cache_scheme": None, + "sparsity_config": {}, + "global_compression_ratio": None, + } + + +def _weight_min_max(weight: torch.Tensor, args: "QuantizationArgs") -> tuple[torch.Tensor, torch.Tensor]: + """Per-tile min/max of a 2D weight, shaped like the scale the strategy expects.""" + from compressed_tensors.quantization import QuantizationStrategy + from compressed_tensors.quantization.utils import maybe_pad_tensor_for_block_quant + + if args.strategy == QuantizationStrategy.BLOCK: + rows, cols = args.block_structure + padded = maybe_pad_tensor_for_block_quant(weight, (rows, cols)) + tiles = padded.unflatten(1, (-1, cols)).unflatten(0, (-1, rows)) + return tiles.amin(dim=(1, 3)), tiles.amax(dim=(1, 3)) + if args.strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): + groups = weight.unflatten(-1, (-1, args.group_size)) + return groups.amin(dim=-1), groups.amax(dim=-1) + if args.strategy == QuantizationStrategy.CHANNEL: + return weight.amin(dim=-1, keepdim=True), weight.amax(dim=-1, keepdim=True) + if args.strategy == QuantizationStrategy.TENSOR: + return weight.amin(), weight.amax() + raise NotImplementedError(f"Trainer-side quantization not implemented for strategy {args.strategy}") + + +def _compress_weight( + weight: torch.Tensor, + scheme: "QuantizationScheme", + *, + global_scale: Optional[torch.Tensor] = None, +) -> dict[str, torch.Tensor]: + """Quantize one module's weight into its compressed-tensors state dict.""" + from compressed_tensors.compressors import BaseCompressor + from compressed_tensors.compressors.format import infer_module_format + from compressed_tensors.quantization import QuantizationStrategy + from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam + + args = scheme.weights + min_vals, max_vals = (t.float() for t in _weight_min_max(weight, args)) + state_dict = {"weight": weight} + if args.strategy == QuantizationStrategy.TENSOR_GROUP: + if global_scale is None: + global_scale = generate_gparam(min_vals.amin(), max_vals.amax()) + state_dict["weight_global_scale"] = global_scale + scale, zero_point = calculate_qparams(min_vals, max_vals, args, global_scale=global_scale) + state_dict["weight_scale"] = scale + if not args.symmetric: + state_dict["weight_zero_point"] = zero_point + compressor = BaseCompressor.get_value_from_registry(infer_module_format(torch.nn.Linear, scheme).value) + return compressor.compress(state_dict, scheme) + + +# HF checkpoints that pack all routed experts into one 3D tensor (e.g. Qwen3.5 MoE). +_PACKED_EXPERTS = re.compile(r"(?P.*\.experts)\.(?Pgate_up_proj|down_proj)(?:\.weight)?$") + + +def _unpack_experts(name: str, tensor: torch.Tensor) -> Iterator[tuple[str, torch.Tensor]]: + """``experts.gate_up_proj`` [E, 2I, H] / ``experts.down_proj`` [E, H, I] -> per-expert 2D ``*.weight``. + + The engines load per-expert names with their scales, while their packed + 3D loaders only map the weight itself. + """ + match = _PACKED_EXPERTS.match(name) + prefix, proj = match["prefix"], match["proj"] + for expert_id, weight in enumerate(tensor): + if proj == "gate_up_proj": + gate, up = weight.chunk(2, dim=0) + yield f"{prefix}.{expert_id}.gate_proj.weight", gate + yield f"{prefix}.{expert_id}.up_proj.weight", up + else: + yield f"{prefix}.{expert_id}.down_proj.weight", weight + + +class CompressedTensorsWeightQuantizer: + """Quantize exported HF weights to match a compressed-tensors rollout config. + + A weight is quantized iff it is a 2D ``*.weight`` whose module is not + ignored and is targeted by a config group, resolved like compressed-tensors + does: a name / "re:" target wins over the ``Linear`` class target. Packed 3D + experts are split per expert first. Everything else passes through unchanged. + """ + + def __init__( + self, + ct_config: Mapping[str, Any], + input_global_scales: Optional[Mapping[str, torch.Tensor]] = None, + ): + self.ignore = list(ct_config.get("ignore") or []) + self.schemes = _group_schemes(ct_config) + self.linear_scheme = next((s for s in self.schemes if "Linear" in s.targets), None) + self.input_global_scales = dict(input_global_scales or {}) + + def set_input_global_scales(self, scales: Mapping[str, torch.Tensor]) -> None: + """Install ModelOpt-observed NVFP4 activation scales keyed by HF module name.""" + self.input_global_scales = dict(scales) + + def refresh_input_global_scales(self, model: torch.nn.Module) -> None: + """Refresh NVFP4 activation scales from ModelOpt before an export.""" + from verl.utils.qat.modelopt import collect_nvfp4_input_global_scales + + scales = collect_nvfp4_input_global_scales(model) + if scales: + self.set_input_global_scales(scales) + + def _input_global_scale(self, module: str, scheme: "QuantizationScheme") -> Optional[torch.Tensor]: + from compressed_tensors.quantization import QuantizationStrategy + + if scheme.input_activations is None or scheme.input_activations.strategy != QuantizationStrategy.TENSOR_GROUP: + return None + scale = self.input_global_scales.get(module) + if scale is None: + match = re.search(r"(?:^|\.)layers\.(\d+)\.mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", module) + if match: + projection = "gate_up" if match.group(2) in ("gate_proj", "up_proj") else "down" + scale = self.input_global_scales.get(f"layers.{match.group(1)}.mlp.experts.{projection}") + if scale is None: + # The first rollout sync precedes any actor forward, so no observer + # value exists yet. One is the neutral CT divisor and leaves the + # dynamic group-16 FP8 local scales their full representable range. + scale = torch.ones(1) + scale = scale.detach().float().reshape(1) + if not torch.isfinite(scale).all() or (scale <= 0).any(): + raise ValueError(f"Invalid input_global_scale for {module!r}: {scale}") + return scale + + def _compress_packed_experts( + self, + name: str, + tensor: torch.Tensor, + scheme: "QuantizationScheme", + ) -> Iterator[tuple[str, torch.Tensor]]: + """Compress packed experts, sharing NVFP4 gate/up tensor scales.""" + from compressed_tensors.quantization import QuantizationStrategy + from compressed_tensors.quantization.utils import generate_gparam + + unpacked = list(_unpack_experts(name, tensor)) + shared_global_scales: dict[tuple[int, str], torch.Tensor] = {} + if scheme.weights.strategy == QuantizationStrategy.TENSOR_GROUP: + by_expert: dict[int, list[torch.Tensor]] = {} + for module_name, weight in unpacked: + match = re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj)\.weight$", module_name) + if match: + by_expert.setdefault(int(match.group(1)), []).append(weight) + for expert_id, weights in by_expert.items(): + amax = torch.stack([weight.float().abs().amax() for weight in weights]).amax() + shared_global_scales[(expert_id, "gate_up")] = generate_gparam(-amax, amax) + + for weight_name, weight in unpacked: + module = weight_name[: -len(".weight")] + match = re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj)$", module) + global_scale = shared_global_scales.get((int(match.group(1)), "gate_up")) if match else None + compressed = _compress_weight(weight, scheme, global_scale=global_scale) + for key, value in compressed.items(): + yield weight_name[: -len("weight")] + key, value + input_scale = self._input_global_scale(module, scheme) + if input_scale is not None: + yield weight_name[: -len("weight")] + "input_global_scale", input_scale + + def _module_scheme(self, module: str) -> Optional["QuantizationScheme"]: + if _matches(module, self.ignore): + return None + for scheme in self.schemes: + if _matches(module, scheme.targets): + return scheme + return self.linear_scheme + + def _scheme_for(self, name: str, tensor: torch.Tensor) -> Optional["QuantizationScheme"]: + if not name.endswith(".weight") or tensor.ndim != 2 or tensor.dtype not in _HIGH_PRECISION_DTYPES: + return None + return self._module_scheme(name[: -len(".weight")]) + + def __call__(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]: + for name, tensor in weights: + packed = _PACKED_EXPERTS.match(name) if tensor.ndim == 3 else None + scheme = self._module_scheme(f"{packed['prefix']}.0.gate_proj") if packed else None + if packed and scheme is not None: + yield from self._compress_packed_experts(name, tensor, scheme) + continue + scheme = self._scheme_for(name, tensor) + if scheme is None: + yield name, tensor + continue + prefix = name[: -len("weight")] + for key, value in _compress_weight(tensor, scheme).items(): + yield prefix + key, value + input_scale = self._input_global_scale(name[: -len(".weight")], scheme) + if input_scale is not None: + yield prefix + "input_global_scale", input_scale + + +class Fp8ConfigWeightQuantizer: + """Export trainer weights in vLLM's native blockwise ``Fp8Config`` layout. + + Unlike :class:`CompressedTensorsWeightQuantizer`, this keeps the rollout's + native ``Fp8Config`` and therefore emits ``weight_scale_inv`` rather than + compressed-tensors' ``weight_scale``. The quantization calculation is + shared with the compressed-tensors FP8_BLOCK exporter. + + DeepSeek-V4 uses this dispatcher entry too, but the quantizer passes through + its private mixed FP8/FP4 layout already exported by Megatron-Bridge. + """ + + def __init__(self, fp8_config: Mapping[str, Any], *, model_types: Iterable[Optional[str]] = ()): + quant_method = str(fp8_config.get("quant_method", "")).lower() + self.passthrough = bool({"deepseek_v4", "deepseek_v41"} & set(model_types)) + if self.passthrough: + return + activation_scheme = fp8_config.get("activation_scheme") + block_size = fp8_config.get("weight_block_size") + store_dtype = fp8_config.get("store_dtype") + if "fp8" not in quant_method: + raise ValueError(f"Expected an FP8 quantization config, got quant_method={quant_method!r}") + if activation_scheme != "dynamic" or block_size != [128, 128] or store_dtype is not None: + raise NotImplementedError( + "Trainer-side native Fp8Config export currently supports only standard " + "FP8_BLOCK (dynamic activations, 128x128 weights, no mixed expert store_dtype); " + f"got activation_scheme={activation_scheme!r}, weight_block_size={block_size!r}, " + f"store_dtype={store_dtype!r}" + ) + self.ignore = list(fp8_config.get("ignored_layers") or fp8_config.get("modules_to_not_convert") or []) + self.ignore_match_mode = fp8_config.get("ignored_layers_match_mode", "exact") + if self.ignore_match_mode not in ("exact", "substring", "suffix"): + raise ValueError(f"Unsupported ignored_layers_match_mode={self.ignore_match_mode!r}") + self.scheme = _scheme_for_preset("FP8_BLOCK", ["Linear"]) + + def _is_ignored(self, module: str) -> bool: + if _matches(module, DEFAULT_IGNORE): + return True + if self.ignore_match_mode == "exact": + return module in self.ignore + if self.ignore_match_mode == "substring": + return any(ignored in module for ignored in self.ignore) + return any(module == ignored or module.endswith(f".{ignored}") for ignored in self.ignore) + + def _compress(self, name: str, weight: torch.Tensor) -> Iterator[tuple[str, torch.Tensor]]: + prefix = name[: -len("weight")] + compressed = _compress_weight(weight, self.scheme) + yield name, compressed["weight"] + yield prefix + "weight_scale_inv", compressed["weight_scale"] + + def __call__(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]: + if self.passthrough: + # Megatron-Bridge already exported DeepSeek-V4's private FP8/FP4 + # layout. vLLM allocated the expert bytes as uint8, so only + # reinterpret the packed values and E8M0 scales before transfer. + for name, tensor in weights: + if ".experts." in name and tensor.dtype in (torch.int8, torch.float8_e8m0fnu): + tensor = tensor.view(torch.uint8) + yield name, tensor + return + for name, tensor in weights: + packed = _PACKED_EXPERTS.match(name) if tensor.ndim == 3 else None + if packed: + probe = f"{packed['prefix']}.0.gate_proj" + if self._is_ignored(probe): + yield name, tensor + else: + for weight_name, weight in _unpack_experts(name, tensor): + yield from self._compress(weight_name, weight) + continue + + if ( + not name.endswith(".weight") + or tensor.ndim != 2 + or tensor.dtype not in _HIGH_PRECISION_DTYPES + or self._is_ignored(name[: -len(".weight")]) + ): + yield name, tensor + continue + yield from self._compress(name, tensor) + + +def _as_quantization_config(config: Any) -> Optional[dict[str, Any]]: + if config is None: + return None + if hasattr(config, "to_dict"): + config = config.to_dict() + if isinstance(config, Mapping): + return dict(config) + raise TypeError(f"quantization_config must be a mapping, got {type(config).__name__}") + + +def dispatch_rollout_weight_quantizer( + *, + model_config: Any, + rollout_config: Any, + qat_config: Any = None, +) -> Optional[CompressedTensorsWeightQuantizer | Fp8ConfigWeightQuantizer]: + """Select the trainer-side exporter for native or overridden rollout quantization. + + The rollout engine keeps its native ``quantization_config`` unless + ``rollout.quantization`` explicitly overrides it. Both native and override + compressed-tensors configs use :class:`CompressedTensorsWeightQuantizer`; + native ``Fp8Config`` uses :class:`Fp8ConfigWeightQuantizer`. DeepSeek-V4 + still gets an FP8 quantizer object, but it passes through the private mixed + FP8/FP4 layout already produced by Megatron-Bridge. + """ + rollout_quantization = getattr(rollout_config, "quantization", None) + linear = getattr(rollout_quantization, "linear", None) + moe = getattr(rollout_quantization, "moe", None) + ignore = getattr(rollout_quantization, "ignore", None) or [] + config_file = getattr(rollout_quantization, "config_file", None) + has_override = bool(linear or moe or config_file) + + hf_config = model_config.hf_config + if has_override: + quantization_config = build_compressed_tensors_config( + linear=linear, + moe=moe, + ignore=ignore, + config_file=config_file, + model_path=model_config.local_path, + ) + else: + native = getattr(hf_config, "quantization_config", None) + if native is None: + native = getattr(getattr(hf_config, "text_config", None), "quantization_config", None) + quantization_config = _as_quantization_config(native) + + qat_mode = qat_config.mode if qat_config is not None and qat_config.enable else None + if quantization_config is None: + if qat_mode is not None: + raise ValueError( + "QAT needs a quantized rollout to sync to: use a model with a native " + "quantization_config or set rollout.quantization.linear / moe / config_file" + ) + return None + + quant_method = str(quantization_config.get("quant_method", "")).lower() + if quant_method == "compressed-tensors": + qat_modes = qat_modes_for_config(quantization_config) if qat_mode is not None else set() + quantizer = CompressedTensorsWeightQuantizer(quantization_config) + elif "fp8" in quant_method: + qat_modes = {"fp8"} + model_types = { + getattr(hf_config, "model_type", None), + getattr(getattr(hf_config, "text_config", None), "model_type", None), + } + quantizer = Fp8ConfigWeightQuantizer(quantization_config, model_types=model_types) + else: + raise NotImplementedError( + f"Trainer-side rollout weight export does not support native quant_method={quant_method!r}" + ) + + if qat_mode is not None and qat_modes != {qat_mode}: + raise ValueError(f"QAT mode {qat_mode!r} does not match rollout quantization (QAT modes {qat_modes})") + return quantizer diff --git a/verl/utils/qat/modelopt.py b/verl/utils/qat/modelopt.py new file mode 100644 index 00000000000..f904cc1aa9e --- /dev/null +++ b/verl/utils/qat/modelopt.py @@ -0,0 +1,355 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ModelOpt fake quantization and Megatron QAT workflow helpers.""" + +import copy + +import modelopt.torch.quantization as mtq +import torch +import torch.nn as nn +import torch.nn.functional as F +from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg +from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + is_registered_quant_backend, + register_quant_backend, +) + +_FP8_BLOCK_BACKEND = "verl_fp8_block" +_MXFP4_BACKEND = "verl_mxfp4" +_NVFP4_WEIGHT_BACKEND = "verl_nvfp4_weight" +_NVFP4_INPUT_BACKEND = "verl_nvfp4_input" +_FP8_MAX = torch.finfo(torch.float8_e4m3fn).max + + +class _FP8BlockFakeQuant(torch.autograd.Function): + """Dynamic FP8 E4M3 fake quantization over blocks of the last two dims, with straight-through grad.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, block_rows: int, block_cols: int) -> torch.Tensor: + shape = x.shape + x2d = x.reshape(-1, shape[-1]).float() + rows, cols = x2d.shape + pad_rows, pad_cols = -rows % block_rows, -cols % block_cols + xp = F.pad(x2d, (0, pad_cols, 0, pad_rows)) + blocks = xp.reshape(xp.shape[0] // block_rows, block_rows, xp.shape[1] // block_cols, block_cols) + scale = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp_min(1e-12) / _FP8_MAX + dq = (blocks / scale).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn).float() * scale + dq = dq.reshape(xp.shape)[:rows, :cols] + return dq.reshape(shape).to(x.dtype) + + @staticmethod + def backward(ctx, grad): + return grad, None, None + + +def _fp8_block_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: + block_rows, block_cols = quantizer.backend_extra_args["block_sizes"] + return _FP8BlockFakeQuant.apply(inputs, block_rows, block_cols) + + +class _Mxfp4FakeQuant(torch.autograd.Function): + """compressed-tensors-compatible MXFP4 fake quantization with STE.""" + + @staticmethod + def forward(ctx, x: torch.Tensor) -> torch.Tensor: + from compressed_tensors.quantization import QuantizationArgs + from compressed_tensors.quantization.lifecycle.forward import fake_quantize + from compressed_tensors.quantization.utils import calculate_qparams + + args = QuantizationArgs( + num_bits=4, + type="float", + symmetric=True, + group_size=32, + strategy="group", + dynamic=False, + scale_dtype="torch.uint8", + ) + shape = x.shape + x2d = x.reshape(-1, shape[-1]) + groups = x2d.unflatten(-1, (-1, 32)) + min_vals, max_vals = groups.amin(dim=-1), groups.amax(dim=-1) + scale, zero_point = calculate_qparams(min_vals.float(), max_vals.float(), args) + return fake_quantize(x2d, scale, zero_point, args).reshape(shape).to(x.dtype) + + @staticmethod + def backward(ctx, grad): + return grad + + +def _mxfp4_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: + return _Mxfp4FakeQuant.apply(inputs) + + +class _Nvfp4FakeQuant(torch.autograd.Function): + """compressed-tensors-compatible NVFP4 QDQ with an explicit global scale.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, global_scale: torch.Tensor) -> torch.Tensor: + from compressed_tensors.quantization import QuantizationArgs + from compressed_tensors.quantization.lifecycle.forward import fake_quantize + from compressed_tensors.quantization.utils import calculate_qparams + + args = QuantizationArgs( + num_bits=4, + type="float", + symmetric=True, + group_size=16, + strategy="tensor_group", + dynamic=False, + scale_dtype="torch.float8_e4m3fn", + ) + shape = x.shape + x2d = x.reshape(-1, shape[-1]) + groups = x2d.unflatten(-1, (-1, 16)) + min_vals, max_vals = groups.amin(dim=-1), groups.amax(dim=-1) + scale, zero_point = calculate_qparams( + min_vals.float(), + max_vals.float(), + args, + global_scale=global_scale.float(), + ) + return ( + fake_quantize( + x2d, + scale, + zero_point, + args, + global_scale=global_scale.float(), + ) + .reshape(shape) + .to(x.dtype) + ) + + @staticmethod + def backward(ctx, grad): + return grad, None + + +def _nvfp4_global_scale(inputs: torch.Tensor) -> torch.Tensor: + from compressed_tensors.quantization.utils import generate_gparam + + amax = inputs.detach().float().abs().amax() + return generate_gparam(-amax, amax).to(device=inputs.device) + + +def _nvfp4_weight_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: + if inputs.ndim > 2: + # TEGroupedLinear stores one expert matrix per leading entry. The + # compressed-tensors checkpoint likewise assigns one global scale per + # expert, so quantize each matrix independently. + return torch.stack( + [_Nvfp4FakeQuant.apply(weight, _nvfp4_global_scale(weight)) for weight in inputs], + dim=0, + ) + return _Nvfp4FakeQuant.apply(inputs, _nvfp4_global_scale(inputs)) + + +def _nvfp4_input_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: + # vLLM's W4A4 kernel needs one static tensor-level scale while the local + # group-16 scales remain dynamic. Keep a running maximum during QAT and use + # that exact scale for both fake quant and checkpoint export. + current = _nvfp4_global_scale(inputs) + previous = getattr(quantizer, "_verl_input_global_scale", None) + if previous is not None: + # Running max of amax == running min of the global divisor. + current = torch.minimum(previous.to(current.device), current) + quantizer._verl_input_global_scale = current.detach() + return _Nvfp4FakeQuant.apply(inputs, current) + + +if not is_registered_quant_backend(_FP8_BLOCK_BACKEND): + register_quant_backend(_FP8_BLOCK_BACKEND, _fp8_block_backend) +if not is_registered_quant_backend(_MXFP4_BACKEND): + register_quant_backend(_MXFP4_BACKEND, _mxfp4_backend) +if not is_registered_quant_backend(_NVFP4_WEIGHT_BACKEND): + register_quant_backend(_NVFP4_WEIGHT_BACKEND, _nvfp4_weight_backend) +if not is_registered_quant_backend(_NVFP4_INPUT_BACKEND): + register_quant_backend(_NVFP4_INPUT_BACKEND, _nvfp4_input_backend) + +# Mirrors the compressed-tensors "fp8" rollout scheme: 128x128 weight blocks and +# dynamic 1x128 activation groups, both with fp32 scales. Only linear GEMMs the +# rollout quantizes may be enabled; anything else (attention BMMs, GDN conv1d, +# vision tower) would make training see a different model than the rollout. +_FP8_BLOCK_QUANTIZER_CFG = { + "*": {"enable": False}, + "*weight_quantizer": { + "num_bits": (4, 3), + "backend": _FP8_BLOCK_BACKEND, + "backend_extra_args": {"block_sizes": [128, 128]}, + "enable": True, + }, + "*input_quantizer": { + "num_bits": (4, 3), + "backend": _FP8_BLOCK_BACKEND, + "backend_extra_args": {"block_sizes": [1, 128]}, + "enable": True, + }, + "*conv1d*": {"enable": False}, + "*vision_model*": {"enable": False}, +} + +_QUANTIZER_CFGS = {"fp8": _FP8_BLOCK_QUANTIZER_CFG} +_SUPPORTED_QAT_MODES = {*_QUANTIZER_CFGS, "mxfp8", "mxfp4", "mxfp4_mxfp8", "nvfp4"} + +_ROUTED_EXPERT_WEIGHT_QUANTIZER = "*mlp.experts*weight_quantizer" +_ROUTED_EXPERT_INPUT_QUANTIZER = "*mlp.experts*input_quantizer" + + +def _modelopt_quantizer_cfg(preset: dict, quantizer_name: str) -> dict: + """Copy one quantizer's numeric config from a native ModelOpt preset.""" + entries = mtq.normalize_quant_cfg_list(copy.deepcopy(preset["quant_cfg"])) + return copy.deepcopy(next(entry["cfg"] for entry in entries if entry["quantizer_name"] == quantizer_name)) + + +def _mxfp4_quantize_config(*, dense_mxfp8: bool) -> tuple[list[dict], None]: + """MXFP4 routed-expert weights with MXFP8 activations. + + ModelOpt has native MXFP4/MXFP4 and MXFP8/MXFP8 presets but no MXFP4/MXFP8 + preset. Their quantizers are independent, so combine the weight numerics + from the former with the input numerics from the latter. + """ + if dense_mxfp8: + quant_cfg = mtq.normalize_quant_cfg_list(copy.deepcopy(mtq.MXFP8_DEFAULT_CFG["quant_cfg"])) + else: + quant_cfg = [{"quantizer_name": "*", "enable": False}] + + weight_cfg = _modelopt_quantizer_cfg(mtq.MXFP4_DEFAULT_CFG, "*weight_quantizer") + # ModelOpt and compressed-tensors use different E2M1 tie handling for a + # small subset of values. Export is compressed-tensors, so QAT must use the + # same QDQ implementation to remain numerically identical. + weight_cfg["backend"] = _MXFP4_BACKEND + quant_cfg.extend( + [ + { + "quantizer_name": _ROUTED_EXPERT_WEIGHT_QUANTIZER, + "cfg": weight_cfg, + }, + { + "quantizer_name": _ROUTED_EXPERT_INPUT_QUANTIZER, + "cfg": _modelopt_quantizer_cfg(mtq.MXFP8_DEFAULT_CFG, "*input_quantizer"), + }, + ] + ) + return quant_cfg, None + + +def _nvfp4_quantize_config() -> tuple[list[dict], None]: + """NVFP4 W4A4 on routed experts with export-compatible running scales.""" + quant_cfg = mtq.normalize_quant_cfg_list(copy.deepcopy(mtq.NVFP4_EXPERTS_ONLY_CFG["quant_cfg"])) + for entry in quant_cfg: + cfg = entry.get("cfg") + if cfg is None or not entry.get("enable", True): + continue + name = entry["quantizer_name"] + if name.endswith("weight_quantizer"): + cfg["backend"] = _NVFP4_WEIGHT_BACKEND + elif name.endswith("input_quantizer"): + cfg["backend"] = _NVFP4_INPUT_BACKEND + return quant_cfg, None + + +def collect_nvfp4_input_global_scales(model: nn.Module | list[nn.Module]) -> dict[str, torch.Tensor]: + """Collect observed expert activation scales under HF-style canonical keys.""" + scales = {} + chunks = model if isinstance(model, list) else [model] + for chunk in chunks: + for name, module in chunk.named_modules(): + scale = getattr(module, "_verl_input_global_scale", None) + if scale is None or not name.endswith("input_quantizer") or ".mlp.experts" not in name: + continue + layer = next((part for part in name.split(".") if part.isdigit()), None) + projection = "gate_up" if "linear_fc1" in name else "down" if "linear_fc2" in name else None + if layer is not None and projection is not None: + scales[f"layers.{layer}.mlp.experts.{projection}"] = scale.detach() + return scales + + +def _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: + cfg = [] + mapping = { + "lm_head": "*output_layer*", + "*mlp.gate": "*router*", + "*self_attn*": "*self_attention*", + } + for pattern in ignore_patterns: + key = mapping.get(pattern, pattern) + cfg.append({"quantizer_name": key, "enable": False}) + return cfg + + +def build_quantize_config( + qat_mode: str, + ignore_patterns: list[str] | None = None, +) -> dict: + """Build a complete ModelOpt quantization config for ``mtq.quantize``.""" + if qat_mode not in _SUPPORTED_QAT_MODES: + raise ValueError(f"Supported QAT modes are {sorted(_SUPPORTED_QAT_MODES)}, got: {qat_mode}") + + if ignore_patterns is None: + ignore_patterns = [] + + ignore_cfg = _ignore_patterns_to_quant_cfg(ignore_patterns) + + if qat_mode == "mxfp8": + # ModelOpt's native preset is exactly the compressed-tensors MXFP8 + # numerics: dynamic groups of 32, E4M3 values and E8M0 scales. + modelopt_config = copy.deepcopy(mtq.MXFP8_DEFAULT_CFG) + quant_cfg = mtq.normalize_quant_cfg_list(modelopt_config["quant_cfg"]) + algorithm = modelopt_config["algorithm"] + elif qat_mode in ("mxfp4", "mxfp4_mxfp8"): + quant_cfg, algorithm = _mxfp4_quantize_config(dense_mxfp8=qat_mode == "mxfp4_mxfp8") + elif qat_mode == "nvfp4": + quant_cfg, algorithm = _nvfp4_quantize_config() + else: + quant_cfg = mtq.normalize_quant_cfg_list(copy.deepcopy(_QUANTIZER_CFGS[qat_mode])) + disabled_cfg = copy.deepcopy(_default_disabled_quantizer_cfg) + if isinstance(disabled_cfg, dict): + disabled_cfg = mtq.normalize_quant_cfg_list(disabled_cfg) + quant_cfg.extend(disabled_cfg) + algorithm = None + quant_cfg.extend(ignore_cfg) + return {"quant_cfg": quant_cfg, "algorithm": algorithm} + + +def apply_qat( + model: nn.Module, + qat_mode: str, + ignore_patterns: list[str] | None = None, +) -> nn.Module: + """Apply Quantization-Aware Training to a Megatron model.""" + config = build_quantize_config(qat_mode, ignore_patterns) + mtq.quantize(model, config) + return model + + +def _get_qat_field(qat_config, key, default=None): + """Extract a field from qat_config, supporting both dict and object-style access.""" + if isinstance(qat_config, dict): + return qat_config.get(key, default) + return getattr(qat_config, key, default) + + +def apply_qat_to_modules(modules, qat_config): + """Apply ModelOpt fake quantization to a list of Megatron module chunks.""" + qat_mode = _get_qat_field(qat_config, "mode", "fp8") + ignore_patterns = _get_qat_field(qat_config, "ignore_patterns", None) + if ignore_patterns is not None: + ignore_patterns = list(ignore_patterns) + + for i in range(len(modules)): + modules[i] = apply_qat(modules[i], qat_mode, ignore_patterns=ignore_patterns) + return modules diff --git a/verl/utils/qat/vllm_patch.py b/verl/utils/qat/vllm_patch.py deleted file mode 100644 index 77d2ec17eeb..00000000000 --- a/verl/utils/qat/vllm_patch.py +++ /dev/null @@ -1,828 +0,0 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -vLLM NVFP4 Patches for Dynamic Weight Updates. - -Enables dynamic weight reloading for NVFP4 quantized models in vLLM. - -Supported schemes: -- Dense: W4A16-FP4, W4A4-FP4 -- MoE: NVFP4-MoE -""" - -import logging -import os -from typing import Optional -from unittest.mock import patch - -import torch -from torch.nn import Parameter - -from verl.utils.device import get_device_name - -logger = logging.getLogger(__name__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - - -class ParamMetaDict(dict): - """ - Dict-like class for parameter management with metadata-based rebuild and tensor swap. - - Supports: - - Rebuild of deleted parameters from saved metadata - - Tensor Swap for parameters with shape changes (address stability for CUDA Graph) - """ - - def __init__(self, model: torch.nn.Module, device: Optional[torch.device] = None): - """ - Initialize ParamMetaDict from a model. - - Args: - model: vLLM model (may be wrapped in ModelRunner) - device: Device for created parameters - """ - super().__init__() - self.device = device - - # Get the actual model (handle vLLM's wrapper structure) - actual_model = model - if hasattr(model, "model"): - actual_model = model.model - self._model = actual_model - - # Build mappings by scanning all modules - self._layer_meta_cache: dict[str, dict] = {} # Cache of _hf_param_meta - self._tensor_swap_layers: dict[str, dict] = {} # Layers needing tensor swap - - self._build_mappings() - - # Initialize with current parameters - for name, param in actual_model.named_parameters(): - self[name] = param - - def _build_mappings(self): - """Build layer metadata cache for rebuild and tensor swap.""" - for layer_name, module in self._model.named_modules(): - # Check for _hf_param_meta which indicates this layer has HF format params - if hasattr(module, "_hf_param_meta"): - self._layer_meta_cache[layer_name] = { - "module": module, - "meta": module._hf_param_meta, - } - - # Check for tensor swap layers (weight_scale with shape change) - if "weight_scale" in module._hf_param_meta: - marlin_refs = getattr(module, "_marlin_tensor_refs", {}) - if "weight_scale" in marlin_refs: - self._tensor_swap_layers[layer_name] = { - "module": module, - "marlin_ref": marlin_refs["weight_scale"], - "hf_meta": module._hf_param_meta["weight_scale"], - } - - # MoE layers (w13_weight_scale, w2_weight_scale) - if "w13_weight_scale" in module._hf_param_meta: - marlin_refs = getattr(module, "_marlin_tensor_refs", {}) - if "w13_weight_scale" in marlin_refs: - self._tensor_swap_layers[f"{layer_name}.w13"] = { - "module": module, - "param_name": "w13_weight_scale", - "marlin_ref": marlin_refs["w13_weight_scale"], - "hf_meta": module._hf_param_meta["w13_weight_scale"], - } - if "w2_weight_scale" in marlin_refs: - self._tensor_swap_layers[f"{layer_name}.w2"] = { - "module": module, - "param_name": "w2_weight_scale", - "marlin_ref": marlin_refs["w2_weight_scale"], - "hf_meta": module._hf_param_meta["w2_weight_scale"], - } - - def _try_rebuild(self, key: str) -> Optional[Parameter]: - """ - Try to rebuild a parameter from metadata if it was deleted. - - Args: - key: Full parameter name - - Returns: - Rebuilt parameter or None if cannot rebuild - """ - # Extract layer name and param name - parts = key.rsplit(".", 1) - if len(parts) != 2: - return None - - layer_name, param_name = parts - - # Check if we have metadata for this layer - if layer_name not in self._layer_meta_cache: - return None - - cache_entry = self._layer_meta_cache[layer_name] - module = cache_entry["module"] - meta = cache_entry["meta"] - - # Check if this param needs rebuild - if param_name not in meta: - return None - - # Already exists on module? - if hasattr(module, param_name): - param = getattr(module, param_name) - if param is not None: - return param - - # Rebuild from metadata - new_param = _create_param_from_meta(module, param_name, meta[param_name], self.device) - module.register_parameter(param_name, new_param) - return new_param - - def prepare_for_reload(self) -> None: - """Replace Marlin-format tensors with HF-shape tensors for reload.""" - for layer_name, swap_info in self._tensor_swap_layers.items(): - module = swap_info["module"] - param_name = swap_info.get("param_name", "weight_scale") - hf_meta = swap_info["hf_meta"] - if hasattr(module, param_name): - new_param = _create_param_from_meta(module, param_name, hf_meta, self.device) - setattr(module, param_name, new_param) - - def __getitem__(self, key: str) -> Parameter: - """Get parameter with rebuild support.""" - # Try standard lookup first - if key in dict.keys(self): - return super().__getitem__(key) - - # Try rebuild from metadata - param = self._try_rebuild(key) - if param is not None: - self[key] = param - return param - - raise KeyError(f"Parameter not found: {key}") - - def __contains__(self, key: str) -> bool: - """Check if parameter exists (with rebuild check).""" - if super().__contains__(key): - return True - - # Check if can rebuild from metadata - parts = key.rsplit(".", 1) - if len(parts) == 2: - layer_name, param_name = parts - if layer_name in self._layer_meta_cache: - meta = self._layer_meta_cache[layer_name]["meta"] - if param_name in meta: - return True - - return False - - def get(self, key: str, default=None): - """Get parameter with default.""" - try: - return self[key] - except KeyError: - return default - - -def _create_param_from_meta( - module: torch.nn.Module, - param_name: str, - meta: dict, - device: Optional[torch.device] = None, -) -> Parameter: - """Create a Parameter from saved metadata. Used by rebuild and tensor swap.""" - shape = meta["shape"] - dtype = meta["dtype"] - dev = device or meta.get("device", get_device_name()) - param_class = meta.get("param_class", Parameter) - - weight_loaders = getattr(module, "_weight_loaders", {}) - weight_loader = weight_loaders.get(param_name) - - data = torch.empty(shape, dtype=dtype, device=dev) - - try: - if param_class is not Parameter and weight_loader is not None: - kwargs = {"data": data, "weight_loader": weight_loader} - if "input_dim" in meta: - kwargs["input_dim"] = meta["input_dim"] - if "output_dim" in meta: - kwargs["output_dim"] = meta["output_dim"] - new_param = param_class(**kwargs) - else: - new_param = Parameter(data, requires_grad=False) - if weight_loader is not None: - new_param.weight_loader = weight_loader - except Exception as e: - logger.warning(f"Failed to create param {param_name} with class {param_class}: {e}, using Parameter") - new_param = Parameter(data, requires_grad=False) - if weight_loader is not None: - new_param.weight_loader = weight_loader - - if "quant_method" in meta: - new_param.quant_method = meta["quant_method"] - - return new_param - - -def save_param_meta(layer: torch.nn.Module, param_name: str): - """Save parameter metadata for rebuild.""" - if not hasattr(layer, "_hf_param_meta"): - layer._hf_param_meta = {} - - param = getattr(layer, param_name, None) - if param is None: - return - - meta = { - "shape": tuple(param.shape), - "dtype": param.dtype, - "device": str(param.device), - "param_class": type(param), # Save the actual parameter class - } - - # Save vLLM-specific attributes needed for reconstruction - if hasattr(param, "_input_dim"): - meta["input_dim"] = param._input_dim - if hasattr(param, "_output_dim"): - meta["output_dim"] = param._output_dim - - # Save MoE-specific attributes (quant_method is required by weight_loader) - if hasattr(param, "quant_method"): - meta["quant_method"] = param.quant_method - - layer._hf_param_meta[param_name] = meta - - -def _check_first_call(layer: torch.nn.Module) -> bool: - """Check if this is the first process_weights call, and increment counter.""" - count = getattr(layer, "_process_weights_call_count", 0) - layer._process_weights_call_count = count + 1 - return count == 0 - - -# Dense W4A16 Patches -def patched_w4a16_process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """Patched process_weights_after_loading for W4A16 Dense layer.""" - import vllm._custom_ops as ops - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - marlin_make_workspace_new, - marlin_permute_scales, - nvfp4_marlin_process_global_scale, - nvfp4_marlin_process_scales, - ) - - is_first_call = _check_first_call(layer) - - group_size = 16 - part_size_n = layer.output_size_per_partition - part_size_k = layer.input_size_per_partition - device = layer.weight_packed.device - param_dtype = getattr(layer, "params_dtype", torch.float16) - - # Save metadata (first call only) - if is_first_call: - save_param_meta(layer, "weight_packed") - save_param_meta(layer, "weight_global_scale") - save_param_meta(layer, "weight_scale") - if not hasattr(layer, "_weight_loaders"): - layer._weight_loaders = {} - for pname in ["weight_packed", "weight_global_scale", "weight_scale"]: - param = getattr(layer, pname, None) - if param is not None and hasattr(param, "weight_loader"): - layer._weight_loaders[pname] = param.weight_loader - - # Get HF format data - weight_packed_hf = layer.weight_packed.data - weight_global_scale_hf = layer.weight_global_scale.data - weight_scale_hf = layer.weight_scale.data - - # Create workspace (first call only) - if is_first_call: - layer.workspace = marlin_make_workspace_new(device) - - # Convert to Marlin format - perm = torch.empty(0, dtype=torch.int, device=device) - qweight = weight_packed_hf.view(torch.int32).T.contiguous() - marlin_weight = ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=part_size_k, - size_n=part_size_n, - num_bits=4, - is_a_8bit=False, - ) - - weight_scale = weight_scale_hf.T.contiguous().to(param_dtype) - weight_scale_permuted = marlin_permute_scales( - s=weight_scale, - size_k=part_size_k, - size_n=part_size_n, - group_size=group_size, - is_a_8bit=False, - ) - marlin_weight_scale = nvfp4_marlin_process_scales(weight_scale_permuted) - - weight_scale_2_raw = (1.0 / weight_global_scale_hf.max()).to(param_dtype) - marlin_weight_scale_2 = nvfp4_marlin_process_global_scale(weight_scale_2_raw) - - # Update compute parameters - if is_first_call: - layer.weight = Parameter(marlin_weight, requires_grad=False) - layer.weight_scale = Parameter(marlin_weight_scale, requires_grad=False) - layer.weight_scale_2 = Parameter(marlin_weight_scale_2, requires_grad=False) - if not hasattr(layer, "_marlin_tensor_refs"): - layer._marlin_tensor_refs = {} - layer._marlin_tensor_refs["weight_scale"] = layer.weight_scale.data - else: - layer.weight.data.copy_(marlin_weight) - layer.weight_scale_2.data.copy_(marlin_weight_scale_2) - marlin_scale_ref = layer._marlin_tensor_refs.get("weight_scale") - if marlin_scale_ref is not None: - marlin_scale_ref.copy_(marlin_weight_scale) - layer.weight_scale = Parameter(marlin_scale_ref, requires_grad=False) - else: - logger.warning("W4A16: _marlin_tensor_refs['weight_scale'] not found") - layer.weight_scale = Parameter(marlin_weight_scale, requires_grad=False) - - # Delete HF parameters - if hasattr(layer, "weight_packed"): - delattr(layer, "weight_packed") - if hasattr(layer, "weight_global_scale"): - delattr(layer, "weight_global_scale") - - -def patched_w4a4_process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """Patched process_weights_after_loading for W4A4 Dense (all backends).""" - from vllm.model_executor.layers.quantization.utils.quant_utils import swizzle_blockscale - - is_first_call = _check_first_call(layer) - - _W4A4_HF_PARAMS = ["weight_packed", "weight_scale", "weight_global_scale", "input_global_scale"] - - if is_first_call: - for pname in _W4A4_HF_PARAMS: - save_param_meta(layer, pname) - if not hasattr(layer, "_weight_loaders"): - layer._weight_loaders = {} - for pname in _W4A4_HF_PARAMS: - param = getattr(layer, pname, None) - if param is not None and hasattr(param, "weight_loader"): - layer._weight_loaders[pname] = param.weight_loader - - weight_packed_data = layer.weight_packed.data - weight_scale_data = layer.weight_scale.data - input_global_scale_data = layer.input_global_scale.data - weight_global_scale_data = layer.weight_global_scale.data - - global_input_scale = input_global_scale_data.max().to(torch.float32) - global_weight_scale = weight_global_scale_data.max().to(torch.float32) - - if self.backend == "flashinfer-trtllm": - from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a - - epilogue_tile_m = 128 - processed_weight = shuffle_matrix_a(weight_packed_data.view(torch.uint8), epilogue_tile_m) - processed_weight_scale = ( - shuffle_matrix_sf_a(weight_scale_data.view(torch.uint8), epilogue_tile_m) - .reshape(weight_scale_data.shape) - .view(torch.float8_e4m3fn) - ) - elif self.backend == "fbgemm": - processed_weight_scale = swizzle_blockscale(weight_scale_data).view(-1).view(torch.uint8) - processed_weight = weight_packed_data - else: - # cutlass / flashinfer-cutlass - processed_weight_scale = swizzle_blockscale(weight_scale_data) - processed_weight = weight_packed_data - - alpha = 1.0 / (global_input_scale * global_weight_scale) - - if is_first_call: - layer.weight_packed = Parameter(processed_weight, requires_grad=False) - layer.weight_scale = Parameter(processed_weight_scale, requires_grad=False) - layer.input_global_scale = Parameter(global_input_scale, requires_grad=False) - layer.weight_global_scale = Parameter(global_weight_scale, requires_grad=False) - layer.alpha = Parameter(alpha, requires_grad=False) - - if not hasattr(layer, "_marlin_tensor_refs"): - layer._marlin_tensor_refs = {} - layer._marlin_tensor_refs["weight_packed"] = layer.weight_packed.data - layer._marlin_tensor_refs["weight_scale"] = layer.weight_scale.data - layer._marlin_tensor_refs["input_global_scale"] = layer.input_global_scale.data - layer._marlin_tensor_refs["weight_global_scale"] = layer.weight_global_scale.data - layer._marlin_tensor_refs["alpha"] = layer.alpha.data - else: - refs = layer._marlin_tensor_refs - for ref_name, new_data in [ - ("weight_packed", processed_weight), - ("weight_scale", processed_weight_scale), - ("input_global_scale", global_input_scale), - ("weight_global_scale", global_weight_scale), - ("alpha", alpha), - ]: - ref = refs.get(ref_name) - if ref is not None: - ref.copy_(new_data) - setattr(layer, ref_name, Parameter(ref, requires_grad=False)) - else: - logger.warning(f"W4A4: _marlin_tensor_refs['{ref_name}'] not found, creating new Parameter") - setattr( - layer, - ref_name, - Parameter( - new_data.clone() if isinstance(new_data, torch.Tensor) else torch.tensor(new_data), - requires_grad=False, - ), - ) - - -def _marlin_repack_experts(packed, perm, size_k, size_n, num_experts): - """Repack weight for each expert into Marlin format and stack.""" - import vllm._custom_ops as ops - - result = [] - for i in range(num_experts): - qweight = packed[i].view(torch.int32).T.contiguous() - result.append( - ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=size_k, - size_n=size_n, - num_bits=4, - is_a_8bit=False, - ) - ) - return torch.stack(result) - - -def _marlin_process_scales_experts(scale_hf, param_dtype, size_k, size_n, group_size, num_experts): - """Process scales for each expert into Marlin format and stack.""" - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - marlin_permute_scales, - nvfp4_marlin_process_scales, - ) - - result = [] - scales = scale_hf.to(param_dtype) - for i in range(num_experts): - s = marlin_permute_scales( - s=scales[i].T, - size_k=size_k, - size_n=size_n, - group_size=group_size, - is_a_8bit=False, - ) - result.append(nvfp4_marlin_process_scales(s)) - return torch.stack(result) - - -def _process_nvfp4_moe_marlin(self, layer: torch.nn.Module, is_first_call: bool) -> None: - """Process MoE layer with MARLIN backend (W4A16).""" - from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import make_nvfp4_moe_kernel - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - marlin_make_workspace_new, - nvfp4_marlin_process_global_scale, - ) - - group_size = 16 - e = layer.num_experts - k = layer.hidden_size - n = layer.intermediate_size_per_partition - device = layer.w13_weight_packed.device - param_dtype = layer.params_dtype - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - - if is_first_call: - layer.workspace = marlin_make_workspace_new(device, 4) - - perm = torch.empty(0, dtype=torch.int, device=device) - - if self.moe.is_act_and_mul and not torch.allclose( - layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1] - ): - logger.warning("w1_weight_global_scale must match w3_weight_global_scale. Accuracy may be affected.") - - size_n_w13, size_k_w13 = n * w13_num_shards, k - size_n_w2, size_k_w2 = k, n - - w13_weight_marlin = _marlin_repack_experts(layer.w13_weight_packed.data, perm, size_k_w13, size_n_w13, e) - w2_weight_marlin = _marlin_repack_experts(layer.w2_weight_packed.data, perm, size_k_w2, size_n_w2, e) - w13_weight_scale_marlin = _marlin_process_scales_experts( - layer.w13_weight_scale.data, param_dtype, size_k_w13, size_n_w13, group_size, e - ) - w2_weight_scale_marlin = _marlin_process_scales_experts( - layer.w2_weight_scale.data, param_dtype, size_k_w2, size_n_w2, group_size, e - ) - - # Process global scales - w13_scale_2 = 1.0 / layer.w13_weight_global_scale[:, 0] - w2_scale_2 = 1.0 / layer.w2_weight_global_scale.data - w13_scale_2_processed = nvfp4_marlin_process_global_scale(w13_scale_2.to(param_dtype)) - w2_scale_2_processed = nvfp4_marlin_process_global_scale(w2_scale_2.to(param_dtype)) - - # Update parameters - if is_first_call: - layer.w13_weight = Parameter(w13_weight_marlin, requires_grad=False) - layer.w2_weight = Parameter(w2_weight_marlin, requires_grad=False) - layer.w13_weight_scale = Parameter(w13_weight_scale_marlin, requires_grad=False) - layer.w2_weight_scale = Parameter(w2_weight_scale_marlin, requires_grad=False) - layer.w13_weight_scale_2 = Parameter(w13_scale_2_processed, requires_grad=False) - layer.w2_weight_scale_2 = Parameter(w2_scale_2_processed, requires_grad=False) - if not hasattr(layer, "_marlin_tensor_refs"): - layer._marlin_tensor_refs = {} - layer._marlin_tensor_refs["w13_weight_scale"] = layer.w13_weight_scale.data - layer._marlin_tensor_refs["w2_weight_scale"] = layer.w2_weight_scale.data - else: - layer.w13_weight.data.copy_(w13_weight_marlin) - layer.w2_weight.data.copy_(w2_weight_marlin) - layer.w13_weight_scale_2.data.copy_(w13_scale_2_processed) - layer.w2_weight_scale_2.data.copy_(w2_scale_2_processed) - w13_marlin_ref = layer._marlin_tensor_refs.get("w13_weight_scale") - w2_marlin_ref = layer._marlin_tensor_refs.get("w2_weight_scale") - if w13_marlin_ref is not None: - w13_marlin_ref.copy_(w13_weight_scale_marlin) - layer.w13_weight_scale = Parameter(w13_marlin_ref, requires_grad=False) - else: - logger.warning("MoE: _marlin_tensor_refs['w13_weight_scale'] not found") - layer.w13_weight_scale.data.copy_(w13_weight_scale_marlin) - if w2_marlin_ref is not None: - w2_marlin_ref.copy_(w2_weight_scale_marlin) - layer.w2_weight_scale = Parameter(w2_marlin_ref, requires_grad=False) - else: - logger.warning("MoE: _marlin_tensor_refs['w2_weight_scale'] not found") - layer.w2_weight_scale.data.copy_(w2_weight_scale_marlin) - - layer.w13_input_scale = None - layer.w2_input_scale = None - - # Initialize kernel - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - if self.moe_quant_config is not None and ( - (not self.moe.moe_parallel_config.use_all2all_kernels) or self.moe.moe_parallel_config.use_naive_all2all_kernels - ): - self.kernel = make_nvfp4_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - ) - - -def _process_nvfp4_moe_flashinfer_cutlass(self, layer: torch.nn.Module, is_first_call: bool) -> None: - """Process MoE layer with FlashInfer/CUTLASS backend (W4A4).""" - from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( - convert_to_nvfp4_moe_kernel_format, - make_nvfp4_moe_kernel, - ) - from vllm.model_executor.utils import replace_parameter - - w13_packed = layer.w13_weight_packed.data - w2_packed = layer.w2_weight_packed.data - w13_scale_hf = layer.w13_weight_scale.data - w2_scale_hf = layer.w2_weight_scale.data - - if self.moe.is_act_and_mul and not torch.allclose( - layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1] - ): - logger.warning("w1_weight_global_scale must match w3_weight_global_scale. Accuracy may be affected.") - w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous() - - w13_temp = Parameter(w13_packed.clone(), requires_grad=False) - w2_temp = Parameter(w2_packed.clone(), requires_grad=False) - - if is_first_call: - layer.w13_weight = w13_temp - layer.w2_weight = w2_temp - - ( - w13, - w13_scale, - w13_scale_2, - a13_scale, - w2, - w2_scale, - w2_scale_2, - a2_scale, - ) = convert_to_nvfp4_moe_kernel_format( - nvfp4_backend=self.nvfp4_backend, - layer=layer, - w13=w13_temp, - w13_scale=w13_scale_hf, - w13_scale_2=(1.0 / w13_weight_global_scale), - a13_scale=(1.0 / layer.w13_input_global_scale), - w2=w2_temp, - w2_scale=w2_scale_hf, - w2_scale_2=(1.0 / layer.w2_weight_global_scale), - a2_scale=(1.0 / layer.w2_input_global_scale), - is_act_and_mul=self.moe.is_act_and_mul, - ) - - # Update parameters - if is_first_call: - replace_parameter(layer, "w13_weight", w13) - replace_parameter(layer, "w2_weight", w2) - layer.w13_weight_scale = Parameter(w13_scale, requires_grad=False) - layer.w2_weight_scale = Parameter(w2_scale, requires_grad=False) - if not hasattr(layer, "_marlin_tensor_refs"): - layer._marlin_tensor_refs = {} - layer._marlin_tensor_refs["w13_weight_scale"] = layer.w13_weight_scale.data - layer._marlin_tensor_refs["w2_weight_scale"] = layer.w2_weight_scale.data - else: - layer.w13_weight.data.copy_(w13.data) - layer.w2_weight.data.copy_(w2.data) - w13_scale_ref = layer._marlin_tensor_refs.get("w13_weight_scale") - w2_scale_ref = layer._marlin_tensor_refs.get("w2_weight_scale") - if w13_scale_ref is not None: - w13_scale_ref.copy_(w13_scale) - layer.w13_weight_scale = Parameter(w13_scale_ref, requires_grad=False) - else: - logger.warning("MoE W4A4: _marlin_tensor_refs['w13_weight_scale'] not found") - layer.w13_weight_scale.data.copy_(w13_scale) - if w2_scale_ref is not None: - w2_scale_ref.copy_(w2_scale) - layer.w2_weight_scale = Parameter(w2_scale_ref, requires_grad=False) - else: - logger.warning("MoE W4A4: _marlin_tensor_refs['w2_weight_scale'] not found") - layer.w2_weight_scale.data.copy_(w2_scale) - - layer.w13_weight_scale_2 = w13_scale_2 - layer.w2_weight_scale_2 = w2_scale_2 - layer.w13_input_scale = a13_scale - layer.w2_input_scale = a2_scale - - # Initialize kernel - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - if self.moe_quant_config is not None and ( - (not self.moe.moe_parallel_config.use_all2all_kernels) or self.moe.moe_parallel_config.use_naive_all2all_kernels - ): - self.kernel = make_nvfp4_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - ) - - -# MoE NVFP4 Patches (entry points) -def patched_nvfp4_moe_process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """Patched process_weights_after_loading for NVFP4 MoE layer.""" - from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend - - is_first_call = _check_first_call(layer) - - # Save metadata (first call only) - if is_first_call: - save_param_meta(layer, "w13_weight_packed") - save_param_meta(layer, "w2_weight_packed") - save_param_meta(layer, "w13_weight_scale") - save_param_meta(layer, "w2_weight_scale") - if not hasattr(layer, "_weight_loaders"): - layer._weight_loaders = {} - for pname in ["w13_weight_packed", "w2_weight_packed", "w13_weight_scale", "w2_weight_scale"]: - param = getattr(layer, pname, None) - if param is not None and hasattr(param, "weight_loader"): - layer._weight_loaders[pname] = param.weight_loader - - is_marlin = self.nvfp4_backend == NvFp4MoeBackend.MARLIN - if is_marlin: - _process_nvfp4_moe_marlin(self, layer, is_first_call) - else: - _process_nvfp4_moe_flashinfer_cutlass(self, layer, is_first_call) - - # Delete HF parameters - if hasattr(layer, "w13_weight_packed"): - delattr(layer, "w13_weight_packed") - if hasattr(layer, "w2_weight_packed"): - delattr(layer, "w2_weight_packed") - - -_PATCH_TARGETS = [ - # Dense W4A16 - ( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w4a16_nvfp4.CompressedTensorsW4A16Fp4.process_weights_after_loading", - patched_w4a16_process_weights_after_loading, - ), - # Dense W4A4 - ( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w4a4_nvfp4.CompressedTensorsW4A4Fp4.process_weights_after_loading", - patched_w4a4_process_weights_after_loading, - ), - # MoE NVFP4 - ( - "vllm.model_executor.layers.quantization.compressed_tensors." - "compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoEMethod.process_weights_after_loading", - patched_nvfp4_moe_process_weights_after_loading, - ), -] - -_applied_patches = [] - - -def apply_qat_patches(): - """Apply NVFP4 patches to support dynamic weight updates. Call before model loading.""" - global _applied_patches - - if _applied_patches: - logger.warning("QAT patches already applied, skipping") - return _applied_patches - - logger.info("Applying NVFP4 patches for dynamic weight loading...") - - for target, replacement in _PATCH_TARGETS: - p = patch(target, replacement) - _applied_patches.append(p) - p.start() - - logger.info(f"Applied {len(_applied_patches)} NVFP4 patches for dynamic weight loading") - return _applied_patches - - -def prepare_qat_for_load_weights(model, device=None): - """ - Prepare QAT model for weight loading. Call ONCE before multi-bucket weight loading. - - Args: - model: vLLM model - device: Device for created parameters - """ - inner_model = model - if hasattr(model, "model"): - inner_model = model.model - - param_meta = ParamMetaDict(inner_model, device=device) - - param_meta.prepare_for_reload() - logger.info(f"[prepare_qat] Tensor swap prepared for {len(param_meta._tensor_swap_layers)} layers") - - # Rebuild deleted (W4A16) or overwritten (W4A4) params back to HF format - rebuilt_count = 0 - for layer_name, cache_entry in param_meta._layer_meta_cache.items(): - module = cache_entry["module"] - for param_name, pm in cache_entry["meta"].items(): - existing = getattr(module, param_name, None) - if existing is not None: - hf_shape = tuple(pm["shape"]) - hf_dtype = pm["dtype"] - if ( - tuple(existing.shape) == hf_shape - and existing.dtype == hf_dtype - and hasattr(existing, "weight_loader") - ): - continue - new_param = _create_param_from_meta(module, param_name, pm, device) - module.register_parameter(param_name, new_param) - rebuilt_count += 1 - - logger.info(f"[prepare_qat] Rebuilt {rebuilt_count} parameters") - inner_model._param_meta_for_restore = param_meta - return param_meta - - -def manual_process_weights_after_loading(model): - """Trigger weight post-processing for all quantized layers after load_weights.""" - dense_count = 0 - moe_count = 0 - - actual_model = model - if hasattr(model, "model"): - actual_model = model.model - - for module in actual_model.modules(): - if hasattr(module, "scheme"): - module.scheme.process_weights_after_loading(module) - dense_count += 1 - - quant_method = getattr(module, "quant_method", None) - if quant_method is not None and not hasattr(module, "scheme"): - if hasattr(quant_method, "process_weights_after_loading"): - # Skip KV cache quantization methods - if "KVCache" in quant_method.__class__.__name__: - continue - quant_method.process_weights_after_loading(module) - moe_count += 1 - - logger.debug(f"Processed {dense_count} dense layers, {moe_count} MoE layers") - return dense_count + moe_count - - -__all__ = [ - "apply_qat_patches", - "prepare_qat_for_load_weights", - "manual_process_weights_after_loading", -] diff --git a/verl/utils/sglang/sglang_fp8_utils.py b/verl/utils/sglang/sglang_fp8_utils.py deleted file mode 100644 index 887d5dadd6e..00000000000 --- a/verl/utils/sglang/sglang_fp8_utils.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -import re -from collections.abc import Iterable -from typing import Any - -from verl.utils.fp8_utils import FP8QuantizerHelper - - -def _get_config_value(config: Any, key: str, default: Any = None) -> Any: - if config is None: - return default - get_value = getattr(config, "get", None) - if callable(get_value): - return get_value(key, default) - return getattr(config, key, default) - - -def _normalize_ignored_layers(ignored_layers: Any) -> list[str]: - if ignored_layers is None: - return [] - if isinstance(ignored_layers, str): - ignored_layers = ignored_layers.split(",") - elif not isinstance(ignored_layers, Iterable): - ignored_layers = [ignored_layers] - - normalized = [] - for layer in ignored_layers: - layer_name = str(layer).strip() - if layer_name: - normalized.append(layer_name) - return normalized - - -def _dedupe_layers(ignored_layers: Iterable[str]) -> list[str]: - seen = set() - deduped = [] - for layer in ignored_layers: - layer_lower = layer.lower() - if layer_lower in seen: - continue - seen.add(layer_lower) - deduped.append(layer) - return deduped - - -def _get_ignored_layers_from_env() -> list[str]: - return _normalize_ignored_layers(os.getenv("SGLANG_FP8_IGNORED_LAYERS")) - - -def get_sglang_fp8_ignored_layers(quant_config: Any = None) -> list[str]: - ignored_layers = [] - ignored_layers.extend(_normalize_ignored_layers(_get_config_value(quant_config, "ignored_layers"))) - ignored_layers.extend(_normalize_ignored_layers(_get_config_value(quant_config, "modules_to_not_convert"))) - ignored_layers.extend(_get_ignored_layers_from_env()) - return _dedupe_layers(ignored_layers) - - -def _matches_ignored_layer(param_name: str, ignored_layer: str) -> bool: - ignored_layer = ignored_layer.strip() - if not ignored_layer: - return False - - name = param_name.strip(".") - module_name = name[: -len(".weight")] if name.lower().endswith(".weight") else name - if ignored_layer.startswith("re:"): - pattern = ignored_layer[3:] - return any(re.match(pattern, candidate) for candidate in (name, module_name)) - - ignored_layer = ignored_layer.lower().strip(".") - name = name.lower() - module_name = module_name.lower() - for candidate in (name, module_name): - if candidate == ignored_layer: - return True - if candidate.startswith(f"{ignored_layer}."): - return True - if candidate.endswith(f".{ignored_layer}"): - return True - if f".{ignored_layer}." in f".{candidate}.": - return True - return False - - -def build_sglang_fp8_quant_config(hf_config: Any = None, ignored_layers: Any = None) -> dict[str, Any]: - """Build SGLang block-wise FP8 config shared by server init and weight sync.""" - fp8_quant_config = { - "activation_scheme": "dynamic", - "fmt": "e4m3", - "quant_method": "fp8", - "weight_block_size": [128, 128], - } - - hf_quant_config = _get_config_value(hf_config, "quantization_config") - merged_ignored_layers = get_sglang_fp8_ignored_layers(hf_quant_config) - merged_ignored_layers.extend(_normalize_ignored_layers(ignored_layers)) - merged_ignored_layers = _dedupe_layers(merged_ignored_layers) - if merged_ignored_layers: - fp8_quant_config["ignored_layers"] = merged_ignored_layers - - return fp8_quant_config - - -class SGLangFP8QuantizerHelper(FP8QuantizerHelper): - def __init__(self, quant_config): - super().__init__(quant_config) - self.ignored_layers = get_sglang_fp8_ignored_layers(quant_config) - - def should_quantize_param(self, param_name): - for ignored_layer in self.ignored_layers: - if _matches_ignored_layer(param_name, ignored_layer): - return False - return super().should_quantize_param(param_name) diff --git a/verl/utils/trtllm/trtllm_fp8_utils.py b/verl/utils/trtllm/trtllm_fp8_utils.py deleted file mode 100644 index 43c15351c06..00000000000 --- a/verl/utils/trtllm/trtllm_fp8_utils.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from verl.utils.fp8_utils import FP8QuantizerHelper - - -class TRTLLMFP8QuantizerHelper(FP8QuantizerHelper): - def __init__(self, quant_config): - super().__init__(quant_config) diff --git a/verl/utils/vllm/vllm_fp4_utils.py b/verl/utils/vllm/vllm_fp4_utils.py index 9120ee07710..184e96191ae 100644 --- a/verl/utils/vllm/vllm_fp4_utils.py +++ b/verl/utils/vllm/vllm_fp4_utils.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Refit support for vLLM's ``Mxfp4MoEMethod`` (DeepSeek V4 routed experts). +"""Refit support for vLLM's MXFP4 and NVFP4 routed experts. Reloading these experts means undoing the kernel's weight conversion: the refit stream carries checkpoint-layout tensors while the live parameters hold whatever @@ -32,6 +32,7 @@ """ import logging +from dataclasses import replace from unittest.mock import patch import torch @@ -40,6 +41,128 @@ _MXFP4_SF_BLOCK = 32 _MXFP4_LIVE_ATTR = "_verl_mxfp4_live_params" +_NVFP4_SF_BLOCK = 16 +_NVFP4_LIVE_ATTR = "_verl_nvfp4_live_params" + + +def install_mxfp4_mxfp8_compressed_tensors_patch(): + """Teach vLLM compressed-tensors to honor MXFP8 expert activations. + + vLLM 0.29 recognizes MXFP4 compressed-tensors weights but ignores their + input-activation scheme and unconditionally constructs its W4A4 method on + Blackwell. Reuse vLLM's canonical MXFP4 method so its oracle can select an + actual MXFP4×MXFP8 backend while retaining the compressed-tensors packed + parameter names expected by the loader. + """ + from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + kMxfp8Dynamic, + select_mxfp4_moe_backend, + ) + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( + CompressedTensorsMoEMethod, + compressed_tensors_moe_w4a4_mxfp4, + ) + from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod + + CompressedTensorsW4A4Mxfp4MoEMethod = compressed_tensors_moe_w4a4_mxfp4.CompressedTensorsW4A4Mxfp4MoEMethod + if getattr(CompressedTensorsMoEMethod, "_verl_mxfp4_mxfp8_patch", False): + return + + class CompressedTensorsW4A8Mxfp4MoEMethod(Mxfp4MoEMethod): + """Compressed-tensors MXFP4 weights with dynamic MXFP8 activations.""" + + _verl_uses_packed_checkpoint_names = True + + def __init__(self, moe): + FusedMoEMethodBase.__init__(self, moe) + self.weight_dtype = "mxfp4" + self.group_size = _MXFP4_SF_BLOCK + # TRTLLM's MXFP4×MXFP8 path is numerically incorrect for Qwen3.5 + # on vLLM 0.29, while FlashInfer CUTLASS matches the BF16 and + # DeepGEMM references. Keep explicit user choices intact, but make + # the verified backend the default for verl's MXFP4 preset. + selection_moe = replace(moe, moe_backend="flashinfer_cutlass") if moe.moe_backend == "auto" else moe + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(selection_moe, activation_key=kMxfp8Dynamic) + self.max_capture_size = moe.max_capture_size + self._cache_permute_indices = {} + self.moe_kernel = None + self.w13_precision_config = None + self.w2_precision_config = None + + def create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ): + CompressedTensorsW4A4Mxfp4MoEMethod.create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ) + self.num_experts = num_experts + self.intermediate_size = intermediate_size_per_partition + self.hidden_size = hidden_size + + def process_weights_after_loading(self, layer): + # The compressed-tensors loader uses ``*_packed`` only for the + # initial load. Refit staging restores the canonical names. + if hasattr(layer, "w13_weight_packed"): + layer.w13_weight = torch.nn.Parameter(layer.w13_weight_packed.data, requires_grad=False) + delattr(layer, "w13_weight_packed") + layer.w2_weight = torch.nn.Parameter(layer.w2_weight_packed.data, requires_grad=False) + delattr(layer, "w2_weight_packed") + return super().process_weights_after_loading(layer) + + original = CompressedTensorsMoEMethod.get_moe_method + + def get_moe_method(quant_config, layer, layer_name): + quant_config._add_fused_moe_to_target_scheme_map() + schemes = [ + quant_config.get_scheme_dict(layer, layer_name + projection) + for projection in (".0.gate_proj", ".0.up_proj", ".0.down_proj") + ] + scheme = schemes[-1] + if not all(candidate == scheme for candidate in schemes): + return original(quant_config, layer, layer_name) + weight_quant = None if scheme is None else scheme.get("weights") + input_quant = None if scheme is None else scheme.get("input_activations") + if quant_config._is_mxfp4(weight_quant) and quant_config._is_mxfp8(input_quant): + logger.info("Using MXFP4 weights with dynamic MXFP8 activations for compressed-tensors MoE") + return CompressedTensorsW4A8Mxfp4MoEMethod(layer.moe_config) + return original(quant_config, layer, layer_name) + + CompressedTensorsMoEMethod.get_moe_method = staticmethod(get_moe_method) + CompressedTensorsMoEMethod._verl_mxfp4_mxfp8_patch = True + + +def install_nvfp4_cutlass_compressed_tensors_patch(): + """Make verl's NVFP4 W4A4 preset select the verified CUTLASS backend.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( + compressed_tensors_moe_w4a4_nvfp4, + ) + + CompressedTensorsW4A4Nvfp4MoEMethod = compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod + if getattr(CompressedTensorsW4A4Nvfp4MoEMethod, "_verl_cutlass_patch", False): + return + original_init = CompressedTensorsW4A4Nvfp4MoEMethod.__init__ + + def __init__(self, moe, layer_name=None, use_a16=False): + if use_a16: + raise NotImplementedError("verl supports NVFP4 only as W4A4 routed experts") + selection_moe = replace(moe, moe_backend="flashinfer_cutlass") if moe.moe_backend == "auto" else moe + original_init(self, selection_moe, layer_name=layer_name, use_a16=False) + + CompressedTensorsW4A4Nvfp4MoEMethod.__init__ = __init__ + CompressedTensorsW4A4Nvfp4MoEMethod._verl_cutlass_patch = True def is_deepseek_v4_model(model): @@ -54,19 +177,6 @@ def is_deepseek_v4_model(model): return getattr(text_config, "model_type", None) == "deepseek_v4" -def iter_deepseek_v4_weights(weights): - """Pass the refit stream through untouched apart from a dtype reinterpret. - - A DSv4 checkpoint already ships quantized experts, so unlike the BF16 path - there is nothing to quantize here. The expert tensors only need to be seen - as the raw ``uint8`` byte layout that ``Mxfp4MoEMethod`` allocated. - """ - for name, weight in weights: - if ".experts." in name and weight.dtype in (torch.int8, torch.float8_e8m0fnu): - weight = weight.view(torch.uint8) - yield name, weight - - def _is_mxfp4_fused_moe_module(module): from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod @@ -90,6 +200,7 @@ def _refittable_mxfp4_backends(): from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend return ( + Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, Mxfp4MoeBackend.DEEPGEMM_MXFP4, Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, @@ -160,6 +271,7 @@ def _stage_mxfp4_moe_params(module): backend's conversion and folds the result back into the live storage. """ live = {} + uses_packed_names = getattr(module.quant_method, "_verl_uses_packed_checkpoint_names", False) for name, (shape, dtype, scale_kind) in _mxfp4_checkpoint_layout(module).items(): param = getattr(module, name, None) if not isinstance(param, torch.nn.Parameter): @@ -170,7 +282,8 @@ def _stage_mxfp4_moe_params(module): staged.weight_loader = module.weight_loader if scale_kind is not None: staged.quant_method = scale_kind - setattr(module, name, staged) + checkpoint_name = f"{name}_packed" if uses_packed_names and name in ("w13_weight", "w2_weight") else name + setattr(module, checkpoint_name, staged) setattr(module, _MXFP4_LIVE_ATTR, live) @@ -205,6 +318,257 @@ def _replace_parameter_in_place(layer, param_name, new_data, prefer_copy=False): setattr(layer, param_name, param) +def _is_nvfp4_fused_moe_module(module): + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( + compressed_tensors_moe_w4a4_nvfp4, + ) + + CompressedTensorsW4A4Nvfp4MoEMethod = compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod + return isinstance(module, RoutedExperts) and isinstance(module.quant_method, CompressedTensorsW4A4Nvfp4MoEMethod) + + +def _nvfp4_checkpoint_layout(module): + """Checkpoint tensors consumed by vLLM's NVFP4 fused-expert loader.""" + num_experts = module.num_experts + intermediate = module.intermediate_size_per_partition + hidden = module.hidden_size + w13_shards = 2 if module.quant_method.moe.is_act_and_mul else 1 + tensor = "tensor" + group = "group" + return { + "w13_weight": ( + (num_experts, w13_shards * intermediate, hidden // 2), + torch.uint8, + None, + "w13_weight_packed", + ), + "w2_weight": ( + (num_experts, hidden, intermediate // 2), + torch.uint8, + None, + "w2_weight_packed", + ), + "w13_weight_scale": ( + (num_experts, w13_shards * intermediate, hidden // _NVFP4_SF_BLOCK), + torch.float8_e4m3fn, + group, + "w13_weight_scale", + ), + "w2_weight_scale": ( + (num_experts, hidden, intermediate // _NVFP4_SF_BLOCK), + torch.float8_e4m3fn, + group, + "w2_weight_scale", + ), + "w13_weight_global_scale": ( + (num_experts, w13_shards), + torch.float32, + tensor, + "w13_weight_global_scale", + ), + "w2_weight_global_scale": ( + (num_experts,), + torch.float32, + tensor, + "w2_weight_global_scale", + ), + "w13_input_global_scale": ( + (num_experts, w13_shards), + torch.float32, + tensor, + "w13_input_global_scale", + ), + "w2_input_global_scale": ( + (num_experts,), + torch.float32, + tensor, + "w2_input_global_scale", + ), + } + + +def _stage_nvfp4_moe_params(module): + live = {} + layout = _nvfp4_checkpoint_layout(module) + for live_name, (shape, dtype, scale_kind, checkpoint_name) in layout.items(): + param = getattr(module, live_name, None) + if not isinstance(param, torch.nn.Parameter): + continue + # Global scales stay in checkpoint layout and are not replaced by + # post-processing, so loading them directly is address-safe. + if live_name.endswith("global_scale"): + continue + live[live_name] = param + staged = torch.nn.Parameter(_mxfp4_staging_data(param, shape, dtype), requires_grad=False) + staged.weight_loader = module.weight_loader + if scale_kind is not None: + staged.quant_method = scale_kind + setattr(module, checkpoint_name, staged) + + # These are derived from the global scales and replaced by the official + # post-process even though they have no checkpoint-side staging tensor. + for name in ("w13_weight_scale_2", "w2_weight_scale_2"): + param = getattr(module, name, None) + if isinstance(param, torch.nn.Parameter): + live[name] = param + # CUDA graphs capture these activation-scale pointers too. The upstream + # post-process assigns freshly derived tensors instead of Parameters, so + # preserve their initial storage and copy each refit's values into it. + for name in ("w13_input_scale", "w2_input_scale"): + tensor = getattr(module, name, None) + if isinstance(tensor, torch.Tensor): + live[name] = tensor + setattr(module, _NVFP4_LIVE_ATTR, live) + + +def _replace_nvfp4_parameter_in_place(layer, param_name, new_data, prefer_copy=False): + from vllm.model_executor.utils import replace_parameter + + live = getattr(layer, _NVFP4_LIVE_ATTR, None) or {} + param = live.pop(param_name, None) + if param is None or new_data is None: + return replace_parameter(layer, param_name, new_data, prefer_copy) + if isinstance(new_data, torch.nn.Parameter): + new_data = new_data.data + if new_data.shape != param.shape or new_data.dtype != param.dtype: + raise RuntimeError( + f"nvfp4 refit re-derived {param_name} as {tuple(new_data.shape)}/{new_data.dtype}, " + f"but the CUDA-graph live parameter is {tuple(param.shape)}/{param.dtype}" + ) + if new_data.data_ptr() != param.data_ptr(): + param.data.copy_(new_data) + setattr(layer, param_name, param) + + +def _restore_nvfp4_input_scales_in_place(module, converted): + converted = list(converted) + live = getattr(module, _NVFP4_LIVE_ATTR, None) or {} + for index, name in ((3, "w13_input_scale"), (7, "w2_input_scale")): + tensor = live.pop(name, None) + new_data = converted[index] + if tensor is None or new_data is None: + continue + if new_data.shape != tensor.shape or new_data.dtype != tensor.dtype: + raise RuntimeError( + f"nvfp4 refit re-derived {name} as {tuple(new_data.shape)}/{new_data.dtype}, " + f"but the CUDA-graph live tensor is {tuple(tensor.shape)}/{tensor.dtype}" + ) + if new_data.data_ptr() != tensor.data_ptr(): + tensor.copy_(new_data) + converted[index] = tensor + return tuple(converted) + + +def _attest_nvfp4_kernel_scale_references(module): + """Verify that the retained CUDA-graph kernel sees the current refit scales.""" + method = module.quant_method + kernel = method.moe_kernel + config = method.moe_quant_config + experts = getattr(kernel, "fused_experts", None) + if experts is None or experts.quant_config is not config: + raise RuntimeError("nvfp4 retained kernel no longer references its retained quant config") + + for config_name, parameter_name in ( + ("g1_alphas", "w13_weight_scale_2"), + ("g2_alphas", "w2_weight_scale_2"), + ("w1_scale", "w13_weight_scale"), + ("w2_scale", "w2_weight_scale"), + ): + actual = getattr(config, config_name, None) + registered = getattr(module, parameter_name, None) + if not isinstance(actual, torch.Tensor) or not isinstance(registered, torch.Tensor): + raise RuntimeError(f"nvfp4 kernel scale reference missing: {config_name}/{parameter_name}") + if actual.device != registered.device or actual.data_ptr() != registered.data_ptr(): + raise RuntimeError(f"nvfp4 kernel has stale scale reference: {config_name}/{parameter_name}") + + for config_name, input_name in (("a1_gscale", "w13_input_scale"), ("a2_gscale", "w2_input_scale")): + actual = getattr(config, config_name, None) + input_scale = getattr(module, input_name, None) + if not isinstance(actual, torch.Tensor) or not isinstance(input_scale, torch.Tensor): + raise RuntimeError(f"nvfp4 activation scale reference missing: {config_name}/{input_name}") + if not torch.equal(actual, 1.0 / input_scale): + raise RuntimeError(f"nvfp4 {config_name} does not match the current activation scale") + + +def _restore_nvfp4_kernel_config_in_place(module, retained_kernel, retained_config): + """Copy newly derived scales into the config storage captured by CUDA graphs.""" + method = module.quant_method + processing_config = method.moe_quant_config + for name in ("a1_gscale", "a2_gscale"): + retained = getattr(retained_config, name, None) + current = getattr(processing_config, name, None) + if not isinstance(retained, torch.Tensor) or not isinstance(current, torch.Tensor): + raise RuntimeError(f"nvfp4 refit cannot restore derived scale {name}") + if retained.shape != current.shape or retained.dtype != current.dtype: + raise RuntimeError( + f"nvfp4 refit re-derived {name} as {tuple(current.shape)}/{current.dtype}, " + f"but the CUDA-graph config has {tuple(retained.shape)}/{retained.dtype}" + ) + retained.copy_(current) + + method.moe_kernel = retained_kernel + method.moe_quant_config = retained_config + _attest_nvfp4_kernel_scale_references(module) + + +def _process_nvfp4_moe_params(module): + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( + compressed_tensors_moe_w4a4_nvfp4, + ) + + original_convert = compressed_tensors_moe_w4a4_nvfp4.convert_to_nvfp4_moe_kernel_format + + def convert_to_nvfp4_moe_kernel_format(*args, **kwargs): + return _restore_nvfp4_input_scales_in_place(module, original_convert(*args, **kwargs)) + + retained_kernel = module.quant_method.moe_kernel + retained_config = module.quant_method.moe_quant_config + with ( + patch.object( + compressed_tensors_moe_w4a4_nvfp4, + "replace_parameter", + _replace_nvfp4_parameter_in_place, + ), + patch.object( + compressed_tensors_moe_w4a4_nvfp4, + "convert_to_nvfp4_moe_kernel_format", + convert_to_nvfp4_moe_kernel_format, + ), + ): + module.quant_method.process_weights_after_loading(module) + _restore_nvfp4_kernel_config_in_place(module, retained_kernel, retained_config) + leftover = getattr(module, _NVFP4_LIVE_ATTR, None) or {} + delattr(module, _NVFP4_LIVE_ATTR) + if leftover: + raise RuntimeError(f"nvfp4 refit left {sorted(leftover)} un-reinstated") + + +def stage_nvfp4_moe_params_for_loading(model): + """Stage NVFP4 W4A4 FlashInfer CUTLASS experts for a bucketed refit.""" + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend + + staged = [] + for module in model.modules(): + if not _is_nvfp4_fused_moe_module(module): + continue + method = module.quant_method + if method.use_a16 or method.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_CUTLASS: + raise NotImplementedError( + "NVFP4 refit supports only W4A4 FlashInfer CUTLASS routed experts; " + f"got use_a16={method.use_a16}, backend={method.nvfp4_backend}" + ) + _stage_nvfp4_moe_params(module) + staged.append(module) + logger.info("Staged %d NVFP4 W4A4 MoE modules for in-place refit", len(staged)) + return staged + + +def process_nvfp4_moe_weights_after_loading(modules): + for module in modules: + _process_nvfp4_moe_params(module) + + def _process_mxfp4_moe_params(module): from vllm.model_executor.layers.quantization import mxfp4 as vllm_mxfp4 diff --git a/verl/utils/vllm/vllm_fp8_utils.py b/verl/utils/vllm/vllm_fp8_utils.py index 4e139b8442d..8cf96d88f67 100644 --- a/verl/utils/vllm/vllm_fp8_utils.py +++ b/verl/utils/vllm/vllm_fp8_utils.py @@ -214,6 +214,25 @@ def _record_pristine_fp8_layout(layer, params): setattr(layer, _FP8_PRISTINE_ATTR, pristine) +def _record_same_layout_repacks(layer, old_params): + """Record direct parameter replacements that shape/dtype checks cannot see.""" + for name in _FP8_REFIT_PARAM_NAMES: + old_param = old_params.get(name) + new_param = getattr(layer, name, None) + if ( + isinstance(old_param, torch.nn.Parameter) + and isinstance(new_param, torch.nn.Parameter) + and tuple(old_param.shape) == tuple(new_param.shape) + and old_param.dtype == new_param.dtype + and old_param.data_ptr() != new_param.data_ptr() + ): + repacked = getattr(layer, _FP8_REPACKED_ATTR, None) + if repacked is None: + repacked = set() + setattr(layer, _FP8_REPACKED_ATTR, repacked) + repacked.add(name) + + def _layer_needs_fp8_staging(layer, pristine) -> bool: repacked = getattr(layer, _FP8_REPACKED_ATTR, None) or () for name, (shape, dtype) in pristine.items(): @@ -313,14 +332,17 @@ def process_fp8_weights_after_loading(layers): delattr(layer, _FP8_LIVE_ATTR) -def _make_process_weights_after_loading_for_vllm20(original_fn): +def _make_process_weights_after_loading_for_vllm20( + original_fn, replace_parameter_path="vllm.model_executor.layers.quantization.fp8.replace_parameter" +): def _patched_process_weights_after_loading(self, layer) -> None: old_params = dict(layer.named_parameters(recurse=False)) _record_pristine_fp8_layout(layer, old_params) - with patch( - "vllm.model_executor.layers.quantization.fp8.replace_parameter", replace_parameter_preserve_subclass - ): + with patch(replace_parameter_path, replace_parameter_preserve_subclass): original_fn(self, layer) + # Some kernels assign ``Parameter`` objects directly rather than calling + # replace_parameter (for example the MXFP8 linear kernels). + _record_same_layout_repacks(layer, old_params) _restore_layer_param_subclass_attrs(layer, old_params) return _patched_process_weights_after_loading @@ -446,6 +468,53 @@ def _create_param_from_subclass_attributes(custom_data, custom_weight): ) +def _build_compressed_tensors_fp8_patchers(): + """Same refit treatment for compressed-tensors FP8/MXFP8 linear and MoE methods.""" + ct = "vllm.model_executor.layers.quantization.compressed_tensors" + ct_moe = f"{ct}.compressed_tensors_moe" + fp8_moe = f"{ct_moe}.compressed_tensors_moe_w8a8_fp8" + mxfp8_moe = f"{ct_moe}.compressed_tensors_moe_w8a8_mxfp8" + try: + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsLinearMethod, + ) + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8 import ( # noqa: E501 + CompressedTensorsW8A8Fp8MoEMethod, + ) + except ImportError: + logger.warning("compressed-tensors FP8 methods not found in this vLLM; skipping their refit patches") + return [] + + wrap = _make_process_weights_after_loading_for_vllm20 + patchers = [ + patch( + f"{ct}.compressed_tensors.CompressedTensorsLinearMethod.process_weights_after_loading", + wrap(CompressedTensorsLinearMethod.process_weights_after_loading), + ), + patch( + f"{fp8_moe}.CompressedTensorsW8A8Fp8MoEMethod.process_weights_after_loading", + wrap(CompressedTensorsW8A8Fp8MoEMethod.process_weights_after_loading, f"{fp8_moe}.replace_parameter"), + ), + ] + try: + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_mxfp8 import ( # noqa: E501 + CompressedTensorsW8A8Mxfp8MoEMethod, + ) + except ImportError: + logger.info("compressed-tensors MXFP8 MoE method not found in this vLLM; skipping its refit patch") + else: + patchers.append( + patch( + f"{mxfp8_moe}.CompressedTensorsW8A8Mxfp8MoEMethod.process_weights_after_loading", + wrap( + CompressedTensorsW8A8Mxfp8MoEMethod.process_weights_after_loading, + f"{mxfp8_moe}.replace_parameter", + ), + ) + ) + return patchers + + def build_fp8_method_patchers(vllm_version): """Patchers that make the FP8 quant methods survive a refit, not yet started. @@ -468,6 +537,7 @@ def build_fp8_method_patchers(vllm_version): return [ patch(linear_path, wrap(Fp8LinearMethod.process_weights_after_loading)), patch(moe_path, wrap(Fp8MoEMethod.process_weights_after_loading)), + *_build_compressed_tensors_fp8_patchers(), ] return [ diff --git a/verl/utils/vllm/vllm_quant_utils.py b/verl/utils/vllm/vllm_quant_utils.py index 702dccfea89..a49105e0b54 100644 --- a/verl/utils/vllm/vllm_quant_utils.py +++ b/verl/utils/vllm/vllm_quant_utils.py @@ -16,9 +16,8 @@ """Entry point for refitting a quantized vLLM engine during RL rollout. Everything the rollout worker calls lives here: startup patching -(``apply_vllm_quant_patches``), the BF16 -> FP8 quantization of the outgoing -weight stream (``quant_weights``), the load itself (``load_quanted_weights``), -and the prepare/finalize pair that brackets a bucketed transfer +(``apply_vllm_quant_patches``), the load itself (``load_quanted_weights``), and +the prepare/finalize pair that brackets a bucketed transfer (``prepare_quanted_weights_for_loading`` / ``process_quanted_weights_after_loading``). The layout surgery each quantization scheme needs is delegated to a sibling @@ -37,22 +36,14 @@ import torch from packaging import version -try: - from vllm.model_executor.layers.linear import LinearBase -except ImportError as e: - raise ImportError("FP8 quantization not available") from e - -try: - from vllm.model_executor.layers.fused_moe.layer import FusedMoE -except ImportError: - FusedMoE = None - -from verl.utils.kernel.fp8_kernel import scaled_fp8_blockwise from verl.utils.vllm.vllm_fp4_utils import ( + install_mxfp4_mxfp8_compressed_tensors_patch, + install_nvfp4_cutlass_compressed_tensors_patch, is_deepseek_v4_model, - iter_deepseek_v4_weights, process_mxfp4_moe_weights_after_loading, + process_nvfp4_moe_weights_after_loading, stage_mxfp4_moe_params_for_loading, + stage_nvfp4_moe_params_for_loading, ) from verl.utils.vllm.vllm_fp8_utils import ( build_fp8_method_patchers, @@ -70,10 +61,6 @@ def _get_vllm_version(): # Ref: https://github.com/NVIDIA-NeMo/RL/commit/bc24887c72a6e1b2699a228bc87c588546dfe6b7 @dataclass() class FP8State: - # A cache of fp8 parameter names, we can check this cache to see if a - # param name corresponds to a fp8 weight - seen_params: set = field(default_factory=lambda: set()) - fp8_param_names: set = field(default_factory=lambda: set()) vllm_patches: list = field(default_factory=lambda: []) @@ -86,205 +73,38 @@ def is_quantized_model(vllm_config): if hasattr(vllm_config, "quant_config"): if isinstance(vllm_config.quant_config, Fp8Config): return True - elif is_vllm_ascend_mx_quant(vllm_config.quant_config): + elif is_compressed_tensors_model(vllm_config): return True return False -# vLLM 0.24.0 (MoE refactor, vllm-project/vllm#41184) removed the ``FusedMoE`` -# ``nn.Module`` class. ``FusedMoE`` is now a factory *function* that builds a -# ``MoERunner``, and the fused expert weight tensors (``w13_weight`` / -# ``w2_weight``) moved onto a ``RoutedExperts`` submodule owned by the runner -# (i.e. ``experts`` -> ``experts.routed_experts``). vLLM 0.26.1 removed the -# ``FusedMoE`` name entirely. Resolve the concrete module classes once so -# ``isinstance`` checks keep working across vLLM versions -- calling -# ``isinstance(x, FusedMoE)`` when ``FusedMoE`` is a function raises -# ``TypeError: isinstance() arg 2 must be a type``, and ``FusedMoE`` may be -# ``None`` when the name no longer exists. -if isinstance(FusedMoE, type): - # vLLM < 0.24: ``FusedMoE`` is itself the expert-weight-holding module. - _MOE_STOP_CLASSES = (FusedMoE,) - _EXPERT_WEIGHT_CLASSES = (FusedMoE,) -else: - _stop_classes: list[type] = [] - _weight_classes: list[type] = [] - try: - from vllm.model_executor.layers.fused_moe import RoutedExperts - - _stop_classes.append(RoutedExperts) - _weight_classes.append(RoutedExperts) - except ImportError: - pass - try: - from vllm.model_executor.layers.fused_moe import MoERunner - - _stop_classes.append(MoERunner) - except ImportError: - pass - _MOE_STOP_CLASSES = tuple(_stop_classes) - _EXPERT_WEIGHT_CLASSES = tuple(_weight_classes) - - -def _is_expert_weight_module(module) -> bool: - """Return whether ``module`` directly owns fused expert weights. - - True for the pre-0.24 ``FusedMoE`` module and the post-0.24 ``RoutedExperts`` - module (both expose ``w13_weight`` / ``w2_weight``). Falls back to duck typing - so the check keeps working even if the concrete classes cannot be imported. - """ - if _EXPERT_WEIGHT_CLASSES and isinstance(module, _EXPERT_WEIGHT_CLASSES): - return True - return hasattr(module, "w13_weight") and hasattr(module, "w2_weight") - - -def _is_fused_moe_block(module) -> bool: - """Return whether ``module`` is the fused-MoE block reached while walking the - module tree, so traversal must stop before descending into per-expert name - parts (e.g. the expert index) that are not real submodules.""" - if _MOE_STOP_CLASSES and isinstance(module, _MOE_STOP_CLASSES): - return True - # ``MoERunner`` owns a ``routed_experts`` child; the weight holder owns the - # expert weight tensors directly. - return hasattr(module, "routed_experts") or _is_expert_weight_module(module) - - -def _resolve_expert_weight_module(module): - """Return the submodule that actually owns the fused expert weights. - - On vLLM >= 0.24 the walked module is a ``MoERunner`` whose ``routed_experts`` - child holds the weights; on older vLLM the module already holds them. - """ - routed_experts = getattr(module, "routed_experts", None) - if routed_experts is not None: - return routed_experts - return module - - -def get_module_from_param_name(model, name: str): - # Split the name into parts (e.g., 'layers', '0', 'self_attn', 'q_proj', 'weight') - # The module path is all but the last part (the parameter's own name) - path_parts = name.split(".") - module_path = path_parts[:-1] - # Replace with the fused model name - packed_modules_mapping = model.packed_modules_mapping - reversed_mapping = { - original_name: fused_name - for fused_name, original_names_list in packed_modules_mapping.items() - for original_name in original_names_list - } - if module_path[-1] in reversed_mapping.keys(): - module_path[-1] = reversed_mapping[module_path[-1]] - - current_module = model - try: - # Traverse the model hierarchy - for part in module_path: - if _is_fused_moe_block(current_module): - # Reached the fused-MoE block: the remaining name parts (expert - # index, per-expert proj name, ...) are handled by vLLM's fused - # weight loader, so return the module owning the expert weights. - return _resolve_expert_weight_module(current_module) - elif isinstance(current_module, torch.nn.ModuleList): - current_module = current_module[int(part)] - else: - current_module = getattr(current_module, part) - except (AttributeError, IndexError, ValueError) as e: - print(f"Warning: Could not find module for parameter '{name}'. Error: {e}") - return _resolve_expert_weight_module(current_module) - - -# Fast-path allowlist: leaf names that can possibly be fp8 weights. -# Dense linear and per-expert-layout weights end in "weight". -# Fused-MoE expert params (all experts in one tensor, no per-expert index) -# use projection names without a ".weight" suffix, e.g. "gate_up_proj" / -# "down_proj". Any other leaf (bias, embedding, norm, ...) is skipped before -# the more expensive module resolution. -_FP8_CANDIDATE_LEAVES: frozenset = frozenset({"weight", "gate_up_proj", "down_proj"}) - - -def is_fp8_weight(name, model): - if name not in fp8_state.seen_params: - fp8_state.seen_params.add(name) - leaf = name.rsplit(".", 1)[-1] - if leaf in _FP8_CANDIDATE_LEAVES: - module = get_module_from_param_name(model, name) - # We currently only quantize linear and fused-MoE expert layers. - is_fp8_linear = isinstance(module, LinearBase) and module.weight.dtype == torch.float8_e4m3fn - is_fp8_moe = ( - _is_expert_weight_module(module) - and module.w13_weight.dtype == torch.float8_e4m3fn - and module.w2_weight.dtype == torch.float8_e4m3fn - ) - if is_fp8_linear or is_fp8_moe: - fp8_state.fp8_param_names.add(name) - return name in fp8_state.fp8_param_names - - -def is_w4a8_weight(module) -> bool: - """Check if a module holds W4A8 (FP4) quantized weights. - - W4A8 stores weights as packed FP4 in uint8 tensors, unlike W8A8 which - uses float8_e4m3fn. This function detects those uint8 weight holders so - that ``quant_weights`` can select the correct destination dtype for the - incoming BF16 weights. - """ - from vllm.model_executor.layers.linear import LinearBase - - is_w4a8_linear = isinstance(module, LinearBase) and module.weight.dtype == torch.uint8 - is_w4a8_moe = ( - _is_expert_weight_module(module) - and module.w13_weight.dtype == torch.uint8 - and module.w2_weight.dtype == torch.uint8 +def is_compressed_tensors_model(vllm_config) -> bool: + """Compressed-tensors rollout, whose refit weights the trainer already quantized.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, ) - return is_w4a8_linear or is_w4a8_moe + return isinstance(getattr(vllm_config, "quant_config", None), CompressedTensorsConfig) -def is_vllm_ascend_mx_quant(quant_config): - try: - from vllm_ascend.quantization import AscendModelSlimConfig - except ImportError: - # vllm_ascend not installed, so this can't be an Ascend MX config - return False - if not isinstance(quant_config, AscendModelSlimConfig): - return False - mx_quant_types = ("W8A8_MXFP8", "W4A8_MXFP") - return any(quant_type in mx_quant_types for quant_type in quant_config.quant_description.values()) +def configure_fp8_ignored_layers_match_mode(vllm_config) -> None: + """Apply verl's optional native-FP8 ignore matching extension.""" + from vllm.model_executor.layers.quantization.fp8 import Fp8Config -def restore_mxfp8_weights_for_loading(model): - for name, module in model.named_modules(): - if ( - (hasattr(module, "_mxfp8_transformed") or hasattr(module, "_mxfp4_transformed")) - and hasattr(module, "quant_method") - and hasattr(module.quant_method, "quant_method") - and hasattr(module.quant_method.quant_method, "restore_weights_for_rl_loading") - ): - module.quant_method.quant_method.restore_weights_for_rl_loading(module) - - -def apply_mxfp8_transformation_after_loading(model): - """Re-apply MXFP8 transformations after weight loading. - - This function iterates through all linear modules in the model and applies - the MXFP8 transformations (transpose, reshape) that are required for NPU - inference. - - Must be called AFTER model.load_weights() in RL training loops. - """ - try: - from vllm.model_executor.layers.linear import LinearBase - except ImportError: - logger.warning("Could not import LinearBase, skipping MXFP8 transformation") + quant_config = getattr(vllm_config, "quant_config", None) + if not isinstance(quant_config, Fp8Config): return - - for name, module in model.named_modules(): - if (isinstance(module, LinearBase) or _is_expert_weight_module(module)) and ( - hasattr(module, "_mxfp8_original_shapes") or hasattr(module, "_mxfp4_original_shapes") - ): - if hasattr(module, "quant_method") and hasattr(module.quant_method, "process_weights_after_loading"): - logger.debug(f"Applying MXFP8 transformation for module: {name}") - module.quant_method.process_weights_after_loading(module) + hf_config = getattr(getattr(vllm_config, "model_config", None), "hf_config", None) + serialized = getattr(hf_config, "quantization_config", None) + if hasattr(serialized, "to_dict"): + serialized = serialized.to_dict() + if not isinstance(serialized, dict): + return + match_mode = serialized.get("ignored_layers_match_mode", "exact") + if match_mode not in ("exact", "substring", "suffix"): + raise ValueError(f"Unsupported ignored_layers_match_mode={match_mode!r}") + quant_config.ignored_layers_match_mode = match_mode def refresh_rocm_attention_weight_caches(model): @@ -325,88 +145,6 @@ def refresh_rocm_attention_weight_caches(model): module._dsv4_wo_a_bf16 = live -def _classify_quant_weight(name, model, is_ascend_mx): - """Determine what kind of quantization a weight parameter needs. - - Returns - ``"mx_fp8"`` Ascend W8A8_MXFP8 or standard FP8 weight on Ascend NPU, - ``"mx_fp4"`` Ascend W4A8_MXFP weight (packed FP4 in uint8), - ``"fp8"`` standard blockwise FP8 weight (non-Ascend), - ``None`` parameter should pass through unquantized. - """ - if is_fp8_weight(name, model): - return "mx_fp8" if is_ascend_mx else "fp8" - if is_ascend_mx and name.endswith("weight"): - module = get_module_from_param_name(model, name) - if is_w4a8_weight(module): - return "mx_fp4" - return None - - -def quant_weights(weights, model, quant_config, dtype=torch.bfloat16): - """Quantize weights to FP8 format using a memory-efficient generator. - - - Args: - weights: Generator or iterable of (name, tensor) pairs - model: The model to check for FP8 weight names - quant_config: Quantization configuration with weight_block_size - dtype: Data type for intermediate computation (default: bfloat16) - - Yields: - Tuples of (name, tensor) for each weight and its scale - """ - - if is_deepseek_v4_model(model): - yield from iter_deepseek_v4_weights(weights) - return - - fp8_state.seen_params.clear() - fp8_state.fp8_param_names.clear() - is_ascend_mx = is_vllm_ascend_mx_quant(quant_config) - if is_ascend_mx: - import torch_npu - for k, v in weights: - quant_type = _classify_quant_weight(k, model, is_ascend_mx) - if quant_type is None: - yield (k, v) - continue - - # Cast the weight into quantized format and its scale factor - if torch.distributed.get_rank() == 0: - logger.debug(f"Quantizing to {quant_type} blockwise: {k}") - if quant_type in ("mx_fp8", "mx_fp4"): - dst_type = torch_npu.float4_e2m1fn_x2 if quant_type == "mx_fp4" else torch_npu.float8_e4m3fn - param_lp, param_scale = torch_npu.npu_dynamic_mx_quant( - v.to(dtype), - axis=-1, - dst_type=dst_type, - ) - param_scale = param_scale.flatten(-2, -1) - if quant_type == "mx_fp4": - # W4A8 params are uint8 holding packed FP4 pairs; the mx quant - # op returns those bytes as float4_e2m1fn_x2, so reinterpret - # them as uint8 to match the parameter dtype. - param_lp = param_lp.view(torch.uint8) - else: - param_lp, param_scale = scaled_fp8_blockwise( - v.to(dtype), - weight_block_size=quant_config.weight_block_size, - ) - param_scale = param_scale.squeeze(-1) - - # Yield the quantized weight - yield (k, param_lp) - - if is_ascend_mx: - yield (k + "_scale", param_scale) - else: - yield (k + "_scale_inv", param_scale) - - # Explicitly delete original tensor reference to help GC - del v, param_lp, param_scale - - def prepare_quanted_weights_for_loading(model): """Restore quantized params to the layout their ``weight_loader`` expects. @@ -417,20 +155,20 @@ def prepare_quanted_weights_for_loading(model): so this is safe to call for any quantization scheme. The returned value is opaque reload state for the paired call. """ - restore_mxfp8_weights_for_loading(model) reload_state: dict[str, Any] = { "fp8_layers": stage_fp8_params_for_loading(model), "mxfp4_moe_modules": stage_mxfp4_moe_params_for_loading(model), + "nvfp4_moe_modules": stage_nvfp4_moe_params_for_loading(model), } return reload_state def process_quanted_weights_after_loading(model, reload_state): """Re-apply the inference layout undone by ``prepare_quanted_weights_for_loading``.""" - apply_mxfp8_transformation_after_loading(model) reload_state = reload_state or {} process_fp8_weights_after_loading(reload_state.get("fp8_layers") or []) process_mxfp4_moe_weights_after_loading(reload_state.get("mxfp4_moe_modules") or []) + process_nvfp4_moe_weights_after_loading(reload_state.get("nvfp4_moe_modules") or []) # Last: the rebuild reads ``wo_a``, which is only back in its inference # layout once the staged FP8 params above have been reinstated. refresh_rocm_attention_weight_caches(model) @@ -446,11 +184,7 @@ def load_quanted_weights(weights, model_runner, is_drafter=False): ) else: model = model_runner.model - quant_config = model_runner.vllm_config.quant_config - vllm_dtype = model_runner.vllm_config.model_config.dtype - weights = list(weights) - weights_quantized = quant_weights(weights, model, quant_config, dtype=vllm_dtype) # Monkey patch the param class to their subclass, as certain models # will check the param type to call the proper weightloader @@ -471,7 +205,7 @@ def load_quanted_weights(weights, model_runner, is_drafter=False): disable_mtp_completeness_check = nullcontext try: with disable_mtp_completeness_check(): - loaded_params = model.load_weights(weights_quantized) + loaded_params = model.load_weights(weights) finally: # Undo the type change above to the original type for name, param in model.named_parameters(): @@ -487,10 +221,12 @@ def apply_vllm_quant_patches(): These have to be in place before the engine builds, because recording each layer's checkpoint layout happens inside the patched hook during the - initial load. The mxfp4 path needs no startup patch -- it derives the - checkpoint layout from the quant method and patches ``replace_parameter`` - only for the duration of its own post-processing call. + initial load. The MXFP4 refit path derives the checkpoint layout from the + quant method; its startup patch only makes compressed-tensors honor MXFP8 + activations. """ + install_mxfp4_mxfp8_compressed_tensors_patch() + install_nvfp4_cutlass_compressed_tensors_patch() if fp8_state.vllm_patches: logger.debug("vLLM quantization patches already applied") return diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index 22bec4a0475..96d318aca4d 100644 --- a/verl/workers/config/engine.py +++ b/verl/workers/config/engine.py @@ -129,7 +129,7 @@ class QATEngineConfig(BaseConfig): Args: enable (bool): Whether to enable QAT, default False - mode (str): Quantization mode, "w4a16" or "w4a4", default "w4a16" + mode (str): Quantization mode, "fp8", "mxfp8", "mxfp4", or "nvfp4", default "fp8" group_size (int): Group size for blockwise quantization, default 16 ignore_patterns (list[str]): Module name patterns to exclude from quantization activation_observer (str): Observer strategy for activation global_scale (W4A4 only) @@ -137,7 +137,7 @@ class QATEngineConfig(BaseConfig): """ enable: bool = False - mode: str = "w4a16" + mode: str = "fp8" group_size: int = 16 ignore_patterns: list[str] = field(default_factory=lambda: ["lm_head", "embed_tokens", "re:.*mlp.gate$"]) activation_observer: str = "static_minmax" diff --git a/verl/workers/config/rollout.py b/verl/workers/config/rollout.py index 4388b177a83..4bfef74d226 100644 --- a/verl/workers/config/rollout.py +++ b/verl/workers/config/rollout.py @@ -13,7 +13,7 @@ # limitations under the License. import warnings from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Optional from omegaconf import MISSING, DictConfig, OmegaConf @@ -31,6 +31,7 @@ "ServerConfig", "PrometheusConfig", "RolloutConfig", + "RolloutQuantizationConfig", "CheckpointEngineConfig", ] @@ -141,6 +142,40 @@ class CheckpointEngineConfig(BaseConfig): custom_backend_module: Optional[str] = None +@dataclass +class RolloutQuantizationConfig(BaseConfig): + """Low-precision rollout as a compressed-tensors ``quantization_config``; the trainer quantizes the weights. + + Either pick compressed-tensors presets per layer kind (``linear`` / ``moe`` / ``ignore``), + or point ``config_file`` at a complete compressed-tensors ``quantization_config`` JSON. + """ + + # Preset for dense linear layers (attention, dense MLP, shared experts), e.g. FP8_BLOCK. + linear: Optional[str] = None + # Preset for routed experts, e.g. FP8_BLOCK. + moe: Optional[str] = None + # Extra "re:" patterns kept in high precision. + ignore: list[str] = field(default_factory=list) + # compressed-tensors quantization_config JSON, used as is; exclusive with the fields above. + config_file: Optional[str] = None + + @property + def enabled(self) -> bool: + return bool(self.linear or self.moe or self.config_file) + + def build_quantization_config(self, model_path: str) -> dict[str, Any]: + """HF ``quantization_config`` the rollout engine starts from and the trainer quantizes refits to.""" + from verl.utils.qat.compressed_tensors import build_compressed_tensors_config + + return build_compressed_tensors_config( + linear=self.linear, + moe=self.moe, + ignore=self.ignore, + config_file=self.config_file, + model_path=model_path, + ) + + @dataclass class RolloutConfig(BaseConfig): _mutable_fields = { @@ -265,9 +300,7 @@ class RolloutConfig(BaseConfig): skip_tokenizer_init: bool = True - quantization: Optional[str] = None - - quantization_config_file: Optional[str] = None + quantization: RolloutQuantizationConfig = field(default_factory=RolloutQuantizationConfig) enable_rollout_routing_replay: bool = False moe_load_balance_metrics_interval: int = 0 @@ -276,8 +309,6 @@ class RolloutConfig(BaseConfig): mtp: MtpConfig = field(default_factory=MtpConfig) - qat: Optional[dict] = None - disaggregation: DisaggregationConfig = field(default_factory=DisaggregationConfig) router_config_path: Optional[str] = None diff --git a/verl/workers/engine/base.py b/verl/workers/engine/base.py index 161faed067c..af1d80d0b3f 100644 --- a/verl/workers/engine/base.py +++ b/verl/workers/engine/base.py @@ -160,6 +160,14 @@ def get_per_tensor_param(self) -> tuple[Generator[tuple[str, torch.Tensor], None """ raise NotImplementedError + def set_rollout_weight_quantizer(self, quantizer) -> None: + """Make :meth:`get_per_tensor_param` emit weights quantized by ``quantizer`` for a low-precision rollout. + + Args: + quantizer: Callable mapping an iterable of ``(name, tensor)`` HF weights to the quantized stream. + """ + raise NotImplementedError(f"{type(self).__name__} does not support trainer-side rollout quantization") + def get_per_tensor_param_shard(self, **kwargs) -> tuple[Generator, Optional[dict]]: """ Like :meth:`get_per_tensor_param`, but yields each rank's *local* parameter shard diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index 354c3a372d8..65e4b89579d 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -205,6 +205,7 @@ def __init__( self._qat_enabled = self._qat_config is not None and getattr(self._qat_config, "enable", False) if self._qat_enabled: logger.info(f"QAT enabled in MegatronEngine: mode={self._qat_config.mode}") + self._rollout_weight_quantizer = None # Router replay configuration for MoE models self.enable_routing_replay = self.engine_config.router_replay.mode != "disabled" @@ -529,7 +530,7 @@ def initialize(self): self.module = self._build_megatron_module() if self._qat_enabled and not self.engine_config.forward_only: - from verl.utils.modelopt import apply_qat_to_modules + from verl.utils.qat.modelopt import apply_qat_to_modules self.module = apply_qat_to_modules(self.module, self._qat_config) @@ -1000,14 +1001,20 @@ def get_per_tensor_param(self, base_sync_done=False, **kwargs): else self.bridge.export_hf_weights(self.module, conversion_tasks=conversion_tasks) ) - # QAT: process weights through QATWeightExporter for quantized weight sync to vLLM 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) + per_tensor_param = ((name, tensor) for name, tensor in per_tensor_param if "_quantizer." not in name) + if self._rollout_weight_quantizer is not None and not adapter_only: + # QAT-trained master weights quantize to the same values the fake quant used in training. + refresh_scales = getattr(self._rollout_weight_quantizer, "refresh_input_global_scales", None) + if refresh_scales is not None: + refresh_scales(self.module) + per_tensor_param = self._rollout_weight_quantizer(per_tensor_param) return per_tensor_param, peft_config + def set_rollout_weight_quantizer(self, quantizer) -> None: + self._rollout_weight_quantizer = quantizer + def _mcore_export_index(self): """Build (once) the per-parameter delta export index: geometry specs and form-B probes derived from the bridge's conversion tasks. See diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index c42ee25ce03..b724eae7c49 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -43,6 +43,7 @@ from verl.utils.metric.utils import Metric from verl.utils.profiler import DistProfiler, DistProfilerExtension, ProfilerConfig, log_gpu_memory_usage from verl.utils.py_functional import append_to_dict +from verl.utils.qat import dispatch_rollout_weight_quantizer from verl.utils.tensordict_utils import maybe_fix_3d_position_ids from verl.utils.torch_functional import allgather_dict_into_dict from verl.workers.config import ( @@ -729,6 +730,15 @@ def init_model(self): backend, is_master=(torch.distributed.get_rank() == 0), bucket_size=bucket_size, **engine_kwargs ) + qat = getattr(getattr(self.actor.engine, "engine_config", None), "qat", None) + weight_quantizer = dispatch_rollout_weight_quantizer( + model_config=model_config, + rollout_config=self.config.rollout, + qat_config=qat, + ) + if weight_quantizer is not None: + self.actor.engine.set_rollout_weight_quantizer(weight_quantizer) + # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo aggressive_empty_cache(force_sync=True) diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index e5b486d07a3..b74af031d89 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -300,14 +300,14 @@ async def launch_server(self, master_address: str = None, master_port: int = Non # (e.g. text "flashinfer" vs vision "flashinfer_cudnn"), so don't mirror it. # Leave None to let sglang's VisionAttention auto-pick per device # (triton_attn on Ada, fa3 on Hopper, fa4 on Blackwell). - quantization = self.config.get("quantization", None) - if quantization is not None: - if quantization == "fp8": - from verl.utils.sglang.sglang_fp8_utils import build_sglang_fp8_quant_config - - fp8_block_quant_kwargs = build_sglang_fp8_quant_config(self.model_config.hf_config) - else: - raise ValueError(f"Currently only support fp8 quantization, got: {quantization}") + model_override_args = {} + quantization_config = None + if self.config.quantization.enabled: + # The trainer quantizes every refit to this config. On Hopper the block-FP8 linear/MoE methods keep + # weights in checkpoint layout after loading, so update_weights_from_tensor can write refits in place. + quantization_config = self.config.quantization.build_quantization_config(self.model_config.local_path) + model_override_args["quantization_config"] = quantization_config + logger.info(f"Rollout compressed-tensors quantization_config: {quantization_config}") infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size args = { "model_path": self.model_config.local_path, @@ -330,15 +330,12 @@ async def launch_server(self, master_address: str = None, master_port: int = Non "attention_backend": attention_backend, "skip_tokenizer_init": self.config.skip_tokenizer_init, "skip_server_warmup": True, - "quantization": quantization, - "json_model_override_args": json.dumps({"quantization_config": fp8_block_quant_kwargs}) - if quantization == "fp8" - else json.dumps({}), + "json_model_override_args": json.dumps(model_override_args), "custom_weight_loader": custom_weight_loader or None, **engine_kwargs, } # bf16 weight updates crash SGLang's default MoE runner on Blackwell (see default_moe_runner_backend) - moe_runner_backend = default_moe_runner_backend(quantization, engine_kwargs, self.model_config.hf_config) + moe_runner_backend = default_moe_runner_backend(quantization_config, engine_kwargs, self.model_config.hf_config) if moe_runner_backend is not None: args["moe_runner_backend"] = moe_runner_backend logger.info( @@ -662,22 +659,21 @@ async def generate( routed_experts = None if self.config.enable_rollout_routing_replay: - if self.config.skip_tokenizer_init: - # convert to numpy - captured = output.get("meta_info", {}).get("routed_experts", None) - routed_experts = captured.numpy() if captured is not None else None - else: - from sglang.srt.layers.moe.routed_experts_capturer import extract_routed_experts_from_meta_info + captured = output.get("meta_info", {}).get("routed_experts", None) + if captured is not None: + from sglang.srt.state_capturer.routed_experts import extract_routed_experts_from_meta_info + + routed_experts = extract_routed_experts_from_meta_info(output) hf_config = self.model_config.hf_config - if not hasattr(hf_config, "num_hidden_layers") or not hasattr(hf_config, "num_experts_per_tok"): + text_config = getattr(hf_config, "text_config", hf_config) + if not hasattr(text_config, "num_hidden_layers") or not hasattr(text_config, "num_experts_per_tok"): raise AttributeError( - "enable_rollout_routing_replay is set, but hf_config is missing " - "'num_hidden_layers' or 'num_experts_per_tok'. This feature requires an MoE model " - "configuration that defines these attributes." + "enable_rollout_routing_replay requires an MoE model config with " + "`num_hidden_layers` and `num_experts_per_tok`" ) - routed_experts = extract_routed_experts_from_meta_info(output).reshape( - -1, hf_config.num_hidden_layers, hf_config.num_experts_per_tok + routed_experts = routed_experts.reshape( + -1, text_config.num_hidden_layers, text_config.num_experts_per_tok ) extra_fields = {"global_steps": self.global_steps} diff --git a/verl/workers/rollout/sglang_rollout/sglang_rollout.py b/verl/workers/rollout/sglang_rollout/sglang_rollout.py index 4ed2a583c81..c80c0a39512 100644 --- a/verl/workers/rollout/sglang_rollout/sglang_rollout.py +++ b/verl/workers/rollout/sglang_rollout/sglang_rollout.py @@ -130,17 +130,6 @@ def __init__( replica_rank: int = -1, ): super().__init__(config, model_config, device_mesh) - if self.config.get("quantization", None) == "fp8": - import sglang - from packaging import version - - from verl.utils.sglang.sglang_fp8_utils import build_sglang_fp8_quant_config - - assert version.parse(sglang.__version__) >= version.parse("0.5.5"), ( - "sglang>=0.5.5 is required for FP8 quantization" - ) - fp8_block_quant_kwargs = build_sglang_fp8_quant_config(self.model_config.hf_config) - self.model_config.hf_config.quantization_config = fp8_block_quant_kwargs self._engine: AsyncHttpServerAdapter = None rank = int(os.environ["RANK"]) @@ -357,18 +346,6 @@ async def update_weights( await self._engine.load_lora_adapter_from_tensor(req) else: update_weights_bucket_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) << 20 - if self.config.get("quantization", None) == "fp8": - from verl.utils.sglang.sglang_fp8_utils import SGLangFP8QuantizerHelper - - logger.info("Convert bf16 weights to fp8 format before loading") - fp8_quantizer_helper = SGLangFP8QuantizerHelper(self.model_config.hf_config.quantization_config) - weights = fp8_quantizer_helper.quant_weights_by_name( - weights, - dtype=self.model_config.hf_config.dtype, - ) - else: - weights = weights - fusion_groups = ( DEEPSEEK_V4_FUSION_GROUPS if getattr(self.model_config.hf_config, "model_type", None) == "deepseek_v4" diff --git a/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py b/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py index bd57233f171..3c42fb63139 100644 --- a/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py +++ b/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py @@ -206,20 +206,8 @@ async def launch_server(self): per_worker_gpu_share = 1.0 / self.max_colocate_count - quantization = self.config.quantization - if quantization is not None: - if quantization == "fp8": - FP8_BLOCK_QUANT_KWARGS = { - "activation_scheme": "dynamic", - "fmt": "e4m3", - "quant_method": "fp8", - "weight_block_size": [128, 128], - } - engine_kwargs["model_kwargs"] = {"quantization_config": FP8_BLOCK_QUANT_KWARGS} - if self.config.load_format != "dummy": - raise ValueError("FP8 quantization is only supported for dummy load format") - else: - raise ValueError(f"Currently only support fp8 quantization, got: {quantization}") + if self.config.quantization.enabled: + raise NotImplementedError("rollout.quantization is only supported by vLLM") llm_kwargs = { "model": self.model_config.local_path, diff --git a/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py b/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py index 0d694937f27..93f5f5cfaf2 100644 --- a/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py +++ b/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py @@ -292,14 +292,6 @@ def __init__( self, config: RolloutConfig, model_config: HFModelConfig, device_mesh: DeviceMesh, replica_rank: int = -1 ): super().__init__(config, model_config, device_mesh) - if self.config.quantization == "fp8" and self.model_config.hf_config is not None: - FP8_BLOCK_QUANT_KWARGS = { - "activation_scheme": "dynamic", - "fmt": "e4m3", - "quant_method": "fp8", - "weight_block_size": [128, 128], - } - self.model_config.hf_config.quantization_config = dict(FP8_BLOCK_QUANT_KWARGS) self._adapter = None self.hybrid_device_mesh = None self.gpu_id = None @@ -467,15 +459,6 @@ async def update_weights( total_available_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) * 1024 * 1024 - if self.config.get("quantization", None) == "fp8": - from verl.utils.trtllm.trtllm_fp8_utils import TRTLLMFP8QuantizerHelper - - fp8_quantizer_helper = TRTLLMFP8QuantizerHelper(self.model_config.hf_config.quantization_config) - weights = fp8_quantizer_helper.quant_weights_by_name( - weights, - dtype=self.model_config.hf_config.dtype, - ) - try: device_uuid = get_device_uuid(int(self.gpu_id)) except Exception as e: diff --git a/verl/workers/rollout/vllm_rollout/utils.py b/verl/workers/rollout/vllm_rollout/utils.py index 213b0727b3b..8112b7e4069 100644 --- a/verl/workers/rollout/vllm_rollout/utils.py +++ b/verl/workers/rollout/vllm_rollout/utils.py @@ -33,6 +33,7 @@ from verl.utils.vllm.rocm_vllm_moe_expert_map import restore_moe_expert_maps from verl.utils.vllm.vllm_quant_utils import ( apply_vllm_quant_patches, + configure_fp8_ignored_layers_match_mode, is_quantized_model, load_quanted_weights, prepare_quanted_weights_for_loading, @@ -186,24 +187,12 @@ def __new__(cls, **kwargs): ) register_verl_delta_weight_transfer_engine() - # 2. patch online fp8 quant. Some models, including DeepSeek-V4, get - # fp8 from the HF config rather than an explicit rollout quantization arg. - if os.environ.get("VERL_VLLM_FP8_QUANT_ENABLED", "0") == "1" or is_quantized_model(vllm_config): - apply_vllm_quant_patches() - # 3. patch QAT (compressed-tensors NVFP4) for dynamic weight loading - quant_config = getattr(vllm_config, "quant_config", None) if vllm_config else None - _is_qat_model = getattr(quant_config, "quant_format", None) == "nvfp4-pack-quantized" - _is_modelopt_qat = type(quant_config).__name__ == "ModelOptNvFp4Config" - if _is_qat_model: - from verl.utils.qat import apply_qat_patches - - apply_qat_patches() - logger.info("Applied QAT (compressed-tensors) patches in vLLM worker subprocess") - elif _is_modelopt_qat: - from verl.utils.modelopt import apply_modelopt_nvfp4_patches - apply_modelopt_nvfp4_patches() - logger.info("Applied ModelOpt NVFP4 patches in vLLM worker subprocess") + # 2. patch the refit of quantized weights, whether the quantization_config comes + # from the checkpoint (e.g. DeepSeek-V4) or from rollout.quantization via hf_overrides. + if is_quantized_model(vllm_config): + configure_fp8_ignored_layers_match_mode(vllm_config) + apply_vllm_quant_patches() # TODO: For ascend NPU, when the corresponding vllm-ascend version is upgraded to v0.13.0, # please remove the VLLM_ASCEND_REQUIRED_ENV_VARS variable replacement action. @@ -214,8 +203,6 @@ def __new__(cls, **kwargs): os.environ[k] = VLLM_ASCEND_REQUIRED_ENV_VARS[k] instance = super().__new__(cls) - instance._is_qat_model = _is_qat_model - instance._is_modelopt_qat = _is_modelopt_qat return instance def _get_drafter_model(self): @@ -276,19 +263,7 @@ def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False for model in self._iter_all_models(): restore_moe_expert_maps(model) - if self._is_qat_model: - # QAT (compressed-tensors): Prepare for weight loading BEFORE receiving any buckets - from verl.utils.qat import prepare_qat_for_load_weights - - for model in self._iter_all_models(): - prepare_qat_for_load_weights(model, device=self.device) - logger.info("QAT: prepare_qat_for_load_weights completed") - elif self._is_modelopt_qat: - from verl.utils.modelopt.vllm_modelopt_patch import prepare_modelopt_for_weight_reload - - prepare_modelopt_for_weight_reload(self.model_runner.model, device=self.device) - logger.info("ModelOpt: prepare_modelopt_for_weight_reload completed") - elif peft_config and base_sync_done: + if peft_config and base_sync_done: # Remove the old LoRA before the new one arrives (applied after is_last below). self.remove_lora(VLLM_LORA_INT_ID) logger.info("LoRA adapter sync: remove old lora and prepare new lora") @@ -339,19 +314,7 @@ def on_bucket_received(weights: list[tuple[str, torch.Tensor]], is_last: bool) - receiver.receive_weights(on_bucket_received=on_bucket_received) # =========================== step 3: process weights after loading =========================== - if self._is_qat_model: - # QAT (compressed-tensors): call process_weights_after_loading AFTER all buckets are received - from verl.utils.qat import manual_process_weights_after_loading - - for model in self._iter_all_models(): - manual_process_weights_after_loading(model) - logger.info("QAT: process_weights_after_loading completed") - elif self._is_modelopt_qat: - from verl.utils.modelopt.vllm_modelopt_patch import modelopt_process_weights_after_loading - - modelopt_process_weights_after_loading(self.model_runner.model) - logger.info("ModelOpt QAT: process_weights_after_loading completed") - elif peft_config and base_sync_done: + if peft_config and base_sync_done: logger.info("LoRA adapter sync, no post-process needed") elif is_quantized_model(self.model_runner.vllm_config): for model, reload_state in quant_reload_states: diff --git a/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/verl/workers/rollout/vllm_rollout/vllm_async_server.py index 21e15c938d3..1e556218b53 100644 --- a/verl/workers/rollout/vllm_rollout/vllm_async_server.py +++ b/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -53,7 +53,6 @@ ) from verl.utils.tokenizer import normalize_token_ids from verl.utils.tracking import RLInsightLogger -from verl.utils.vllm.vllm_quant_utils import apply_vllm_quant_patches from verl.workers.config import HFModelConfig, RolloutConfig from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput from verl.workers.rollout.utils import ( @@ -318,7 +317,7 @@ async def launch_server(self, master_address: str = None, master_port: int = Non set_expandable_segments(True) - quantization, hf_overrides = self._apply_quantization() + hf_overrides = self._apply_quantization() compilation_config = engine_kwargs.pop("compilation_config", None) or {} if isinstance(compilation_config, str): @@ -358,7 +357,6 @@ async def launch_server(self, master_address: str = None, master_port: int = Non "tensor_parallel_size": self.config.tensor_model_parallel_size, "seed": self.replica_rank + self.config.seed, "override_generation_config": json.dumps(override_generation_config), - "quantization": quantization, "hf_overrides": hf_overrides, "scheduling_policy": self.config.scheduling_policy, "compilation_config": compilation_config, @@ -1392,69 +1390,30 @@ def _get_override_generation_config(self) -> dict: max_new_tokens=self.config.response_length, ) - def _apply_quantization(self) -> tuple[Optional[str], dict]: - """Process quantization config. Returns (quantization_str, hf_overrides).""" - quantization = self.config.quantization - hf_overrides = {} - - # Handle QAT (Quantization-Aware Training) configuration - qat_config_dict = getattr(self.config, "qat", {}) or {} - if qat_config_dict.get("enable", False): - from verl.utils.qat import QATConfig, load_quantization_config - - qat_config = QATConfig(**qat_config_dict) - quantization_config_dict = load_quantization_config(qat_config) - quant_method = quantization_config_dict.get("quant_method", None) - - if quant_method == "modelopt": - from verl.utils.modelopt import apply_modelopt_nvfp4_patches + def _apply_quantization(self) -> dict: + """hf_overrides for a rollout quantization set in the rollout config. - apply_modelopt_nvfp4_patches() - quantization = "modelopt" - elif quant_method == "compressed-tensors": - from verl.utils.qat import apply_qat_patches - - apply_qat_patches() - quantization = "compressed-tensors" - else: - raise ValueError(f"Unsupported quant_method: {quant_method}") - - logger.info(f"QAT quantization config injected (quant_method={quant_method})") - hf_overrides["quantization_config"] = quantization_config_dict - elif quantization is not None: - # Handle other quantization methods (fp8, torchao) - _SUPPORTED_QUANTIZATION = ["fp8", "torchao", "ascend"] - if quantization not in _SUPPORTED_QUANTIZATION: - raise ValueError(f"Currently only support {_SUPPORTED_QUANTIZATION} quantization, got: {quantization}") - - if quantization == "fp8": - # Ignore MoE router layers for FP8 quantization - all_mlp_gate_layers = [] - for layer in range(self.model_config.hf_config.num_hidden_layers): - all_mlp_gate_layers.append(f"model.layers.{layer}.mlp.gate") - - FP8_BLOCK_QUANT_KWARGS = { - "activation_scheme": "dynamic", - "fmt": "e4m3", - "quant_method": "fp8", - "weight_block_size": [128, 128], - "ignored_layers": all_mlp_gate_layers, - } - hf_overrides["quantization_config"] = dict(FP8_BLOCK_QUANT_KWARGS) - # Will remove the patch after vllm support on-the-fly quant for rollout natively. - apply_vllm_quant_patches() - # for subprocesses patching - os.environ["VERL_VLLM_FP8_QUANT_ENABLED"] = "1" - - model_quantization_config = getattr(self.model_config.hf_config, "quantization_config", {}) or {} - if quantization is None and model_quantization_config.get("quant_method") == "fp8": - apply_vllm_quant_patches() - os.environ["VERL_VLLM_FP8_QUANT_ENABLED"] = "1" - - if quantization is not None and self.config.quantization_config_file is not None: - hf_overrides["quantization_config_file"] = self.config.quantization_config_file - - return quantization, hf_overrides + Models shipping their own ``quantization_config`` need nothing here; either way the + workers install the refit patches from the resulting vLLM ``quant_config``. + """ + quantization = self.config.quantization + if not quantization.enabled: + return {} + + # The trainer quantizes every refit to this config; vLLM picks the method up from hf_overrides. + quantization_config = quantization.build_quantization_config(self.model_config.local_path) + # FP8_BLOCK carries FP32 block scales. DeepGEMM's E8M0 mode (vLLM's default) would round them + # to powers of two and drift from the trainer. MX formats carry E8M0 scales natively and must + # retain vLLM's E8M0 path. + uses_fp32_block_scales = any( + group.get("weights", {}).get("strategy") == "block" + and group.get("weights", {}).get("scale_dtype") in {None, "torch.float32"} + for group in quantization_config.get("config_groups", {}).values() + ) + if uses_fp32_block_scales: + os.environ.setdefault("VLLM_USE_DEEP_GEMM_E8M0", "0") + logger.info(f"Rollout compressed-tensors quantization_config: {quantization_config}") + return {"quantization_config": quantization_config} def _get_worker_extension_cls(self) -> str: """Return the fully-qualified colocate worker extension class name."""