feat: 新增 MXFP4 QAT 伪量化 Linear（带 STE）-amct-AtomGit

Mamct_pytorch/experimental/fakequant/README.md+7-1

	@@ -14,9 +14,15 @@ fakequant/
14	│ ├── python/mxfp4/# Python 包装	14	│ ├── python/mxfp4/# Python 包装
15	│ ├── reference/ # 纯 PyTorch 参考实现	15	│ ├── reference/ # 纯 PyTorch 参考实现
16	│ └── tests/	16	│ └── tests/
17	-└── README.md	17	+├── mxfp4_qat/ # MXFP4 量化感知训练（QAT）：带 STE 的 nn.Linear 替代实现
		18	+│ ├── fake_quant.py# MXFP4 QDQ + STE autograd.Function + 量化器模块
		19	+│ └── linear.py # MXFP4QATLinear + convert_to_mxfp4_qat
		20	+├── README.md
		21	+└── README_en.md
18	```	22	```
19		23
		24	+`mxfp4_ascendc/` 面向推理侧精度验证（快速跑出伪量化数值），`mxfp4_qat/` 面向训练侧（让模型在训练中适应 MXFP4 误差）；后者在 NPU 上会自动复用前者的算子加速。
		25	+
20	## 说明	26	## 说明
21		27
22	- 本模块属于试验特性（`experimental`），接口与实现可能随硬件能力演进而调整。	28	- 本模块属于试验特性（`experimental`），接口与实现可能随硬件能力演进而调整。

Aamct_pytorch/experimental/fakequant/README_en.md+30-0

@@ -0,0 +1,30 @@
	1	+# FakeQuant Simulation Toolkit
	2	+
	3	+This directory provides a fake quantization toolkit. It simulates the numerical behavior of quantization formats that the hardware does not yet natively support, so that accuracy can be verified and algorithms can be evaluated in software.
	4	+
	5	+Typical cases include formats that have not yet landed in the NPU operator stack, or environments where the corresponding low-bit compute units cannot be enabled. The fake-quant path still reproduces an approximate quantized numerical behavior.
	6	+
	7	+## Directory Layout
	8	+
	9	+```
	10	+fakequant/
	11	+├── mxfp4_ascendc/ # MXFP4 Ascend-C fake-quant operator (layout aligned with amct_ops/hifloat8_cast)
	12	+│ ├── op_kernel/ # device kernel + tiling
	13	+│ ├── op_extension/# Torch host + TORCH_LIBRARY registration
	14	+│ ├── python/mxfp4/# Python wrapper
	15	+│ ├── reference/ # pure PyTorch reference implementation
	16	+│ └── tests/
	17	+├── mxfp4_qat/ # MXFP4 quantization-aware training (QAT): STE-based nn.Linear replacement
	18	+│ ├── fake_quant.py# MXFP4 QDQ + STE autograd.Function + quantizer module
	19	+│ └── linear.py # MXFP4QATLinear + convert_to_mxfp4_qat
	20	+├── README.md
	21	+└── README_en.md
	22	+```
	23	+
	24	+`mxfp4_ascendc/` targets inference-side accuracy verification (producing fake-quant values quickly). `mxfp4_qat/` targets training (letting the model adapt to MXFP4 error during training). On NPU, the latter automatically reuses the former's operator for acceleration.
	25	+
	26	+## Notes
	27	+
	28	+- This module is experimental (`experimental`). Interfaces and implementations may change as hardware capability evolves.
	29	+- Fake-quant results are for accuracy alignment and scheme validation. They are not equivalent to the performance of real low-bit operators on the target hardware.
	30	+- The MXFP4 operator follows the three-layer layout of `amct_ops/hifloat8_cast`, but stays under `amct_pytorch/experimental/` during the experimental stage and is not moved into `amct_ops/` yet.

Mamct_pytorch/experimental/fakequant/mxfp4_ascendc/README.md+4-1

	@@ -36,7 +36,8 @@ mxfp4_ascendc/
36	├── CMakeLists.txt # 构建入口	36	├── CMakeLists.txt # 构建入口
37	├── build.sh # 一键编译并 stage .so 到 python/mxfp4/	37	├── build.sh # 一键编译并 stage .so 到 python/mxfp4/
38	├── tests/ # 正确性 / inv_scale / benchmark	38	├── tests/ # 正确性 / inv_scale / benchmark
39	-└── README.md	39	+├── README.md
		40	+└── README_en.md
