已开启
feat: graph-based PTQ/QAT 支持 A8W4 量化 #269
feat: graph-based PTQ/QAT 支持 A8W4 量化 #269
已开启
QQR创建于 15 天前
QQR成员
15 天前

描述

支持 graph-based PTQ/QAT A8W4 量化,主要改动如下:

  • PTQ/QAT 接口模式支持 activation INT8 per-tensor、weight INT4 per-tensor/per-channel 配置。
  • QAT 单算子仅为 Conv2dQAT、LinearQAT 开放 INT4 weight;A16W4 和其他 QAT 单算子 INT4 配置显式报错,已有 A8W8/A16W8 行为保持兼容。
  • 用户调用 torch.onnx.export(..., opset_version=21) 导出 QAT 单算子标准 Q/DQ,AMCT 校验版本要求:INT4 weight 使用 output_dtype=INT4 且省略 zero-point;Conv2d per-channel 使用 axis=1,Linear 通过互逆 Transpose 使用 axis=1。
  • PTQ/QAT 接口模式按最终 deploy pack 轴检查权重 shape;Gemm transB=1 检查轴 0,其他目标算子检查最后一轴。
  • pack 轴为奇数时整层跳过 A8W4,不降级为 A8W8;QAT 单算子构造时对 Conv2d W、Linear out_features 奇数配置显式报错。
  • deploy 权重使用原生 ONNX INT4、保持原始逻辑 shape,并沿最后一轴将两个有符号 INT4 按低 nibble 在前写入 raw_data
  • ONNX 不提供原生 INT4 枚举时,非 A8W4 路径保持可用,A8W4 路径给出明确能力错误。

关联的Issue

https://gitcode.com/annqr/amct_open/issues/2

如何测试

运行相关单元测试,检查导出、权重 pack 和 QAT 转换行为。

执行命令:

python3 -m pytest tests/amct_pytorch/testcase_python/configuration/test_check.py tests/amct_pytorch/testcase_python/configuration/test_config_base_checks.py tests/amct_pytorch/testcase_python/configuration/test_retrain_config_base_unit.py tests/amct_pytorch/testcase_python/custom_op/test_arq_retrain_symbolic.py tests/amct_pytorch/testcase_python/custom_op/test_ulq_scale_retrain_symbolic.py tests/amct_pytorch/testcase_python/nn/test_qat_op.py tests/amct_pytorch/testcase_python/optimizer/test_pack_int4_weight_pass.py -q

文档更新

无文档更新。

类型标签

likedislike
合并受阻
QQQR成员
15 天前 创建了 pull request,commit 6b8faa84
QQQR成员
15 天前 关联了issue:[Requirement|需求建议]: 支持cv网络int8*int4数据格式量化
atomgit-bot
atomgit-bot
15 天前 评论:

变更摘要

本 PR 为 graph-based PTQ/QAT 增加 A8W4(INT4 权重 + INT8 激活)量化支持,目标算子为 conv2d/linear。核心改动包括:将 arq_retrain/ulq_scale_retrain 的 ONNX 导出收敛到新增的 add_qdq 助手,以原生 ONNX TensorProto.INT4(要求 opset 21)导出 4bit 权重;INT4 可 pack 性校验由原来的 Cin 轴改为最终 Deploy pack 轴(末轴,GemmtransB=1 时为 axis 0);原先无法 pack 的层由「降级 INT8(A8W8)」改为「跳过该层、保持浮点」;Conv2dQAT/LinearQAT 开放 INT4 支持(LinearQAT 同时放开 channel_wise 限制),并在 qat_base.py 增加激活仅支持 per-tensor、INT4 权重须配 INT8 激活等约束;PackInt4WeightPass 改为直接写出原生 ONNX INT4,删除基于 numpy 的 nibble 手动打包逻辑。

主要改动

  • 新增统一 QDQ 导出助手 add_qdq(新文件 qdq_symbolic.py:将 arq_retrain.py/ulq_scale_retrain.py 中直接生成 QUANTIZE_LINEAR/DEQUANTIZE_LINEAR 的分支收敛为调用 add_qdq;4bit 路径经 _check_int4_export 校验 opset 21 与原生 TensorProto.INT4,per-channel 时校验 wts_scales 数量等于 out_channelsLinear 在 channel_wise 下用两次 Transpose(perm=[1,0]) 将量化轴落于 axis 1。
  • INT4 pack 轴校验从 Cin 轴改为最终 Deploy pack 轴check.pycheck_int4_cin_pack_supported 重构为 check_int4_weight_pack_supported(原方法保留为兼容别名),不再因 grouped conv 或 Cin 为奇数拒绝,而是检查权重张量末轴(GemmtransB=1 时为 axis 0)长度是否为偶数,并在奇数时输出告警日志。
  • 无法 pack 的 INT4 层由降级 INT8 改为跳过、保持浮点config_base.pycheck_int4_weight_quant_axisretrain_config_base.pyskip_int4_unpackable_layers(原 downgrade_int4_unpackable_layers,保留兼容包装)不再将 num_bits/dst_type 从 4 改回 8,而是置 quant_enable/retrain_enable 为 False 并将该层从 supported_layersPARAM_POOL 量化层列表剔除。
  • QAT 模块开放 INT4 并增加配置约束Conv2dQAT/LinearQAT 新增 _supported_weight_dst_types = (INT8, INT4)LinearQAT.check_quantifiable 移除原先「不支持 channel_wise」的 RuntimeErrorqat_base.py 新增激活 channel_wise 报错、INT4 权重须搭配 INT8 激活的校验,且 Conv2dQAT/LinearQAT 分别在权重末轴(W)或 out_features 为奇数时抛出 ValueError
  • PackInt4WeightPass 改为写出原生 ONNX INT4:删除基于 numpy 的 pack_along_axis 两两 nibble 打包逻辑,pack_int4_weight_node 仅接收权重节点,校验末轴为偶数后通过 TensorProtoHelper.set_dataINT4 类型、保持原始 shape 写回权重(RNN 的 recurrence_weight 同样处理)。
likedislike
不准确?
atomgit-bot
atomgit-bot
15 天前 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CANN-robotCANN-robot成员
15 天前 添加了label:cann-cla/yes
此处折叠了195条消息 查看更多
CANN-robotCANN-robot成员
6 小时前 添加了label:ci-pipeline-running
CANN-robotCANN-robot成员
6 小时前 删除了label:ci-pipeline-running
CANN-robotCANN-robot成员
6 小时前 添加了label:ci-pipeline-passed
fujun19成员
5 小时前 评论:

/lgtm

likedislike
QQR成员
3 小时前 评论:
likedislike