`quantizer.py` — 这个文件实现了 QATQuantizer -- 用于将 QAT 训练后的 FP16/BF16...¶

模块路径: verl.utils.qat.quantizer

文件概述¶

这个文件实现了 QATQuantizer -- 用于将 QAT 训练后的 FP16/BF16 权重真正量化为 NVFP4 格式的量化器。与 linear.py 的假量化不同，这里是真正的量化，产生打包后的低精度权重和 scale 参数。

使用 compressed_tensors 库的 API 进行实际的量化和打包。

关键代码讲解¶

1. Blockwise Scale 计算¶

def compute_blockwise_scale(weight, global_scale, group_size=16):
    """计算分块 scale（FP8 E4M3 格式）"""
    out_features, in_features = weight.shape
    weight_reshaped = weight.view(out_features, num_groups, group_size)

    # 每个分块的最大绝对值
    block_max = torch.amax(torch.abs(weight_reshaped), dim=-1)

    # 计算 local scale 并用 global scale 归一化
    local_scale = block_max / FP4_E2M1_DATA.max
    blockwise_scale = torch.clamp(
        global_scale * local_scale,
        min=-FP8_E4M3_DATA.max, max=FP8_E4M3_DATA.max
    ).to(torch.float8_e4m3fn)
    return blockwise_scale

NVFP4 的两级 scale 体系：global_scale（整个权重矩阵）+ blockwise_scale（每个分组），确保量化精度。

2. Scale 融合¶

FUSE_PATTERNS = {
    "qkv": ["q_proj", "k_proj", "v_proj"],
    "gate_up": ["gate_proj", "up_proj"],
}

def fuse_global_scales(layer_global_scales, strategy="min"):
    """将 QKV/GateUp 组的 global scale 融合（取最小值）"""
    for _, patterns in FUSE_PATTERNS.items():
        matched = [children[p] for p in patterns if p in children]
        if len(matched) == len(patterns):
            group_scales = [layer_global_scales[n] for n in matched]
            fused_scale = torch.min(torch.cat(group_scales)).reshape([1])
            for layer_name in matched:
                fused_scales[layer_name] = fused_scale.clone()

取最小值（min）策略确保量化范围足以覆盖所有融合层。

3. QATQuantizer 类¶

class QATQuantizer:
    def __init__(self, mode="w4a16", group_size=16, ignore_patterns=None, device=None):
        self._compressor = NVFP4PackedCompressor()
        self._quant_args = QuantizationArgs(
            num_bits=4, type=QuantizationType.FLOAT, symmetric=True,
            strategy=QuantizationStrategy.TENSOR_GROUP, group_size=group_size
        )

4. 流式量化¶

def quantize_with_fusion(self, params, target_device=None):
    """流式量化：逐层消费输入，逐层产出量化结果"""
    current_layer_idx = _sentinel
    layer_buffer = {}

    for name, tensor in params:
        layer_idx = self._extract_layer_idx(name)

        # 层边界：刷新上一层的缓冲
        if layer_idx != current_layer_idx and layer_buffer:
            yield from self._process_layer_group(current_layer_idx, layer_buffer, ...)
            layer_buffer = {}

        layer_buffer[name] = tensor_cpu

    # 刷新最后一层
    if layer_buffer:
        yield from self._process_layer_group(...)

流式处理的优势：不需要将所有参数同时加载到内存，适合大模型。每次只处理一个 decoder layer 的参数。

5. 层级量化处理¶

def _process_layer_group(self, layer_idx, layer_params, input_global_scales, output_device):
    # 1. 将权重移到 GPU
    weight_gpu = tensor.to(device=self.device, dtype=self.param_dtype)
    # 2. 计算 global scale
    amax = torch.amax(torch.abs(weight_gpu))
    layer_global_scales[layer_name] = generate_gparam(...)
    # 3. 融合 global scale
    fused_global_scales = fuse_global_scales(layer_global_scales)
    # 4. 计算 blockwise scale
    weight_scale = compute_blockwise_scale(weight_gpu, fused_global_scale)
    # 5. 压缩打包
    weight_packed = self._compressor.compress_weight(weight=weight_gpu, scale=weight_scale, ...)

    results.append((f"{layer_name}.weight_packed", weight_packed))
    results.append((f"{layer_name}.weight_scale", weight_scale))
    results.append((f"{layer_name}.weight_global_scale", fused_global_scale))

核心类/函数列表¶

类/函数名	作用
`QATQuantizer`	主量化器类
`compute_blockwise_scale`	计算分块 scale
`fuse_global_scales`	融合相关层的 global scale
`quantize_with_fusion`	流式量化（generator）
`_process_layer_group`	处理一层的量化

与其他模块的关系¶

在训练完成后使用，将训练好的权重转换为量化格式
使用 compressed_tensors 库的 NVFP4PackedCompressor
产出的量化权重可以被 vLLM 加载（通过 vllm_patch.py 的补丁）

小结¶

QATQuantizer 是训练到推理的桥梁。它将 QAT 训练后的全精度权重转换为 NVFP4 打包格式，同时计算两级 scale（global + blockwise）并执行 QKV/GateUp 融合。流式处理设计使其能够处理超大模型。

quantizer.py — 这个文件实现了 QATQuantizer -- 用于将 QAT 训练后的 FP16/BF16...¶