40	```	41	```
41		42
42	## 从源码编译	43	## 从源码编译
	@@ -86,6 +87,8 @@ quant_dequant_mxfp4(
86		87
87	AIV 核数由 host 侧 `PlatformAscendC::GetCoreNumAiv()` 运行时查询，无需手动指定。	88	AIV 核数由 host 侧 `PlatformAscendC::GetCoreNumAiv()` 运行时查询，无需手动指定。
88		89
		90	+需要在训练中使用（带 STE 的可微伪量化）见 `../mxfp4_qat/`，其 `backend="auto"` 会在 NPU 张量上自动调用本算子。
		91	+
89	### 性能	92	### 性能
90		93
91	\| Shape \| torch_npu \| Ascend-C \| 加速比 \|	94	\| Shape \| torch_npu \| Ascend-C \| 加速比 \|

Aamct_pytorch/experimental/fakequant/mxfp4_ascendc/README_en.md+98-0

@@ -0,0 +1,98 @@
	1	+# MXFP4 Ascend-C Accelerated Operator (experimental)
	2	+
	3	+An Ascend-C custom kernel that implements the MXFP4 (Microscaling FP4 E2M1) fake-quant operator on Ascend NPU. The directory layout is aligned with `amct_ops/hifloat8_cast` (`op_kernel` / `op_extension` / `python`). The implementation remains under `amct_pytorch/experimental/fakequant/` (experimental; not part of `amct_ops`).
	4	+
	5	+Compared with the torch_npu software path: 3.3x speedup (large matrices), 18x speedup (small matrices). Correctness is bit-exact against the PyTorch reference implementation.
	6	+
	7	+## Runtime Environment
	8	+
	9	+\| Component \| Version \|
	10	+\|-----------\|---------\|
	11	+\| Hardware \| Ascend 910B3 / compatible SoC \|
	12	+\| CANN \| 8.2.RC1+ \|
	13	+\| Python \| 3.10 (aarch64) \|
	14	+\| PyTorch \| 2.6.0 \|
	15	+\| torch_npu \| 2.6.0.post4 \|
	16	+
	17	+> The open-source repository does not ship prebuilt `.so` files. Compile locally against your SoC / CANN / Python ABI.
	18	+
	19	+## Directory Structure
	20	+
	21	+```
	22	+mxfp4_ascendc/
	23	+├── op_kernel/
	24	+│ ├── mxfp4_kernel.cpp # Ascend-C device kernel
	25	+│ └── mxfp4_tiling.h # Host/device shared tiling constants and structs
	26	+├── op_extension/
	27	+│ ├── mxfp4_torch.cpp # PyTorch host: tiling + ACLRT_LAUNCH_KERNEL
	28	+│ ├── ops.h # C++ host API declaration (namespace AscendKernel)
	29	+│ └── register.cpp # TORCH_LIBRARY_FRAGMENT(amct, ...) + Meta
	30	+├── python/
	31	+│ └── mxfp4/
	32	+│ ├── __init__.py # load .so, self-check, re-export
	33	+│ └── ops.py # thin Python wrapper (pad / dtype)
	34	+├── reference/
	35	+│ └── mxfp4_ref.py # pure PyTorch reference implementation
	36	+├── CMakeLists.txt # build entry
	37	+├── build.sh # one-shot build and stage .so into python/mxfp4/
	38	+├── tests/ # correctness / inv_scale / benchmark
	39	+├── README.md
	40	+└── README_en.md
	41	+```
	42	+
	43	+## Build from Source
	44	+
	45	+```bash
	46	+cd /path/to/mxfp4_ascendc
	47	+
	48	+# Build (takes a few minutes); on success, .so files are copied to python/mxfp4/
	49	+bash build.sh
	50	+
	51	+# Correctness + performance tests
	52	+python tests/test_mxfp4.py
	53	+# python tests/test_inv_scale.py # inv_scale parameter correctness
	54	+# python tests/bench_qdq.py # extra performance comparison
	55	+```
	56	+
	57	+Specify SoC:
	58	+
	59	+```bash
	60	+SOC_VERSION=Ascend910_9392 bash build.sh
	61	+```
	62	+
	63	+## Quick Start
	64	+
	65	+```python
	66	+import sys
	67	+sys.path.insert(0, "/path/to/mxfp4_ascendc/python")
	68	+
	69	+from mxfp4 import quant_dequant_mxfp4
	70	+
	71	+x_npu = x.npu()
	72	+result = quant_dequant_mxfp4(x_npu)
	73	+
	74	+# Equivalent low-level call (input must already be float32 flat, numel a multiple of 32)
	75	+# result = torch.ops.amct.quant_dequant_mxfp4(x_flat, 1.0)
	76	+```
	77	+
	78	+### API
	79	+
	80	+```python
	81	+quant_dequant_mxfp4(
	82	+ x: torch.Tensor, # any shape, float32 recommended, on NPU
	83	+ block_size: int = 32, # quantization block width (must be 32)
	84	+ inv_scale_factor_scale: float = 1.0,
	85	+) -> torch.Tensor # same shape / dtype / device
	86	+```
	87	+
	88	+AIV core count is queried at runtime on the host via `PlatformAscendC::GetCoreNumAiv()`; you do not need to set it manually.
	89	+
	90	+For training use (differentiable fake quant with STE), see `../mxfp4_qat/`. Its `backend="auto"` automatically calls this operator on NPU tensors.
	91	+
	92	+### Performance
	93	+
	94	+\| Shape \| torch_npu \| Ascend-C \| Speedup \|
	95	+\|-------\|-----------\|----------\|---------\|
	96	+\| (64, 4096) \| 0.69 ms \| 0.038 ms \| 18.1x \|
	97	+\| (256, 4096) \| 0.72 ms \| 0.059 ms \| 12.3x \|
	98	+\| (1024, 4096) \| 0.72 ms \| 0.219 ms \| 3.28x \|

Aamct_pytorch/experimental/fakequant/mxfp4_qat/README.md+180-0

@@ -0,0 +1,180 @@
	1	+# MXFP4 QAT Linear（experimental）
	2	+
	3	+一个最小可用的 MXFP4 量化感知训练（QAT）实现：`MXFP4QATLinear` 可直接替换 `torch.nn.Linear`，前向按 MXFP4 数值伪量化、反向通过 STE（straight-through estimator，直通估计器）把梯度回传到高精度主权重，从而让模型在训练阶段就适应 MXFP4 的量化误差。
	4	+
	5	+本目录只提供训练侧的量化算子与层，不含训练框架。第 5 节说明如何把它接入你自己的框架（Megatron / DeepSpeed / HuggingFace Trainer 等）。
	6	+
	7	+## 1. MXFP4 与 STE
	8	+
	9	+MXFP4 = 逐元素 FP4 E2M1 尾数 + 逐 block 的 E8M0（2 的幂）共享缩放：
	10	+
	11	+- 最后一维每 `block_size`（默认 32）个相邻元素共享一个 scale；
	12	+- `scale = 2^round(log2(max_abs / scale_factor))`，`scale_factor` 默认 `6.0`；
	13	+- 元素码本为 `{0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}`（乘以 scale）。
	14	+
	15	+量化是分段常量函数，几乎处处导数为 0，无法直接反传。STE 的做法是：前向用量化值、反向把量化算子当作恒等映射，因此梯度可以穿过量化点抵达高精度权重：
	16	+
	17	+```
	18	+forward : y = Q(x)
	19	+backward: dL/dx = dL/dy # 普通 STE
	20	+ dL/dx = dL/dy * (\|x\| <= 6*scale) # clipped STE（clip_grad=True）
	21	+```
	22	+
	23	+`clip_grad=True` 时会把发生截断（saturation）的元素梯度置零。因为 scale 被强制取整到 2 的幂（可能向下取整），block 内最大值有约一半概率超出 `6*scale` 而被截断，这些位置的梯度方向具有误导性，屏蔽后训练通常更稳。
	24	+
	25	+> 实现上等价于 MindSpeed-LLM 的 `x + (x_q - x).detach()` 写法，本目录改用显式的
	26	+> `torch.autograd.Function`，便于在 `backward` 里做梯度屏蔽，也更直观。
	27	+
	28	+## 2. 目录结构
	29	+
	30	+```
	31	+mxfp4_qat/
	32	+├── fake_quant.py # MXFP4 QDQ、STE autograd.Function、量化器模块、配置
	33	+├── linear.py # MXFP4QATLinear + convert_to_mxfp4_qat
	34	+├── README.md
	35	+└── README_en.md
	36	+```
	37	+
	38	+`fake_quant.py` 只依赖 `torch`，可以单文件拷进任意训练仓使用。
	39	+
	40	+## 3. 快速开始
	41	+
	42	+```python
	43	+import sys
	44	+sys.path.insert(0, ".../amct_pytorch/experimental/fakequant")
	45	+
	46	+from mxfp4_qat import MXFP4QATConfig, convert_to_mxfp4_qat
	47	+
	48	+model.load_state_dict(torch.load(ckpt)) # 从 float 预训练权重出发
	49	+convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_input=True))
	50	+# 其余训练代码不变
	51	+```
	52	+
	53	+## 4. API
	54	+
	55	+```python
	56	+import sys
	57	+sys.path.insert(0, ".../amct_pytorch/experimental/fakequant")
	58	+
	59	+from mxfp4_qat import MXFP4QATConfig, MXFP4QATLinear, convert_to_mxfp4_qat
	60	+```
	61	+
	62	+### `MXFP4QATConfig`
	63	+
	64	+\| 字段 \| 默认 \| 说明 \|
	65	+\|------\|------\|------\|
	66	+\| `quantize_weight` \| `True` \| 是否伪量化权重 \|
	67	+\| `quantize_input` \| `False` \| 是否伪量化层输入。`False` → W4A16（推荐起点），`True` → W4A4 \|
	68	+\| `block_size` \| `32` \| 共享 scale 的元素数；Ascend-C 算子只支持 32 \|
	69	+\| `scale_factor` \| `6.0` \| 增大 → scale 变小，inlier 分辨率更高但截断更多；减小则相反 \|
	70	+\| `clip_grad` \| `False` \| `True` 使用 clipped STE \|
	71	+\| `backend` \| `"auto"` \| `"auto"`（NPU 上自动用 Ascend-C 算子，否则纯 PyTorch）/ `"torch"` / `"npu"` \|
	72	+
	73	+### `MXFP4QATLinear`
	74	+
	75	+`nn.Linear` 的子类，前向为 `F.linear(Q(x), Q(W), b)`（bias 不量化）。
	76	+
	77	+由于量化器无参数、无 buffer，`state_dict` 与 float 层完全一致：float 权重可以直接 load 进 QAT 模型，QAT 训练完的权重也可以 load 回 float 模型或交给 AMCT 的 deploy 流程导出。
	78	+
	79	+```python
	80	+layer = MXFP4QATLinear(in_features, out_features, config=MXFP4QATConfig())
	81	+layer = MXFP4QATLinear.from_linear(existing_linear, config) # 复用原 Parameter，不额外占显存
	82	+```
	83	+
	84	+### `convert_to_mxfp4_qat(module, config=None, skip_names=())`
	85	+
	86	+原地递归替换 `module` 下所有 `nn.Linear`。`skip_names` 按模块点分路径做子串匹配，命中则跳过该子树：
	87	+
	88	+```python
	89	+convert_to_mxfp4_qat(
	90	+ model,
	91	+ MXFP4QATConfig(quantize_input=True, clip_grad=True),
	92	+ skip_names=("lm_head", "embed_tokens"), # 敏感层保持 float
	93	+)
	94	+```
	95	+
	96	+### 底层函数
	97	+
	98	+```python
	99	+mxfp4_quant_dequant(x, block_size=32, scale_factor=6.0, backend="auto") # 无梯度，与 mxfp4_ascendc 参考实现 bit-exact
	100	+mxfp4_fake_quant(x, block_size=32, scale_factor=6.0, clip_grad=False, backend="auto") # 带 STE
	101	+mxfp4_saturation_mask(x, block_size=32, scale_factor=6.0) # 截断位置掩码
	102	+MXFP4FakeQuantizer(block_size=32, scale_factor=6.0, clip_grad=False, backend="auto") # nn.Module 形态
	103	+```
	104	+
	105	+## 5. 接入自己的训练框架
	106	+
	107	+### 方式一：模型里是标准 `nn.Linear`
	108	+
	109	+建模完成、加载完预训练权重之后，`optimizer` 创建之前插入一行即可，其余训练代码不用改：
	110	+
	111	+```python
	112	+model = build_model()
	113	+model.load_state_dict(torch.load(ckpt)) # 从 float 预训练权重出发
	114	+
	115	+convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_input=True))
	116	+
	117	+optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
	118	+# ... 正常训练循环 ...
	119	+```
	120	+
	121	+`from_linear` 复用原 `Parameter` 对象，所以在 optimizer 之后转换也不会失效；但放在 optimizer 之前更保险（避免 parameter group 引用悬空）。
	122	+
	123	+### 方式二：框架有自定义 Linear（Megatron `ColumnParallelLinear` 等）
	124	+
	125	+无法用继承替换时，把 `MXFP4FakeQuantizer` 作为量化器挂到层上，在 `forward` 里手动调用即可 —— 这正是 MindSpeed-LLM 的做法：
	126	+
	127	+```python
	128	+from mxfp4_qat import MXFP4FakeQuantizer
	129	+
	130	+
	131	+class FakeQuantColumnParallelLinear(ColumnParallelLinear):
	132	+ def __init__(self, args, *kwargs):
	133	+ super().__init__(args, *kwargs)
	134	+ self.weight_quantizer = MXFP4FakeQuantizer()
	135	+ self.input_quantizer = MXFP4FakeQuantizer()
	136	+
	137	+ def forward(self, input_, weight=None, **kwargs):
	138	+ input_ = self.input_quantizer(input_)
	139	+ # 父类 forward 读取 self.weight，因此临时替换 .data 后再恢复。
	140	+ # 前向已在量化值上完成，梯度经 STE 正确回传到原始高精度权重。
	141	+ original = self.weight.data
	142	+ self.weight.data = self.weight_quantizer(original)
	143	+ try:
	144	+ return super().forward(input_, weight=weight, **kwargs)
	145	+ finally:
	146	+ self.weight.data = original
	147	+```
	148	+
	149	+MoE 的 GroupedMatmul 专家权重同理：在调用 GMM 之前对 `w1` / `w2` 和 permute 后的专家输入各过一次量化器。
	150	+
	151	+### 方式三：只想复用量化算子
	152	+
	153	+直接调用 `mxfp4_fake_quant(x)`，它就是一个可微的 `Tensor -> Tensor` 函数，放在任何位置（KV cache、logits、residual 等）都可以。
	154	+
	155	+### 训练建议
	156	+
	157	+- 从 float 预训练权重出发做 QAT 微调，不要从随机初始化开始训。
	158	+- 先 W4A16 再 W4A4：激活量化掉点通常明显大于权重量化，先确认 `quantize_input=False` 能收敛。
	159	+- 学习率取预训练的 1/10 左右并配 cosine 衰减。
	160	+- 敏感层保持 float：`lm_head`、embedding、第一/最后一层通常通过 `skip_names` 排除。
	161	+- NPU 上务必用 Ascend-C 后端：纯 PyTorch 路径每次 QDQ 有十几个 elementwise kernel，大模型训练开销不可忽略；`backend="auto"` 会在 NPU 张量上自动切换。
	162	+- 训练完成后用 `amct_pytorch` 的 deploy 流程导出真实低比特权重；QAT 只是让权重"适应"MXFP4，导出仍需常规量化链路。
	163	+
	164	+### Ascend-C 算子的定位
	165	+
	166	+`backend="auto"` 时按以下顺序查找 `mxfp4` 包：环境变量 `MXFP4_ASCENDC_PATH` → 同级 `../mxfp4_ascendc/python`。算子需先自行编译：
	167	+
	168	+```bash
	169	+cd ../mxfp4_ascendc && bash build.sh
	170	+```
	171	+
	172	+未编译或不在 NPU 上时自动退回纯 PyTorch 路径（结果 bit-exact 一致，仅速度不同）；显式指定 `backend="npu"` 而算子不可用时会抛出带修复指引的 `RuntimeError`。
	173	+
	174	+## 6. 限制
	175	+
	176	+- 属于试验特性（`experimental`），接口可能调整。
	177	+- 只覆盖 `nn.Linear`；卷积、Embedding、Attention 内部的 matmul 未处理。
	178	+- scale 与截断阈值均由数据静态推导，未实现可学习的 scale / clipping（LSQ、PACT 等）。
	179	+- `block_size != 32` 只有纯 PyTorch 路径支持。
	180	+- 伪量化仅复现 MXFP4 的数值行为，不代表目标硬件上真实低比特算子的性能。

Aamct_pytorch/experimental/fakequant/mxfp4_qat/README_en.md+180-0

@@ -0,0 +1,180 @@
	1	+# MXFP4 QAT Linear (experimental)
	2	+
	3	+A minimal usable MXFP4 quantization-aware training (QAT) implementation: `MXFP4QATLinear` can replace `torch.nn.Linear` directly. The forward pass fake-quantizes values with MXFP4 numerics; the backward pass uses STE (straight-through estimator) to send gradients back to the high-precision master weights, so the model adapts to MXFP4 quantization error during training.
	4	+
	5	+This directory only provides training-side quantization operators and layers. It does not include a training framework. Section 5 describes how to plug them into your own stack (Megatron / DeepSpeed / HuggingFace Trainer, etc.).
	6	+
	7	+## 1. MXFP4 and STE
	8	+
	9	+MXFP4 = per-element FP4 E2M1 mantissa + per-block E8M0 (power-of-two) shared scale:
	10	+
	11	+- Every `block_size` (default 32) adjacent elements along the last dimension share one scale;
	12	+- `scale = 2^round(log2(max_abs / scale_factor))`, with `scale_factor` defaulting to `6.0`;
	13	+- The element codebook is `{0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}` (multiplied by scale).
	14	+
	15	+Quantization is a piecewise-constant function whose derivative is zero almost everywhere, so it cannot back-propagate directly. STE uses the quantized value in the forward pass and treats the quantizer as identity in the backward pass, so gradients can pass through the quantization point to the high-precision weights:
	16	+
	17	+```
	18	+forward : y = Q(x)
	19	+backward: dL/dx = dL/dy # plain STE
	20	+ dL/dx = dL/dy * (\|x\| <= 6*scale) # clipped STE (clip_grad=True)
	21	+```
	22	+
	23	+When `clip_grad=True`, gradients of saturated (clipped) elements are zeroed. Because scale is forced onto a power of two (and may round down), the block maximum has about a 50% chance of exceeding `6*scale` and being clipped. Gradients at those positions are misleading; masking them usually makes training more stable.
	24	+
	25	+> This is numerically equivalent to MindSpeed-LLM's `x + (x_q - x).detach()` form. This directory uses an explicit
	26	+> `torch.autograd.Function` instead, which makes gradient masking in `backward` straightforward and easier to follow.
	27	+
	28	+## 2. Directory Structure
	29	+
	30	+```
	31	+mxfp4_qat/
	32	+├── fake_quant.py # MXFP4 QDQ, STE autograd.Function, quantizer module, config
	33	+├── linear.py # MXFP4QATLinear + convert_to_mxfp4_qat
	34	+├── README.md
	35	+└── README_en.md
	36	+```
	37	+
	38	+`fake_quant.py` depends only on `torch` and can be copied as a single file into any training repository.
	39	+
	40	+## 3. Quick Start
	41	+
	42	+```python
	43	+import sys
	44	+sys.path.insert(0, ".../amct_pytorch/experimental/fakequant")
	45	+
	46	+from mxfp4_qat import MXFP4QATConfig, convert_to_mxfp4_qat
	47	+
	48	+model.load_state_dict(torch.load(ckpt)) # start from float pretrained weights
	49	+convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_input=True))
	50	+# remaining training code is unchanged
	51	+```
	52	+
	53	+## 4. API
	54	+
	55	+```python
	56	+import sys
	57	+sys.path.insert(0, ".../amct_pytorch/experimental/fakequant")
	58	+
	59	+from mxfp4_qat import MXFP4QATConfig, MXFP4QATLinear, convert_to_mxfp4_qat
	60	+```
	61	+
	62	+### `MXFP4QATConfig`
	63	+
	64	+\| Field \| Default \| Description \|
	65	+\|-------\|---------\|-------------\|
	66	+\| `quantize_weight` \| `True` \| Whether to fake-quantize weights \|
	67	+\| `quantize_input` \| `False` \| Whether to fake-quantize layer inputs. `False` → W4A16 (recommended starting point), `True` → W4A4 \|
	68	+\| `block_size` \| `32` \| Number of elements sharing a scale; the Ascend-C operator only supports 32 \|
	69	+\| `scale_factor` \| `6.0` \| Larger → smaller scale, higher inlier resolution but more clipping; smaller is the opposite \|
	70	+\| `clip_grad` \| `False` \| `True` uses clipped STE \|
	71	+\| `backend` \| `"auto"` \| `"auto"` (use the Ascend-C operator on NPU tensors, otherwise pure PyTorch) / `"torch"` / `"npu"` \|
	72	+
	73	+### `MXFP4QATLinear`
	74	+
	75	+A subclass of `nn.Linear`. Forward is `F.linear(Q(x), Q(W), b)` (bias is not quantized).
	76	+
	77	+Because the quantizer has no parameters and no buffers, the `state_dict` matches a float layer exactly: float weights can be loaded into a QAT model, and QAT-trained weights can be loaded back into a float model or handed to AMCT's deploy flow for export.
	78	+
	79	+```python
	80	+layer = MXFP4QATLinear(in_features, out_features, config=MXFP4QATConfig())
	81	+layer = MXFP4QATLinear.from_linear(existing_linear, config) # reuse the original Parameter, no extra GPU memory
	82	+```
	83	+
	84	+### `convert_to_mxfp4_qat(module, config=None, skip_names=())`
	85	+
	86	+Recursively replace all `nn.Linear` modules under `module` in place. `skip_names` does substring matching on dotted module paths; a hit skips that subtree:
	87	+
	88	+```python
	89	+convert_to_mxfp4_qat(
	90	+ model,
	91	+ MXFP4QATConfig(quantize_input=True, clip_grad=True),
	92	+ skip_names=("lm_head", "embed_tokens"), # keep sensitive layers in float
	93	+)
	94	+```
	95	+
	96	+### Low-level functions
	97	+
	98	+```python
	99	+mxfp4_quant_dequant(x, block_size=32, scale_factor=6.0, backend="auto") # no grad, bit-exact with mxfp4_ascendc reference
	100	+mxfp4_fake_quant(x, block_size=32, scale_factor=6.0, clip_grad=False, backend="auto") # with STE
	101	+mxfp4_saturation_mask(x, block_size=32, scale_factor=6.0) # saturation-position mask
	102	+MXFP4FakeQuantizer(block_size=32, scale_factor=6.0, clip_grad=False, backend="auto") # nn.Module form
	103	+```
	104	+
	105	+## 5. Integrating with Your Training Framework
	106	+
	107	+### Option 1: The model uses standard `nn.Linear`
	108	+
	109	+After the model is built and pretrained weights are loaded, insert one line before creating the `optimizer`. The rest of the training code stays unchanged:
	110	+
	111	+```python
	112	+model = build_model()
	113	+model.load_state_dict(torch.load(ckpt)) # start from float pretrained weights
	114	+
	115	+convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_input=True))
	116	+
	117	+optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
	118	+# ... normal training loop ...
	119	+```
	120	+
	121	+`from_linear` reuses the original `Parameter` objects, so converting after the optimizer is created still works; converting before the optimizer is safer (avoids dangling parameter-group references).
	122	+
	123	+### Option 2: The framework has a custom Linear (e.g. Megatron `ColumnParallelLinear`)
	124	+
	125	+When inheritance-based replacement is not possible, attach `MXFP4FakeQuantizer` as a quantizer on the layer and call it manually in `forward` — this is how MindSpeed-LLM does it:
	126	+
	127	+```python
	128	+from mxfp4_qat import MXFP4FakeQuantizer
	129	+
	130	+
	131	+class FakeQuantColumnParallelLinear(ColumnParallelLinear):
	132	+ def __init__(self, args, *kwargs):
	133	+ super().__init__(args, *kwargs)
	134	+ self.weight_quantizer = MXFP4FakeQuantizer()
	135	+ self.input_quantizer = MXFP4FakeQuantizer()
	136	+
	137	+ def forward(self, input_, weight=None, **kwargs):
	138	+ input_ = self.input_quantizer(input_)
	139	+ # The parent forward reads self.weight, so temporarily replace .data then restore.
	140	+ # Forward already ran on quantized values; STE sends gradients back to the original high-precision weights.
	141	+ original = self.weight.data
	142	+ self.weight.data = self.weight_quantizer(original)
	143	+ try:
	144	+ return super().forward(input_, weight=weight, **kwargs)
	145	+ finally:
	146	+ self.weight.data = original
	147	+```
	148	+
	149	+The same pattern applies to MoE GroupedMatmul expert weights: run the quantizer on `w1` / `w2` and on the permuted expert inputs before calling GMM.
	150	+
	151	+### Option 3: Reuse only the quantization operator
	152	+
	153	+Call `mxfp4_fake_quant(x)` directly. It is a differentiable `Tensor -> Tensor` function and can be placed anywhere (KV cache, logits, residual, etc.).
	154	+
	155	+### Training tips
	156	+
	157	+- Start QAT fine-tuning from float pretrained weights; do not train from random initialization.
	158	+- W4A16 first, then W4A4: activation quantization usually drops accuracy more than weight quantization. Confirm `quantize_input=False` converges first.
	159	+- Use about 1/10 of the pretraining learning rate with cosine decay.
	160	+- Keep sensitive layers in float: `lm_head`, embeddings, and the first/last layers are usually excluded via `skip_names`.
	161	+- On NPU, use the Ascend-C backend: the pure PyTorch path launches more than ten elementwise kernels per QDQ, which is not negligible for large-model training. `backend="auto"` switches automatically on NPU tensors.
	162	+- After training, export real low-bit weights with the `amct_pytorch` deploy flow. QAT only makes weights "adapt" to MXFP4; export still needs the regular quantization pipeline.
	163	+
	164	+### Role of the Ascend-C operator
	165	+
	166	+When `backend="auto"`, the `mxfp4` package is looked up in this order: environment variable `MXFP4_ASCENDC_PATH` → sibling `../mxfp4_ascendc/python`. Compile the operator first:
	167	+
	168	+```bash
	169	+cd ../mxfp4_ascendc && bash build.sh
	170	+```
	171	+
	172	+If it is not compiled or the tensor is not on NPU, the code falls back to the pure PyTorch path (bit-exact results, different speed only). Explicitly setting `backend="npu"` when the operator is unavailable raises a `RuntimeError` with fix instructions.
	173	+
	174	+## 6. Limitations
	175	+
	176	+- This is experimental (`experimental`); interfaces may change.
	177	+- Only `nn.Linear` is covered; convolution, Embedding, and matmul inside Attention are not handled.
	178	+- Scale and clip thresholds are derived statically from data. Learnable scale / clipping (LSQ, PACT, etc.) is not implemented.
	179	+- `block_size != 32` is supported only on the pure PyTorch path.
	180	+- Fake quant only reproduces MXFP4 numerical behavior. It does not represent the performance of real low-bit operators on the target hardware.

Aamct_pytorch/experimental/fakequant/mxfp4_qat/__init__.py+60-0

@@ -0,0 +1,60 @@
	1	+# -- coding: UTF-8 --
	2	+# ----------------------------------------------------------------------------
	3	+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+# ----------------------------------------------------------------------------
	17	+"""Minimal MXFP4 quantisation-aware training (QAT) building blocks.
	18	+
	19	+Usage:
	20	+ import sys
	21	+ sys.path.insert(0, ".../amct_pytorch/experimental/fakequant")
	22	+
	23	+ from mxfp4_qat import MXFP4QATConfig, convert_to_mxfp4_qat
	24	+
	25	+ convert_to_mxfp4_qat(model, MXFP4QATConfig(quantize_input=True))
	26	+ # ... train as usual; nothing else changes ...
	27	+"""
	28	+
	29	+from __future__ import annotations
	30	+
	31	+__all__ = [
	32	+ "BACKEND_AUTO",
	33	+ "BACKEND_NPU",
	34	+ "BACKEND_TORCH",
	35	+ "BLOCK_SIZE",
	36	+ "MXFP4_E2M1_MAX",
	37	+ "MXFP4FakeQuantizer",
	38	+ "MXFP4QATConfig",
	39	+ "MXFP4QATLinear",
	40	+ "SCALE_FACTOR",
	41	+ "convert_to_mxfp4_qat",
	42	+ "mxfp4_fake_quant",
	43	+ "mxfp4_quant_dequant",
	44	+ "mxfp4_saturation_mask",
	45	+]
	46	+
	47	+from .fake_quant import (
	48	+ BACKEND_AUTO,
	49	+ BACKEND_NPU,
	50	+ BACKEND_TORCH,
	51	+ BLOCK_SIZE,
	52	+ MXFP4_E2M1_MAX,
	53	+ SCALE_FACTOR,
	54	+ MXFP4FakeQuantizer,
	55	+ MXFP4QATConfig,
	56	+ mxfp4_fake_quant,
	57	+ mxfp4_quant_dequant,
	58	+ mxfp4_saturation_mask,
	59	+)
	60	+from .linear import MXFP4QATLinear, convert_to_mxfp4_qat

Aamct_pytorch/experimental/fakequant/mxfp4_qat/fake_quant.py+373-0

@@ -0,0 +1,373 @@
	1	+# -- coding: UTF-8 --
	2	+# ----------------------------------------------------------------------------
	3	+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+# ----------------------------------------------------------------------------
	17	+"""MXFP4 fake quantisation with a straight-through estimator (STE).
	18	+
	19	+MXFP4 = per-element FP4 E2M1 mantissa plus a per-block E8M0 (power-of-two)
	20	+shared scale, with ``block_size`` consecutive elements of the last dimension
	21	+sharing one scale.
	22	+
	23	+Quantisation is not differentiable (it is piecewise constant, so its true
	24	+derivative is zero almost everywhere). QAT therefore uses the straight-through
	25	+estimator: the forward pass sees quantised values while the backward pass
	26	+treats the quantiser as an identity, letting the gradient reach the underlying
	27	+high-precision weights and activations.
	28	+
	29	+This module is intentionally self-contained (torch only) so it can be copied
	30	+into a third-party training framework as-is. The Ascend-C kernel under
	31	+``../mxfp4_ascendc`` is used automatically when it is available and the tensor
	32	+lives on an NPU; otherwise a pure-PyTorch path runs on any device.
	33	+"""
	34	+
	35	+from __future__ import annotations
	36	+
	37	+__all__ = [
	38	+ "BACKEND_AUTO",
	39	+ "BACKEND_NPU",
	40	+ "BACKEND_TORCH",
	41	+ "BLOCK_SIZE",
	42	+ "MXFP4_E2M1_MAX",
	43	+ "MXFP4FakeQuantizer",
	44	+ "MXFP4QATConfig",
	45	+ "SCALE_FACTOR",
	46	+ "mxfp4_fake_quant",
	47	+ "mxfp4_quant_dequant",
	48	+ "mxfp4_saturation_mask",
	49	+]
	50	+
	51	+import dataclasses
	52	+import os
	53	+import sys
	54	+from typing import Callable
	55	+
	56	+import torch
	57	+import torch.nn.functional as F
	58	+from torch import nn
	59	+
	60	+BLOCK_SIZE = 32
	61	+SCALE_FACTOR = 6.0
	62	+MXFP4_E2M1_MAX = 6.0
	63	+
	64	+# Smallest representable block scale, kept in sync with MXFP4_MIN_SCALE_RAW in
	65	+# ../mxfp4_ascendc/op_kernel/mxfp4_tiling.h.
	66	+_MIN_SCALE_RAW = 2.0**-30
	67	+
	68	+# Positive FP4 E2M1 codebook and the midpoints between its adjacent entries.
	69	+_E2M1_CODES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
	70	+_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00)
	71	+_codebook_cache: dict[tuple[torch.device, torch.dtype], tuple[torch.Tensor, ...]] = {}
	72	+
	73	+BACKEND_AUTO = "auto"
	74	+BACKEND_TORCH = "torch"
	75	+BACKEND_NPU = "npu"
	76	+_BACKENDS = (BACKEND_AUTO, BACKEND_TORCH, BACKEND_NPU)
	77	+
	78	+_PKG_DIR = os.path.dirname(os.path.abspath(__file__))
	79	+# Ascend-C kernel shipped next to this package; overridable via env var so a
	80	+# framework can point at an out-of-tree build.
	81	+_VENDORED_KERNEL_PATH = os.path.join(_PKG_DIR, os.pardir, "mxfp4_ascendc", "python")
	82	+
	83	+_npu_qdq: Callable[..., torch.Tensor] \| None = None
	84	+_npu_load_error: str \| None = None
	85	+
	86	+
	87	+def _validate(block_size: int, scale_factor: float) -> None:
	88	+ if block_size <= 0:
	89	+ raise ValueError(f"block_size must be positive, got {block_size}")
	90	+ if scale_factor <= 0:
	91	+ raise ValueError(f"scale_factor must be positive, got {scale_factor}")
	92	+
	93	+
	94	+def _block_view(x: torch.Tensor, block_size: int) -> tuple[torch.Tensor, int]:
	95	+ """Reshape x to ``(..., n_block, block_size)``, zero-padding the last dim.
	96	+
	97	+ Padding is applied on the last dimension rather than on a flattened view so
	98	+ that adjacent rows never end up sharing an MXFP4 block.
	99	+ """
	100	+ if x.ndim == 0:
	101	+ raise ValueError("x must have at least 1 dimension")
	102	+ last_dim = x.shape[-1]
	103	+ pad = (block_size - last_dim % block_size) % block_size
	104	+ x_fp = x.to(torch.float32)
	105	+ if pad:
	106	+ x_fp = F.pad(x_fp, (0, pad))
	107	+ return x_fp.reshape(*x_fp.shape[:-1], -1, block_size), pad
	108	+
	109	+
	110	+def _unpad(blocked: torch.Tensor, last_dim: int, pad: int) -> torch.Tensor:
	111	+ flat = blocked.reshape(*blocked.shape[:-2], -1)
	112	+ return flat[..., :last_dim] if pad else flat
	113	+
	114	+
	115	+def _block_scale(x_blocked: torch.Tensor, scale_factor: float) -> torch.Tensor:
	116	+ """Per-block E8M0 scale: the power of two nearest to ``max_abs / scale_factor``."""
	117	+ max_abs = x_blocked.abs().amax(dim=-1, keepdim=True)
	118	+ raw_scale = torch.clamp(max_abs / scale_factor, min=_MIN_SCALE_RAW)
	119	+ return torch.exp2(torch.round(torch.log2(raw_scale)))
	120	+
	121	+
	122	+def _codebook(like: torch.Tensor) -> tuple[torch.Tensor, ...]:
	123	+ key = (like.device, like.dtype)
	124	+ if key not in _codebook_cache:
	125	+ _codebook_cache[key] = tuple(
	126	+ torch.tensor(values, device=like.device, dtype=like.dtype)
	127	+ for values in (_E2M1_MIDPOINTS, _E2M1_CODES)
	128	+ )
	129	+ return _codebook_cache[key]
	130	+
	131	+
	132	+def _round_to_e2m1(y_abs: torch.Tensor) -> torch.Tensor:
	133	+ """Round each magnitude to the nearest FP4 E2M1 code.
	134	+
	135	+ Since the thresholds are the codebook midpoints, the number of thresholds a
	136	+ value exceeds is exactly the index of its nearest code.
	137	+ """
	138	+ midpoints, codes = _codebook(y_abs)
	139	+ return codes[torch.bucketize(y_abs, midpoints, right=True)]
	140	+
	141	+
	142	+def _quant_dequant_torch(
	143	+ x: torch.Tensor, block_size: int, scale_factor: float
	144	+) -> torch.Tensor:
	145	+ x_blocked, pad = _block_view(x, block_size)
	146	+ scale = _block_scale(x_blocked, scale_factor)
	147	+ y = x_blocked / scale
	148	+ q = torch.sign(y) * _round_to_e2m1(y.abs())
	149	+ return _unpad(q * scale, x.shape[-1], pad).to(x.dtype)
	150	+
	151	+
	152	+def _load_npu_qdq() -> Callable[..., torch.Tensor]:
	153	+ """Import the Ascend-C MXFP4 kernel wrapper, caching success and failure."""
	154	+ global _npu_qdq, _npu_load_error
	155	+
	156	+ if _npu_qdq is not None:
	157	+ return _npu_qdq
	158	+ if _npu_load_error is not None:
	159	+ raise RuntimeError(_npu_load_error)
	160	+
	161	+ search_path = os.environ.get("MXFP4_ASCENDC_PATH") or os.path.realpath(
	162	+ _VENDORED_KERNEL_PATH
	163	+ )
	164	+ if search_path not in sys.path:
	165	+ sys.path.insert(0, search_path)
	166	+
	167	+ try:
	168	+ import mxfp4
	169	+
	170	+ _npu_qdq = mxfp4.quant_dequant_mxfp4
	171	+ except Exception as e:
	172	+ _npu_load_error = (
	173	+ f"cannot load the Ascend-C MXFP4 kernel from '{search_path}': {e}. "
	174	+ "Build it with `bash build.sh` in mxfp4_ascendc/, or point "
	175	+ "MXFP4_ASCENDC_PATH at a directory containing the built `mxfp4` "
	176	+ f"package, or use backend='{BACKEND_TORCH}'."
	177	+ )
	178	+ raise RuntimeError(_npu_load_error) from e
	179	+
	180	+ return _npu_qdq
	181	+
	182	+
	183	+def _npu_qdq_available() -> bool:
	184	+ try:
	185	+ _load_npu_qdq()
	186	+ except RuntimeError:
	187	+ return False
	188	+ return True
	189	+
	190	+
	191	+def mxfp4_quant_dequant(
	192	+ x: torch.Tensor,
	193	+ block_size: int = BLOCK_SIZE,
	194	+ scale_factor: float = SCALE_FACTOR,
	195	+ backend: str = BACKEND_AUTO,
	196	+) -> torch.Tensor:
	197	+ """Quantise x to MXFP4 and dequantise it back, without any autograd hook.
	198	+
	199	+ Bit-exactly matches ``mxfp4_ascendc/reference/mxfp4_ref.quant_dequant_mxfp4``.
	200	+ Use :func:`mxfp4_fake_quant` instead when training.
	201	+
	202	+ Args:
	203	+ x: Tensor with at least one dimension, any dtype/device.
	204	+ block_size: Elements per shared scale. The NPU kernel only supports 32.
	205	+ scale_factor: Divisor applied to each block maximum before rounding the
	206	+ scale to a power of two. ``6.0`` maps the block maximum onto the
	207	+ largest E2M1 code.
	208	+ backend: ``"auto"`` picks the Ascend-C kernel for NPU tensors when it is
	209	+ available and falls back to PyTorch; ``"torch"`` / ``"npu"`` force
	210	+ one path.
	211	+
	212	+ Returns:
	213	+ Tensor with the same shape, dtype and device as x.
	214	+
	215	+ Raises:
	216	+ ValueError: On a non-positive ``block_size`` / ``scale_factor``, a 0-dim
	217	+ input, or an unknown ``backend``.
	218	+ RuntimeError: If ``backend="npu"`` but the kernel cannot be loaded.
	219	+ """
	220	+ _validate(block_size, scale_factor)
	221	+ if backend not in _BACKENDS:
	222	+ raise ValueError(f"backend must be one of {_BACKENDS}, got {backend!r}")
	223	+
	224	+ use_npu = backend == BACKEND_NPU or (
	225	+ backend == BACKEND_AUTO and x.device.type == "npu" and _npu_qdq_available()
	226	+ )
	227	+ if not use_npu:
	228	+ return _quant_dequant_torch(x, block_size, scale_factor)
	229	+
	230	+ # The kernel hard-codes 1 / 6.0 and multiplies it by inv_scale_factor_scale,
	231	+ # so scale_factor == SCALE_FACTOR / inv_scale_factor_scale.
	232	+ return _load_npu_qdq()(
	233	+ x,
	234	+ block_size=block_size,
	235	+ inv_scale_factor_scale=SCALE_FACTOR / scale_factor,
	236	+ )
	237	+
	238	+
	239	+def mxfp4_saturation_mask(
	240	+ x: torch.Tensor,
	241	+ block_size: int = BLOCK_SIZE,
	242	+ scale_factor: float = SCALE_FACTOR,
	243	+) -> torch.Tensor:
	244	+ """Return a bool mask that is ``True`` where MXFP4 quantisation clips x.
	245	+
	246	+ An element saturates when its magnitude exceeds ``6 * block_scale``, which
	247	+ happens because the block scale is rounded to a power of two and may round
	248	+ down. Gradients of clipped elements are misleading under a plain STE, so
	249	+ :class:`MXFP4QATConfig` can mask them out (clipped STE).
	250	+ """
	251	+ _validate(block_size, scale_factor)
	252	+ x_blocked, pad = _block_view(x, block_size)
	253	+ limit = MXFP4_E2M1_MAX * _block_scale(x_blocked, scale_factor)
	254	+ return _unpad(x_blocked.abs() > limit, x.shape[-1], pad)
	255	+
	256	+
	257	+class _MXFP4FakeQuantSTE(torch.autograd.Function):
	258	+ """Fake-quantise in forward; pass the gradient through unchanged in backward."""
	259	+
	260	+ @staticmethod
	261	+ def forward(ctx, x, block_size, scale_factor, clip_grad, backend):
	262	+ ctx.clip_grad = clip_grad
	263	+ if clip_grad:
	264	+ ctx.save_for_backward(mxfp4_saturation_mask(x, block_size, scale_factor))
	265	+ return mxfp4_quant_dequant(x, block_size, scale_factor, backend)
	266	+
	267	+ @staticmethod
	268	+ def backward(ctx, grad_output):
	269	+ if ctx.clip_grad:
	270	+ (saturated,) = ctx.saved_tensors
	271	+ grad_output = grad_output.masked_fill(saturated, 0.0)
	272	+ return grad_output, None, None, None, None
	273	+
	274	+
	275	+def mxfp4_fake_quant(
	276	+ x: torch.Tensor,
	277	+ block_size: int = BLOCK_SIZE,
	278	+ scale_factor: float = SCALE_FACTOR,
	279	+ clip_grad: bool = False,
	280	+ backend: str = BACKEND_AUTO,
	281	+) -> torch.Tensor:
	282	+ """Differentiable MXFP4 fake quantisation (STE).
	283	+
	284	+ Args:
	285	+ x: Tensor to fake-quantise, typically a weight or an activation.
	286	+ block_size: Elements per shared scale.
	287	+ scale_factor: See :func:`mxfp4_quant_dequant`.
	288	+ clip_grad: If ``True``, zero the gradient of elements that saturated
	289	+ (clipped STE) instead of passing everything through.
	290	+ backend: See :func:`mxfp4_quant_dequant`.
	291	+
	292	+ Returns:
	293	+ Fake-quantised tensor, same shape/dtype/device as x, whose gradient
	294	+ flows back to x unchanged (or masked when ``clip_grad``).
	295	+ """
	296	+ return _MXFP4FakeQuantSTE.apply(x, block_size, scale_factor, clip_grad, backend)
	297	+
	298	+
	299	+class MXFP4FakeQuantizer(nn.Module):
	300	+ """Stateless module wrapper around :func:`mxfp4_fake_quant`.
	301	+
	302	+ Holds no parameters or buffers, so inserting it into a model leaves the
	303	+ ``state_dict`` untouched and checkpoints stay interchangeable with the
	304	+ float model.
	305	+ """
	306	+
	307	+ def __init__(
	308	+ self,
	309	+ block_size: int = BLOCK_SIZE,
	310	+ scale_factor: float = SCALE_FACTOR,
	311	+ clip_grad: bool = False,
	312	+ backend: str = BACKEND_AUTO,
	313	+ ) -> None:
	314	+ super().__init__()
	315	+ _validate(block_size, scale_factor)
	316	+ if backend not in _BACKENDS:
	317	+ raise ValueError(f"backend must be one of {_BACKENDS}, got {backend!r}")
	318	+ self.block_size = block_size
	319	+ self.scale_factor = scale_factor
	320	+ self.clip_grad = clip_grad
	321	+ self.backend = backend
	322	+
	323	+ def forward(self, x: torch.Tensor) -> torch.Tensor:
	324	+ return mxfp4_fake_quant(
	325	+ x, self.block_size, self.scale_factor, self.clip_grad, self.backend
	326	+ )
	327	+
	328	+ def extra_repr(self) -> str:
	329	+ return (
	330	+ f"block_size={self.block_size}, scale_factor={self.scale_factor}, "
	331	+ f"clip_grad={self.clip_grad}, backend={self.backend!r}"
	332	+ )
	333	+
	334	+
	335	+@dataclasses.dataclass
	336	+class MXFP4QATConfig:
	337	+ """MXFP4 QAT settings shared by every converted layer.
	338	+
	339	+ Attributes:
	340	+ quantize_weight: Fake-quantise the weight. Disabling it turns the layer
	341	+ into an activation-only experiment.
	342	+ quantize_input: Fake-quantise the layer input as well. ``False`` gives
	343	+ W4A16 (the recommended starting point); ``True`` gives W4A4.
	344	+ block_size: Elements per shared scale. The Ascend-C kernel needs 32.
	345	+ scale_factor: See :func:`mxfp4_quant_dequant`. Raising it shrinks the
	346	+ block scale, resolving inliers more finely but clipping outliers
	347	+ harder; lowering it does the opposite.
	348	+ clip_grad: Use a clipped STE instead of a plain one.
	349	+ backend: ``"auto"`` / ``"torch"`` / ``"npu"``.
	350	+ """
	351	+
	352	+ quantize_weight: bool = True
	353	+ quantize_input: bool = False
	354	+ block_size: int = BLOCK_SIZE
	355	+ scale_factor: float = SCALE_FACTOR
	356	+ clip_grad: bool = False
	357	+ backend: str = BACKEND_AUTO
	358	+
	359	+ def __post_init__(self) -> None:
	360	+ _validate(self.block_size, self.scale_factor)
	361	+ if self.backend not in _BACKENDS:
	362	+ raise ValueError(
	363	+ f"backend must be one of {_BACKENDS}, got {self.backend!r}"
	364	+ )
	365	+
	366	+ def make_quantizer(self) -> MXFP4FakeQuantizer:
	367	+ """Build a fresh quantizer module carrying this configuration."""
	368	+ return MXFP4FakeQuantizer(
	369	+ block_size=self.block_size,
	370	+ scale_factor=self.scale_factor,
	371	+ clip_grad=self.clip_grad,
	372	+ backend=self.backend,
	373	+ )

