feat: graph-based PTQ/QAT 支持 A8W4 量化-amct-AtomGit

已开启

feat: graph-based PTQ/QAT 支持 A8W4 量化 #269

QQR创建于 11 天前从annqr/amct_open： feature/a8w4-graph-ptq-qat 合入到 cann/amct： master

feat: graph-based PTQ/QAT 支持 A8W4 量化 #269

已开启

QQR创建于 11 天前

从annqr/amct_open feature/a8w4-graph-ptq-qat合入到cann/amct master

QQR成员

11 天前

描述

支持 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 行为保持兼容。
QAT symbolic 按 ONNX opset 21 导出标准 Q/DQ：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 路径给出明确能力错误。

Diff 摘要：

.../amct_pytorch/configuration/check.py            |  51 ++++----
.../custom_op/arq_retrain/arq_retrain.py           |  39 ++++---
.../amct_pytorch/custom_op/qdq_symbolic.py         |  61 ++++++++++
.../ulq_scale_retrain/ulq_scale_retrain.py         |  39 ++++---
.../amct_pytorch/nn/module/quantization/conv2d.py  |  12 ++
.../amct_pytorch/nn/module/quantization/linear.py  |  16 ++-
.../nn/module/quantization/qat_base.py             |  22 +++-
.../optimizer/pack_int4_weight_pass.py             |  55 +++------
.../graph_based/common/config/config_base.py       |  38 +++---
.../common/retrain_config/retrain_config_base.py   |  27 +++--
.../testcase_python/configuration/test_check.py    | 116 +++++++++++++-----
.../configuration/test_config_base_checks.py       |  55 +++++++++
.../configuration/test_retrain_config_base_unit.py |  70 ++++++++---
.../custom_op/test_arq_retrain_symbolic.py         | 125 ++++++++++++++++++--
.../custom_op/test_ulq_scale_retrain_symbolic.py   | 103 ++++++++++++++--
.../amct_pytorch/testcase_python/nn/test_qat_op.py | 129 ++++++++++++++++++++-
.../optimizer/test_pack_int4_weight_pass.py        |  85 ++++++--------
17 files changed, 813 insertions(+), 230 deletions(-)

关联的Issue

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

如何测试

增量 codecheck：PASS
git diff --check：PASS
相关 UT：196 passed, 4 subtests passed

执行命令：

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

文档更新

无文档更新。

类型标签

Bug修复
新特性
性能优化
文档更新
代码重构
其他，请描述：

合并受阻

Q QQR成员

11 天前创建了 pull request，commit 6b8faa84

Q QQR成员

11 天前关联了issue：[Requirement|需求建议]: 支持cv网络int8*int4数据格式量化

atomgit-bot

11 天前评论：

变更摘要

本 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 轴（末轴，Gemm 且 transB=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_channels，Linear 在 channel_wise 下用两次 Transpose(perm=[1,0]) 将量化轴落于 axis 1。
INT4 pack 轴校验从 Cin 轴改为最终 Deploy pack 轴：check.py 中 check_int4_cin_pack_supported 重构为 check_int4_weight_pack_supported（原方法保留为兼容别名），不再因 grouped conv 或 Cin 为奇数拒绝，而是检查权重张量末轴（Gemm 且 transB=1 时为 axis 0）长度是否为偶数，并在奇数时输出告警日志。
无法 pack 的 INT4 层由降级 INT8 改为跳过、保持浮点：config_base.py 的 check_int4_weight_quant_axis 与 retrain_config_base.py 的 skip_int4_unpackable_layers（原 downgrade_int4_unpackable_layers，保留兼容包装）不再将 num_bits/dst_type 从 4 改回 8，而是置 quant_enable/retrain_enable 为 False 并将该层从 supported_layers 与 PARAM_POOL 量化层列表剔除。
QAT 模块开放 INT4 并增加配置约束：Conv2dQAT/LinearQAT 新增 _supported_weight_dst_types = (INT8, INT4)，LinearQAT.check_quantifiable 移除原先「不支持 channel_wise」的 RuntimeError；qat_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_data 以 INT4 类型、保持原始 shape 写回权重（RNN 的 recurrence_weight 同样处理）。

atomgit-bot

11 天前评论：

代码审查

✅ 未发现问题

CANN-robot成员

11 天前添加了label：cann-cla/yes

此处折叠了178条消息查看更多

yaoguangxiu成员18 小时前进行代码检视1

amct_pytorch/classic/graph_based/amct_pytorch/configuration/check.py

358356 dims = QuantOpInfo.get_node_tensor(wnode).dims

359- if cin_axis < len(dims) and dims[cin_axis] % 2 == 1:

357+ if dims and dims[pack_axis] % 2 == 1:

358+ actual_axis = pack_axis if pack_axis >= 0 else len(dims) + pack_axis

yaoguangxiu18 小时前评论：

【一般】这里不需要处理负数，日志可以直接打印负数

yaoguangxiu成员18 小时前进行代码检视1

amct_pytorch/classic/graph_based/amct_pytorch/configuration/check.py

357+ if dims and dims[pack_axis] % 2 == 1:

358+ actual_axis = pack_axis if pack_axis >= 0 else len(dims) + pack_axis

359+ LOGGER.logw(

360+ "Skip A8W4 layer '{}': ONNX weight shape {} has odd Deploy "

yaoguangxiu18 小时前评论：

【一般】这里不一定是A8W4，不需要描述activation信息

yaoguangxiu成员18 小时前进行代码检视1

amct_pytorch/classic/graph_based/amct_pytorch/configuration/check.py

337- """

338- # layer_name 预期能取到 node，取不到属于异常，交由 get_node_by_name 抛出

328+ def is_int4_weight_pack_axis_even(graph, layer_name):

329+ """Return whether all weights have an even final Deploy pack axis."""

yaoguangxiu18 小时前评论：

【建议】补充函数功能描述，尤其是轴判断的信息

yaoguangxiu成员17 小时前进行代码检视1

amct_pytorch/classic/graph_based/amct_pytorch/custom_op/arq_retrain/arq_retrain.py

52- Function: ArqRetrain foward funtion.

57+ Function: ArqRetrain forward function.

5358 """

59+ if is_dynamo_export():

yaoguangxiu17 小时前评论：

AMCT内部没有效用dynamo=True的选项，也没有暴露接口给用户，这里的处理是不会被触发的？

yaoguangxiu成员17 小时前进行代码检视1

amct_pytorch/classic/graph_based/amct_pytorch/custom_op/qdq_symbolic.py

45+def check_int4_export(wts_param):

46+ if not hasattr(TensorProto, 'INT4'):

47+ raise RuntimeError('INT4 export requires ONNX native TensorProto.INT4.')

48+ if _globals.GLOBALS.export_onnx_opset_version != 21:

yaoguangxiu17 小时前评论：

torch.onnx.export 是否使用 opset 21 是AMCT自己控制的吧？
这里只是针对QAT单算子场景吗？