From fa1f6d2f794cc3bb5cd1631bd953670e5828953b Mon Sep 17 00:00:00 2001 From: wuxibin Date: Wed, 14 Oct 2026 21:29:49 +0800 Subject: [PATCH 01/18] [1/N][BREAKING][megatron,vllm] feat: refactor megatron QAT training --- docs/low_precision/qat.md | 219 +++++++ tests/experimental/agent_loop/agent_utils.py | 9 +- .../test_megatron_vllm_logprob_diff.py | 192 ++++++ ...t_vllm_weight_name_normalization_on_cpu.py | 10 +- verl/trainer/config/engine/megatron.yaml | 3 +- verl/trainer/config/rollout/rollout.yaml | 4 +- verl/utils/modelopt/__init__.py | 15 +- verl/utils/modelopt/qat_utils.py | 8 - verl/utils/modelopt/qat_weight_exporter.py | 275 -------- verl/utils/modelopt/quantize.py | 74 ++- verl/utils/modelopt/vllm_modelopt_patch.py | 619 ------------------ verl/utils/qat/__init__.py | 23 +- verl/utils/qat/compressed_tensors.py | 317 +++++++++ verl/utils/vllm/vllm_fp8_utils.py | 37 +- .../vllm_nvfp4_utils.py} | 0 verl/utils/vllm/vllm_quant_utils.py | 16 +- verl/workers/config/engine.py | 2 +- verl/workers/config/rollout.py | 6 +- verl/workers/engine/base.py | 8 + .../engine/megatron/transformer_impl.py | 12 +- verl/workers/engine_workers.py | 20 + .../sglang_rollout/async_sglang_server.py | 3 + .../trtllm_rollout/trtllm_async_server.py | 3 + verl/workers/rollout/vllm_rollout/utils.py | 53 +- .../rollout/vllm_rollout/vllm_async_server.py | 63 +- 25 files changed, 958 insertions(+), 1033 deletions(-) create mode 100644 docs/low_precision/qat.md create mode 100644 tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py delete mode 100644 verl/utils/modelopt/qat_weight_exporter.py delete mode 100644 verl/utils/modelopt/vllm_modelopt_patch.py create mode 100644 verl/utils/qat/compressed_tensors.py rename verl/utils/{qat/vllm_patch.py => vllm/vllm_nvfp4_utils.py} (100%) diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md new file mode 100644 index 00000000000..d802c6e656a --- /dev/null +++ b/docs/low_precision/qat.md @@ -0,0 +1,219 @@ +# Quantization Aware Training + +## Low Precision Inference + +### 背景 + +RL 训练中 rollout(推理)通常占据大部分时间,而大规模 MoE 模型的 rollout 往往受限于显存容量和访存带宽。使用低精度(FP8 / MXFP8 / MXFP4 / NVFP4)推理可以: + +- 降低权重显存占用,留出更多 KV cache,提高并发; +- 降低 decode 阶段的访存量,提高吞吐; +- 在 Hopper/Blackwell 上使用 FP8/FP4 Tensor Core 提升 prefill 算力。 + +代价是训练(高精度)与推理(低精度)之间的数值差异,即 training-inference mismatch,它会影响 importance sampling 和 off-policy 修正的准确性,严重时会导致训练不稳定。 + +### 模型权重格式 + +HuggingFace checkpoint 通常有两种精度格式: + +1. **高精度**:权重以 BF16 存储,`config.json` 中没有 `quantization_config`。 +2. **低精度**:`config.json` 中有 `quantization_config`,描述各层的量化方式。不同模型对 linear(attention、dense MLP、shared experts)和 MoE(routed experts)往往采用不同的格式。 + +| 模型 | 结构 | `quant_method` | Linear | MoE (routed experts) | 不量化的层 | +|---|---|---|---|---|---| +| [Qwen/Qwen3.8-2.4T-A95B](https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B/blob/main/config.json) | 92 层,512 experts,top-10 | 无 | BF16 | BF16 | - | +| [zai-org/GLM-5.2](https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json) | 78 层,256 experts,top-8,前 3 层 dense | 无 | BF16 | BF16 | - | +| [zai-org/GLM-5.3](https://huggingface.co/zai-org/GLM-5.3/blob/main/config.json) | 与 GLM-5.2 相同 | `fp8` | FP8 E4M3,128x128 block,FP32 scale(`weight_scale_inv`),activation 动态量化 | 同 Linear | router `mlp.gate`、`indexers_proj`、`embed_tokens`、`lm_head`、MTP `eh_proj`、各类 norm(`modules_to_not_convert`) | +| [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3/blob/main/config.json) | 93 层,896 experts,第 1 层 dense | `compressed-tensors` | BF16 | MXFP4 W4A16:FP4 E2M1,1x32 group,UE8M0 scale(`weight_packed` + `weight_scale`),activation 不量化 | `self_attn`、`shared_experts`、dense `mlp`、`lm_head`、`vision_tower`、`mm_projector`(`ignore` 正则) | +| [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/main/config.json) | 40 层,384 experts,top-6 | `fp8` | FP8 E4M3,32x32 block(`wo_a` 可能是 128x128),UE8M0 scale(`scale_fmt: ue8m0`) | FP4 E2M1,1x32 group,UE8M0 scale(`expert_dtype: fp4`),格式与 MXFP4 相同 | router `ffn.gate`、indexer `wk`/`weights_proj`、`compressor` | + +说明: + +- GLM-5.3 是标准的 HF FineGrainedFP8 格式(与 DeepSeek-V3 相同),vLLM 和 SGLang 的 `Fp8Config.from_config` 都能直接解析。 +- Kimi-K3 使用 compressed-tensors 格式:`config_groups` 以 `Linear` 为 target,再用 `ignore` 排除非 expert 层,从而实现 "linear BF16 + MoE MXFP4"。 +- DeepSeek-V4.1-Flash 的 `fp8` 配置带有 `scale_fmt`、`expert_dtype` 等私有字段,需要引擎提供专门的支持(如 vLLM 的 `deepseek_v4_fp8`),不属于通用格式。 + +### 训练与推理精度 + +- **rollout 侧**:低精度权重直接使用其自带的 `quantization_config` 推理;高精度权重既可以高精度推理,也可以通过 override `quantization_config` 的方式使用低精度推理,以获取更高的吞吐。 +- **trainer 侧**:总是使用高精度训练。低精度权重会被反量化到高精度后再训练,如 DeepSeek-V4.1-Flash、Kimi-K3。为了减少 training-inference mismatch,低精度推理时训练侧可以开启 Quantization Aware Training(QAT)。 + +### 低精度推理方案:override quantization_config + +rollout 统一采用 **override HF `quantization_config` + `load_format=dummy`** 的方式: + +1. 引擎启动时用 `dummy` 加载,并把 `quantization_config` override 成目标量化格式。这样引擎会按"已量化 checkpoint"创建参数,例如 FP8 的 `weight` + `weight_scale_inv`、MXFP4 的 `weight_packed` + `weight_scale`。 +2. 每次权重同步时,把 BF16 权重量化成 checkpoint 格式(量化可以在 trainer 侧或 rollout worker 侧完成),再交给引擎的 `load_weights`,由引擎完成融合、TP 切分以及 kernel 布局转换。 + +对于原生低精度的模型(如 GLM-5.3、Kimi-K3),不需要 override,权重同步时按其自带的 `quantization_config` 量化即可。 + +#### 为什么选择 compressed-tensors + +目标是灵活表达 (linear, moe) x (fp8, mxfp8, mxfp4, nvfp4) 的任意组合。各候选格式的能力如下: + +| 格式 | 按层混合 | 覆盖的精度 | 备注 | +|---|---|---|---| +| `fp8`(HF FineGrainedFP8) | 否 | FP8 block;SGLang 支持 `quant_method: mxfp8` | verl 现有方案,两个引擎都支持 | +| `modelopt_fp8` / `modelopt_fp4` | 否 | FP8 / NVFP4 | | +| `modelopt_mixed` | 是 | FP8 + NVFP4 | 不支持 MX 格式 | +| `quark` | 是 | MXFP4 + FP8 | 主要面向 AMD | +| **`compressed-tensors`** | **是**(`config_groups`) | **FP8 / MXFP8 / MXFP4 / NVFP4** | Kimi-K3 和 verl QAT 已在使用 | + +compressed-tensors 的优势: + +- 用 `config_groups` 为不同 target 指定不同方案。target 匹配时先按层名(支持 `re:` 正则)、再按类名,因此 `re:.*mlp\.experts\..*` 会优先于 `Linear` 命中 routed experts。 +- 是开放的格式规范(`compressed_tensors.QuantizationArgs`),vLLM 和 SGLang 都支持。 +- verl QAT 已经产出 compressed-tensors 格式,NVFP4 W4A4 所需的 `input_global_scale` 由 QAT observer 维护,可以直接复用。 + +#### 引擎支持矩阵(compressed-tensors) + +| 精度 | GPU | vLLM Linear | vLLM MoE | SGLang Linear | SGLang MoE | +|---|---|---|---|---|---| +| FP8(block / channel) | Hopper / Blackwell | ✅ `W8A8Fp8` | ✅ `W8A8Fp8MoE` | ✅ `W8A8Fp8` | ✅ `W8A8Fp8MoE` | +| MXFP8 | Blackwell | ✅ `W8A8Mxfp8` | ✅ `W8A8Mxfp8MoE` | ❌(只能用 `Fp8Config(quant_method="mxfp8")`,整模型统一) | ❌(同左) | +| MXFP4 | Blackwell | ✅ `W4A4Mxfp4` | ✅ `W4A4Mxfp4MoE` | ❌ | ✅ `Mxfp4MoEMethod` | +| NVFP4 | Blackwell | ✅ W4A4 / W4A16 | ✅ W4A4 / W4A16 | ✅ 仅 W4A4 | ✅ W4A4 | + +vLLM 可以覆盖全部组合;SGLang 只支持其中一部分,缺失部分需要 SGLang 上游补齐(可参考 vLLM 对应的 scheme 移植)。 + +硬件方面:MXFP8 / MXFP4 / NVFP4 只在 Blackwell(SM100)上支持。Hopper 上 vLLM 虽然会退化为 Marlin weight-only kernel(激活保持 bf16),但数值与 Blackwell 上的 W8A8 / W4A4 不同,verl 不支持这种用法。 + +#### 配置 + +```yaml +actor_rollout_ref: + rollout: + quantization: # null 表示不做低精度推理 + linear: mxfp8 # fp8 | mxfp8 | mxfp4 | nvfp4 | null | {weight, activation} + moe: mxfp4 + ignore: [] # 额外跳过的层,支持精确层名或 "re:" + quantization_config_file: null # 直接提供 compressed-tensors/HF JSON 时优先使用(QAT、torchao) +``` + +为了向后兼容,`quantization: fp8` 等价于 `{linear: fp8, moe: fp8}`。router(`mlp.gate`)和 `lm_head` 默认加入 ignore。 + +#### 配置翻译规则 + +verl 内部根据 `(engine, linear, moe, ignore)` 生成 override 用的 `quantization_config`: + +1. **linear 和 moe 都是 fp8**:生成 HF FineGrainedFP8,沿用现有路径。 + + ```json + {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [128, 128], + "ignored_layers": ["model.layers.0.mlp.gate", "..."]} + ``` + +2. **linear 和 moe 都是 mxfp8**:SGLang 生成 `{"quant_method": "mxfp8", ...}`,vLLM 生成 compressed-tensors。 +3. **其他组合**:生成 compressed-tensors。 + + ```json + { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "ignore": ["re:.*mlp\\.gate$", "lm_head"], + "config_groups": { + "moe": { + "targets": ["re:.*mlp\\.experts\\..*"], + "format": "mxfp4-pack-quantized", + "weights": {"type": "float", "num_bits": 4, "strategy": "group", "group_size": 32, + "symmetric": true, "dynamic": false, "scale_dtype": "torch.uint8"}, + "input_activations": null + }, + "linear": { + "targets": ["Linear"], + "weights": {"type": "float", "num_bits": 8, "strategy": "group", "group_size": 32, + "symmetric": true, "dynamic": false, "scale_dtype": "torch.uint8"}, + "input_activations": {"type": "float", "num_bits": 8, "strategy": "group", "group_size": 32, + "symmetric": true, "dynamic": true, "scale_dtype": "torch.uint8"} + } + } + } + ``` + + 上例为 "linear MXFP8 + MoE MXFP4 (W4A16)",MoE 部分与 Kimi-K3 自带的配置一致。各字段的具体取值需要与所用的 `compressed-tensors` 版本核对。 + +4. **SGLang 校验**:生成的配置如果包含上述支持矩阵中 SGLang 不支持的组合,直接报错。 + +#### 权重同步 + +- 量化器输出 **checkpoint 格式** 的张量:按 expert 拆分、不做 TP 切分,参数名与 checkpoint 一致,例如: + - FP8:`weight` + `weight_scale_inv`; + - compressed-tensors FP4:`weight_packed` + `weight_scale`,NVFP4 另有 `weight_global_scale`、`input_global_scale`。 +- vLLM 侧**不使用** vLLM 自带的 layerwise reload(`initialize_layerwise_reload` / `finalize_layerwise_reload`)。它会把一个 layer 的所有权重都攒齐之后才开始处理,对大规模 MoE 模型来说,一个 MoE layer 的全部 expert 权重会同时驻留在显存里,容易 OOM。 +- vLLM 侧沿用 `verl/utils/vllm/vllm_fp8_utils.py`(`Fp8LinearMethod` / `Fp8MoEMethod`)和 `verl/utils/vllm/vllm_fp4_utils.py`(`Mxfp4MoEMethod`)的 **stage → load → process** 做法,入口在 `verl/utils/vllm/vllm_quant_utils.py`。它要解决的问题是:kernel 后处理会把 weight/scale 改写成推理布局(如 DeepGEMM 把 block scale 打包成 int32 UE8M0,Marlin 对 FP4 权重做 repack,FlashInfer 交换 w13 的 gate/up 并做 tile shuffle),checkpoint 格式的张量无法直接写回。具体分三步: + 1. **记录 checkpoint 布局**:首次加载时,通过 patch `process_weights_after_loading` 记下各参数的 checkpoint shape/dtype(只记布局,不保留张量),或者像 MXFP4 那样直接按 `create_weights` 推导出来。 + 2. **stage(权重同步前)**:把活跃参数暂存起来,为每个参数换上一个 checkpoint 布局的 staging buffer,并带上 `weight_loader` 等参数元数据,供 `load_weights` 写入。只要字节数相同,就直接 reinterpret 原有存储(例如 MXFP4 expert 权重在 DeepGEMM 和 Marlin 下都是这样),不额外分配显存;只有元素个数会变化的参数(主要是 block scale,大小约为权重的 1/16)才需要新分配。权重按 bucket 流式写入,不会把整层攒在一起。 + 3. **process(全部 bucket 加载完之后)**:重新执行原有的 `process_weights_after_loading`,同时 patch `replace_parameter`,把后处理结果 `copy_` 回原来的活跃参数存储,CUDA graph 捕获的地址保持不变。这一步必须在 `replace_parameter` 调用时就写回,因为 MoE 的 `_setup_kernel` 末尾会用当时 layer 上的张量重建 `moe_quant_config` 和 `moe_kernel`。 +- 新增 compressed-tensors scheme(MXFP8、MXFP4、NVFP4 的 linear 和 MoE)时,按同样的模式实现各自的 stage/process:从 `create_weights` 推导出 checkpoint 布局,拦截对应模块里的 `replace_parameter`,并对不能原地写回的后端(如参数不以 Parameter 形式存放的 Triton MXFP4 后端)显式报错。注意 compressed-tensors MXFP4 MoE 会在后处理中把 `w13_weight_packed` 改名为 `w13_weight`,stage 时需要恢复成 checkpoint 中的参数名。 +- SGLang 侧沿用 trainer 侧量化后 `update_weights_from_tensor` 的方式。compressed-tensors 方案下是否会重新执行后处理,需要验证。 + +#### 注意事项 + +1. **NVFP4 W4A4 需要静态的 `input_global_scale`**:使用 `dummy` 加载时没有校准数据。可以使用 W4A16(仅 vLLM 支持),或开启 QAT 由 observer 维护该 scale,或额外进行一次校准。MX 格式的 activation 为动态量化,不需要校准。 +2. **数值格式必须逐 bit 对齐**:UE8M0 scale 的编码、FP4 nibble 的打包顺序、NVFP4 两级 scale 的约定,都必须与引擎期望的一致。建议为每种格式编写往返测试:verl 量化,引擎反量化,再与 BF16 对比。 +3. **activation 精度由 scheme 决定**:例如 compressed-tensors 的 MXFP4 linear 是 W4A4,而 Kimi-K3 的 MoE 是 W4A16(`input_activations: null`)。配置中的 `activation` 字段需要准确映射到 `input_activations`。 + +## Quantization Aware Training + +QAT 在训练时对权重(W4A4 时还包括 activation)做 fake quantization:前向使用"量化后再反量化"的值,反向通过 STE 直接回传梯度,使模型学会容忍低精度推理带来的误差,从而缩小 training-inference mismatch。权重同步时,trainer 把 BF16 权重真正量化成 checkpoint 格式,再发给 rollout。 + +QAT 由两部分组成: + +1. **训练侧 fake quantization**:与训练后端强相关,需要把量化器插入后端实际使用的 Linear 层。 +2. **导出格式**:权重同步时输出的张量名和 scale 约定,它决定 rollout 侧使用哪个 `quantization_config`。 + +这两部分是相互独立的。训练侧可以按后端选择不同的实现,导出格式则应该统一为 compressed-tensors,与上文的 low precision inference 方案保持一致。 + +### FSDP/VeOmni + +- **fake quantization**:使用 verl 自己实现的 `QATLinear`(`verl/utils/qat/linear.py`),替换模型中的 `nn.Linear`。W4A4 模式下,activation 的 `input_global_scale` 由 observer 在训练中维护。 +- **导出**:`verl/utils/qat/quantizer.py` 输出 compressed-tensors NVFP4 格式,即 `weight_packed`、`weight_scale`、`weight_global_scale`,W4A4 时还有 `input_global_scale`。 +- **global scale 融合**:vLLM 会把 q/k/v、gate/up 融合成一个层,融合后的层只能有一个 global scale,因此导出时需要用 `fuse_global_scales` 对齐各分片的 global scale(取 min)。 +- **rollout 配置**:`recipe/qat/config/nvfp4_w4a16.json`、`nvfp4_w4a4.json`(`quant_method: compressed-tensors`)。 + +### Megatron-LM + +Megatron-LM + TransformerEngine + ModelOpt + +- **fake quantization**:Megatron 中实际计算的是 TE 实现的层(带 TP 切分的 `ColumnParallelLinear` / `RowParallelLinear`,以及 expert 使用的 `TEGroupedLinear`),无法沿用 FSDP 替换 `nn.Linear` 的方式。Megatron-Bridge 官方的做法是使用 `modelopt_transformer_layer_spec` 构建模型,再调用 `mtq.quantize` 插入 ModelOpt 量化器(`verl/utils/modelopt/quantize.py`)。因此训练侧继续使用 ModelOpt。 +- **导出(现状)**:`verl/utils/modelopt/qat_weight_exporter.py` 输出 ModelOpt checkpoint 格式,rollout 使用 `quant_method: modelopt`(`recipe/qat/config/nvfp4_w4a16_megatron.json`),vLLM 侧需要单独维护一套 patch(`verl/utils/modelopt/vllm_modelopt_patch.py`)。 + +### 两个后端的导出差异 + +两个后端导出的都是 NVFP4,量化方式完全相同:FP4 E2M1,每 16 个元素一个 FP8 E4M3 block scale,外加一个 FP32 global scale。区别只在参数名和 global scale 的表示方式: + +| | FSDP/VeOmni(compressed-tensors) | Megatron(ModelOpt) | +|---|---|---| +| FP4 打包后的权重 | `weight_packed` | `weight` | +| block scale(FP8 E4M3,每 16 个元素一个) | `weight_scale` | `weight_scale` | +| weight global scale | `weight_global_scale` = `6 * 448 / amax` | `weight_scale_2` = `amax / (6 * 448)` | +| activation global scale(W4A4) | `input_global_scale` | `input_scale`(同为倒数关系) | +| `ignore_patterns` 语法 | `re:` 正则,匹配 HF 层名 | fnmatch,匹配 Megatron 模块名,并手写映射(`lm_head` 对应 `*output_layer*`,`*mlp.gate` 对应 `*router*`) | +| vLLM patch | `verl/utils/qat/vllm_patch.py` | `verl/utils/modelopt/vllm_modelopt_patch.py` | + +可以看出,`weight_global_scale` 与 `weight_scale_2` 互为倒数,`input_global_scale` 与 `input_scale` 同理。 + +### 统一方案:训练侧保留 ModelOpt,导出改为 compressed-tensors + +在 `QATWeightExporter._quantize_nvfp4` 中改为输出 compressed-tensors 格式: + +```python +yield (name.replace(".weight", ".weight_packed"), quantized) +yield (_derive_scale_name(name, "weight_scale"), w_scale) +yield (_derive_scale_name(name, "weight_global_scale"), 1.0 / w_scale_2) +if input_scale is not None: + yield (_derive_scale_name(name, "input_global_scale"), 1.0 / input_scale) +``` + +这样改之后: + +- **rollout 配置统一**:两个后端共用 `nvfp4_w4a16.json` / `nvfp4_w4a4.json`,`nvfp4_w4a16_megatron.json` 可以删除。rollout 侧不需要关心训练后端,与上文 "override quantization_config" 的方案一致。 +- **少维护一套 patch**:`vllm_modelopt_patch.py` 可以下线。`qat/vllm_patch.py` 也可以改写成与 `vllm_fp8_utils.py`、`vllm_fp4_utils.py` 一致的 stage/process 模式,统一维护。 +- **global scale 天然一致**:Megatron 的 `linear_qkv`、`linear_fc1` 本身就是融合层,只有一个量化器和一个 amax,导出后 q/k/v、gate/up 的 global scale 自然相同,不需要额外的 `fuse_global_scales`。MoE 的 `w1`/`w3` 同理,这也满足 vLLM compressed-tensors NVFP4 MoE 要求 `w13` 两个 global scale 相同的约束。 +- **ignore 统一用 HF 层名和 `re:` 语法**:Megatron 侧通过 bridge 的 `mapping_registry` 把 HF 层名翻译到 Megatron 模块名(可复用 `_iter_hf_to_megatron_matches`),不再需要手写映射。 + +### 待验证 + +1. **FP4 打包逐 bit 一致**:ModelOpt 的 `to_quantized_weight` 与 compressed-tensors 的 `compress_weight` 的 nibble 打包顺序、block scale 排布是否一致。可以对同一个 BF16 权重分别过两条导出路径,比较输出张量或 vLLM 反量化后的结果。 +2. **activation scale 的倒数关系**:确认 `NVFP4QTensor.get_activation_scaling_factor` 的定义为 `amax / (6 * 448)`,这决定了 W4A4 时 `input_global_scale` 取倒数是否正确。Megatron 目前仅支持 W4A16(`build_quantize_config` 中写死),暂不阻塞。 +3. **MX 格式 fake quantization**:若 Megatron QAT 要扩展到 MXFP8 / MXFP4,需要确认 ModelOpt 是否提供对应的量化配置。导出侧同样按 compressed-tensors 的 MX scheme 输出即可。 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/experimental/agent_loop/test_megatron_vllm_logprob_diff.py b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py new file mode 100644 index 00000000000..78ef5807654 --- /dev/null +++ b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py @@ -0,0 +1,192 @@ +# 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, recompute old_log_probs with Megatron, and check the rollout-vs-actor log-prob gap. + +Usage (8 GPUs): + MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 \ + .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py + MODEL_PATH=~/models/Qwen3.5-35B-A3B EP=8 \ + .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py + # fp8 blockwise rollout, weights quantized on the trainer side + MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK \ + .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py + # same, with Megatron ModelOpt QAT fake-quantizing to the rollout scheme + MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK QAT_MODE=fp8 \ + .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py +""" + +import json +import os + +import numpy as np +import pytest +import ray +import torch +from omegaconf import DictConfig +from hydra import compose, initialize_config_dir + +from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager_and_workers +from verl.protocol import DataProto +from verl.trainer.ppo.rollout_corr_helper import compute_offpolicy_metrics +from verl.utils import tensordict_utils as tu +from verl.utils.debug.metrics import calculate_debug_metrics +from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding + +TP = int(os.getenv("TP", 2)) +PP = int(os.getenv("PP", 1)) +CP = int(os.getenv("CP", 1)) +EP = int(os.getenv("EP", int(os.environ.get("EP", 1)))) +ETP = int(os.getenv("ETP", 1)) +GEN_TP = int(os.getenv("GEN_TP", 2)) +N_GPUS = int(os.getenv("N_GPUS", 8)) +MAX_K3_KL = float(os.getenv("MAX_K3_KL", 1e-2)) +# Rollout compressed-tensors presets, e.g. QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK; unset keeps the rollout in bf16. +QUANT_LINEAR = os.getenv("QUANT_LINEAR") +QUANT_MOE = os.getenv("QUANT_MOE") +# Megatron ModelOpt QAT mode matching the rollout scheme, e.g. QAT_MODE=fp8; unset disables QAT. +QAT_MODE = os.getenv("QAT_MODE") + + +@pytest.fixture +def init_config() -> DictConfig: + + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose( + config_name="ppo_trainer", + overrides=[ + "model_engine=megatron", + f"trainer.n_gpus_per_node={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={TP}", + f"actor_rollout_ref.actor.megatron.pipeline_model_parallel_size={PP}", + f"actor_rollout_ref.actor.megatron.context_parallel_size={CP}", + f"actor_rollout_ref.actor.megatron.expert_model_parallel_size={EP}", + f"actor_rollout_ref.actor.megatron.expert_tensor_parallel_size={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", + "actor_rollout_ref.rollout.name=vllm", + "actor_rollout_ref.rollout.mode=async", + f"actor_rollout_ref.rollout.tensor_model_parallel_size={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", + ], + ) + + config.actor_rollout_ref.model.path = os.path.expanduser(os.getenv("MODEL_PATH", "~/models/Qwen3.5-35B-A3B")) + if QUANT_LINEAR or QUANT_MOE: + config.actor_rollout_ref.rollout.quantization = {"linear": QUANT_LINEAR, "moe": QUANT_MOE, "ignore": []} + if QAT_MODE: + config.actor_rollout_ref.actor.megatron.qat.enable = True + config.actor_rollout_ref.actor.megatron.qat.mode = QAT_MODE + return config + + +def compute_old_log_prob(actor_rollout_wg, 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) + entropy = no_padding_2_padding(tu.get(output, "entropy"), batch_td) + return DataProto.from_tensordict( + tu.get_tensordict({"old_log_probs": log_probs.float(), "entropys": entropy.float()}) + ) + + +def test_megatron_vllm_logprob_diff(init_config): + 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_agent_loop_manager_and_workers(init_config) + + 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)?", + ] + 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)), + }, + ) + n = init_config.actor_rollout_ref.rollout.n + batch = batch.repeat(n) + + # rollout + result = agent_loop_manager.generate_sequences(prompts=batch) + checkpoint_manager.sleep_replicas() + assert len(result) == len(raw_prompts) * n + assert "rollout_log_probs" in result.batch + result.meta_info.pop("timing", None) + + # recompute old_log_probs with megatron + result.meta_info["temperature"] = init_config.actor_rollout_ref.rollout.temperature + result.meta_info["global_token_num"] = torch.sum(result.batch["attention_mask"], dim=-1).tolist() + old_log_prob = compute_old_log_prob(actor_rollout_wg, result) + old_log_prob.batch.pop("entropys") + result = result.union(old_log_prob) + + old_log_probs = result.batch["old_log_probs"] + rollout_log_probs = result.batch["rollout_log_probs"] + response_mask = result.batch["response_mask"].bool() + assert old_log_probs.shape == rollout_log_probs.shape == response_mask.shape + + 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()) + print(json.dumps(metrics, indent=2)) + + assert np.isfinite(metrics["kl"]) and np.isfinite(metrics["k3_kl"]) + assert metrics["k3_kl"] < MAX_K3_KL, f"k3_kl {metrics['k3_kl']} exceeds {MAX_K3_KL}" + + print("Test passed!") + ray.shutdown() 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..5f991403d59 100644 --- a/tests/utils/test_vllm_weight_name_normalization_on_cpu.py +++ b/tests/utils/test_vllm_weight_name_normalization_on_cpu.py @@ -296,6 +296,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 +732,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 +765,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 +908,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/verl/trainer/config/engine/megatron.yaml b/verl/trainer/config/engine/megatron.yaml index 0ed6528e80b..a49f8848f0d 100644 --- a/verl/trainer/config/engine/megatron.yaml +++ b/verl/trainer/config/engine/megatron.yaml @@ -140,7 +140,8 @@ qat: _target_: verl.workers.config.QATEngineConfig # Whether to enable QAT enable: false - # Quantization mode: "w4a16" (weight-only FP4) or "w4a4" (weight + activation FP4) + # Quantization mode: "w4a16" (weight-only FP4), or "fp8" (128x128 block weights, 1x128 dynamic + # activations) to train against a rollout.quantization={linear: FP8_BLOCK, moe: FP8_BLOCK} rollout mode: "w4a16" # Group size for blockwise quantization (NVFP4 requires 16) group_size: 16 diff --git a/verl/trainer/config/rollout/rollout.yaml b/verl/trainer/config/rollout/rollout.yaml index dbf59f10ded..a7750555d01 100644 --- a/verl/trainer/config/rollout/rollout.yaml +++ b/verl/trainer/config/rollout/rollout.yaml @@ -481,7 +481,9 @@ 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 +# type of quantization in vllm, currently support fp8 and torchao. +# Or a mapping of compressed-tensors presets {linear: FP8_BLOCK, moe: FP8_BLOCK, ignore: []} for a rollout +# whose weights are quantized on the trainer side (vLLM + Megatron only for now). quantization: null # extra quantization information serialized in a config file, e.g. torchao_config.json diff --git a/verl/utils/modelopt/__init__.py b/verl/utils/modelopt/__init__.py index cc151f98f20..faf96665dd4 100644 --- a/verl/utils/modelopt/__init__.py +++ b/verl/utils/modelopt/__init__.py @@ -13,24 +13,13 @@ # 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.""" +"""ModelOpt fake quantization for Megatron QAT training.""" -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.qat_utils import apply_qat_to_modules 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 index ca002c5132c..b5fc0b1afe8 100644 --- a/verl/utils/modelopt/qat_utils.py +++ b/verl/utils/modelopt/qat_utils.py @@ -35,11 +35,3 @@ def apply_qat_to_modules(modules, qat_config): 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 index dee7274099c..9e56af81f23 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -13,13 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ModelOpt NVFP4 quantization config and application for Megatron QAT.""" +"""ModelOpt quantization config and application for Megatron QAT.""" 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, +) _NVFP4_W4A16_QUANTIZER_CFG = { "*weight_quantizer": { @@ -31,6 +37,63 @@ "*input_quantizer": {"enable": False}, } +_FP8_BLOCK_BACKEND = "verl_fp8_block" +_FP8_MAX = torch.finfo(torch.float8_e4m3fn).max + + +class _FP8BlockFakeQuant(torch.autograd.Function): + """Dynamic FP8 E4M3 fake quantization over ``(rows, cols)`` blocks of the last two dims, 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) + + +if not is_registered_quant_backend(_FP8_BLOCK_BACKEND): + register_quant_backend(_FP8_BLOCK_BACKEND, _fp8_block_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 = {"w4a16": _NVFP4_W4A16_QUANTIZER_CFG, "fp8": _FP8_BLOCK_QUANTIZER_CFG} + def _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: cfg = [] @@ -52,21 +115,22 @@ def build_quantize_config( 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 qat_mode not in _QUANTIZER_CFGS: + raise ValueError(f"Supported QAT modes are {sorted(_QUANTIZER_CFGS)}, 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) + 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) quant_cfg.extend(ignore_cfg) - return {"quant_cfg": quant_cfg, "algorithm": "max"} + # The fp8 backend derives its scales from each forward's tensor, so there is nothing to calibrate. + return {"quant_cfg": quant_cfg, "algorithm": "max" if qat_mode == "w4a16" else None} def apply_qat( 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..9c86b810435 100644 --- a/verl/utils/qat/__init__.py +++ b/verl/utils/qat/__init__.py @@ -21,7 +21,9 @@ - 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 + +The vLLM side (NVFP4 refit patches) lives in ``verl.utils.vllm.vllm_nvfp4_utils``. Usage: from verl.utils.qat import apply_qat, QATConfig @@ -37,21 +39,22 @@ 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, + +from verl.utils.qat.compressed_tensors import ( + CompressedTensorsWeightQuantizer, + build_compressed_tensors_config, + is_compressed_tensors_quantization, + qat_mode_for_scheme, ) __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", + "build_compressed_tensors_config", + "is_compressed_tensors_quantization", + "qat_mode_for_scheme", ] diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py new file mode 100644 index 00000000000..b1dcc53702b --- /dev/null +++ b/verl/utils/qat/compressed_tensors.py @@ -0,0 +1,317 @@ +# 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 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 + +``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 os +import re +import struct +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", + "build_compressed_tensors_config", + "is_compressed_tensors_quantization", + "qat_mode_for_scheme", +] + +# 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", +} + +# 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 is_compressed_tensors_quantization(quantization: Any) -> bool: + """Whether the rollout ``quantization`` field selects this path (a mapping).""" + return isinstance(quantization, Mapping) + + +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 + + +def qat_mode_for_scheme(scheme: str) -> Optional[str]: + """ModelOpt QAT mode whose fake quant matches the given rollout preset.""" + return SUPPORTED_SCHEMES.get(scheme.upper()) + + +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 _iter_checkpoint_shapes(model_path: str) -> Iterator[tuple[str, list[int]]]: + """Yield ``(name, shape)`` from safetensors headers without reading tensor data.""" + index_file = os.path.join(model_path, "model.safetensors.index.json") + if os.path.exists(index_file): + with open(index_file) as f: + files = sorted(set(json.load(f)["weight_map"].values())) + else: + files = [f for f in sorted(os.listdir(model_path)) if f.endswith(".safetensors")] + for file in files: + with open(os.path.join(model_path, file), "rb") as f: + (header_len,) = struct.unpack(" str: + """``model.layers.3.linear_attn.in_proj_a`` -> ``re:.*linear_attn\\.in_proj_a$``.""" + suffix = re.split(r"\.\d+\.", module_name)[-1] + return "re:.*" + re.escape(suffix) + "$" + + +def _weight_tile(args: "QuantizationArgs") -> Optional[tuple[int, int]]: + """(rows, cols) a weight's shape must be a multiple of, if the strategy tiles it.""" + from compressed_tensors.quantization import QuantizationStrategy + + if args.strategy == QuantizationStrategy.BLOCK: + return tuple(args.block_structure) + if args.strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): + return 1, args.group_size + return None + + +def _incompatible_shape_ignores(model_path: str, ignore: list[str], tile: tuple[int, int]) -> list[str]: + """Patterns for 2D weights whose shape is not a multiple of the weight tile.""" + patterns: list[str] = [] + for name, shape in _iter_checkpoint_shapes(model_path): + if not name.endswith(".weight") or len(shape) != 2: + continue + if shape[0] % tile[0] == 0 and shape[1] % tile[1] == 0: + continue + module = name[: -len(".weight")] + if _matches(module, ignore + patterns): + continue + patterns.append(_module_pattern(module)) + return patterns + + +def _config_group(targets: list[str], preset: str) -> dict[str, Any]: + from compressed_tensors.compressors.format import infer_module_format + from compressed_tensors.quantization import preset_name_to_scheme + + scheme = preset_name_to_scheme(preset, targets) + group = scheme.model_dump(mode="json", exclude_none=True) + group["format"] = infer_module_format(torch.nn.Linear, scheme).value + return group + + +def build_compressed_tensors_config(quantization: Mapping, model_path: str) -> dict[str, Any]: + """Build the HF compressed-tensors ``quantization_config`` for the rollout engine.""" + from compressed_tensors.quantization import preset_name_to_scheme + + linear = _check_scheme("linear", quantization.get("linear")) + moe = _check_scheme("moe", quantization.get("moe")) + if linear is None and moe is None: + raise ValueError("rollout quantization needs at least one of `linear` / `moe`") + + ignore = DEFAULT_IGNORE + list(quantization.get("ignore") or []) + for preset in {linear, moe} - {None}: + tile = _weight_tile(preset_name_to_scheme(preset, ["Linear"]).weights) + if tile is not None: + ignore += _incompatible_shape_ignores(model_path, ignore, tile) + + 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") -> 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} + global_scale = None + if args.strategy == QuantizationStrategy.TENSOR_GROUP: + 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; routed experts take the MoE group's scheme when present, the rest + the ``Linear`` group's. Packed 3D experts are split per expert first. + Everything else passes through unchanged. + """ + + def __init__(self, ct_config: Mapping[str, Any]): + from compressed_tensors.quantization import QuantizationScheme + + self.ignore = list(ct_config["ignore"]) + self.moe_scheme: Optional[QuantizationScheme] = None + self.linear_scheme: Optional[QuantizationScheme] = None + for group in ct_config["config_groups"].values(): + scheme = QuantizationScheme.model_validate({k: v for k, v in group.items() if k != "format"}) + if MOE_TARGET in scheme.targets: + self.moe_scheme = scheme + if "Linear" in scheme.targets: + self.linear_scheme = scheme + if self.moe_scheme is None and MOE_TARGET not in self.ignore: + self.moe_scheme = 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 + module = name[: -len(".weight")] + if _matches(module, self.ignore): + return None + if _matches(module, [MOE_TARGET]): + return self.moe_scheme + return self.linear_scheme + + def __call__(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]: + for name, tensor in weights: + if tensor.ndim == 3 and self.moe_scheme is not None and _PACKED_EXPERTS.match(name): + yield from self(_unpack_experts(name, tensor)) + 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 diff --git a/verl/utils/vllm/vllm_fp8_utils.py b/verl/utils/vllm/vllm_fp8_utils.py index 4e139b8442d..8155da0c432 100644 --- a/verl/utils/vllm/vllm_fp8_utils.py +++ b/verl/utils/vllm/vllm_fp8_utils.py @@ -313,13 +313,13 @@ 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) _restore_layer_param_subclass_attrs(layer, old_params) @@ -446,6 +446,34 @@ def _create_param_from_subclass_attributes(custom_data, custom_weight): ) +def _build_compressed_tensors_fp8_patchers(): + """Same refit treatment for the compressed-tensors linear and FP8 MoE methods.""" + ct = "vllm.model_executor.layers.quantization.compressed_tensors" + ct_moe = f"{ct}.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8" + 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 + return [ + patch( + f"{ct}.compressed_tensors.CompressedTensorsLinearMethod.process_weights_after_loading", + wrap(CompressedTensorsLinearMethod.process_weights_after_loading), + ), + patch( + f"{ct_moe}.CompressedTensorsW8A8Fp8MoEMethod.process_weights_after_loading", + wrap(CompressedTensorsW8A8Fp8MoEMethod.process_weights_after_loading, f"{ct_moe}.replace_parameter"), + ), + ] + + def build_fp8_method_patchers(vllm_version): """Patchers that make the FP8 quant methods survive a refit, not yet started. @@ -468,6 +496,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/qat/vllm_patch.py b/verl/utils/vllm/vllm_nvfp4_utils.py similarity index 100% rename from verl/utils/qat/vllm_patch.py rename to verl/utils/vllm/vllm_nvfp4_utils.py diff --git a/verl/utils/vllm/vllm_quant_utils.py b/verl/utils/vllm/vllm_quant_utils.py index 702dccfea89..a12bcdeae12 100644 --- a/verl/utils/vllm/vllm_quant_utils.py +++ b/verl/utils/vllm/vllm_quant_utils.py @@ -88,10 +88,21 @@ def is_quantized_model(vllm_config): return True elif is_vllm_ascend_mx_quant(vllm_config.quant_config): return True + elif is_compressed_tensors_model(vllm_config): + return True return False +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 isinstance(getattr(vllm_config, "quant_config", None), CompressedTensorsConfig) + + # 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`` / @@ -450,7 +461,10 @@ def load_quanted_weights(weights, model_runner, is_drafter=False): vllm_dtype = model_runner.vllm_config.model_config.dtype weights = list(weights) - weights_quantized = quant_weights(weights, model, quant_config, dtype=vllm_dtype) + if is_compressed_tensors_model(model_runner.vllm_config): + weights_quantized = weights + else: + 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 diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index 22bec4a0475..7f45d3c4628 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, "w4a16", "w4a4" or "fp8" (Megatron only), default "w4a16" 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) diff --git a/verl/workers/config/rollout.py b/verl/workers/config/rollout.py index 4388b177a83..96733581cd5 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 @@ -265,7 +265,9 @@ class RolloutConfig(BaseConfig): skip_tokenizer_init: bool = True - quantization: Optional[str] = None + # str: legacy engine-side quantization ("fp8", "torchao", ...); + # mapping {linear, moe, ignore} of compressed-tensors presets (e.g. FP8_BLOCK): rollout quantized by the trainer. + quantization: Optional[Any] = None quantization_config_file: 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..3a8fa9889a9 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" @@ -1000,14 +1001,17 @@ 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. + 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..2cc4aa1fc20 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -43,6 +43,12 @@ 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 ( + CompressedTensorsWeightQuantizer, + build_compressed_tensors_config, + is_compressed_tensors_quantization, + qat_mode_for_scheme, +) 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 +735,20 @@ def init_model(self): backend, is_master=(torch.distributed.get_rank() == 0), bucket_size=bucket_size, **engine_kwargs ) + rollout_quantization = self.config.rollout.get("quantization", None) + qat = self.config.rollout.get("qat", None) or {} + if qat.get("enable", False) and not is_compressed_tensors_quantization(rollout_quantization): + raise ValueError( + "QAT needs a quantized rollout to sync to: set rollout.quantization, " + "e.g. {linear: FP8_BLOCK, moe: FP8_BLOCK}" + ) + if is_compressed_tensors_quantization(rollout_quantization): + schemes = {rollout_quantization.get("linear"), rollout_quantization.get("moe")} - {None} + if qat.get("enable", False) and {qat_mode_for_scheme(s) for s in schemes} != {qat.get("mode")}: + raise ValueError(f"QAT mode {qat.get('mode')!r} does not match the rollout quantization {schemes}") + ct_config = build_compressed_tensors_config(rollout_quantization, model_config.local_path) + self.actor.engine.set_rollout_weight_quantizer(CompressedTensorsWeightQuantizer(ct_config)) + # 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..32aeb44aeeb 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -17,6 +17,7 @@ import logging import os import secrets +from collections.abc import Mapping from pathlib import Path from typing import Any, Optional @@ -301,6 +302,8 @@ async def launch_server(self, master_address: str = None, master_port: int = Non # 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 isinstance(quantization, Mapping): + raise NotImplementedError("rollout.quantization as a {linear, moe} mapping is only supported by vLLM") if quantization is not None: if quantization == "fp8": from verl.utils.sglang.sglang_fp8_utils import build_sglang_fp8_quant_config diff --git a/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py b/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py index bd57233f171..646de2e7515 100644 --- a/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py +++ b/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py @@ -14,6 +14,7 @@ import asyncio import logging import os +from collections.abc import Mapping from typing import Any, Optional import ray @@ -207,6 +208,8 @@ async def launch_server(self): per_worker_gpu_share = 1.0 / self.max_colocate_count quantization = self.config.quantization + if isinstance(quantization, Mapping): + raise NotImplementedError("rollout.quantization as a {linear, moe} mapping is only supported by vLLM") if quantization is not None: if quantization == "fp8": FP8_BLOCK_QUANT_KWARGS = { diff --git a/verl/workers/rollout/vllm_rollout/utils.py b/verl/workers/rollout/vllm_rollout/utils.py index 213b0727b3b..db9bc8daa4c 100644 --- a/verl/workers/rollout/vllm_rollout/utils.py +++ b/verl/workers/rollout/vllm_rollout/utils.py @@ -186,24 +186,20 @@ def __new__(cls, **kwargs): ) register_verl_delta_weight_transfer_engine() + quant_config = getattr(vllm_config, "quant_config", None) if vllm_config else None + _is_nvfp4_packed = getattr(quant_config, "quant_format", None) == "nvfp4-pack-quantized" # 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): + if os.environ.get("VERL_VLLM_FP8_QUANT_ENABLED", "0") == "1" or ( + is_quantized_model(vllm_config) and not _is_nvfp4_packed + ): 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 + # 3. patch compressed-tensors NVFP4 for dynamic weight loading + if _is_nvfp4_packed: + from verl.utils.vllm.vllm_nvfp4_utils 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") + logger.info("Applied compressed-tensors NVFP4 patches in vLLM worker subprocess") # 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 +210,7 @@ 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 + instance._is_nvfp4_packed = _is_nvfp4_packed return instance def _get_drafter_model(self): @@ -276,18 +271,13 @@ 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 + if self._is_nvfp4_packed: + # Compressed-tensors NVFP4: prepare for weight loading BEFORE receiving any buckets + from verl.utils.vllm.vllm_nvfp4_utils 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") + logger.info("NVFP4: prepare_qat_for_load_weights completed") elif 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) @@ -339,18 +329,13 @@ 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 + if self._is_nvfp4_packed: + # Compressed-tensors NVFP4: call process_weights_after_loading AFTER all buckets are received + from verl.utils.vllm.vllm_nvfp4_utils 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") + logger.info("NVFP4: process_weights_after_loading completed") elif peft_config and base_sync_done: logger.info("LoRA adapter sync, no post-process needed") elif is_quantized_model(self.model_runner.vllm_config): @@ -399,7 +384,7 @@ def _update_weights( param_updates, buffer_updates, named_buffers = split_buffer_updates(self.model_runner.model, weights) # Add the FP8 related logic here as sharding manager has been deprecated. # Check if FP8 quantization is enabled and apply appropriate weight loading - if is_quantized_model(self.model_runner.vllm_config): + if is_quantized_model(self.model_runner.vllm_config) and not self._is_nvfp4_packed: logger.info(f"FP8 model detected (async): {self.model_runner.vllm_config.quant_config}") # Convert bf16 weights to fp8 format before loading model_updates = drop_tied_alias_updates(self.model_runner.model, param_updates) diff --git a/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/verl/workers/rollout/vllm_rollout/vllm_async_server.py index 21e15c938d3..65d75ae6cd6 100644 --- a/verl/workers/rollout/vllm_rollout/vllm_async_server.py +++ b/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -51,6 +51,7 @@ relocate_rollout_traces, rollout_profiler_global_ranks, ) +from verl.utils.qat import build_compressed_tensors_config, is_compressed_tensors_quantization 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 @@ -1397,54 +1398,20 @@ def _apply_quantization(self) -> tuple[Optional[str], dict]: 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 - - 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" + if is_compressed_tensors_quantization(quantization): + # The trainer quantizes every refit to this config; vLLM picks the method up from hf_overrides. + hf_overrides["quantization_config"] = build_compressed_tensors_config( + quantization, self.model_config.local_path + ) + apply_vllm_quant_patches() + os.environ["VERL_VLLM_FP8_QUANT_ENABLED"] = "1" + # The schemes carry fp32 block scales. DeepGEMM's E8M0 mode (vLLM's default) rounds weight and + # activation scales to powers of two, so the rollout would drift from what the trainer quantized. + os.environ.setdefault("VLLM_USE_DEEP_GEMM_E8M0", "0") + logger.info(f"Rollout compressed-tensors quantization_config: {hf_overrides['quantization_config']}") + return None, hf_overrides + if quantization == "fp8": + raise ValueError("rollout.quantization=fp8 is replaced by the mapping form, e.g. {linear: FP8_BLOCK, moe: FP8_BLOCK}") model_quantization_config = getattr(self.model_config.hf_config, "quantization_config", {}) or {} if quantization is None and model_quantization_config.get("quant_method") == "fp8": From 10c4b50b1ffeead3effee7cb76ffa9ac7c6976d4 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Thu, 15 Oct 2026 16:00:06 +0800 Subject: [PATCH 02/18] dev --- .../workflows/e2e_ppo_grpo_trainer_trtllm.yml | 28 ---- .../e2e_ppo_trainer_megatron_sglang.yml | 5 - .../e2e_ppo_trainer_megatron_vllm.yml | 2 +- .../e2e_ppo_trainer_megatron_vllm_2.yml | 2 +- .../e2e_ppo_trainer_megatron_vllm_rocm.yml | 2 +- .../model_dev/parameter_and_metrics.md | 4 +- .../ascend_backend_features.md | 4 +- .../feature_support/npu_advance_features.md | 4 +- docs/low_precision/fp8.md | 61 ++------ docs/low_precision/qat.md | 26 ++-- .../run_nemotron_nano_v3_30b_a3b_megatron.sh | 2 +- .../run_qwen3_30b_a3b_megatron.sh | 18 ++- .../test_megatron_vllm_logprob_diff.py | 7 +- tests/special_e2e/run_ppo_trainer_megatron.sh | 3 +- tests/utils/test_sglang_fp8_utils.py | 107 -------------- ...b_a3b_base_math_megatron_8_8_mis_trtllm.sh | 7 - .../_generated_ppo_megatron_trainer.yaml | 16 ++- .../_generated_ppo_torchtitan_trainer.yaml | 16 ++- .../config/_generated_ppo_trainer.yaml | 16 ++- .../config/_generated_ppo_veomni_trainer.yaml | 16 ++- verl/trainer/config/reward/reward.yaml | 7 +- verl/trainer/config/rollout/rollout.yaml | 26 ++-- verl/utils/qat/__init__.py | 15 +- verl/utils/qat/compressed_tensors.py | 136 ++++++++++++------ verl/utils/sglang/sglang_fp8_utils.py | 127 ---------------- verl/utils/trtllm/trtllm_fp8_utils.py | 21 --- verl/workers/config/rollout.py | 39 ++++- verl/workers/engine_workers.py | 30 ++-- .../sglang_rollout/async_sglang_server.py | 18 +-- .../rollout/sglang_rollout/sglang_rollout.py | 23 --- .../trtllm_rollout/trtllm_async_server.py | 19 +-- .../rollout/trtllm_rollout/trtllm_rollout.py | 17 --- verl/workers/rollout/vllm_rollout/utils.py | 10 +- .../rollout/vllm_rollout/vllm_async_server.py | 49 +++---- 34 files changed, 288 insertions(+), 595 deletions(-) delete mode 100644 tests/utils/test_sglang_fp8_utils.py delete mode 100644 verl/utils/sglang/sglang_fp8_utils.py delete mode 100644 verl/utils/trtllm/trtllm_fp8_utils.py 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..fffac95d850 100644 --- a/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml +++ b/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml @@ -200,11 +200,6 @@ jobs: run: | $UV_RUN ray stop --force ALL_OFFLOAD=True VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 LR_WARMUP_STEPS=1 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh - - 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 - 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/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/low_precision/fp8.md b/docs/low_precision/fp8.md index d91a0020447..12987b9821a 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 rollout with the Megatron training backend; SGLang, 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/qat.md b/docs/low_precision/qat.md index d802c6e656a..ee9bae72c30 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -83,28 +83,24 @@ vLLM 可以覆盖全部组合;SGLang 只支持其中一部分,缺失部分 ```yaml actor_rollout_ref: rollout: - quantization: # null 表示不做低精度推理 - linear: mxfp8 # fp8 | mxfp8 | mxfp4 | nvfp4 | null | {weight, activation} - moe: mxfp4 + quantization: # linear / moe / config_file 全为 null 表示不做低精度推理 + linear: FP8_BLOCK # compressed-tensors preset,目前支持 FP8_BLOCK + moe: FP8_BLOCK ignore: [] # 额外跳过的层,支持精确层名或 "re:" - quantization_config_file: null # 直接提供 compressed-tensors/HF JSON 时优先使用(QAT、torchao) + config_file: null # 直接提供 compressed-tensors quantization_config JSON,与 linear/moe/ignore 互斥 ``` -为了向后兼容,`quantization: fp8` 等价于 `{linear: fp8, moe: fp8}`。router(`mlp.gate`)和 `lm_head` 默认加入 ignore。 +只支持 compressed-tensors,不再支持 `quantization: fp8` 字符串和 torchao。router(`mlp.gate`)、`lm_head`、embedding、视觉塔和 MTP 默认加入 ignore,形状不能被 block 整除的层(如 Qwen3.5 `in_proj_a/b`)自动加入 ignore。 #### 配置翻译规则 -verl 内部根据 `(engine, linear, moe, ignore)` 生成 override 用的 `quantization_config`: +verl 内部(`verl/utils/qat/compressed_tensors.py`)生成 override 用的 compressed-tensors `quantization_config`,trainer 侧量化器使用同一份配置: -1. **linear 和 moe 都是 fp8**:生成 HF FineGrainedFP8,沿用现有路径。 +1. **linear 和 moe 相同**:一个 `targets: ["Linear"]` 的 config group。 +2. **linear 和 moe 不同**:MoE 用 `re:.*mlp\.experts\..*` 单独成组,其余 `Linear` 一组;某一项为 null 时对应层保持高精度。 +3. **`config_file`**:JSON 可以是 `quantization_config` 本身,也可以是带 `quantization_config` 字段的 HF `config.json`,原样使用(不追加默认 ignore)。每个 config group 的数值配置必须与某个已支持的 preset 一致,否则报错;QAT mode 也据此校验。 - ```json - {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [128, 128], - "ignored_layers": ["model.layers.0.mlp.gate", "..."]} - ``` - -2. **linear 和 moe 都是 mxfp8**:SGLang 生成 `{"quant_method": "mxfp8", ...}`,vLLM 生成 compressed-tensors。 -3. **其他组合**:生成 compressed-tensors。 + 多精度组合(目前未支持)的配置示例: ```json { @@ -132,7 +128,7 @@ verl 内部根据 `(engine, linear, moe, ignore)` 生成 override 用的 `quanti 上例为 "linear MXFP8 + MoE MXFP4 (W4A16)",MoE 部分与 Kimi-K3 自带的配置一致。各字段的具体取值需要与所用的 `compressed-tensors` 版本核对。 -4. **SGLang 校验**:生成的配置如果包含上述支持矩阵中 SGLang 不支持的组合,直接报错。 +4. **引擎**:目前只有 vLLM + Megatron 支持;SGLang / TRT-LLM 设置 `rollout.quantization` 会直接报错。 #### 权重同步 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/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py index 78ef5807654..14a0458f459 100644 --- a/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py +++ b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py @@ -33,8 +33,8 @@ import pytest import ray import torch -from omegaconf import DictConfig from hydra import compose, initialize_config_dir +from omegaconf import DictConfig from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager_and_workers from verl.protocol import DataProto @@ -60,7 +60,6 @@ @pytest.fixture def init_config() -> DictConfig: - with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): config = compose( config_name="ppo_trainer", @@ -94,8 +93,8 @@ def init_config() -> DictConfig: ) config.actor_rollout_ref.model.path = os.path.expanduser(os.getenv("MODEL_PATH", "~/models/Qwen3.5-35B-A3B")) - if QUANT_LINEAR or QUANT_MOE: - config.actor_rollout_ref.rollout.quantization = {"linear": QUANT_LINEAR, "moe": QUANT_MOE, "ignore": []} + config.actor_rollout_ref.rollout.quantization.linear = QUANT_LINEAR + config.actor_rollout_ref.rollout.quantization.moe = QUANT_MOE if QAT_MODE: config.actor_rollout_ref.actor.megatron.qat.enable = True config.actor_rollout_ref.actor.megatron.qat.mode = QAT_MODE 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/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/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..8f36fb143dd 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -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 @@ -864,8 +867,11 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + 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..a4c649a7ccc 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,11 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + 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..787b46d660e 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,11 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + 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..1a9e539fb5c 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,11 @@ reward: 'n': 1 prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + linear: null + moe: null + ignore: [] + config_file: null mtp: null disaggregation: enabled: false diff --git a/verl/trainer/config/reward/reward.yaml b/verl/trainer/config/reward/reward.yaml index 6553e56737e..9614422429f 100644 --- a/verl/trainer/config/reward/reward.yaml +++ b/verl/trainer/config/reward/reward.yaml @@ -72,8 +72,11 @@ reward_model: prompt_length: 2048 response_length: 2048 - quantization: null - quantization_config_file: null + quantization: + 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 a7750555d01..05032607599 100644 --- a/verl/trainer/config/rollout/rollout.yaml +++ b/verl/trainer/config/rollout/rollout.yaml @@ -481,20 +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. -# Or a mapping of compressed-tensors presets {linear: FP8_BLOCK, moe: FP8_BLOCK, ignore: []} for a rollout -# whose weights are quantized on the trainer side (vLLM + Megatron only for now). -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 + 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/qat/__init__.py b/verl/utils/qat/__init__.py index 9c86b810435..71afd229bb2 100644 --- a/verl/utils/qat/__init__.py +++ b/verl/utils/qat/__init__.py @@ -32,6 +32,11 @@ model = apply_qat(model, config) # Before FSDP wrapping """ +from verl.utils.qat.compressed_tensors import ( + CompressedTensorsWeightQuantizer, + build_compressed_tensors_config, + qat_modes_for_config, +) from verl.utils.qat.core import ( QATConfig, apply_qat, @@ -40,13 +45,6 @@ load_quantization_config, ) -from verl.utils.qat.compressed_tensors import ( - CompressedTensorsWeightQuantizer, - build_compressed_tensors_config, - is_compressed_tensors_quantization, - qat_mode_for_scheme, -) - __all__ = [ "QATConfig", "apply_qat", @@ -55,6 +53,5 @@ "invalidate_all_scales", "CompressedTensorsWeightQuantizer", "build_compressed_tensors_config", - "is_compressed_tensors_quantization", - "qat_mode_for_scheme", + "qat_modes_for_config", ] diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py index b1dcc53702b..6bc59439637 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -13,14 +13,17 @@ # limitations under the License. """Low-precision rollout expressed as a compressed-tensors ``quantization_config``. -The rollout config selects a compressed-tensors preset scheme per layer kind:: +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 -``build_compressed_tensors_config`` turns that into the HF ``quantization_config`` +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 @@ -46,8 +49,7 @@ "SUPPORTED_SCHEMES", "CompressedTensorsWeightQuantizer", "build_compressed_tensors_config", - "is_compressed_tensors_quantization", - "qat_mode_for_scheme", + "qat_modes_for_config", ] # Matched against vLLM's unfused expert names (``...mlp.experts.0.gate_proj``), @@ -76,25 +78,48 @@ _HIGH_PRECISION_DTYPES = (torch.bfloat16, torch.float16, torch.float32) -def is_compressed_tensors_quantization(quantization: Any) -> bool: - """Whether the rollout ``quantization`` field selects this path (a mapping).""" - return isinstance(quantization, Mapping) - - 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)}" - ) + raise ValueError(f"Unsupported rollout quantization {kind}={scheme!r}, choose from {sorted(SUPPORTED_SCHEMES)}") return scheme -def qat_mode_for_scheme(scheme: str) -> Optional[str]: - """ModelOpt QAT mode whose fake quant matches the given rollout preset.""" - return SUPPORTED_SCHEMES.get(scheme.upper()) +_NUMERIC_FIELDS = {"num_bits", "type", "symmetric", "group_size", "strategy", "block_structure", "dynamic"} + + +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.""" + from compressed_tensors.quantization import preset_name_to_scheme + + for preset in SUPPORTED_SCHEMES: + reference = preset_name_to_scheme(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``.""" + return {SUPPORTED_SCHEMES[_preset_of(scheme)] for scheme in _group_schemes(ct_config)} def _matches(name: str, targets: Iterable[str]) -> bool: @@ -166,16 +191,47 @@ def _config_group(targets: list[str], preset: str) -> dict[str, Any]: return group -def build_compressed_tensors_config(quantization: Mapping, model_path: str) -> dict[str, Any]: - """Build the HF compressed-tensors ``quantization_config`` for the rollout engine.""" +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: str, +) -> dict[str, Any]: + """Build the HF compressed-tensors ``quantization_config`` for the rollout engine. + + Prefer ``RolloutQuantizationConfig.build_quantization_config``; the arguments mirror its fields. + """ from compressed_tensors.quantization import preset_name_to_scheme - linear = _check_scheme("linear", quantization.get("linear")) - moe = _check_scheme("moe", quantization.get("moe")) + 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 is None and moe is None: raise ValueError("rollout quantization needs at least one of `linear` / `moe`") - ignore = DEFAULT_IGNORE + list(quantization.get("ignore") or []) + ignore = DEFAULT_IGNORE + ignore for preset in {linear, moe} - {None}: tile = _weight_tile(preset_name_to_scheme(preset, ["Linear"]).weights) if tile is not None: @@ -273,39 +329,33 @@ 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; routed experts take the MoE group's scheme when present, the rest - the ``Linear`` group's. Packed 3D experts are split per expert first. - Everything else passes through unchanged. + 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]): - from compressed_tensors.quantization import QuantizationScheme - - self.ignore = list(ct_config["ignore"]) - self.moe_scheme: Optional[QuantizationScheme] = None - self.linear_scheme: Optional[QuantizationScheme] = None - for group in ct_config["config_groups"].values(): - scheme = QuantizationScheme.model_validate({k: v for k, v in group.items() if k != "format"}) - if MOE_TARGET in scheme.targets: - self.moe_scheme = scheme - if "Linear" in scheme.targets: - self.linear_scheme = scheme - if self.moe_scheme is None and MOE_TARGET not in self.ignore: - self.moe_scheme = self.linear_scheme + 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) - 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 - module = name[: -len(".weight")] + def _module_scheme(self, module: str) -> Optional["QuantizationScheme"]: if _matches(module, self.ignore): return None - if _matches(module, [MOE_TARGET]): - return self.moe_scheme + 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: - if tensor.ndim == 3 and self.moe_scheme is not None and _PACKED_EXPERTS.match(name): + packed = _PACKED_EXPERTS.match(name) if tensor.ndim == 3 else None + if packed and self._module_scheme(f"{packed['prefix']}.0.gate_proj") is not None: yield from self(_unpack_experts(name, tensor)) continue scheme = self._scheme_for(name, tensor) 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/workers/config/rollout.py b/verl/workers/config/rollout.py index 96733581cd5..79b5a0080d1 100644 --- a/verl/workers/config/rollout.py +++ b/verl/workers/config/rollout.py @@ -31,6 +31,7 @@ "ServerConfig", "PrometheusConfig", "RolloutConfig", + "RolloutQuantizationConfig", "CheckpointEngineConfig", ] @@ -141,6 +142,36 @@ 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,11 +296,7 @@ class RolloutConfig(BaseConfig): skip_tokenizer_init: bool = True - # str: legacy engine-side quantization ("fp8", "torchao", ...); - # mapping {linear, moe, ignore} of compressed-tensors presets (e.g. FP8_BLOCK): rollout quantized by the trainer. - quantization: Optional[Any] = 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 @@ -278,8 +305,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_workers.py b/verl/workers/engine_workers.py index 2cc4aa1fc20..8d0a80f27b3 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -45,9 +45,7 @@ from verl.utils.py_functional import append_to_dict from verl.utils.qat import ( CompressedTensorsWeightQuantizer, - build_compressed_tensors_config, - is_compressed_tensors_quantization, - qat_mode_for_scheme, + qat_modes_for_config, ) from verl.utils.tensordict_utils import maybe_fix_3d_position_ids from verl.utils.torch_functional import allgather_dict_into_dict @@ -58,6 +56,7 @@ McoreEngineConfig, MtpConfig, RolloutConfig, + RolloutQuantizationConfig, TrainingWorkerConfig, ) from verl.workers.rollout.base import BaseRollout, get_rollout_class @@ -735,18 +734,23 @@ def init_model(self): backend, is_master=(torch.distributed.get_rank() == 0), bucket_size=bucket_size, **engine_kwargs ) - rollout_quantization = self.config.rollout.get("quantization", None) - qat = self.config.rollout.get("qat", None) or {} - if qat.get("enable", False) and not is_compressed_tensors_quantization(rollout_quantization): + rollout_quantization: RolloutQuantizationConfig = omega_conf_to_dataclass( + self.config.rollout.quantization, RolloutQuantizationConfig + ) + qat = getattr(getattr(self.actor.engine, "engine_config", None), "qat", None) + qat_mode = qat.mode if qat is not None and qat.enable else None + if qat_mode is not None and not rollout_quantization.enabled: raise ValueError( - "QAT needs a quantized rollout to sync to: set rollout.quantization, " - "e.g. {linear: FP8_BLOCK, moe: FP8_BLOCK}" + "QAT needs a quantized rollout to sync to: set rollout.quantization.linear / moe " + "(e.g. FP8_BLOCK) or rollout.quantization.config_file" ) - if is_compressed_tensors_quantization(rollout_quantization): - schemes = {rollout_quantization.get("linear"), rollout_quantization.get("moe")} - {None} - if qat.get("enable", False) and {qat_mode_for_scheme(s) for s in schemes} != {qat.get("mode")}: - raise ValueError(f"QAT mode {qat.get('mode')!r} does not match the rollout quantization {schemes}") - ct_config = build_compressed_tensors_config(rollout_quantization, model_config.local_path) + if rollout_quantization.enabled: + ct_config = rollout_quantization.build_quantization_config(model_config.local_path) + rollout_qat_modes = qat_modes_for_config(ct_config) + if qat_mode is not None and rollout_qat_modes != {qat_mode}: + raise ValueError( + f"QAT mode {qat_mode!r} does not match the rollout quantization (QAT modes {rollout_qat_modes})" + ) self.actor.engine.set_rollout_weight_quantizer(CompressedTensorsWeightQuantizer(ct_config)) # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 32aeb44aeeb..40843994590 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -17,7 +17,6 @@ import logging import os import secrets -from collections.abc import Mapping from pathlib import Path from typing import Any, Optional @@ -301,16 +300,8 @@ 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 isinstance(quantization, Mapping): - raise NotImplementedError("rollout.quantization as a {linear, moe} mapping is only supported by vLLM") - 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}") + if self.config.quantization.enabled: + raise NotImplementedError("rollout.quantization is only supported by vLLM") infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size args = { "model_path": self.model_config.local_path, @@ -333,10 +324,7 @@ 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({}), "custom_weight_loader": custom_weight_loader or None, **engine_kwargs, } 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 646de2e7515..3c42fb63139 100644 --- a/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py +++ b/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py @@ -14,7 +14,6 @@ import asyncio import logging import os -from collections.abc import Mapping from typing import Any, Optional import ray @@ -207,22 +206,8 @@ async def launch_server(self): per_worker_gpu_share = 1.0 / self.max_colocate_count - quantization = self.config.quantization - if isinstance(quantization, Mapping): - raise NotImplementedError("rollout.quantization as a {linear, moe} mapping is only supported by vLLM") - 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 db9bc8daa4c..f8792bc1667 100644 --- a/verl/workers/rollout/vllm_rollout/utils.py +++ b/verl/workers/rollout/vllm_rollout/utils.py @@ -186,14 +186,14 @@ def __new__(cls, **kwargs): ) register_verl_delta_weight_transfer_engine() + + # 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. quant_config = getattr(vllm_config, "quant_config", None) if vllm_config else None _is_nvfp4_packed = getattr(quant_config, "quant_format", None) == "nvfp4-pack-quantized" - # 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) and not _is_nvfp4_packed - ): + if is_quantized_model(vllm_config) and not _is_nvfp4_packed: apply_vllm_quant_patches() + # 3. patch compressed-tensors NVFP4 for dynamic weight loading if _is_nvfp4_packed: from verl.utils.vllm.vllm_nvfp4_utils import apply_qat_patches diff --git a/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/verl/workers/rollout/vllm_rollout/vllm_async_server.py index 65d75ae6cd6..e7774c6df4a 100644 --- a/verl/workers/rollout/vllm_rollout/vllm_async_server.py +++ b/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -51,10 +51,8 @@ relocate_rollout_traces, rollout_profiler_global_ranks, ) -from verl.utils.qat import build_compressed_tensors_config, is_compressed_tensors_quantization 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 ( @@ -319,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): @@ -359,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, @@ -1393,35 +1390,23 @@ 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 = {} + def _apply_quantization(self) -> dict: + """hf_overrides for a rollout quantization set in the rollout config. - if is_compressed_tensors_quantization(quantization): - # The trainer quantizes every refit to this config; vLLM picks the method up from hf_overrides. - hf_overrides["quantization_config"] = build_compressed_tensors_config( - quantization, self.model_config.local_path - ) - apply_vllm_quant_patches() - os.environ["VERL_VLLM_FP8_QUANT_ENABLED"] = "1" - # The schemes carry fp32 block scales. DeepGEMM's E8M0 mode (vLLM's default) rounds weight and - # activation scales to powers of two, so the rollout would drift from what the trainer quantized. - os.environ.setdefault("VLLM_USE_DEEP_GEMM_E8M0", "0") - logger.info(f"Rollout compressed-tensors quantization_config: {hf_overrides['quantization_config']}") - return None, hf_overrides - if quantization == "fp8": - raise ValueError("rollout.quantization=fp8 is replaced by the mapping form, e.g. {linear: FP8_BLOCK, moe: FP8_BLOCK}") - - 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) + # The schemes carry fp32 block scales. DeepGEMM's E8M0 mode (vLLM's default) rounds weight and + # activation scales to powers of two, so the rollout would drift from what the trainer quantized. + 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.""" From 409117ed0ee8a0490cca829c5159f34d1c2c6b36 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Thu, 15 Oct 2026 16:40:30 +0800 Subject: [PATCH 03/18] sglang --- docs/low_precision/fp8.md | 2 +- docs/low_precision/qat.md | 4 ++-- .../agent_loop/test_megatron_vllm_logprob_diff.py | 8 ++++++-- verl/trainer/config/rollout/rollout.yaml | 2 +- .../rollout/sglang_rollout/async_sglang_server.py | 9 +++++++-- 5 files changed, 17 insertions(+), 8 deletions(-) diff --git a/docs/low_precision/fp8.md b/docs/low_precision/fp8.md index 12987b9821a..a9dda1021f6 100644 --- a/docs/low_precision/fp8.md +++ b/docs/low_precision/fp8.md @@ -28,7 +28,7 @@ The rollout engine starts from a compressed-tensors `quantization_config` with d - Used in Deepseek, which is 1x128 quantization for activations and 128x128 quantization for model weights - Dense models and MoE models - Async rollout interfaces -- vLLM rollout with the Megatron training backend; SGLang, TRT-LLM and FSDP are not supported +- vLLM or SGLang rollout (SGLang validated on Hopper only) with the Megatron training backend; TRT-LLM and FSDP are not supported ### Usage diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index ee9bae72c30..67d4b28c21c 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -128,7 +128,7 @@ verl 内部(`verl/utils/qat/compressed_tensors.py`)生成 override 用的 co 上例为 "linear MXFP8 + MoE MXFP4 (W4A16)",MoE 部分与 Kimi-K3 自带的配置一致。各字段的具体取值需要与所用的 `compressed-tensors` 版本核对。 -4. **引擎**:目前只有 vLLM + Megatron 支持;SGLang / TRT-LLM 设置 `rollout.quantization` 会直接报错。 +4. **引擎**:vLLM 和 SGLang(训练侧目前只有 Megatron)支持;TRT-LLM 设置 `rollout.quantization` 会直接报错。SGLang 目前只验证了 Hopper 上的 FP8 block。 #### 权重同步 @@ -141,7 +141,7 @@ verl 内部(`verl/utils/qat/compressed_tensors.py`)生成 override 用的 co 2. **stage(权重同步前)**:把活跃参数暂存起来,为每个参数换上一个 checkpoint 布局的 staging buffer,并带上 `weight_loader` 等参数元数据,供 `load_weights` 写入。只要字节数相同,就直接 reinterpret 原有存储(例如 MXFP4 expert 权重在 DeepGEMM 和 Marlin 下都是这样),不额外分配显存;只有元素个数会变化的参数(主要是 block scale,大小约为权重的 1/16)才需要新分配。权重按 bucket 流式写入,不会把整层攒在一起。 3. **process(全部 bucket 加载完之后)**:重新执行原有的 `process_weights_after_loading`,同时 patch `replace_parameter`,把后处理结果 `copy_` 回原来的活跃参数存储,CUDA graph 捕获的地址保持不变。这一步必须在 `replace_parameter` 调用时就写回,因为 MoE 的 `_setup_kernel` 末尾会用当时 layer 上的张量重建 `moe_quant_config` 和 `moe_kernel`。 - 新增 compressed-tensors scheme(MXFP8、MXFP4、NVFP4 的 linear 和 MoE)时,按同样的模式实现各自的 stage/process:从 `create_weights` 推导出 checkpoint 布局,拦截对应模块里的 `replace_parameter`,并对不能原地写回的后端(如参数不以 Parameter 形式存放的 Triton MXFP4 后端)显式报错。注意 compressed-tensors MXFP4 MoE 会在后处理中把 `w13_weight_packed` 改名为 `w13_weight`,stage 时需要恢复成 checkpoint 中的参数名。 -- SGLang 侧沿用 trainer 侧量化后 `update_weights_from_tensor` 的方式。compressed-tensors 方案下是否会重新执行后处理,需要验证。 +- SGLang 侧沿用 trainer 侧量化后 `update_weights_from_tensor` 的方式,它只调用 `load_weights`,不会重新执行 `process_weights_after_loading`。Hopper 上 compressed-tensors FP8 block 的 linear 和 Triton MoE 后处理不改写 weight/scale 布局,因此可以直接写回,不需要 patch。Blackwell 上 DeepGEMM 会把 block scale 改写成 UE8M0、FlashInfer TRTLLM MoE 会交换 w13,届时需要像 vLLM 一样实现 stage/process。 #### 注意事项 diff --git a/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py index 14a0458f459..2cc978b9da1 100644 --- a/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py +++ b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py @@ -11,7 +11,7 @@ # 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, recompute old_log_probs with Megatron, and check the rollout-vs-actor log-prob gap. +"""Rollout with vLLM / SGLang, recompute old_log_probs with Megatron, and check the rollout-vs-actor log-prob gap. Usage (8 GPUs): MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 \ @@ -24,6 +24,9 @@ # same, with Megatron ModelOpt QAT fake-quantizing to the rollout scheme MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK QAT_MODE=fp8 \ .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py + # SGLang rollout + ENGINE=sglang MODEL_PATH=~/models/Qwen3.5-35B-A3B EP=8 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK \ + .venv-sglang/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py """ import json @@ -49,6 +52,7 @@ EP = int(os.getenv("EP", int(os.environ.get("EP", 1)))) ETP = int(os.getenv("ETP", 1)) GEN_TP = int(os.getenv("GEN_TP", 2)) +ENGINE = os.getenv("ENGINE", "vllm") N_GPUS = int(os.getenv("N_GPUS", 8)) MAX_K3_KL = float(os.getenv("MAX_K3_KL", 1e-2)) # Rollout compressed-tensors presets, e.g. QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK; unset keeps the rollout in bf16. @@ -80,7 +84,7 @@ def init_config() -> DictConfig: "++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", - "actor_rollout_ref.rollout.name=vllm", + f"actor_rollout_ref.rollout.name={ENGINE}", "actor_rollout_ref.rollout.mode=async", f"actor_rollout_ref.rollout.tensor_model_parallel_size={GEN_TP}", "actor_rollout_ref.rollout.gpu_memory_utilization=0.6", diff --git a/verl/trainer/config/rollout/rollout.yaml b/verl/trainer/config/rollout/rollout.yaml index 05032607599..d503665c4ba 100644 --- a/verl/trainer/config/rollout/rollout.yaml +++ b/verl/trainer/config/rollout/rollout.yaml @@ -482,7 +482,7 @@ prometheus: served_model_name: ${oc.select:actor_rollout_ref.model.path,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 + Megatron only for now). Off when all are null. +# weights and the trainer quantizes every weight sync (vLLM / SGLang + Megatron only for now). Off when all are null. quantization: # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 40843994590..341cbe712f7 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -300,8 +300,13 @@ 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). + model_override_args = {} if self.config.quantization.enabled: - raise NotImplementedError("rollout.quantization is only supported by vLLM") + # 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, @@ -324,7 +329,7 @@ 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, - "json_model_override_args": json.dumps({}), + "json_model_override_args": json.dumps(model_override_args), "custom_weight_loader": custom_weight_loader or None, **engine_kwargs, } From 97aeb4570dbc75a815473f4ad4305b1a7be26653 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Thu, 15 Oct 2026 18:15:56 +0800 Subject: [PATCH 04/18] docs --- .../e2e_ppo_trainer_megatron_sglang.yml | 4 + docs/index.rst | 1 + docs/low_precision/qat.md | 289 ++++++++---------- scripts/rollout_logprob_diff.py | 247 +++++++++++++++ .../test_megatron_vllm_logprob_diff.py | 195 ------------ 5 files changed, 373 insertions(+), 363 deletions(-) create mode 100644 scripts/rollout_logprob_diff.py delete mode 100644 tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py diff --git a/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml b/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml index fffac95d850..30929b926a3 100644 --- a/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml +++ b/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml @@ -200,6 +200,10 @@ jobs: run: | $UV_RUN ray stop --force ALL_OFFLOAD=True VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 LR_WARMUP_STEPS=1 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout + run: | + $UV_RUN ray stop --force + 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/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/qat.md b/docs/low_precision/qat.md index 67d4b28c21c..0ed70b6d9ca 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -1,215 +1,168 @@ -# Quantization Aware Training +# Low-Precision Rollout and Quantization-Aware Training -## Low Precision Inference +Last updated: 10/15/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: -RL 训练中 rollout(推理)通常占据大部分时间,而大规模 MoE 模型的 rollout 往往受限于显存容量和访存带宽。使用低精度(FP8 / MXFP8 / MXFP4 / NVFP4)推理可以: +- 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. -- 降低权重显存占用,留出更多 KV cache,提高并发; -- 降低 decode 阶段的访存量,提高吞吐; -- 在 Hopper/Blackwell 上使用 FP8/FP4 Tensor Core 提升 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. -代价是训练(高精度)与推理(低精度)之间的数值差异,即 training-inference mismatch,它会影响 importance sampling 和 off-policy 修正的准确性,严重时会导致训练不稳定。 +## Background: model weight formats -### 模型权重格式 +HuggingFace checkpoints come in two flavors: -HuggingFace checkpoint 通常有两种精度格式: +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). -1. **高精度**:权重以 BF16 存储,`config.json` 中没有 `quantization_config`。 -2. **低精度**:`config.json` 中有 `quantization_config`,描述各层的量化方式。不同模型对 linear(attention、dense MLP、shared experts)和 MoE(routed experts)往往采用不同的格式。 - -| 模型 | 结构 | `quant_method` | Linear | 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 层,512 experts,top-10 | 无 | BF16 | BF16 | - | -| [zai-org/GLM-5.2](https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json) | 78 层,256 experts,top-8,前 3 层 dense | 无 | BF16 | BF16 | - | -| [zai-org/GLM-5.3](https://huggingface.co/zai-org/GLM-5.3/blob/main/config.json) | 与 GLM-5.2 相同 | `fp8` | FP8 E4M3,128x128 block,FP32 scale(`weight_scale_inv`),activation 动态量化 | 同 Linear | router `mlp.gate`、`indexers_proj`、`embed_tokens`、`lm_head`、MTP `eh_proj`、各类 norm(`modules_to_not_convert`) | -| [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3/blob/main/config.json) | 93 层,896 experts,第 1 层 dense | `compressed-tensors` | BF16 | MXFP4 W4A16:FP4 E2M1,1x32 group,UE8M0 scale(`weight_packed` + `weight_scale`),activation 不量化 | `self_attn`、`shared_experts`、dense `mlp`、`lm_head`、`vision_tower`、`mm_projector`(`ignore` 正则) | -| [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/main/config.json) | 40 层,384 experts,top-6 | `fp8` | FP8 E4M3,32x32 block(`wo_a` 可能是 128x128),UE8M0 scale(`scale_fmt: ue8m0`) | FP4 E2M1,1x32 group,UE8M0 scale(`expert_dtype: fp4`),格式与 MXFP4 相同 | router `ffn.gate`、indexer `wk`/`weights_proj`、`compressor` | - -说明: - -- GLM-5.3 是标准的 HF FineGrainedFP8 格式(与 DeepSeek-V3 相同),vLLM 和 SGLang 的 `Fp8Config.from_config` 都能直接解析。 -- Kimi-K3 使用 compressed-tensors 格式:`config_groups` 以 `Linear` 为 target,再用 `ignore` 排除非 expert 层,从而实现 "linear BF16 + MoE MXFP4"。 -- DeepSeek-V4.1-Flash 的 `fp8` 配置带有 `scale_fmt`、`expert_dtype` 等私有字段,需要引擎提供专门的支持(如 vLLM 的 `deepseek_v4_fp8`),不属于通用格式。 - -### 训练与推理精度 +| [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` | -- **rollout 侧**:低精度权重直接使用其自带的 `quantization_config` 推理;高精度权重既可以高精度推理,也可以通过 override `quantization_config` 的方式使用低精度推理,以获取更高的吞吐。 -- **trainer 侧**:总是使用高精度训练。低精度权重会被反量化到高精度后再训练,如 DeepSeek-V4.1-Flash、Kimi-K3。为了减少 training-inference mismatch,低精度推理时训练侧可以开启 Quantization Aware Training(QAT)。 +- 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. -### 低精度推理方案:override quantization_config +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`. -rollout 统一采用 **override HF `quantization_config` + `load_format=dummy`** 的方式: +## How it works -1. 引擎启动时用 `dummy` 加载,并把 `quantization_config` override 成目标量化格式。这样引擎会按"已量化 checkpoint"创建参数,例如 FP8 的 `weight` + `weight_scale_inv`、MXFP4 的 `weight_packed` + `weight_scale`。 -2. 每次权重同步时,把 BF16 权重量化成 checkpoint 格式(量化可以在 trainer 侧或 rollout worker 侧完成),再交给引擎的 `load_weights`,由引擎完成融合、TP 切分以及 kernel 布局转换。 +- **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 an overridden `quantization_config`, so it allocates quantized parameters as if it were loading a quantized checkpoint (e.g. FP8 `weight` + `weight_scale`). On every weight sync, the trainer quantizes its BF16 weights to that checkpoint format, and the engine's weight loader handles fusion, tensor-parallel sharding and kernel-specific layouts. -对于原生低精度的模型(如 GLM-5.3、Kimi-K3),不需要 override,权重同步时按其自带的 `quantization_config` 量化即可。 +Your HuggingFace checkpoint stays in BF16; no offline quantization or calibration step is needed. The rollout format is described in the open [compressed-tensors](https://github.com/neuralmagic/compressed-tensors) format, which both vLLM and SGLang understand and which can express different schemes for linear and MoE layers. -#### 为什么选择 compressed-tensors +## Support matrix -目标是灵活表达 (linear, moe) x (fp8, mxfp8, mxfp4, nvfp4) 的任意组合。各候选格式的能力如下: +The trainer must be Megatron (FSDP is not yet supported for this feature; see [NVFP4 QAT](nvfp4_qat.md) for FSDP). 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. -| 格式 | 按层混合 | 覆盖的精度 | 备注 | -|---|---|---|---| -| `fp8`(HF FineGrainedFP8) | 否 | FP8 block;SGLang 支持 `quant_method: mxfp8` | verl 现有方案,两个引擎都支持 | -| `modelopt_fp8` / `modelopt_fp4` | 否 | FP8 / NVFP4 | | -| `modelopt_mixed` | 是 | FP8 + NVFP4 | 不支持 MX 格式 | -| `quark` | 是 | MXFP4 + FP8 | 主要面向 AMD | -| **`compressed-tensors`** | **是**(`config_groups`) | **FP8 / MXFP8 / MXFP4 / NVFP4** | Kimi-K3 和 verl QAT 已在使用 | +| Precision | Preset | QAT mode | GPU | vLLM Linear | vLLM MoE | SGLang Linear | SGLang MoE | Status in verl | +|---|---|---|---|---|---|---|---|---| +| FP8 block (E4M3, 128x128 weight blocks, dynamic 1x128 activations) | `FP8_BLOCK` | `fp8` | Hopper / Blackwell | `W8A8Fp8` | `W8A8Fp8MoE` | `W8A8Fp8` | `W8A8Fp8MoE` | ✅ Hopper (vLLM, SGLang); Blackwell not yet verified | +| MXFP8 | - | - | Blackwell | `W8A8Mxfp8` | `W8A8Mxfp8MoE` | ❌ (only whole-model `quant_method: mxfp8`) | ❌ | Planned | +| MXFP4 | - | - | Blackwell | `W4A4Mxfp4` | `W4A4Mxfp4MoE` | ❌ | `Mxfp4MoEMethod` | Planned | +| NVFP4 | - | - | Blackwell | W4A4 / W4A16 | W4A4 / W4A16 | W4A4 only | W4A4 | Planned | -compressed-tensors 的优势: - -- 用 `config_groups` 为不同 target 指定不同方案。target 匹配时先按层名(支持 `re:` 正则)、再按类名,因此 `re:.*mlp\.experts\..*` 会优先于 `Linear` 命中 routed experts。 -- 是开放的格式规范(`compressed_tensors.QuantizationArgs`),vLLM 和 SGLang 都支持。 -- verl QAT 已经产出 compressed-tensors 格式,NVFP4 W4A4 所需的 `input_global_scale` 由 QAT observer 维护,可以直接复用。 - -#### 引擎支持矩阵(compressed-tensors) - -| 精度 | GPU | vLLM Linear | vLLM MoE | SGLang Linear | SGLang MoE | -|---|---|---|---|---|---| -| FP8(block / channel) | Hopper / Blackwell | ✅ `W8A8Fp8` | ✅ `W8A8Fp8MoE` | ✅ `W8A8Fp8` | ✅ `W8A8Fp8MoE` | -| MXFP8 | Blackwell | ✅ `W8A8Mxfp8` | ✅ `W8A8Mxfp8MoE` | ❌(只能用 `Fp8Config(quant_method="mxfp8")`,整模型统一) | ❌(同左) | -| MXFP4 | Blackwell | ✅ `W4A4Mxfp4` | ✅ `W4A4Mxfp4MoE` | ❌ | ✅ `Mxfp4MoEMethod` | -| NVFP4 | Blackwell | ✅ W4A4 / W4A16 | ✅ W4A4 / W4A16 | ✅ 仅 W4A4 | ✅ W4A4 | +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. -vLLM 可以覆盖全部组合;SGLang 只支持其中一部分,缺失部分需要 SGLang 上游补齐(可参考 vLLM 对应的 scheme 移植)。 +## Configuration -硬件方面:MXFP8 / MXFP4 / NVFP4 只在 Blackwell(SM100)上支持。Hopper 上 vLLM 虽然会退化为 Marlin weight-only kernel(激活保持 bf16),但数值与 Blackwell 上的 W8A8 / W4A4 不同,verl 不支持这种用法。 - -#### 配置 +### Rollout quantization ```yaml actor_rollout_ref: rollout: - quantization: # linear / moe / config_file 全为 null 表示不做低精度推理 - linear: FP8_BLOCK # compressed-tensors preset,目前支持 FP8_BLOCK - moe: FP8_BLOCK - ignore: [] # 额外跳过的层,支持精确层名或 "re:" - config_file: null # 直接提供 compressed-tensors quantization_config JSON,与 linear/moe/ignore 互斥 + 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 ``` -只支持 compressed-tensors,不再支持 `quantization: fp8` 字符串和 torchao。router(`mlp.gate`)、`lm_head`、embedding、视觉塔和 MTP 默认加入 ignore,形状不能被 block 整除的层(如 Qwen3.5 `in_proj_a/b`)自动加入 ignore。 - -#### 配置翻译规则 - -verl 内部(`verl/utils/qat/compressed_tensors.py`)生成 override 用的 compressed-tensors `quantization_config`,trainer 侧量化器使用同一份配置: - -1. **linear 和 moe 相同**:一个 `targets: ["Linear"]` 的 config group。 -2. **linear 和 moe 不同**:MoE 用 `re:.*mlp\.experts\..*` 单独成组,其余 `Linear` 一组;某一项为 null 时对应层保持高精度。 -3. **`config_file`**:JSON 可以是 `quantization_config` 本身,也可以是带 `quantization_config` 字段的 HF `config.json`,原样使用(不追加默认 ignore)。每个 config group 的数值配置必须与某个已支持的 preset 一致,否则报错;QAT mode 也据此校验。 - - 多精度组合(目前未支持)的配置示例: - - ```json - { - "quant_method": "compressed-tensors", - "format": "mixed-precision", - "ignore": ["re:.*mlp\\.gate$", "lm_head"], - "config_groups": { - "moe": { - "targets": ["re:.*mlp\\.experts\\..*"], - "format": "mxfp4-pack-quantized", - "weights": {"type": "float", "num_bits": 4, "strategy": "group", "group_size": 32, - "symmetric": true, "dynamic": false, "scale_dtype": "torch.uint8"}, - "input_activations": null - }, - "linear": { - "targets": ["Linear"], - "weights": {"type": "float", "num_bits": 8, "strategy": "group", "group_size": 32, - "symmetric": true, "dynamic": false, "scale_dtype": "torch.uint8"}, - "input_activations": {"type": "float", "num_bits": 8, "strategy": "group", "group_size": 32, - "symmetric": true, "dynamic": true, "scale_dtype": "torch.uint8"} - } - } - } - ``` - - 上例为 "linear MXFP8 + MoE MXFP4 (W4A16)",MoE 部分与 Kimi-K3 自带的配置一致。各字段的具体取值需要与所用的 `compressed-tensors` 版本核对。 +- 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. +- The following layers are always kept in BF16: `lm_head`, embeddings, MoE routers (`mlp.gate`, `shared_expert_gate`), vision towers and MTP layers. Layers whose shapes are not divisible by the block size (e.g. Qwen3.5 `linear_attn.in_proj_a/b`) are skipped automatically. +- `ignore` adds more layers on top of these, using HuggingFace layer names. -4. **引擎**:vLLM 和 SGLang(训练侧目前只有 Megatron)支持;TRT-LLM 设置 `rollout.quantization` 会直接报错。SGLang 目前只验证了 Hopper 上的 FP8 block。 +From the command line: -#### 权重同步 - -- 量化器输出 **checkpoint 格式** 的张量:按 expert 拆分、不做 TP 切分,参数名与 checkpoint 一致,例如: - - FP8:`weight` + `weight_scale_inv`; - - compressed-tensors FP4:`weight_packed` + `weight_scale`,NVFP4 另有 `weight_global_scale`、`input_global_scale`。 -- vLLM 侧**不使用** vLLM 自带的 layerwise reload(`initialize_layerwise_reload` / `finalize_layerwise_reload`)。它会把一个 layer 的所有权重都攒齐之后才开始处理,对大规模 MoE 模型来说,一个 MoE layer 的全部 expert 权重会同时驻留在显存里,容易 OOM。 -- vLLM 侧沿用 `verl/utils/vllm/vllm_fp8_utils.py`(`Fp8LinearMethod` / `Fp8MoEMethod`)和 `verl/utils/vllm/vllm_fp4_utils.py`(`Mxfp4MoEMethod`)的 **stage → load → process** 做法,入口在 `verl/utils/vllm/vllm_quant_utils.py`。它要解决的问题是:kernel 后处理会把 weight/scale 改写成推理布局(如 DeepGEMM 把 block scale 打包成 int32 UE8M0,Marlin 对 FP4 权重做 repack,FlashInfer 交换 w13 的 gate/up 并做 tile shuffle),checkpoint 格式的张量无法直接写回。具体分三步: - 1. **记录 checkpoint 布局**:首次加载时,通过 patch `process_weights_after_loading` 记下各参数的 checkpoint shape/dtype(只记布局,不保留张量),或者像 MXFP4 那样直接按 `create_weights` 推导出来。 - 2. **stage(权重同步前)**:把活跃参数暂存起来,为每个参数换上一个 checkpoint 布局的 staging buffer,并带上 `weight_loader` 等参数元数据,供 `load_weights` 写入。只要字节数相同,就直接 reinterpret 原有存储(例如 MXFP4 expert 权重在 DeepGEMM 和 Marlin 下都是这样),不额外分配显存;只有元素个数会变化的参数(主要是 block scale,大小约为权重的 1/16)才需要新分配。权重按 bucket 流式写入,不会把整层攒在一起。 - 3. **process(全部 bucket 加载完之后)**:重新执行原有的 `process_weights_after_loading`,同时 patch `replace_parameter`,把后处理结果 `copy_` 回原来的活跃参数存储,CUDA graph 捕获的地址保持不变。这一步必须在 `replace_parameter` 调用时就写回,因为 MoE 的 `_setup_kernel` 末尾会用当时 layer 上的张量重建 `moe_quant_config` 和 `moe_kernel`。 -- 新增 compressed-tensors scheme(MXFP8、MXFP4、NVFP4 的 linear 和 MoE)时,按同样的模式实现各自的 stage/process:从 `create_weights` 推导出 checkpoint 布局,拦截对应模块里的 `replace_parameter`,并对不能原地写回的后端(如参数不以 Parameter 形式存放的 Triton MXFP4 后端)显式报错。注意 compressed-tensors MXFP4 MoE 会在后处理中把 `w13_weight_packed` 改名为 `w13_weight`,stage 时需要恢复成 checkpoint 中的参数名。 -- SGLang 侧沿用 trainer 侧量化后 `update_weights_from_tensor` 的方式,它只调用 `load_weights`,不会重新执行 `process_weights_after_loading`。Hopper 上 compressed-tensors FP8 block 的 linear 和 Triton MoE 后处理不改写 weight/scale 布局,因此可以直接写回,不需要 patch。Blackwell 上 DeepGEMM 会把 block scale 改写成 UE8M0、FlashInfer TRTLLM MoE 会交换 w13,届时需要像 vLLM 一样实现 stage/process。 - -#### 注意事项 - -1. **NVFP4 W4A4 需要静态的 `input_global_scale`**:使用 `dummy` 加载时没有校准数据。可以使用 W4A16(仅 vLLM 支持),或开启 QAT 由 observer 维护该 scale,或额外进行一次校准。MX 格式的 activation 为动态量化,不需要校准。 -2. **数值格式必须逐 bit 对齐**:UE8M0 scale 的编码、FP4 nibble 的打包顺序、NVFP4 两级 scale 的约定,都必须与引擎期望的一致。建议为每种格式编写往返测试:verl 量化,引擎反量化,再与 BF16 对比。 -3. **activation 精度由 scheme 决定**:例如 compressed-tensors 的 MXFP4 linear 是 W4A4,而 Kimi-K3 的 MoE 是 W4A16(`input_activations: null`)。配置中的 `activation` 字段需要准确映射到 `input_activations`。 - -## Quantization Aware Training - -QAT 在训练时对权重(W4A4 时还包括 activation)做 fake quantization:前向使用"量化后再反量化"的值,反向通过 STE 直接回传梯度,使模型学会容忍低精度推理带来的误差,从而缩小 training-inference mismatch。权重同步时,trainer 把 BF16 权重真正量化成 checkpoint 格式,再发给 rollout。 - -QAT 由两部分组成: +```bash +actor_rollout_ref.rollout.quantization.linear=FP8_BLOCK \ +actor_rollout_ref.rollout.quantization.moe=FP8_BLOCK +``` -1. **训练侧 fake quantization**:与训练后端强相关,需要把量化器插入后端实际使用的 Linear 层。 -2. **导出格式**:权重同步时输出的张量名和 scale 约定,它决定 rollout 侧使用哪个 `quantization_config`。 +### 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$" + ] +} +``` -这两部分是相互独立的。训练侧可以按后端选择不同的实现,导出格式则应该统一为 compressed-tensors,与上文的 low precision inference 方案保持一致。 +### QAT (Megatron) -### FSDP/VeOmni +QAT is optional. Enable it on the actor and pick the mode that matches the rollout preset: -- **fake quantization**:使用 verl 自己实现的 `QATLinear`(`verl/utils/qat/linear.py`),替换模型中的 `nn.Linear`。W4A4 模式下,activation 的 `input_global_scale` 由 observer 在训练中维护。 -- **导出**:`verl/utils/qat/quantizer.py` 输出 compressed-tensors NVFP4 格式,即 `weight_packed`、`weight_scale`、`weight_global_scale`,W4A4 时还有 `input_global_scale`。 -- **global scale 融合**:vLLM 会把 q/k/v、gate/up 融合成一个层,融合后的层只能有一个 global scale,因此导出时需要用 `fuse_global_scales` 对齐各分片的 global scale(取 min)。 -- **rollout 配置**:`recipe/qat/config/nvfp4_w4a16.json`、`nvfp4_w4a4.json`(`quant_method: compressed-tensors`)。 +```yaml +actor_rollout_ref: + actor: + megatron: + qat: + enable: true + mode: fp8 # must match rollout.quantization (FP8_BLOCK -> fp8) +``` -### Megatron-LM +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. -Megatron-LM + TransformerEngine + ModelOpt +### Engine notes -- **fake quantization**:Megatron 中实际计算的是 TE 实现的层(带 TP 切分的 `ColumnParallelLinear` / `RowParallelLinear`,以及 expert 使用的 `TEGroupedLinear`),无法沿用 FSDP 替换 `nn.Linear` 的方式。Megatron-Bridge 官方的做法是使用 `modelopt_transformer_layer_spec` 构建模型,再调用 `mtq.quantize` 插入 ModelOpt 量化器(`verl/utils/modelopt/quantize.py`)。因此训练侧继续使用 ModelOpt。 -- **导出(现状)**:`verl/utils/modelopt/qat_weight_exporter.py` 输出 ModelOpt checkpoint 格式,rollout 使用 `quant_method: modelopt`(`recipe/qat/config/nvfp4_w4a16_megatron.json`),vLLM 侧需要单独维护一套 patch(`verl/utils/modelopt/vllm_modelopt_patch.py`)。 +- **vLLM**: `VLLM_USE_DEEP_GEMM_E8M0` defaults to `0` when rollout quantization is enabled, so block scales stay in FP32. +- **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 -两个后端导出的都是 NVFP4,量化方式完全相同:FP4 E2M1,每 16 个元素一个 FP8 E4M3 block scale,外加一个 FP32 global scale。区别只在参数名和 global scale 的表示方式: +Qwen3.5-35B-A3B on 8 GPUs with FP8 block rollout and FP8 QAT: -| | FSDP/VeOmni(compressed-tensors) | Megatron(ModelOpt) | -|---|---|---| -| FP4 打包后的权重 | `weight_packed` | `weight` | -| block scale(FP8 E4M3,每 16 个元素一个) | `weight_scale` | `weight_scale` | -| weight global scale | `weight_global_scale` = `6 * 448 / amax` | `weight_scale_2` = `amax / (6 * 448)` | -| activation global scale(W4A4) | `input_global_scale` | `input_scale`(同为倒数关系) | -| `ignore_patterns` 语法 | `re:` 正则,匹配 HF 层名 | fnmatch,匹配 Megatron 模块名,并手写映射(`lm_head` 对应 `*output_layer*`,`*mlp.gate` 对应 `*router*`) | -| vLLM patch | `verl/utils/qat/vllm_patch.py` | `verl/utils/modelopt/vllm_modelopt_patch.py` | +```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 \ + ... +``` -可以看出,`weight_global_scale` 与 `weight_scale_2` 互为倒数,`input_global_scale` 与 `input_scale` 同理。 +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). -### 统一方案:训练侧保留 ModelOpt,导出改为 compressed-tensors +## Measuring the mismatch -在 `QATWeightExporter._quantize_nvfp4` 中改为输出 compressed-tensors 格式: +`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: -```python -yield (name.replace(".weight", ".weight_packed"), quantized) -yield (_derive_scale_name(name, "weight_scale"), w_scale) -yield (_derive_scale_name(name, "weight_global_scale"), 1.0 / w_scale_2) -if input_scale is not None: - yield (_derive_scale_name(name, "input_global_scale"), 1.0 / input_scale) +```bash +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 ``` -这样改之后: +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, H-series GPUs): -- **rollout 配置统一**:两个后端共用 `nvfp4_w4a16.json` / `nvfp4_w4a4.json`,`nvfp4_w4a16_megatron.json` 可以删除。rollout 侧不需要关心训练后端,与上文 "override quantization_config" 的方案一致。 -- **少维护一套 patch**:`vllm_modelopt_patch.py` 可以下线。`qat/vllm_patch.py` 也可以改写成与 `vllm_fp8_utils.py`、`vllm_fp4_utils.py` 一致的 stage/process 模式,统一维护。 -- **global scale 天然一致**:Megatron 的 `linear_qkv`、`linear_fc1` 本身就是融合层,只有一个量化器和一个 amax,导出后 q/k/v、gate/up 的 global scale 自然相同,不需要额外的 `fuse_global_scales`。MoE 的 `w1`/`w3` 同理,这也满足 vLLM compressed-tensors NVFP4 MoE 要求 `w13` 两个 global scale 相同的约束。 -- **ignore 统一用 HF 层名和 `re:` 语法**:Megatron 侧通过 bridge 的 `mapping_registry` 把 HF 层名翻译到 Megatron 模块名(可复用 `_iter_hf_to_megatron_matches`),不再需要手写映射。 +| Rollout | vLLM | SGLang | +|---|---|---| +| BF16 | ~0.0014 | - | +| FP8_BLOCK | 0.0039 | 0.0036 | +| FP8_BLOCK + FP8 QAT | 0.0033 | 0.0033 | -### 待验证 +## Limitations -1. **FP4 打包逐 bit 一致**:ModelOpt 的 `to_quantized_weight` 与 compressed-tensors 的 `compress_weight` 的 nibble 打包顺序、block scale 排布是否一致。可以对同一个 BF16 权重分别过两条导出路径,比较输出张量或 vLLM 反量化后的结果。 -2. **activation scale 的倒数关系**:确认 `NVFP4QTensor.get_activation_scaling_factor` 的定义为 `amax / (6 * 448)`,这决定了 W4A4 时 `input_global_scale` 取倒数是否正确。Megatron 目前仅支持 W4A16(`build_quantize_config` 中写死),暂不阻塞。 -3. **MX 格式 fake quantization**:若 Megatron QAT 要扩展到 MXFP8 / MXFP4,需要确认 ModelOpt 是否提供对应的量化配置。导出侧同样按 compressed-tensors 的 MX scheme 输出即可。 +- Only the `FP8_BLOCK` preset on Hopper is supported today. +- Checkpoints that already ship a `quantization_config` (e.g. FP8 or compressed-tensors checkpoints) are not supported as a training target together with low-precision rollout yet; use a BF16 checkpoint. +- Megatron QAT and the rollout use separate ignore lists. ModelOpt may fake-quantize a few layers that the rollout keeps in BF16 (e.g. Qwen3.5 `in_proj_a/b`); this slightly overestimates the quantization error for those layers but does not affect correctness. diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py new file mode 100644 index 00000000000..00eec619ee6 --- /dev/null +++ b/scripts/rollout_logprob_diff.py @@ -0,0 +1,247 @@ +# 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 + # 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, e.g. FP8_BLOCK") + parser.add_argument("--quant-moe", default=None, help="rollout.quantization.moe preset, e.g. FP8_BLOCK") + parser.add_argument("--qat-mode", default=None, help="Megatron ModelOpt QAT mode matching the rollout, e.g. fp8") + 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 + if args.qat_mode: + config.actor_rollout_ref.actor.megatron.qat.enable = True + config.actor_rollout_ref.actor.megatron.qat.mode = args.qat_mode + 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() + 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) + + 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()) + 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/test_megatron_vllm_logprob_diff.py b/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py deleted file mode 100644 index 2cc978b9da1..00000000000 --- a/tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py +++ /dev/null @@ -1,195 +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. -"""Rollout with vLLM / SGLang, recompute old_log_probs with Megatron, and check the rollout-vs-actor log-prob gap. - -Usage (8 GPUs): - MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 \ - .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py - MODEL_PATH=~/models/Qwen3.5-35B-A3B EP=8 \ - .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py - # fp8 blockwise rollout, weights quantized on the trainer side - MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK \ - .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py - # same, with Megatron ModelOpt QAT fake-quantizing to the rollout scheme - MODEL_PATH=~/models/Qwen3.5-0.8B EP=1 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK QAT_MODE=fp8 \ - .venv/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py - # SGLang rollout - ENGINE=sglang MODEL_PATH=~/models/Qwen3.5-35B-A3B EP=8 QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK \ - .venv-sglang/bin/python -m pytest -s tests/experimental/agent_loop/test_megatron_vllm_logprob_diff.py -""" - -import json -import os - -import numpy as np -import pytest -import ray -import torch -from hydra import compose, initialize_config_dir -from omegaconf import DictConfig - -from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager_and_workers -from verl.protocol import DataProto -from verl.trainer.ppo.rollout_corr_helper import compute_offpolicy_metrics -from verl.utils import tensordict_utils as tu -from verl.utils.debug.metrics import calculate_debug_metrics -from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding - -TP = int(os.getenv("TP", 2)) -PP = int(os.getenv("PP", 1)) -CP = int(os.getenv("CP", 1)) -EP = int(os.getenv("EP", int(os.environ.get("EP", 1)))) -ETP = int(os.getenv("ETP", 1)) -GEN_TP = int(os.getenv("GEN_TP", 2)) -ENGINE = os.getenv("ENGINE", "vllm") -N_GPUS = int(os.getenv("N_GPUS", 8)) -MAX_K3_KL = float(os.getenv("MAX_K3_KL", 1e-2)) -# Rollout compressed-tensors presets, e.g. QUANT_LINEAR=FP8_BLOCK QUANT_MOE=FP8_BLOCK; unset keeps the rollout in bf16. -QUANT_LINEAR = os.getenv("QUANT_LINEAR") -QUANT_MOE = os.getenv("QUANT_MOE") -# Megatron ModelOpt QAT mode matching the rollout scheme, e.g. QAT_MODE=fp8; unset disables QAT. -QAT_MODE = os.getenv("QAT_MODE") - - -@pytest.fixture -def init_config() -> DictConfig: - with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): - config = compose( - config_name="ppo_trainer", - overrides=[ - "model_engine=megatron", - f"trainer.n_gpus_per_node={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={TP}", - f"actor_rollout_ref.actor.megatron.pipeline_model_parallel_size={PP}", - f"actor_rollout_ref.actor.megatron.context_parallel_size={CP}", - f"actor_rollout_ref.actor.megatron.expert_model_parallel_size={EP}", - f"actor_rollout_ref.actor.megatron.expert_tensor_parallel_size={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={ENGINE}", - "actor_rollout_ref.rollout.mode=async", - f"actor_rollout_ref.rollout.tensor_model_parallel_size={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", - ], - ) - - config.actor_rollout_ref.model.path = os.path.expanduser(os.getenv("MODEL_PATH", "~/models/Qwen3.5-35B-A3B")) - config.actor_rollout_ref.rollout.quantization.linear = QUANT_LINEAR - config.actor_rollout_ref.rollout.quantization.moe = QUANT_MOE - if QAT_MODE: - config.actor_rollout_ref.actor.megatron.qat.enable = True - config.actor_rollout_ref.actor.megatron.qat.mode = QAT_MODE - return config - - -def compute_old_log_prob(actor_rollout_wg, 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) - entropy = no_padding_2_padding(tu.get(output, "entropy"), batch_td) - return DataProto.from_tensordict( - tu.get_tensordict({"old_log_probs": log_probs.float(), "entropys": entropy.float()}) - ) - - -def test_megatron_vllm_logprob_diff(init_config): - 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_agent_loop_manager_and_workers(init_config) - - 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)?", - ] - 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)), - }, - ) - n = init_config.actor_rollout_ref.rollout.n - batch = batch.repeat(n) - - # rollout - result = agent_loop_manager.generate_sequences(prompts=batch) - checkpoint_manager.sleep_replicas() - assert len(result) == len(raw_prompts) * n - assert "rollout_log_probs" in result.batch - result.meta_info.pop("timing", None) - - # recompute old_log_probs with megatron - result.meta_info["temperature"] = init_config.actor_rollout_ref.rollout.temperature - result.meta_info["global_token_num"] = torch.sum(result.batch["attention_mask"], dim=-1).tolist() - old_log_prob = compute_old_log_prob(actor_rollout_wg, result) - old_log_prob.batch.pop("entropys") - result = result.union(old_log_prob) - - old_log_probs = result.batch["old_log_probs"] - rollout_log_probs = result.batch["rollout_log_probs"] - response_mask = result.batch["response_mask"].bool() - assert old_log_probs.shape == rollout_log_probs.shape == response_mask.shape - - 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()) - print(json.dumps(metrics, indent=2)) - - assert np.isfinite(metrics["kl"]) and np.isfinite(metrics["k3_kl"]) - assert metrics["k3_kl"] < MAX_K3_KL, f"k3_kl {metrics['k3_kl']} exceeds {MAX_K3_KL}" - - print("Test passed!") - ray.shutdown() From 3b73bf4b05d5049aa54f9ee49205bc4fc2936506 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 16 Oct 2026 13:11:17 +0800 Subject: [PATCH 05/18] mxfp8 --- .pre-commit-config.yaml | 6 +- docs/low_precision/qat.md | 44 ++++++++---- scripts/rollout_logprob_diff.py | 32 ++++++++- ...test_vllm_fp8_restage_same_shape_on_cpu.py | 20 ++++++ verl/trainer/config/engine/megatron.yaml | 4 +- verl/utils/modelopt/quantize.py | 64 ++++++++++++++++- verl/utils/qat/compressed_tensors.py | 71 ++++--------------- verl/utils/vllm/vllm_fp8_utils.py | 51 +++++++++++-- verl/workers/config/engine.py | 2 +- verl/workers/config/rollout.py | 6 +- .../rollout/vllm_rollout/vllm_async_server.py | 13 +++- 11 files changed, 223 insertions(+), 90 deletions(-) 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/low_precision/qat.md b/docs/low_precision/qat.md index 0ed70b6d9ca..679c133752f 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -1,6 +1,6 @@ # Low-Precision Rollout and Quantization-Aware Training -Last updated: 10/15/2026 +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: @@ -42,12 +42,12 @@ Your HuggingFace checkpoint stays in BF16; no offline quantization or calibratio The trainer must be Megatron (FSDP is not yet supported for this feature; see [NVFP4 QAT](nvfp4_qat.md) for FSDP). 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 | Status in verl | -|---|---|---|---|---|---|---|---|---| -| FP8 block (E4M3, 128x128 weight blocks, dynamic 1x128 activations) | `FP8_BLOCK` | `fp8` | Hopper / Blackwell | `W8A8Fp8` | `W8A8Fp8MoE` | `W8A8Fp8` | `W8A8Fp8MoE` | ✅ Hopper (vLLM, SGLang); Blackwell not yet verified | -| MXFP8 | - | - | Blackwell | `W8A8Mxfp8` | `W8A8Mxfp8MoE` | ❌ (only whole-model `quant_method: mxfp8`) | ❌ | Planned | -| MXFP4 | - | - | Blackwell | `W4A4Mxfp4` | `W4A4Mxfp4MoE` | ❌ | `Mxfp4MoEMethod` | Planned | -| NVFP4 | - | - | Blackwell | W4A4 / W4A16 | W4A4 / W4A16 | W4A4 only | W4A4 | Planned | +| 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; Blackwell under verification | ✅ Hopper; Blackwell not yet verified | +| MXFP8 | `MXFP8` | `mxfp8` | Blackwell | `W8A8Mxfp8` | `W8A8Mxfp8MoE` | ❌ (only whole-model `quant_method: mxfp8`) | ❌ | ✅ Blackwell | Blocked by engine support | +| MXFP4 | - | - | Blackwell | `W4A4Mxfp4` | `W4A4Mxfp4MoE` | ❌ | `Mxfp4MoEMethod` | Planned | Planned | +| NVFP4 | - | - | Blackwell | W4A4 / W4A16 | W4A4 / W4A16 | W4A4 only | W4A4 | Planned | Planned | 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. @@ -67,8 +67,8 @@ actor_rollout_ref: - 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. -- The following layers are always kept in BF16: `lm_head`, embeddings, MoE routers (`mlp.gate`, `shared_expert_gate`), vision towers and MTP layers. Layers whose shapes are not divisible by the block size (e.g. Qwen3.5 `linear_attn.in_proj_a/b`) are skipped automatically. -- `ignore` adds more layers on top of these, using HuggingFace layer names. +- 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: @@ -77,6 +77,14 @@ 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. @@ -115,14 +123,15 @@ actor_rollout_ref: megatron: qat: enable: true - mode: fp8 # must match rollout.quantization (FP8_BLOCK -> fp8) + mode: fp8 # must match rollout.quantization (FP8_BLOCK -> fp8, MXFP8 -> mxfp8) ``` 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. +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` when rollout quantization is enabled, so block scales stay in FP32. +- **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. - **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 @@ -161,8 +170,17 @@ Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arg | FP8_BLOCK | 0.0039 | 0.0036 | | FP8_BLOCK + FP8 QAT | 0.0033 | 0.0033 | +Blackwell reference `k3_kl` for vLLM MXFP8: + +| Model | Rollout | k3_kl | +|---|---|---| +| Qwen3.5-0.8B | BF16 | 0.00066 | +| | MXFP8 linear | 0.01391 | +| Qwen3.5-35B-A3B (EP=8, rollout TP=2) | MXFP8 linear + MoE | 0.00148 | +| | MXFP8 linear + MoE + MXFP8 QAT | 0.00288 | + ## Limitations -- Only the `FP8_BLOCK` preset on Hopper is supported today. +- `FP8_BLOCK` is supported on Hopper. `MXFP8` with vLLM is supported on Blackwell; SGLang has no per-layer MXFP8 support. - Checkpoints that already ship a `quantization_config` (e.g. FP8 or compressed-tensors checkpoints) are not supported as a training target together with low-precision rollout yet; use a BF16 checkpoint. -- Megatron QAT and the rollout use separate ignore lists. ModelOpt may fake-quantize a few layers that the rollout keeps in BF16 (e.g. Qwen3.5 `in_proj_a/b`); this slightly overestimates the quantization error for those layers but does not affect correctness. +- Megatron QAT and the rollout use separate ignore lists and naming conventions. Users must keep them aligned. diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py index 00eec619ee6..610ac3e8f8b 100644 --- a/scripts/rollout_logprob_diff.py +++ b/scripts/rollout_logprob_diff.py @@ -21,6 +21,13 @@ # 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' # 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 @@ -83,9 +90,23 @@ def parse_args() -> argparse.Namespace: 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, e.g. FP8_BLOCK") - parser.add_argument("--quant-moe", default=None, help="rollout.quantization.moe preset, e.g. FP8_BLOCK") - parser.add_argument("--qat-mode", default=None, help="Megatron ModelOpt QAT mode matching the rollout, e.g. fp8") + parser.add_argument( + "--quant-linear", default=None, help="rollout.quantization.linear preset, e.g. FP8_BLOCK or MXFP8" + ) + parser.add_argument("--quant-moe", default=None, help="rollout.quantization.moe preset, e.g. FP8_BLOCK or MXFP8") + 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("--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() @@ -129,9 +150,14 @@ def build_config(args: argparse.Namespace) -> DictConfig: 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, + ] return config 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/verl/trainer/config/engine/megatron.yaml b/verl/trainer/config/engine/megatron.yaml index a49f8848f0d..8033ba5e0ba 100644 --- a/verl/trainer/config/engine/megatron.yaml +++ b/verl/trainer/config/engine/megatron.yaml @@ -140,8 +140,8 @@ qat: _target_: verl.workers.config.QATEngineConfig # Whether to enable QAT enable: false - # Quantization mode: "w4a16" (weight-only FP4), or "fp8" (128x128 block weights, 1x128 dynamic - # activations) to train against a rollout.quantization={linear: FP8_BLOCK, moe: FP8_BLOCK} rollout + # Quantization mode: "w4a16" (weight-only FP4), "fp8" (FP8_BLOCK rollout), or "mxfp8" + # (MXFP8 rollout with group-of-32 E8M0-scaled weights and activations) mode: "w4a16" # Group size for blockwise quantization (NVFP4 requires 16) group_size: 16 diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py index 9e56af81f23..50cec72b1df 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -38,6 +38,7 @@ } _FP8_BLOCK_BACKEND = "verl_fp8_block" +_MXFP8_BACKEND = "verl_mxfp8" _FP8_MAX = torch.finfo(torch.float8_e4m3fn).max @@ -67,8 +68,45 @@ def _fp8_block_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: return _FP8BlockFakeQuant.apply(inputs, block_rows, block_cols) +class _MXFP8FakeQuant(torch.autograd.Function): + """Dynamic MXFP8 fake quantization over groups on the last dimension, with E8M0 scales.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, group_size: int) -> torch.Tensor: + if x.shape[-1] % group_size: + raise ValueError( + f"MXFP8 requires the last dimension ({x.shape[-1]}) to be divisible by group_size={group_size}; " + "add this layer to the QAT and rollout ignore lists" + ) + + groups = x.float().unflatten(-1, (-1, group_size)) + amax = groups.abs().amax(dim=-1, keepdim=True) + + # compressed-tensors rounds each group amax to E8M0 before applying the + # FP8 E4M3 element-format offset (2**8). Its power-of-two rounding moves + # to the next exponent at a 1.75 significand. + base_exp = torch.floor(torch.log2(amax)) + base = torch.exp2(base_exp) + rounded_exp = base_exp + (amax >= 1.75 * base) + encoded_exp = (127 + rounded_exp - 8).clamp(0, 255) + scale = torch.exp2(encoded_exp - 127) + + dq = (groups / scale).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn).float() * scale + return dq.flatten(start_dim=-2).to(x.dtype) + + @staticmethod + def backward(ctx, grad): + return grad, None + + +def _mxfp8_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: + return _MXFP8FakeQuant.apply(inputs, quantizer.backend_extra_args["group_size"]) + + if not is_registered_quant_backend(_FP8_BLOCK_BACKEND): register_quant_backend(_FP8_BLOCK_BACKEND, _fp8_block_backend) +if not is_registered_quant_backend(_MXFP8_BACKEND): + register_quant_backend(_MXFP8_BACKEND, _mxfp8_backend) # Mirrors the compressed-tensors "fp8" rollout scheme: 128x128 weight blocks and # dynamic 1x128 activation groups, both with fp32 scales. Only linear GEMMs the @@ -92,7 +130,31 @@ def _fp8_block_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: "*vision_model*": {"enable": False}, } -_QUANTIZER_CFGS = {"w4a16": _NVFP4_W4A16_QUANTIZER_CFG, "fp8": _FP8_BLOCK_QUANTIZER_CFG} +# Mirrors compressed-tensors MXFP8: FP8 E4M3 values in groups of 32 with +# dynamic E8M0 scales for both weights and activations. +_MXFP8_QUANTIZER_CFG = { + "*": {"enable": False}, + "*weight_quantizer": { + "num_bits": (4, 3), + "backend": _MXFP8_BACKEND, + "backend_extra_args": {"group_size": 32}, + "enable": True, + }, + "*input_quantizer": { + "num_bits": (4, 3), + "backend": _MXFP8_BACKEND, + "backend_extra_args": {"group_size": 32}, + "enable": True, + }, + "*conv1d*": {"enable": False}, + "*vision_model*": {"enable": False}, +} + +_QUANTIZER_CFGS = { + "w4a16": _NVFP4_W4A16_QUANTIZER_CFG, + "fp8": _FP8_BLOCK_QUANTIZER_CFG, + "mxfp8": _MXFP8_QUANTIZER_CFG, +} def _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py index 6bc59439637..e842d69d780 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -31,9 +31,7 @@ """ import json -import os import re -import struct from collections.abc import Iterable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Optional @@ -60,6 +58,7 @@ # has been validated, mapped to the ModelOpt QAT mode that fake-quantizes identically. SUPPORTED_SCHEMES: dict[str, Optional[str]] = { "FP8_BLOCK": "fp8", + "MXFP8": "mxfp8", } # Layers that stay in high precision regardless of model. Patterns must start @@ -87,7 +86,16 @@ def _check_scheme(kind: str, scheme: Optional[str]) -> Optional[str]: return scheme -_NUMERIC_FIELDS = {"num_bits", "type", "symmetric", "group_size", "strategy", "block_structure", "dynamic"} +_NUMERIC_FIELDS = { + "num_bits", + "type", + "symmetric", + "group_size", + "strategy", + "block_structure", + "dynamic", + "scale_dtype", +} def _numerics(args: Optional["QuantizationArgs"]) -> Optional[dict[str, Any]]: @@ -132,55 +140,6 @@ def _matches(name: str, targets: Iterable[str]) -> bool: return False -def _iter_checkpoint_shapes(model_path: str) -> Iterator[tuple[str, list[int]]]: - """Yield ``(name, shape)`` from safetensors headers without reading tensor data.""" - index_file = os.path.join(model_path, "model.safetensors.index.json") - if os.path.exists(index_file): - with open(index_file) as f: - files = sorted(set(json.load(f)["weight_map"].values())) - else: - files = [f for f in sorted(os.listdir(model_path)) if f.endswith(".safetensors")] - for file in files: - with open(os.path.join(model_path, file), "rb") as f: - (header_len,) = struct.unpack(" str: - """``model.layers.3.linear_attn.in_proj_a`` -> ``re:.*linear_attn\\.in_proj_a$``.""" - suffix = re.split(r"\.\d+\.", module_name)[-1] - return "re:.*" + re.escape(suffix) + "$" - - -def _weight_tile(args: "QuantizationArgs") -> Optional[tuple[int, int]]: - """(rows, cols) a weight's shape must be a multiple of, if the strategy tiles it.""" - from compressed_tensors.quantization import QuantizationStrategy - - if args.strategy == QuantizationStrategy.BLOCK: - return tuple(args.block_structure) - if args.strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): - return 1, args.group_size - return None - - -def _incompatible_shape_ignores(model_path: str, ignore: list[str], tile: tuple[int, int]) -> list[str]: - """Patterns for 2D weights whose shape is not a multiple of the weight tile.""" - patterns: list[str] = [] - for name, shape in _iter_checkpoint_shapes(model_path): - if not name.endswith(".weight") or len(shape) != 2: - continue - if shape[0] % tile[0] == 0 and shape[1] % tile[1] == 0: - continue - module = name[: -len(".weight")] - if _matches(module, ignore + patterns): - continue - patterns.append(_module_pattern(module)) - return patterns - - def _config_group(targets: list[str], preset: str) -> dict[str, Any]: from compressed_tensors.compressors.format import infer_module_format from compressed_tensors.quantization import preset_name_to_scheme @@ -212,14 +171,12 @@ def build_compressed_tensors_config( moe: Optional[str], ignore: Iterable[str], config_file: Optional[str], - model_path: 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. """ - from compressed_tensors.quantization import preset_name_to_scheme - ignore = list(ignore or []) if config_file: if linear or moe or ignore: @@ -232,10 +189,6 @@ def build_compressed_tensors_config( raise ValueError("rollout quantization needs at least one of `linear` / `moe`") ignore = DEFAULT_IGNORE + ignore - for preset in {linear, moe} - {None}: - tile = _weight_tile(preset_name_to_scheme(preset, ["Linear"]).weights) - if tile is not None: - ignore += _incompatible_shape_ignores(model_path, ignore, tile) config_groups: dict[str, Any] = {} if linear == moe: diff --git a/verl/utils/vllm/vllm_fp8_utils.py b/verl/utils/vllm/vllm_fp8_utils.py index 8155da0c432..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(): @@ -321,6 +340,9 @@ def _patched_process_weights_after_loading(self, layer) -> None: _record_pristine_fp8_layout(layer, old_params) 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 @@ -447,9 +469,11 @@ def _create_param_from_subclass_attributes(custom_data, custom_weight): def _build_compressed_tensors_fp8_patchers(): - """Same refit treatment for the compressed-tensors linear and FP8 MoE methods.""" + """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.compressed_tensors_moe_w8a8_fp8" + 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, @@ -462,16 +486,33 @@ def _build_compressed_tensors_fp8_patchers(): return [] wrap = _make_process_weights_after_loading_for_vllm20 - return [ + patchers = [ patch( f"{ct}.compressed_tensors.CompressedTensorsLinearMethod.process_weights_after_loading", wrap(CompressedTensorsLinearMethod.process_weights_after_loading), ), patch( - f"{ct_moe}.CompressedTensorsW8A8Fp8MoEMethod.process_weights_after_loading", - wrap(CompressedTensorsW8A8Fp8MoEMethod.process_weights_after_loading, f"{ct_moe}.replace_parameter"), + 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): diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index 7f45d3c4628..38d7aba396d 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", "w4a4" or "fp8" (Megatron only), default "w4a16" + mode (str): Quantization mode, "w4a16", "w4a4", "fp8" or "mxfp8" (Megatron only), default "w4a16" 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) diff --git a/verl/workers/config/rollout.py b/verl/workers/config/rollout.py index 79b5a0080d1..4bfef74d226 100644 --- a/verl/workers/config/rollout.py +++ b/verl/workers/config/rollout.py @@ -168,7 +168,11 @@ def build_quantization_config(self, model_path: str) -> dict[str, Any]: 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 + linear=self.linear, + moe=self.moe, + ignore=self.ignore, + config_file=self.config_file, + model_path=model_path, ) diff --git a/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/verl/workers/rollout/vllm_rollout/vllm_async_server.py index e7774c6df4a..1e556218b53 100644 --- a/verl/workers/rollout/vllm_rollout/vllm_async_server.py +++ b/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -1402,9 +1402,16 @@ def _apply_quantization(self) -> dict: # 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) - # The schemes carry fp32 block scales. DeepGEMM's E8M0 mode (vLLM's default) rounds weight and - # activation scales to powers of two, so the rollout would drift from what the trainer quantized. - os.environ.setdefault("VLLM_USE_DEEP_GEMM_E8M0", "0") + # 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} From 700839ecf1895cd593753e033af34076ebcc8146 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 16 Oct 2026 14:22:24 +0800 Subject: [PATCH 06/18] fix mxfp8 QAT --- docs/low_precision/qat.md | 4 +- verl/utils/modelopt/quantize.py | 85 +++++++-------------------------- 2 files changed, 19 insertions(+), 70 deletions(-) diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index 679c133752f..512fe9e84f7 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -126,7 +126,7 @@ actor_rollout_ref: mode: fp8 # must match rollout.quantization (FP8_BLOCK -> fp8, MXFP8 -> mxfp8) ``` -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. +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. 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 @@ -177,7 +177,7 @@ Blackwell reference `k3_kl` for vLLM MXFP8: | Qwen3.5-0.8B | BF16 | 0.00066 | | | MXFP8 linear | 0.01391 | | Qwen3.5-35B-A3B (EP=8, rollout TP=2) | MXFP8 linear + MoE | 0.00148 | -| | MXFP8 linear + MoE + MXFP8 QAT | 0.00288 | +| | MXFP8 linear + MoE + MXFP8 QAT | 0.00315 | ## Limitations diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py index 50cec72b1df..ba8dbdac2ce 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -38,7 +38,6 @@ } _FP8_BLOCK_BACKEND = "verl_fp8_block" -_MXFP8_BACKEND = "verl_mxfp8" _FP8_MAX = torch.finfo(torch.float8_e4m3fn).max @@ -68,45 +67,8 @@ def _fp8_block_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: return _FP8BlockFakeQuant.apply(inputs, block_rows, block_cols) -class _MXFP8FakeQuant(torch.autograd.Function): - """Dynamic MXFP8 fake quantization over groups on the last dimension, with E8M0 scales.""" - - @staticmethod - def forward(ctx, x: torch.Tensor, group_size: int) -> torch.Tensor: - if x.shape[-1] % group_size: - raise ValueError( - f"MXFP8 requires the last dimension ({x.shape[-1]}) to be divisible by group_size={group_size}; " - "add this layer to the QAT and rollout ignore lists" - ) - - groups = x.float().unflatten(-1, (-1, group_size)) - amax = groups.abs().amax(dim=-1, keepdim=True) - - # compressed-tensors rounds each group amax to E8M0 before applying the - # FP8 E4M3 element-format offset (2**8). Its power-of-two rounding moves - # to the next exponent at a 1.75 significand. - base_exp = torch.floor(torch.log2(amax)) - base = torch.exp2(base_exp) - rounded_exp = base_exp + (amax >= 1.75 * base) - encoded_exp = (127 + rounded_exp - 8).clamp(0, 255) - scale = torch.exp2(encoded_exp - 127) - - dq = (groups / scale).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn).float() * scale - return dq.flatten(start_dim=-2).to(x.dtype) - - @staticmethod - def backward(ctx, grad): - return grad, None - - -def _mxfp8_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: - return _MXFP8FakeQuant.apply(inputs, quantizer.backend_extra_args["group_size"]) - - if not is_registered_quant_backend(_FP8_BLOCK_BACKEND): register_quant_backend(_FP8_BLOCK_BACKEND, _fp8_block_backend) -if not is_registered_quant_backend(_MXFP8_BACKEND): - register_quant_backend(_MXFP8_BACKEND, _mxfp8_backend) # Mirrors the compressed-tensors "fp8" rollout scheme: 128x128 weight blocks and # dynamic 1x128 activation groups, both with fp32 scales. Only linear GEMMs the @@ -130,31 +92,11 @@ def _mxfp8_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: "*vision_model*": {"enable": False}, } -# Mirrors compressed-tensors MXFP8: FP8 E4M3 values in groups of 32 with -# dynamic E8M0 scales for both weights and activations. -_MXFP8_QUANTIZER_CFG = { - "*": {"enable": False}, - "*weight_quantizer": { - "num_bits": (4, 3), - "backend": _MXFP8_BACKEND, - "backend_extra_args": {"group_size": 32}, - "enable": True, - }, - "*input_quantizer": { - "num_bits": (4, 3), - "backend": _MXFP8_BACKEND, - "backend_extra_args": {"group_size": 32}, - "enable": True, - }, - "*conv1d*": {"enable": False}, - "*vision_model*": {"enable": False}, -} - _QUANTIZER_CFGS = { "w4a16": _NVFP4_W4A16_QUANTIZER_CFG, "fp8": _FP8_BLOCK_QUANTIZER_CFG, - "mxfp8": _MXFP8_QUANTIZER_CFG, } +_SUPPORTED_QAT_MODES = {*_QUANTIZER_CFGS, "mxfp8"} def _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: @@ -177,22 +119,29 @@ def build_quantize_config( ignore_patterns: list[str] | None = None, ) -> dict: """Build a complete ModelOpt quantization config for ``mtq.quantize``.""" - if qat_mode not in _QUANTIZER_CFGS: - raise ValueError(f"Supported QAT modes are {sorted(_QUANTIZER_CFGS)}, got: {qat_mode}") + 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) - 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) + 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"] + 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 = "max" if qat_mode == "w4a16" else None quant_cfg.extend(ignore_cfg) - # The fp8 backend derives its scales from each forward's tensor, so there is nothing to calibrate. - return {"quant_cfg": quant_cfg, "algorithm": "max" if qat_mode == "w4a16" else None} + return {"quant_cfg": quant_cfg, "algorithm": algorithm} def apply_qat( From f6714c1f52ffe0dcc5a79141d56df3c76ee1b758 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 16 Oct 2026 15:13:59 +0800 Subject: [PATCH 07/18] enable R3 --- scripts/rollout_logprob_diff.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py index 610ac3e8f8b..c50509a093d 100644 --- a/scripts/rollout_logprob_diff.py +++ b/scripts/rollout_logprob_diff.py @@ -27,7 +27,8 @@ .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' + --qat-ignore '*linear_attn.in_proj_a' --qat-ignore '*linear_attn.in_proj_b' \\ + --r3-router-replay # 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 @@ -107,6 +108,11 @@ def parse_args() -> argparse.Namespace: 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("--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() @@ -158,6 +164,9 @@ def build_config(args: argparse.Namespace) -> DictConfig: *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 From d5a28a2376e89b0fcd1496a2b52dd8ccd26f91c8 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 16 Oct 2026 17:10:58 +0800 Subject: [PATCH 08/18] fix sglang R3 --- docs/low_precision/qat.md | 16 ++++++++----- .../sglang_rollout/async_sglang_server.py | 23 +++++++++---------- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index 512fe9e84f7..2b3f7dd0a6e 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -159,16 +159,20 @@ Use `actor_rollout_ref.rollout.name=sglang` for SGLang. For complete scripts, se ```bash 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 + --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, H-series GPUs): +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, H-series GPUs), with R3 router replay enabled: -| Rollout | vLLM | SGLang | +| Rollout | vLLM + R3 | SGLang + R3 | |---|---|---| -| BF16 | ~0.0014 | - | -| FP8_BLOCK | 0.0039 | 0.0036 | -| FP8_BLOCK + FP8 QAT | 0.0033 | 0.0033 | +| BF16 | 0.00031 | 0.00031 | +| FP8_BLOCK | 0.00225 | 0.00231 | +| FP8_BLOCK + FP8 QAT | 0.00187 | 0.00188 | Blackwell reference `k3_kl` for vLLM MXFP8: diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 341cbe712f7..02afee29fb5 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -658,22 +658,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} From 7711e029b202b089ded13fa06fc8088e874dc17a Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 16 Oct 2026 23:35:57 +0800 Subject: [PATCH 09/18] mxfp4 --- docs/low_precision/qat.md | 42 ++++++----- scripts/rollout_logprob_diff.py | 9 ++- verl/utils/modelopt/quantize.py | 82 ++++++++++++++++++++- verl/utils/qat/compressed_tensors.py | 29 ++++++-- verl/utils/vllm/vllm_fp4_utils.py | 105 ++++++++++++++++++++++++++- verl/utils/vllm/vllm_quant_utils.py | 8 +- 6 files changed, 241 insertions(+), 34 deletions(-) diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index 2b3f7dd0a6e..eeaf756c311 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -44,13 +44,15 @@ The trainer must be Megatron (FSDP is not yet supported for this feature; see [N | 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; Blackwell under verification | ✅ Hopper; Blackwell not yet verified | +| 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 | - | - | Blackwell | `W4A4Mxfp4` | `W4A4Mxfp4MoE` | ❌ | `Mxfp4MoEMethod` | Planned | Planned | -| NVFP4 | - | - | Blackwell | W4A4 / W4A16 | W4A4 / W4A16 | W4A4 only | W4A4 | Planned | Planned | +| 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` (planned) | `nvfp4` (planned) | Blackwell | ❌ | W4A4 / W4A16 | ❌ | W4A4 | Planned (MoE only) | Planned (MoE only) | 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 @@ -67,6 +69,7 @@ actor_rollout_ref: - 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` is valid only for `moe`. Use `linear: null, moe: MXFP4` for BF16 dense layers or `linear: MXFP8, moe: MXFP4` 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. @@ -123,15 +126,16 @@ actor_rollout_ref: megatron: qat: enable: true - mode: fp8 # must match rollout.quantization (FP8_BLOCK -> fp8, MXFP8 -> mxfp8) + mode: fp8 # FP8_BLOCK -> fp8, MXFP8 -> mxfp8, MoE MXFP4 -> mxfp4 ``` -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. +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. For `mxfp4`, verl uses compressed-tensors-compatible MXFP4 weight QDQ plus ModelOpt's MXFP8 input quantizer, enabled only on routed experts. The custom weight QDQ is necessary because ModelOpt's native E2M1 tie handling differs from compressed-tensors export for a small subset of values. 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. - **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 @@ -166,22 +170,20 @@ python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --e --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, H-series GPUs), with R3 router replay enabled: - -| Rollout | vLLM + R3 | SGLang + R3 | -|---|---|---| -| BF16 | 0.00031 | 0.00031 | -| FP8_BLOCK | 0.00225 | 0.00231 | -| FP8_BLOCK + FP8 QAT | 0.00187 | 0.00188 | +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): -Blackwell reference `k3_kl` for vLLM MXFP8: - -| Model | Rollout | k3_kl | -|---|---|---| -| Qwen3.5-0.8B | BF16 | 0.00066 | -| | MXFP8 linear | 0.01391 | -| Qwen3.5-35B-A3B (EP=8, rollout TP=2) | MXFP8 linear + MoE | 0.00148 | -| | MXFP8 linear + MoE + MXFP8 QAT | 0.00315 | +| 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 | ## Limitations diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py index c50509a093d..5da3d1e9774 100644 --- a/scripts/rollout_logprob_diff.py +++ b/scripts/rollout_logprob_diff.py @@ -29,6 +29,9 @@ --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 # 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 @@ -91,10 +94,8 @@ def parse_args() -> argparse.Namespace: 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, e.g. FP8_BLOCK or MXFP8" - ) - parser.add_argument("--quant-moe", default=None, help="rollout.quantization.moe preset, e.g. FP8_BLOCK or MXFP8") + 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, or MXFP4") parser.add_argument( "--quant-ignore", action="append", diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py index ba8dbdac2ce..e216c14eab0 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -38,6 +38,7 @@ } _FP8_BLOCK_BACKEND = "verl_fp8_block" +_MXFP4_BACKEND = "verl_mxfp4" _FP8_MAX = torch.finfo(torch.float8_e4m3fn).max @@ -67,8 +68,44 @@ def _fp8_block_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: 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) + + 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) # Mirrors the compressed-tensors "fp8" rollout scheme: 128x128 weight blocks and # dynamic 1x128 activation groups, both with fp32 scales. Only linear GEMMs the @@ -96,7 +133,48 @@ def _fp8_block_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: "w4a16": _NVFP4_W4A16_QUANTIZER_CFG, "fp8": _FP8_BLOCK_QUANTIZER_CFG, } -_SUPPORTED_QAT_MODES = {*_QUANTIZER_CFGS, "mxfp8"} +_SUPPORTED_QAT_MODES = {*_QUANTIZER_CFGS, "mxfp8", "mxfp4", "mxfp4_mxfp8"} + +_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 _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: @@ -133,6 +211,8 @@ def build_quantize_config( 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") else: quant_cfg = mtq.normalize_quant_cfg_list(copy.deepcopy(_QUANTIZER_CFGS[qat_mode])) disabled_cfg = copy.deepcopy(_default_disabled_quantizer_cfg) diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py index e842d69d780..d17c10674bd 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -59,6 +59,11 @@ 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", } # Layers that stay in high precision regardless of model. Patterns must start @@ -104,10 +109,9 @@ def _numerics(args: Optional["QuantizationArgs"]) -> Optional[dict[str, Any]]: def _preset_of(scheme: "QuantizationScheme") -> Optional[str]: """The supported preset ``scheme`` quantizes like, ignoring observers and targets.""" - from compressed_tensors.quantization import preset_name_to_scheme for preset in SUPPORTED_SCHEMES: - reference = preset_name_to_scheme(preset, scheme.targets) + 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") @@ -127,7 +131,10 @@ def _group_schemes(ct_config: Mapping[str, Any]) -> list["QuantizationScheme"]: 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``.""" - return {SUPPORTED_SCHEMES[_preset_of(scheme)] for scheme in _group_schemes(ct_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: @@ -140,11 +147,21 @@ def _matches(name: str, targets: Iterable[str]) -> bool: return False -def _config_group(targets: list[str], preset: str) -> dict[str, Any]: - from compressed_tensors.compressors.format import infer_module_format +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 @@ -185,6 +202,8 @@ def build_compressed_tensors_config( linear = _check_scheme("linear", linear) moe = _check_scheme("moe", moe) + if linear == "MXFP4": + raise ValueError("MXFP4 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`") diff --git a/verl/utils/vllm/vllm_fp4_utils.py b/verl/utils/vllm/vllm_fp4_utils.py index 9120ee07710..09050cf22aa 100644 --- a/verl/utils/vllm/vllm_fp4_utils.py +++ b/verl/utils/vllm/vllm_fp4_utils.py @@ -32,6 +32,7 @@ """ import logging +from dataclasses import replace from unittest.mock import patch import torch @@ -42,6 +43,105 @@ _MXFP4_LIVE_ATTR = "_verl_mxfp4_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 is_deepseek_v4_model(model): if model is None: return False @@ -90,6 +190,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 +261,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 +272,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) diff --git a/verl/utils/vllm/vllm_quant_utils.py b/verl/utils/vllm/vllm_quant_utils.py index a12bcdeae12..98fee59f768 100644 --- a/verl/utils/vllm/vllm_quant_utils.py +++ b/verl/utils/vllm/vllm_quant_utils.py @@ -49,6 +49,7 @@ from verl.utils.kernel.fp8_kernel import scaled_fp8_blockwise from verl.utils.vllm.vllm_fp4_utils import ( + install_mxfp4_mxfp8_compressed_tensors_patch, is_deepseek_v4_model, iter_deepseek_v4_weights, process_mxfp4_moe_weights_after_loading, @@ -501,10 +502,11 @@ 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() if fp8_state.vllm_patches: logger.debug("vLLM quantization patches already applied") return From ae0b8c088154588d640795c1bbf1952286385804 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 13:36:32 +0800 Subject: [PATCH 10/18] nvfp4 --- docs/low_precision/nvfp4_qat.md | 109 +-- docs/low_precision/qat.md | 20 +- scripts/rollout_logprob_diff.py | 9 +- .../_generated_ppo_megatron_trainer.yaml | 6 +- verl/trainer/config/engine/megatron.yaml | 6 +- verl/utils/modelopt/qat_utils.py | 2 +- verl/utils/modelopt/quantize.py | 137 ++- verl/utils/qat/__init__.py | 6 +- verl/utils/qat/compressed_tensors.py | 99 ++- verl/utils/vllm/vllm_fp4_utils.py | 221 ++++- verl/utils/vllm/vllm_nvfp4_utils.py | 828 ------------------ verl/utils/vllm/vllm_quant_utils.py | 6 + verl/workers/config/engine.py | 4 +- .../engine/megatron/transformer_impl.py | 3 + verl/workers/rollout/vllm_rollout/utils.py | 32 +- 15 files changed, 519 insertions(+), 969 deletions(-) delete mode 100644 verl/utils/vllm/vllm_nvfp4_utils.py 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 index eeaf756c311..cc671358311 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -40,14 +40,14 @@ Your HuggingFace checkpoint stays in BF16; no offline quantization or calibratio ## Support matrix -The trainer must be Megatron (FSDP is not yet supported for this feature; see [NVFP4 QAT](nvfp4_qat.md) for FSDP). 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. +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` (planned) | `nvfp4` (planned) | Blackwell | ❌ | W4A4 / W4A16 | ❌ | W4A4 | Planned (MoE only) | Planned (MoE only) | +| 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. @@ -69,7 +69,7 @@ actor_rollout_ref: - 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` is valid only for `moe`. Use `linear: null, moe: MXFP4` for BF16 dense layers or `linear: MXFP8, moe: MXFP4` for MXFP8 dense layers. +- `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. @@ -126,16 +126,17 @@ actor_rollout_ref: megatron: qat: enable: true - mode: fp8 # FP8_BLOCK -> fp8, MXFP8 -> mxfp8, MoE MXFP4 -> mxfp4 + 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. For `mxfp4`, verl uses compressed-tensors-compatible MXFP4 weight QDQ plus ModelOpt's MXFP8 input quantizer, enabled only on routed experts. The custom weight QDQ is necessary because ModelOpt's native E2M1 tie handling differs from compressed-tensors export for a small subset of values. Use `mxfp4_mxfp8` when dense linear layers also use MXFP8. +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 @@ -184,6 +185,15 @@ Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arg | 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.25778 | +| Blackwell | vLLM | ❌ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS | 0.32112 | + +The NVFP4 path is functionally validated with CUDA Graph, but its current +Qwen3.5-35B-A3B numerical gap is too large for training. The QAT result is not +an accuracy acceptance result; it records the current implementation state and +shows that additional kernel/QDQ parity work is required. Repeated runs varied +because FlashInfer's fused finalize is nondeterministic (`0.219–0.258` without +QAT and `0.263–0.380` with QAT). ## Limitations diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py index 5da3d1e9774..2cc5f2f37f8 100644 --- a/scripts/rollout_logprob_diff.py +++ b/scripts/rollout_logprob_diff.py @@ -32,6 +32,9 @@ # 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 @@ -95,7 +98,11 @@ def parse_args() -> argparse.Namespace: 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, or MXFP4") + 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", diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index 8f36fb143dd..ef22213ab04 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 @@ -634,7 +634,7 @@ critic: qat: _target_: verl.workers.config.QATEngineConfig enable: false - mode: w4a16 + mode: fp8 group_size: 16 ignore_patterns: - lm_head diff --git a/verl/trainer/config/engine/megatron.yaml b/verl/trainer/config/engine/megatron.yaml index 8033ba5e0ba..0ee54390a31 100644 --- a/verl/trainer/config/engine/megatron.yaml +++ b/verl/trainer/config/engine/megatron.yaml @@ -140,9 +140,9 @@ qat: _target_: verl.workers.config.QATEngineConfig # Whether to enable QAT enable: false - # Quantization mode: "w4a16" (weight-only FP4), "fp8" (FP8_BLOCK rollout), or "mxfp8" - # (MXFP8 rollout with group-of-32 E8M0-scaled weights and activations) - 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/utils/modelopt/qat_utils.py b/verl/utils/modelopt/qat_utils.py index b5fc0b1afe8..c83b2e01eb8 100644 --- a/verl/utils/modelopt/qat_utils.py +++ b/verl/utils/modelopt/qat_utils.py @@ -27,7 +27,7 @@ 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") + 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) diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py index e216c14eab0..a43c8f92ebc 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -27,18 +27,10 @@ register_quant_backend, ) -_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}, -} - _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 @@ -102,10 +94,93 @@ 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 None: + # A neutral divisor keeps the full FP8 local-scale range available and + # also gives the first pre-rollout weight sync a deterministic value. + previous = torch.ones_like(current) + 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 @@ -129,11 +204,8 @@ def _mxfp4_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: "*vision_model*": {"enable": False}, } -_QUANTIZER_CFGS = { - "w4a16": _NVFP4_W4A16_QUANTIZER_CFG, - "fp8": _FP8_BLOCK_QUANTIZER_CFG, -} -_SUPPORTED_QAT_MODES = {*_QUANTIZER_CFGS, "mxfp8", "mxfp4", "mxfp4_mxfp8"} +_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" @@ -177,6 +249,37 @@ def _mxfp4_quantize_config(*, dense_mxfp8: bool) -> tuple[list[dict], None]: 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 = { @@ -213,13 +316,15 @@ def build_quantize_config( 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 = "max" if qat_mode == "w4a16" else None + algorithm = None quant_cfg.extend(ignore_cfg) return {"quant_cfg": quant_cfg, "algorithm": algorithm} diff --git a/verl/utils/qat/__init__.py b/verl/utils/qat/__init__.py index 71afd229bb2..3488abe004c 100644 --- a/verl/utils/qat/__init__.py +++ b/verl/utils/qat/__init__.py @@ -15,7 +15,7 @@ """ 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) @@ -23,12 +23,10 @@ - quantizer.py: QATQuantizer for true quantization + scale computation utilities - compressed_tensors.py: compressed-tensors rollout config and trainer-side weight quantizer -The vLLM side (NVFP4 refit patches) lives in ``verl.utils.vllm.vllm_nvfp4_utils``. - 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 """ diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py index d17c10674bd..ad8b24383da 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -64,6 +64,9 @@ # 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 @@ -202,8 +205,8 @@ def build_compressed_tensors_config( linear = _check_scheme("linear", linear) moe = _check_scheme("moe", moe) - if linear == "MXFP4": - raise ValueError("MXFP4 is supported only for routed MoE experts; use BF16 or MXFP8 dense linear layers") + 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`") @@ -254,7 +257,12 @@ def _weight_min_max(weight: torch.Tensor, args: "QuantizationArgs") -> tuple[tor raise NotImplementedError(f"Trainer-side quantization not implemented for strategy {args.strategy}") -def _compress_weight(weight: torch.Tensor, scheme: "QuantizationScheme") -> dict[str, torch.Tensor]: +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 @@ -264,9 +272,9 @@ def _compress_weight(weight: torch.Tensor, scheme: "QuantizationScheme") -> dict args = scheme.weights min_vals, max_vals = (t.float() for t in _weight_min_max(weight, args)) state_dict = {"weight": weight} - global_scale = None if args.strategy == QuantizationStrategy.TENSOR_GROUP: - global_scale = generate_gparam(min_vals.amin(), max_vals.amax()) + 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 @@ -306,10 +314,81 @@ class CompressedTensorsWeightQuantizer: experts are split per expert first. Everything else passes through unchanged. """ - def __init__(self, ct_config: Mapping[str, Any]): + 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.modelopt.quantize 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): @@ -327,8 +406,9 @@ def _scheme_for(self, name: str, tensor: torch.Tensor) -> Optional["Quantization 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 - if packed and self._module_scheme(f"{packed['prefix']}.0.gate_proj") is not None: - yield from self(_unpack_experts(name, tensor)) + 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: @@ -337,3 +417,6 @@ def __call__(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tupl 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 diff --git a/verl/utils/vllm/vllm_fp4_utils.py b/verl/utils/vllm/vllm_fp4_utils.py index 09050cf22aa..795cf8c6f93 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 @@ -41,6 +41,8 @@ _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(): @@ -142,6 +144,27 @@ def get_moe_method(quant_config, layer, layer_name): 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): if model is None: return False @@ -308,6 +331,202 @@ 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 _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)) + + 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) + 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_nvfp4_utils.py b/verl/utils/vllm/vllm_nvfp4_utils.py deleted file mode 100644 index 77d2ec17eeb..00000000000 --- a/verl/utils/vllm/vllm_nvfp4_utils.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/vllm/vllm_quant_utils.py b/verl/utils/vllm/vllm_quant_utils.py index 98fee59f768..2d7a0d58716 100644 --- a/verl/utils/vllm/vllm_quant_utils.py +++ b/verl/utils/vllm/vllm_quant_utils.py @@ -50,10 +50,13 @@ 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, @@ -433,6 +436,7 @@ def prepare_quanted_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 @@ -443,6 +447,7 @@ def process_quanted_weights_after_loading(model, reload_state): 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) @@ -507,6 +512,7 @@ def apply_vllm_quant_patches(): 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 38d7aba396d..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", "w4a4", "fp8" or "mxfp8" (Megatron only), 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/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index 3a8fa9889a9..84b0cfa1dc3 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -1005,6 +1005,9 @@ def get_per_tensor_param(self, base_sync_done=False, **kwargs): 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 diff --git a/verl/workers/rollout/vllm_rollout/utils.py b/verl/workers/rollout/vllm_rollout/utils.py index f8792bc1667..92b5da77628 100644 --- a/verl/workers/rollout/vllm_rollout/utils.py +++ b/verl/workers/rollout/vllm_rollout/utils.py @@ -189,18 +189,9 @@ def __new__(cls, **kwargs): # 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. - quant_config = getattr(vllm_config, "quant_config", None) if vllm_config else None - _is_nvfp4_packed = getattr(quant_config, "quant_format", None) == "nvfp4-pack-quantized" - if is_quantized_model(vllm_config) and not _is_nvfp4_packed: + if is_quantized_model(vllm_config): apply_vllm_quant_patches() - # 3. patch compressed-tensors NVFP4 for dynamic weight loading - if _is_nvfp4_packed: - from verl.utils.vllm.vllm_nvfp4_utils import apply_qat_patches - - apply_qat_patches() - logger.info("Applied compressed-tensors NVFP4 patches in vLLM worker subprocess") - # 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. # This is only a fix for vllm version < v0.13.0. @@ -210,7 +201,6 @@ def __new__(cls, **kwargs): os.environ[k] = VLLM_ASCEND_REQUIRED_ENV_VARS[k] instance = super().__new__(cls) - instance._is_nvfp4_packed = _is_nvfp4_packed return instance def _get_drafter_model(self): @@ -271,14 +261,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_nvfp4_packed: - # Compressed-tensors NVFP4: prepare for weight loading BEFORE receiving any buckets - from verl.utils.vllm.vllm_nvfp4_utils import prepare_qat_for_load_weights - - for model in self._iter_all_models(): - prepare_qat_for_load_weights(model, device=self.device) - logger.info("NVFP4: prepare_qat_for_load_weights 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") @@ -329,14 +312,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_nvfp4_packed: - # Compressed-tensors NVFP4: call process_weights_after_loading AFTER all buckets are received - from verl.utils.vllm.vllm_nvfp4_utils import manual_process_weights_after_loading - - for model in self._iter_all_models(): - manual_process_weights_after_loading(model) - logger.info("NVFP4: 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: @@ -384,7 +360,7 @@ def _update_weights( param_updates, buffer_updates, named_buffers = split_buffer_updates(self.model_runner.model, weights) # Add the FP8 related logic here as sharding manager has been deprecated. # Check if FP8 quantization is enabled and apply appropriate weight loading - if is_quantized_model(self.model_runner.vllm_config) and not self._is_nvfp4_packed: + if is_quantized_model(self.model_runner.vllm_config): logger.info(f"FP8 model detected (async): {self.model_runner.vllm_config.quant_config}") # Convert bf16 weights to fp8 format before loading model_updates = drop_tied_alias_updates(self.model_runner.model, param_updates) From c7690670f279636bcc8b571d69e87ebac05301eb Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 15:34:00 +0800 Subject: [PATCH 11/18] nvfp4 --- docs/low_precision/qat.md | 11 ++----- verl/utils/vllm/vllm_fp4_utils.py | 55 +++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 9 deletions(-) diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index cc671358311..a2b18938f48 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -185,15 +185,8 @@ Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arg | 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.25778 | -| Blackwell | vLLM | ❌ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS | 0.32112 | - -The NVFP4 path is functionally validated with CUDA Graph, but its current -Qwen3.5-35B-A3B numerical gap is too large for training. The QAT result is not -an accuracy acceptance result; it records the current implementation state and -shows that additional kernel/QDQ parity work is required. Repeated runs varied -because FlashInfer's fused finalize is nondeterministic (`0.219–0.258` without -QAT and `0.263–0.380` with QAT). +| Blackwell | vLLM | ✅ | NVFP4 MoE | FlashInfer CUTLASS | 0.00813 | +| Blackwell | vLLM | ✅ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS | 0.00574 | ## Limitations diff --git a/verl/utils/vllm/vllm_fp4_utils.py b/verl/utils/vllm/vllm_fp4_utils.py index 795cf8c6f93..a7c9102d4b5 100644 --- a/verl/utils/vllm/vllm_fp4_utils.py +++ b/verl/utils/vllm/vllm_fp4_utils.py @@ -473,6 +473,58 @@ def _restore_nvfp4_input_scales_in_place(module, converted): 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, @@ -483,6 +535,8 @@ def _process_nvfp4_moe_params(module): 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, @@ -496,6 +550,7 @@ def convert_to_nvfp4_moe_kernel_format(*args, **kwargs): ), ): 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: From cbe9c7d8b95f2df4bf47a5b67842db572f4847c5 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 17:56:24 +0800 Subject: [PATCH 12/18] fp8 --- docs/low_precision/qat.md | 34 +- tests/utils/test_modelopt_mxfp8_on_cpu.py | 407 ++++++++++++++++++ tests/utils/test_vllm_mxfp4_restage_on_cpu.py | 153 +++++++ .../utils/test_vllm_quant_utils_moe_on_cpu.py | 240 ----------- ...t_vllm_weight_name_normalization_on_cpu.py | 1 + .../test_vllm_weight_update_utils_on_cpu.py | 1 + verl/utils/qat/__init__.py | 4 + verl/utils/qat/compressed_tensors.py | 159 +++++++ verl/utils/vllm/vllm_fp4_utils.py | 13 - verl/utils/vllm/vllm_quant_utils.py | 322 +------------- verl/workers/engine_workers.py | 30 +- verl/workers/rollout/vllm_rollout/utils.py | 2 + 12 files changed, 783 insertions(+), 583 deletions(-) create mode 100644 tests/utils/test_modelopt_mxfp8_on_cpu.py create mode 100644 tests/utils/test_vllm_mxfp4_restage_on_cpu.py delete mode 100644 tests/utils/test_vllm_quant_utils_moe_on_cpu.py diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index a2b18938f48..51e4767def9 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -34,9 +34,30 @@ The trainer always trains in high precision; low-precision checkpoints are dequa ## 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 an overridden `quantization_config`, so it allocates quantized parameters as if it were loading a quantized checkpoint (e.g. FP8 `weight` + `weight_scale`). On every weight sync, the trainer quantizes its BF16 weights to that checkpoint format, and the engine's weight loader handles fusion, tensor-parallel sharding and kernel-specific layouts. - -Your HuggingFace checkpoint stays in BF16; no offline quantization or calibration step is needed. The rollout format is described in the open [compressed-tensors](https://github.com/neuralmagic/compressed-tensors) format, which both vLLM and SGLang understand and which can express different schemes for linear and MoE layers. +- **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. + +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 @@ -191,5 +212,10 @@ Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arg ## Limitations - `FP8_BLOCK` is supported on Hopper. `MXFP8` with vLLM is supported on Blackwell; SGLang has no per-layer MXFP8 support. -- Checkpoints that already ship a `quantization_config` (e.g. FP8 or compressed-tensors checkpoints) are not supported as a training target together with low-precision rollout yet; use a BF16 checkpoint. +- 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/tests/utils/test_modelopt_mxfp8_on_cpu.py b/tests/utils/test_modelopt_mxfp8_on_cpu.py new file mode 100644 index 00000000000..8ff200170bf --- /dev/null +++ b/tests/utils/test_modelopt_mxfp8_on_cpu.py @@ -0,0 +1,407 @@ +# 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.modelopt.quantize import ( + _Mxfp4FakeQuant, + _Nvfp4FakeQuant, + _nvfp4_global_scale, + _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_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_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_vllm_mxfp4_restage_on_cpu.py b/tests/utils/test_vllm_mxfp4_restage_on_cpu.py new file mode 100644 index 00000000000..c3a931bf0fb --- /dev/null +++ b/tests/utils/test_vllm_mxfp4_restage_on_cpu.py @@ -0,0 +1,153 @@ +# 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 torch + +from verl.utils.vllm.vllm_fp4_utils import ( + _MXFP4_LIVE_ATTR, + _NVFP4_LIVE_ATTR, + _restore_nvfp4_input_scales_in_place, + _restore_nvfp4_kernel_config_in_place, + _stage_mxfp4_moe_params, + _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 5f991403d59..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 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/verl/utils/qat/__init__.py b/verl/utils/qat/__init__.py index 3488abe004c..6d7b5f0aecf 100644 --- a/verl/utils/qat/__init__.py +++ b/verl/utils/qat/__init__.py @@ -32,7 +32,9 @@ 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 ( @@ -50,6 +52,8 @@ "enable_qat_fuse", "invalidate_all_scales", "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 index ad8b24383da..8911d0ccf43 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -46,6 +46,8 @@ "MOE_TARGET", "SUPPORTED_SCHEMES", "CompressedTensorsWeightQuantizer", + "Fp8ConfigWeightQuantizer", + "dispatch_rollout_weight_quantizer", "build_compressed_tensors_config", "qat_modes_for_config", ] @@ -420,3 +422,160 @@ def __call__(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tupl 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 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/vllm/vllm_fp4_utils.py b/verl/utils/vllm/vllm_fp4_utils.py index a7c9102d4b5..184e96191ae 100644 --- a/verl/utils/vllm/vllm_fp4_utils.py +++ b/verl/utils/vllm/vllm_fp4_utils.py @@ -177,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 diff --git a/verl/utils/vllm/vllm_quant_utils.py b/verl/utils/vllm/vllm_quant_utils.py index 2d7a0d58716..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,10 @@ 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, @@ -74,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: []) @@ -90,8 +73,6 @@ 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): - return True elif is_compressed_tensors_model(vllm_config): return True @@ -107,199 +88,23 @@ def is_compressed_tensors_model(vllm_config) -> bool: return isinstance(getattr(vllm_config, "quant_config", None), CompressedTensorsConfig) -# 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 - ) - return is_w4a8_linear or is_w4a8_moe - - -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 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. +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 - 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): @@ -340,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. @@ -432,7 +155,6 @@ 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), @@ -443,7 +165,6 @@ def prepare_quanted_weights_for_loading(model): 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 []) @@ -463,14 +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) - if is_compressed_tensors_model(model_runner.vllm_config): - weights_quantized = weights - else: - 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 @@ -491,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(): diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index 8d0a80f27b3..b724eae7c49 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -43,10 +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 ( - CompressedTensorsWeightQuantizer, - qat_modes_for_config, -) +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 ( @@ -56,7 +53,6 @@ McoreEngineConfig, MtpConfig, RolloutConfig, - RolloutQuantizationConfig, TrainingWorkerConfig, ) from verl.workers.rollout.base import BaseRollout, get_rollout_class @@ -734,24 +730,14 @@ def init_model(self): backend, is_master=(torch.distributed.get_rank() == 0), bucket_size=bucket_size, **engine_kwargs ) - rollout_quantization: RolloutQuantizationConfig = omega_conf_to_dataclass( - self.config.rollout.quantization, RolloutQuantizationConfig - ) qat = getattr(getattr(self.actor.engine, "engine_config", None), "qat", None) - qat_mode = qat.mode if qat is not None and qat.enable else None - if qat_mode is not None and not rollout_quantization.enabled: - raise ValueError( - "QAT needs a quantized rollout to sync to: set rollout.quantization.linear / moe " - "(e.g. FP8_BLOCK) or rollout.quantization.config_file" - ) - if rollout_quantization.enabled: - ct_config = rollout_quantization.build_quantization_config(model_config.local_path) - rollout_qat_modes = qat_modes_for_config(ct_config) - if qat_mode is not None and rollout_qat_modes != {qat_mode}: - raise ValueError( - f"QAT mode {qat_mode!r} does not match the rollout quantization (QAT modes {rollout_qat_modes})" - ) - self.actor.engine.set_rollout_weight_quantizer(CompressedTensorsWeightQuantizer(ct_config)) + 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/vllm_rollout/utils.py b/verl/workers/rollout/vllm_rollout/utils.py index 92b5da77628..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, @@ -190,6 +191,7 @@ def __new__(cls, **kwargs): # 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, From 5c4e63528debfa89c72e8fc8bbe0f1eb13caa303 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 20:14:49 +0800 Subject: [PATCH 13/18] refactor --- scripts/rollout_logprob_diff.py | 51 ++++++++++++------- .../special_sanity/check_device_api_usage.py | 1 - tests/utils/test_modelopt_mxfp8_on_cpu.py | 2 +- verl/utils/modelopt/__init__.py | 25 --------- verl/utils/modelopt/qat_utils.py | 37 -------------- verl/utils/qat/__init__.py | 1 + verl/utils/qat/compressed_tensors.py | 2 +- .../{modelopt/quantize.py => qat/modelopt.py} | 31 ++++++++--- .../engine/megatron/transformer_impl.py | 2 +- 9 files changed, 62 insertions(+), 90 deletions(-) delete mode 100644 verl/utils/modelopt/__init__.py delete mode 100644 verl/utils/modelopt/qat_utils.py rename verl/utils/{modelopt/quantize.py => qat/modelopt.py} (92%) diff --git a/scripts/rollout_logprob_diff.py b/scripts/rollout_logprob_diff.py index 2cc5f2f37f8..103dcdde4dd 100644 --- a/scripts/rollout_logprob_diff.py +++ b/scripts/rollout_logprob_diff.py @@ -121,6 +121,12 @@ def parse_args() -> argparse.Namespace: 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() @@ -229,6 +235,8 @@ def compute_old_log_prob(actor_rollout_wg: RayWorkerGroup, batch: DataProto) -> 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( @@ -255,26 +263,35 @@ def main() -> int: }, ).repeat(config.actor_rollout_ref.rollout.n) - result = agent_loop_manager.generate_sequences(prompts=batch) - checkpoint_manager.sleep_replicas() - result.meta_info.pop("timing", None) + 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)) + 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() + 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()) + 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() 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 index 8ff200170bf..fc0bb000eb6 100644 --- a/tests/utils/test_modelopt_mxfp8_on_cpu.py +++ b/tests/utils/test_modelopt_mxfp8_on_cpu.py @@ -24,7 +24,7 @@ # isort: off import verl.models.mcore # noqa: F401 import modelopt.torch.quantization as mtq -from verl.utils.modelopt.quantize import ( +from verl.utils.qat.modelopt import ( _Mxfp4FakeQuant, _Nvfp4FakeQuant, _nvfp4_global_scale, diff --git a/verl/utils/modelopt/__init__.py b/verl/utils/modelopt/__init__.py deleted file mode 100644 index faf96665dd4..00000000000 --- a/verl/utils/modelopt/__init__.py +++ /dev/null @@ -1,25 +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 fake quantization for Megatron QAT training.""" - -from verl.utils.modelopt.qat_utils import apply_qat_to_modules -from verl.utils.modelopt.quantize import apply_qat, build_quantize_config - -__all__ = [ - "build_quantize_config", - "apply_qat", - "apply_qat_to_modules", -] diff --git a/verl/utils/modelopt/qat_utils.py b/verl/utils/modelopt/qat_utils.py deleted file mode 100644 index c83b2e01eb8..00000000000 --- a/verl/utils/modelopt/qat_utils.py +++ /dev/null @@ -1,37 +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", "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/__init__.py b/verl/utils/qat/__init__.py index 6d7b5f0aecf..fc62bfac3e7 100644 --- a/verl/utils/qat/__init__.py +++ b/verl/utils/qat/__init__.py @@ -22,6 +22,7 @@ - linear.py: QATLinear layer with Triton kernels for fake quantization - quantizer.py: QATQuantizer for true quantization + scale computation utilities - 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 diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py index 8911d0ccf43..44e940a2821 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -332,7 +332,7 @@ def set_input_global_scales(self, scales: Mapping[str, torch.Tensor]) -> None: def refresh_input_global_scales(self, model: torch.nn.Module) -> None: """Refresh NVFP4 activation scales from ModelOpt before an export.""" - from verl.utils.modelopt.quantize import collect_nvfp4_input_global_scales + from verl.utils.qat.modelopt import collect_nvfp4_input_global_scales scales = collect_nvfp4_input_global_scales(model) if scales: diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/qat/modelopt.py similarity index 92% rename from verl/utils/modelopt/quantize.py rename to verl/utils/qat/modelopt.py index a43c8f92ebc..48821864b30 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/qat/modelopt.py @@ -5,7 +5,7 @@ # 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 +# 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, @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ModelOpt quantization config and application for Megatron QAT.""" +"""ModelOpt fake quantization and Megatron QAT workflow helpers.""" import copy @@ -35,7 +35,7 @@ class _FP8BlockFakeQuant(torch.autograd.Function): - """Dynamic FP8 E4M3 fake quantization over ``(rows, cols)`` blocks of the last two dims, straight-through grad.""" + """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: @@ -148,7 +148,7 @@ def _nvfp4_global_scale(inputs: torch.Tensor) -> torch.Tensor: def _nvfp4_weight_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: if inputs.ndim > 2: - # TEGroupedLinear stores one expert matrix per leading entry. The + # 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( @@ -288,9 +288,7 @@ def _ignore_patterns_to_quant_cfg(ignore_patterns: list[str]) -> list[dict]: "*self_attn*": "*self_attention*", } for pattern in ignore_patterns: - key = pattern - if key in mapping: - key = mapping[key] + key = mapping.get(pattern, pattern) cfg.append({"quantizer_name": key, "enable": False}) return cfg @@ -338,3 +336,22 @@ def apply_qat( 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/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index 84b0cfa1dc3..65e4b89579d 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -530,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) From 7d2217742833d3e2c5306b38a9794561f48163f2 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 20:39:47 +0800 Subject: [PATCH 14/18] fix sglang --- tests/utils/test_modelopt_mxfp8_on_cpu.py | 31 +++++++++++++++++++ verl/utils/qat/compressed_tensors.py | 2 ++ .../sglang_rollout/async_sglang_server.py | 3 +- 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/utils/test_modelopt_mxfp8_on_cpu.py b/tests/utils/test_modelopt_mxfp8_on_cpu.py index fc0bb000eb6..f88c652a7dd 100644 --- a/tests/utils/test_modelopt_mxfp8_on_cpu.py +++ b/tests/utils/test_modelopt_mxfp8_on_cpu.py @@ -173,6 +173,37 @@ def test_native_fp8_config_export_uses_vllm_scale_names(): 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( { diff --git a/verl/utils/qat/compressed_tensors.py b/verl/utils/qat/compressed_tensors.py index 44e940a2821..3408132742c 100644 --- a/verl/utils/qat/compressed_tensors.py +++ b/verl/utils/qat/compressed_tensors.py @@ -460,6 +460,8 @@ def __init__(self, fp8_config: Mapping[str, Any], *, model_types: Iterable[Optio 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": diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 02afee29fb5..b74af031d89 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -301,6 +301,7 @@ async def launch_server(self, master_address: str = None, master_port: int = Non # Leave None to let sglang's VisionAttention auto-pick per device # (triton_attn on Ada, fa3 on Hopper, fa4 on Blackwell). 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. @@ -334,7 +335,7 @@ async def launch_server(self, master_address: str = None, master_port: int = Non **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( From 95638197660477195d2f8fdbe82a8c3dcf913564 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 20:41:04 +0800 Subject: [PATCH 15/18] update doc --- docs/low_precision/qat.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index 51e4767def9..211448c07b6 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -47,6 +47,8 @@ The trainer always trains in high precision; low-precision checkpoints are dequa 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 From d68a50e8699474d24c82962ad3d53d65196d3515 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 20:59:59 +0800 Subject: [PATCH 16/18] [test] fix: isolate FP4 CPU tests from vLLM Co-authored-by: Cursor --- tests/utils/test_vllm_mxfp4_restage_on_cpu.py | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/tests/utils/test_vllm_mxfp4_restage_on_cpu.py b/tests/utils/test_vllm_mxfp4_restage_on_cpu.py index c3a931bf0fb..54ab034b610 100644 --- a/tests/utils/test_vllm_mxfp4_restage_on_cpu.py +++ b/tests/utils/test_vllm_mxfp4_restage_on_cpu.py @@ -12,18 +12,33 @@ # 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 -from verl.utils.vllm.vllm_fp4_utils import ( - _MXFP4_LIVE_ATTR, - _NVFP4_LIVE_ATTR, - _restore_nvfp4_input_scales_in_place, - _restore_nvfp4_kernel_config_in_place, - _stage_mxfp4_moe_params, - _stage_nvfp4_moe_params, -) +_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: From 05152e618062b309a00695b3d3d149a2418589bc Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 21:05:54 +0800 Subject: [PATCH 17/18] [cfg] fix: instantiate reward rollout quantization Co-authored-by: Cursor --- verl/trainer/config/_generated_ppo_megatron_trainer.yaml | 1 + verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml | 1 + verl/trainer/config/_generated_ppo_trainer.yaml | 1 + verl/trainer/config/_generated_ppo_veomni_trainer.yaml | 1 + verl/trainer/config/reward/reward.yaml | 1 + 5 files changed, 5 insertions(+) diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index ef22213ab04..c4d6ef0b80f 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -868,6 +868,7 @@ reward: prompt_length: 2048 response_length: 2048 quantization: + _target_: verl.workers.config.RolloutQuantizationConfig linear: null moe: null ignore: [] diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index a4c649a7ccc..ef5dbc68306 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -774,6 +774,7 @@ reward: prompt_length: 2048 response_length: 2048 quantization: + _target_: verl.workers.config.RolloutQuantizationConfig linear: null moe: null ignore: [] diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index 787b46d660e..bf41fa143a6 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -837,6 +837,7 @@ reward: prompt_length: 2048 response_length: 2048 quantization: + _target_: verl.workers.config.RolloutQuantizationConfig linear: null moe: null ignore: [] diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index 1a9e539fb5c..481c245f1ff 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -808,6 +808,7 @@ reward: prompt_length: 2048 response_length: 2048 quantization: + _target_: verl.workers.config.RolloutQuantizationConfig linear: null moe: null ignore: [] diff --git a/verl/trainer/config/reward/reward.yaml b/verl/trainer/config/reward/reward.yaml index 9614422429f..afb4bd0ed16 100644 --- a/verl/trainer/config/reward/reward.yaml +++ b/verl/trainer/config/reward/reward.yaml @@ -73,6 +73,7 @@ reward_model: response_length: 2048 quantization: + _target_: verl.workers.config.RolloutQuantizationConfig linear: null moe: null ignore: [] From b0025e35a1b530f530636a604f614dc0e14764e6 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Mon, 19 Oct 2026 21:58:12 +0800 Subject: [PATCH 18/18] [test, model] fix: initialize NVFP4 scale from observations Co-authored-by: Cursor --- tests/utils/test_modelopt_mxfp8_on_cpu.py | 17 +++++++++++++++++ .../test_megatron_liger_fused_kernels_on_cpu.py | 7 ++++++- verl/utils/qat/modelopt.py | 8 +++----- 3 files changed, 26 insertions(+), 6 deletions(-) diff --git a/tests/utils/test_modelopt_mxfp8_on_cpu.py b/tests/utils/test_modelopt_mxfp8_on_cpu.py index f88c652a7dd..d0e63a05679 100644 --- a/tests/utils/test_modelopt_mxfp8_on_cpu.py +++ b/tests/utils/test_modelopt_mxfp8_on_cpu.py @@ -28,6 +28,7 @@ _Mxfp4FakeQuant, _Nvfp4FakeQuant, _nvfp4_global_scale, + _nvfp4_input_backend, _nvfp4_weight_backend, build_quantize_config, ) @@ -409,6 +410,22 @@ def test_nvfp4_grouped_expert_qdq_uses_one_global_scale_per_expert(): 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, 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/utils/qat/modelopt.py b/verl/utils/qat/modelopt.py index 48821864b30..f904cc1aa9e 100644 --- a/verl/utils/qat/modelopt.py +++ b/verl/utils/qat/modelopt.py @@ -164,11 +164,9 @@ def _nvfp4_input_backend(inputs: torch.Tensor, quantizer) -> torch.Tensor: # 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 None: - # A neutral divisor keeps the full FP8 local-scale range available and - # also gives the first pre-rollout weight sync a deterministic value. - previous = torch.ones_like(current) - current = torch.minimum(previous.to(current.device), current) + 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)