Aamct_pytorch/experimental/fakequant/mxfp4_qat/linear.py+145-0

@@ -0,0 +1,145 @@
	1	+# -- coding: UTF-8 --
	2	+# ----------------------------------------------------------------------------
	3	+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+# ----------------------------------------------------------------------------
	17	+"""A drop-in ``torch.nn.Linear`` replacement performing MXFP4 QAT."""
	18	+
	19	+from __future__ import annotations
	20	+
	21	+__all__ = [
	22	+ "MXFP4QATLinear",
	23	+ "convert_to_mxfp4_qat",
	24	+]
	25	+
	26	+import torch
	27	+import torch.nn.functional as F
	28	+from torch import nn
	29	+
	30	+from .fake_quant import MXFP4QATConfig
	31	+
	32	+
	33	+class MXFP4QATLinear(nn.Linear):
	34	+ """``nn.Linear`` whose weight (and optionally input) is MXFP4 fake-quantised.
	35	+
	36	+ Master weights stay in high precision and are what the optimiser updates;
	37	+ every forward pass quantises a throwaway copy so the loss reflects MXFP4
	38	+ numerics, and the STE routes the gradient back to the master weight.
	39	+
	40	+ Because it subclasses ``nn.Linear`` and the quantizers are stateless, the
	41	+ ``state_dict`` is identical to the float layer's, so float checkpoints load
	42	+ into a QAT model and QAT checkpoints load back into the float model.
	43	+
	44	+ Args:
	45	+ in_features: Size of each input sample.
	46	+ out_features: Size of each output sample.
	47	+ bias: Whether to learn an additive bias. The bias is never quantised.
	48	+ device: Device of the created parameters.
	49	+ dtype: Dtype of the created parameters.
	50	+ config: MXFP4 QAT settings; defaults to ``MXFP4QATConfig()`` (W4A16).
	51	+ """
	52	+
	53	+ def __init__(
	54	+ self,
	55	+ in_features: int,
	56	+ out_features: int,
	57	+ bias: bool = True,
	58	+ device: torch.device \| str \| None = None,
	59	+ dtype: torch.dtype \| None = None,
	60	+ config: MXFP4QATConfig \| None = None,
	61	+ ) -> None:
	62	+ super().__init__(in_features, out_features, bias, device=device, dtype=dtype)
	63	+ self.config = config if config is not None else MXFP4QATConfig()
	64	+ self.weight_quantizer = (
	65	+ self.config.make_quantizer() if self.config.quantize_weight else None
	66	+ )
	67	+ self.input_quantizer = (
	68	+ self.config.make_quantizer() if self.config.quantize_input else None
	69	+ )
	70	+
	71	+ @classmethod
	72	+ def from_linear(
	73	+ cls, linear: nn.Linear, config: MXFP4QATConfig \| None = None
	74	+ ) -> MXFP4QATLinear:
	75	+ """Wrap an existing ``nn.Linear``, reusing its parameter objects.
	76	+
	77	+ The weight and bias tensors are adopted rather than copied, so the
	78	+ conversion costs no extra memory and any optimiser state or parameter
	79	+ group already referencing them stays valid.
	80	+ """
	81	+ qat_linear = cls(
	82	+ linear.in_features,
	83	+ linear.out_features,
	84	+ bias=linear.bias is not None,
	85	+ device="meta",
	86	+ dtype=linear.weight.dtype,
	87	+ config=config,
	88	+ )
	89	+ qat_linear.weight = linear.weight
	90	+ if linear.bias is not None:
	91	+ qat_linear.bias = linear.bias
	92	+ return qat_linear
	93	+
	94	+ def forward(self, input: torch.Tensor) -> torch.Tensor:
	95	+ if self.input_quantizer is not None:
	96	+ input = self.input_quantizer(input)
	97	+ weight = self.weight
	98	+ if self.weight_quantizer is not None:
	99	+ weight = self.weight_quantizer(weight)
	100	+ return F.linear(input, weight, self.bias)
	101	+
	102	+
	103	+def convert_to_mxfp4_qat(
	104	+ module: nn.Module,
	105	+ config: MXFP4QATConfig \| None = None,
	106	+ skip_names: tuple[str, ...] = (),
	107	+) -> nn.Module:
	108	+ """Recursively replace every ``nn.Linear`` in module with an MXFP4 QAT one.
	109	+
	110	+ The replacement happens in place; the returned reference is the same object,
	111	+ provided only for chaining. Subclasses of ``nn.Linear`` are included so
	112	+ framework wrappers (Megatron-style Linear subclasses, etc.) are converted
	113	+ rather than left in float. Already converted ``MXFP4QATLinear`` layers are
	114	+ skipped, so calling this twice is a no-op.
	115	+
	116	+ Args:
	117	+ module: Root module to rewrite.
	118	+ config: Settings applied to every replaced layer.
	119	+ skip_names: Substrings matched against the dotted module path. A match
	120	+ skips that submodule and everything below it, which is the usual way
	121	+ to keep sensitive layers (``"lm_head"``, ``"embed"``) in float.
	122	+
	123	+ Returns:
	124	+ The same module, with its linear layers replaced.
	125	+ """
	126	+ return _convert(module, config, skip_names, prefix="")
	127	+
	128	+
	129	+def _convert(
	130	+ module: nn.Module,
	131	+ config: MXFP4QATConfig \| None,
	132	+ skip_names: tuple[str, ...],
	133	+ prefix: str,
	134	+) -> nn.Module:
	135	+ for name, child in list(module.named_children()):
	136	+ path = f"{prefix}{name}"
	137	+ if any(pattern in path for pattern in skip_names):
	138	+ continue
	139	+ if isinstance(child, MXFP4QATLinear):
	140	+ continue
	141	+ if isinstance(child, nn.Linear):
	142	+ setattr(module, name, MXFP4QATLinear.from_linear(child, config))
	143	+ else:
	144	+ _convert(child, config, skip_names, prefix=f"{path}.")
	145	+ return module