AMCT FakeQuant 模拟伪量化工具包MXFP4 推理精度验证与 QAT 训练实践【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct本文以 AMCT昇腾 AI 处理器亲和的模型压缩工具中的experimental/fakequant模块为主体讲解当 NPU 尚不原生支持某类低比特格式时如何在软件侧用“伪量化”复现 MXFP4Microscaling FP4的量化数值行为既包括基于 Ascend-C 自定义 kernel 的推理侧精度验证算子也包括基于 STE直通估计器的 MXFP4 量化感知训练QAT替换层。读完本文读者可以独立完成 Ascend-C 算子的编译与正确性验证并掌握在自有训练框架中接入 MXFP4 QAT Linear 的完整流程与调参要点。1. 工具包定位为什么需要模拟伪量化在低比特格式落地过程中常出现这样的时间差算法侧已经需要评估 MXFP4 这类新格式的精度上限但目标格式的 NPU 算子栈尚未合入或当前环境无法启用对应的低比特计算单元。AMCT 的 FakeQuant 工具包README正是为此设计的在硬件尚不原生支持某些量化格式时在软件侧模拟对应格式的量化精度以便进行精度验证与算法评估。需要明确它的适用边界原文档即有说明本模块属于试验特性experimental接口与实现可能随硬件能力演进而调整伪量化结果用于精度对齐与方案验证不等同于目标硬件上真实低比特算子的性能表现MXFP4 算子实现参考 amct_ops/hifloat8_cast 的三层结构op_kernel/op_extension/python但因处于试验阶段暂不迁入amct_ops/。工具包由两个职责互补的子模块组成fakequant/ ├── mxfp4_ascendc/ # MXFP4 Ascend-C 伪量化算子目录对齐 amct_ops/hifloat8_cast │ ├── op_kernel/ # device kernel tiling │ ├── op_extension/# Torch host TORCH_LIBRARY 注册 │ ├── python/mxfp4/# Python 包装 │ ├── reference/ # 纯 PyTorch 参考实现 │ └── tests/ ├── mxfp4_qat/ # MXFP4 量化感知训练QAT带 STE 的 nn.Linear 替代实现 │ ├── fake_quant.py# MXFP4 QDQ STE autograd.Function 量化器模块 │ └── linear.py # MXFP4QATLinear convert_to_mxfp4_qat ├── README.md └── README_en.md其中mxfp4_ascendc/面向推理侧精度验证快速跑出伪量化数值mxfp4_qat/面向训练侧让模型在训练中适应 MXFP4 误差后者在 NPU 上会自动复用前者的算子加速。二者共享同一套 MXFP4 数值定义保证“推理时看到的误差”和“训练时适应的误差”完全一致。2. MXFP4 数值模型E2M1 码本 E8M0 逐块缩放理解整个工具包的前提是 MXFP4 的数值模型逐元素 FP4 E2M1 尾数 逐 block 的 E8M02 的幂共享缩放。具体规则为最后一维每block_size默认 32个相邻元素共享一个 scalescale 2^round(log2(max_abs / scale_factor))scale_factor默认6.0即把 block 内最大值映射到码本最大码值 6 附近元素码本为{0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}乘以 scale。这三条规则在仓库中有三份一一对应的实现可以交叉印证1纯 PyTorch 参考实现reference/mxfp4_ref.pyL42-L61先用F.pad对最后一维补齐到 block 整数倍注意是按最后一维 padding 而非 flatten 后 padding否则相邻行会错误地共享同一个 MXFP4 block再按abs().amax求块内最大值经clamp(min2**-30)防下溢后用exp2(round(log2(...)))取 2 的幂作为 scale元素舍入则通过一组torch.where(y_abs 中点, 码值, 原值)级联完成——中点序列(0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00)恰好是 E2M1 相邻码值的中点。2QAT 侧的向量化实现mxfp4_qat/fake_quant.pyL115-L149_block_scale()与参考实现逐步对应元素舍入改用torch.bucketize(y_abs, midpoints, rightTrue)一步索引到最近码值L132-L139并缓存了按(device, dtype)组织的码本张量文件头部常量BLOCK_SIZE 32、SCALE_FACTOR 6.0、MXFP4_E2M1_MAX 6.0L60-L62定义了默认数值配置最小块尺度_MIN_SCALE_RAW 2**-30明确注释为“与mxfp4_tiling.h中的MXFP4_MIN_SCALE_RAW保持同步”。3Ascend-C kernel 的宿主/设备共享常量op_kernel/mxfp4_tiling.hL22-L27constexpr int32_t MXFP4_BLOCK_SIZE 32; constexpr float MXFP4_SCALE_FACTOR 6.0f; constexpr float MXFP4_INV_SCALE_FACTOR 1.0f / 6.0f; constexpr float MXFP4_MIN_SCALE_RAW 9.313225746e-10f; // 2^-30 constexpr int32_t MXFP4_BLOCKS_PER_TILE 208; constexpr int32_t MXFP4_TILE_ELEMS MXFP4_BLOCK_SIZE * MXFP4_BLOCKS_PER_TILE;kernel 内部以1/6.0为内置倒数因子运行时再乘以一个可配置的invScaleMulBitstiling 结构中保存的是 IEEE-754 位模式这一设计正是 Python 包装层inv_scale_factor_scale参数的底层来源见第 4 节。3. mxfp4_ascendcNPU 侧 Ascend-C 伪量化算子3.1 运行环境子目录 README 给出的环境基线组件版本硬件Ascend 910B3 / 兼容 SoCCANN8.2.RC1Python3.10 (aarch64)PyTorch2.6.0torch_npu2.6.0.post4开源仓不附带预编译.so需要按本机 SoC / CANN / Python ABI 自行编译。构建入口是 mxfp4_ascendc/CMakeLists.txt其中定义了三个关键 CMake 缓存变量L25-L28SOC_VERSION默认Ascend910B3可用SOC_VERSIONAscend910_9392 bash build.sh形式覆盖RUN_MODEnpu / cpu / simASCEND_CANN_PACKAGE_PATH默认/usr/local/Ascend/ascend-toolkit/latest。从源码结构看CMake 构建分两段先经 CANN 的ascendc_kernel_cmake/ascendc.cmake用 ccec 编译 device kernelop_kernel/mxfp4_kernel.cpp为ascendc_kernels_${RUN_MODE}共享库再编译 Torch 扩展libmxfp4_ops.soL125-L149后者链接 Torch、kernel 库与ascendcl并显式探测torch._C._GLIBCXX_USE_CXX11_ABI以保持 ABI 一致——这就是“必须按本机 ABI 自行编译”的原因。3.2 三层结构与 PyTorch 算子注册目录布局对齐amct_ops/hifloat8_castmxfp4_ascendc/ ├── op_kernel/ │ ├── mxfp4_kernel.cpp # Ascend-C device kernel │ └── mxfp4_tiling.h # Host/device 共用 tiling 常量与结构体 ├── op_extension/ │ ├── mxfp4_torch.cpp # PyTorch hosttiling ACLRT_LAUNCH_KERNEL │ ├── ops.h # C host 接口声明namespace AscendKernel │ └── register.cpp # TORCH_LIBRARY_FRAGMENT(amct, ...) Meta ├── python/ │ └── mxfp4/ │ ├── __init__.py # 加载 .so、自检、re-export │ └── ops.py # 薄 Python 包装pad / dtype ├── reference/ │ └── mxfp4_ref.py # 纯 PyTorch 参考实现 ├── CMakeLists.txt # 构建入口 ├── tests/ # 正确性 / inv_scale / benchmark └── README.md算子注册在 op_extension/register.cppL23-L45中完成值得注意的细节是注册方式TORCH_LIBRARY_FRAGMENT(amct, m) { m.def(quant_dequant_mxfp4(Tensor x, float inv_scale_factor_scale1.0) - Tensor); }使用TORCH_LIBRARY_FRAGMENT而非TORCH_LIBRARY允许与 AMCT 其他扩展在同一命名空间amct下增量注册通过TORCH_LIBRARY_IMPL(amct, PrivateUse1, m)提供 NPU 设备实现委托给AscendKernel::Mxfp4QuantDequantTorch即 host 侧完成 tiling 并ACLRT_LAUNCH_KERNEL启动通过TORCH_LIBRARY_IMPL(amct, Meta, m)提供仅返回 shape 的 Meta 实现at::empty(x.sizes(), x.options())使该算子可安全穿过符号追踪/dynamo 类路径的 shape 推导。因此底层原始调用形式为torch.ops.amct.quant_dequant_mxfp4(x_flat, 1.0)约定输入是float32、连续contiguous且 numel 为 32 的倍数的 flat 张量AIV 核数由 host 侧PlatformAscendC::GetCoreNumAiv()运行时查询无需手动指定。3.3 Python 包装层的 padding 约定python/mxfp4/ops.pyL25-L87是面向用户的正式入口mxfp4.quant_dequant_mxfp4(x, block_size32, inv_scale_factor_scale1.0)它做了四件“脏活”校验block_size 32NPU kernel 硬约束与输入必须位于 NPU 设备转为 float32 后先在最后一维 padding到 block 整数倍再 reshape 成 flat——源码注释明确指出若先 flatten 再 pad当最后一维不是 32 的倍数时相邻行会合并进同一个 MXFP4 block与参考实现不一致调用torch.ops.amct.quant_dequant_mxfp4(x_flat, float(inv_scale_factor_scale))恢复原 shape、切掉 padding、还原输入 dtype。其中inv_scale_factor_scale与 QAT 侧的scale_factor参数是倒数关系kernel 内置1/6.0包装层用它乘上运行时乘子故scale_factor SCALE_FACTOR / inv_scale_factor_scale。这一换算在 fake_quant.pyL230-L236中完成用户无论走哪条路径面对的语义参数都是scale_factor。3.4 编译、正确性验证与性能按子目录 README 的指引编译并测试编译成功后.so会被 stage 到python/mxfp4/cd /path/to/mxfp4_ascendc # 编译约数分钟成功后自动将 .so 拷到 python/mxfp4/ bash build.sh # 测试正确性 性能 python tests/test_mxfp4.py # python tests/test_inv_scale.py # inv_scale 参数正确性 # python tests/bench_qdq.py # 额外性能对比正确性测试 tests/test_mxfp4.pyL35-L72将 Ascend-C 输出与mxfp4_ref纯 PyTorch 参考实现逐元素比对atol1e-6测试形状刻意覆盖了最后一维不是 32 倍数的场景——(2, 33)、(3, 17)、(4, 48, 33)——专门验证“按行 block”而非“flatten 后 pad”的语义文档口径为与 PyTorch 参考实现 bit-exact 一致。性能对比同一文档口径torch_npu 软件路径 vs Ascend-C 自定义 kernelShapetorch_npuAscend-C加速比(64, 4096)0.69 ms0.038 ms18.1x(256, 4096)0.72 ms0.059 ms12.3x(1024, 4096)0.72 ms0.219 ms3.28x即小矩阵约 18x、大矩阵约 3.3x 加速。该收益的解释与 QAT 侧的训练建议直接相关torch_npu 软件路径每次 QDQ 要发射十几个 elementwise kernelkernel 启动开销在小 shape 下占比极高这正是自定义单 kernel 方案收益最大的区间。3.5 快速使用import sys sys.path.insert(0, /path/to/mxfp4_ascendc/python) from mxfp4 import quant_dequant_mxfp4 x_npu x.npu() result quant_dequant_mxfp4(x_npu) # 等价底层调用输入需已是 float32 flatnumel 为 32 的倍数 # result torch.ops.amct.quant_dequant_mxfp4(x_flat, 1.0)API 签名quant_dequant_mxfp4( x: torch.Tensor, # 任意 shape建议 float32在 NPU 上 block_size: int 32, # 量化 block 宽度必须为 32 inv_scale_factor_scale: float 1.0, ) - torch.Tensor # 同 shape / dtype / device4. mxfp4_qatMXFP4 量化感知训练4.1 STE 与 clipped STE量化是分段常量函数几乎处处导数为 0无法直接反传。STEstraight-through estimator直通估计器的做法是前向用量化值、反向把量化算子当作恒等映射从而让梯度穿过量化点抵达高精度权重forward : y Q(x) backward: dL/dx dL/dy # 普通 STE dL/dx dL/dy * (|x| 6*scale) # clipped STEclip_gradTrueclip_gradTrue时把**发生截断saturation**的元素梯度置零。其动机来自 MXFP4 的数值特性scale 被强制取整到 2 的幂round(log2(...))可能向下取整block 内最大值约有半数概率超出6*scale而被截断这些位置的梯度方向具有误导性屏蔽后训练通常更稳。源码中该机制对应 fake_quant.py 的三处实现mxfp4_saturation_mask()L239-L254返回|x| 6 * block_scale的 bool 掩码_MXFP4FakeQuantSTEL257-L272显式torch.autograd.Functionforward中按需保存截断掩码backward中grad_output.masked_fill(saturated, 0.0)后原样透传普通路径mxfp4_fake_quant()L275-L296可微的Tensor - Tensor函数。子目录 README 还说明了实现取向它等价于 MindSpeed-LLM 的x (x_q - x).detach()写法但改用显式autograd.Function以便在backward里做梯度屏蔽。该模块只依赖torchfake_quant.py的 import 仅torch与torch.nn.functional可以单文件拷进任意训练仓使用。4.2 MXFP4QATConfig 全参数from mxfp4_qat import MXFP4QATConfig, MXFP4QATLinear, convert_to_mxfp4_qat字段默认说明quantize_weightTrue是否伪量化权重置False则退化为纯激活实验quantize_inputFalse是否伪量化层输入。False→ W4A16推荐起点True→ W4A4block_size32共享 scale 的元素数Ascend-C 算子只支持 32scale_factor6.0增大 → scale 变小inlier 分辨率更高但截断更多减小则相反clip_gradFalseTrue使用 clipped STEbackendautoautoNPU 上自动用 Ascend-C 算子否则纯 PyTorch/torch/npudataclass 的__post_init__会对block_size、scale_factor的正数约束与backend取值做校验fake_quant.py L335-L373。4.3 MXFP4QATLinear无状态的 nn.Linear 替换linear.py 中的MXFP4QATLinear是nn.Linear的子类L33-L100设计上有两个关键性质1主权重保持高精度前向量化的是“一次性副本”。forward为F.linear(Q(x), Q(W), b)bias 不量化每次前向把当前高精度权重过一遍量化器得到伪量化值参与计算optimizer 更新的始终是高精度“master weight”loss 反映的却是 MXFP4 数值。2state_dict与 float 层完全一致。由于量化器无参数、无 bufferfloat 权重可以直接 load 进 QAT 模型QAT 训练完的权重也可以 load 回 float 模型或交给 AMCT 的 deploy 流程导出真实低比特权重。构造方式layer MXFP4QATLinear(in_features, out_features, configMXFP4QATConfig()) layer MXFP4QATLinear.from_linear(existing_linear, config) # 复用原 Parameter不额外占显存from_linearL71-L92在meta设备上创建壳层后直接采用adopt原nn.Linear的weight/biasParameter 对象而非拷贝因此不占额外显存已引用这些 Parameter 的 optimizer 状态与 parameter group 也不受影响。4.4 convert_to_mxfp4_qat原地批量替换convert_to_mxfp4_qat(module, configNone, skip_names())linear.py L103-L145原地递归替换模块树下所有nn.Linear含子类因此 Megatron 风格的 Linear 子类也会被转换而非遗留 float。要点skip_names按模块点分路径做子串匹配命中则跳过该子树连同其下所有层已转换的MXFP4QATLinear会被识别并跳过因此重复调用是幂等 no-op返回的是同一个 module 对象仅用于链式调用。convert_to_mxfp4_qat( model, MXFP4QATConfig(quantize_inputTrue, clip_gradTrue), skip_names(lm_head, embed_tokens), # 敏感层保持 float )4.5 快速开始import sys sys.path.insert(0, .../amct_pytorch/experimental/fakequant) from mxfp4_qat import MXFP4QATConfig, convert_to_mxfp4_qat model.load_state_dict(torch.load(ckpt)) # 从 float 预训练权重出发 convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_inputTrue)) # 其余训练代码不变4.6 接入自有训练框架的三种方式方式一模型里是标准nn.Linear。建模完成、加载完预训练权重之后optimizer创建之前插入一行即可其余训练代码不用改model build_model() model.load_state_dict(torch.load(ckpt)) # 从 float 预训练权重出发 convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_inputTrue)) optimizer torch.optim.AdamW(model.parameters(), lr2e-5) # ... 正常训练循环 ...from_linear复用原Parameter对象所以在 optimizer 之后转换也不会失效但放在 optimizer 之前更保险避免 parameter group 引用悬空。方式二框架有自定义 LinearMegatronColumnParallelLinear等。无法用继承替换时把MXFP4FakeQuantizernn.Module形态同样无状态挂到层上、在forward里手动调用——这正是 MindSpeed-LLM 的做法from mxfp4_qat import MXFP4FakeQuantizer class FakeQuantColumnParallelLinear(ColumnParallelLinear): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.weight_quantizer MXFP4FakeQuantizer() self.input_quantizer MXFP4FakeQuantizer() def forward(self, input_, weightNone, **kwargs): input_ self.input_quantizer(input_) # 父类 forward 读取 self.weight因此临时替换 .data 后再恢复。 # 前向已在量化值上完成梯度经 STE 正确回传到原始高精度权重。 original self.weight.data self.weight.data self.weight_quantizer(original) try: return super().forward(input_, weightweight, **kwargs) finally: self.weight.data originalMoE 的 GroupedMatmul 专家权重同理在调用 GMM 之前对w1/w2和 permute 后的专家输入各过一次量化器。方式三只想复用量化算子。直接调用mxfp4_fake_quant(x)它是可微的Tensor - Tensor函数可放在任何位置KV cache、logits、residual 等。底层函数全集mxfp4_quant_dequant(x, block_size32, scale_factor6.0, backendauto) # 无梯度与 mxfp4_ascendc 参考实现 bit-exact mxfp4_fake_quant(x, block_size32, scale_factor6.0, clip_gradFalse, backendauto) # 带 STE mxfp4_saturation_mask(x, block_size32, scale_factor6.0) # 截断位置掩码 MXFP4FakeQuantizer(block_size32, scale_factor6.0, clip_gradFalse, backendauto) # nn.Module 形态4.7 训练建议继承原文档从 float 预训练权重出发做 QAT 微调不要从随机初始化开始训先 W4A16 再 W4A4激活量化掉点通常明显大于权重量化先确认quantize_inputFalse能收敛学习率取预训练的 1/10 左右并配 cosine 衰减敏感层保持 floatlm_head、embedding、第一/最后一层通常通过skip_names排除NPU 上务必用 Ascend-C 后端纯 PyTorch 路径每次 QDQ 有十几个 elementwise kernel大模型训练开销不可忽略训练完成后用amct_pytorch的 deploy 流程导出真实低比特权重QAT 只是让权重“适应”MXFP4导出仍需常规量化链路。4.8 backend 解析机制QAT 与 Ascend-C 算子的衔接backendauto时的后端选择在 fake_quant.pyL224-L236中实现仅当张量位于 NPU 设备且Ascend-C kernel 可加载时才走 NPU 路径。kernel 包的查找顺序为环境变量MXFP4_ASCENDC_PATH→ 同级../mxfp4_ascendc/python目录_VENDORED_KERNEL_PATHL81成功与失败都会缓存避免重复探测。算子需先自行编译cd ../mxfp4_ascendc bash build.sh未编译或不在 NPU 上时自动退回纯 PyTorch 路径——两条路径结果 bit-exact 一致仅速度不同显式指定backendnpu而算子不可用时会抛出带修复指引重新编译、设置MXFP4_ASCENDC_PATH或改用backendtorch的RuntimeError。5. 限制与使用前提以下限制均继承自子目录 README使用本工具包前应逐条确认属于试验特性experimental接口可能调整QAT 侧只覆盖nn.Linear卷积、Embedding、Attention 内部的 matmul 未处理scale 与截断阈值均由数据静态推导未实现可学习的 scale / clippingLSQ、PACT 等block_size ! 32只有纯 PyTorch 路径支持Ascend-C kernel 与 Python 包装层均硬约束 32伪量化仅复现 MXFP4 的数值行为不代表目标硬件上真实低比特算子的性能Ascend-C 算子需要 CANN 8.2.RC1、aarch64 环境与对应 SoC且开源仓不附带预编译.so需自行编译。6. 小结AMCT 的experimental/fakequant用一套统一的 MXFP4 数值定义串起了压缩工具链上的两个关键环节mxfp4_ascendc提供与 PyTorch 参考实现 bit-exact 一致、并在小 shape 下最高约 18x 加速的 NPU 侧 QDQ 算子支撑推理精度快速验证mxfp4_qat在其上叠加 STE/clipped-STE 可微伪量化与MXFP4QATLinear以“state_dict 与 float 完全兼容”为设计底线让 QAT 可以低成本插入任意训练框架。由于二者共享同一数值模型E2M1 码本 E8M0 逐块 2 的幂 scale常量在 mxfp4_tiling.h、mxfp4_ref.py、fake_quant.py 三处保持一致训练阶段适应的误差与部署验证时观测的误差天然对齐这正是“模拟伪量化”作为低比特格式落地前验证手段的价值所在。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考