M06-3 | 訓練時的圖 ≠ convert 後的圖？XNNPACK 預設 QAT 設定的一個陷阱

Open In Colab

QAT 的前提是「訓練時模擬的量化誤差 = 推論時的量化誤差」。這本 notebook 檢查這個前提：把訓練好的 QAT 圖（fake quant 圖，eval 模式、observer 關）和它 convert 後的 Q/DQ 圖餵同一批資料，輸出應該完全一樣。

fused fake quant 是什麼：這裡的 fused 和 M04 的 conv + BN 折疊無關。FakeQuantize 模組內部有三個步驟：

observer 更新 min／max。
由 min／max 算 scale／zero_point。
quantize → clamp → dequantize。

FusedMovingAvgObsFakeQuantize 把三步合成一顆 ATen kernel（torch.fused_moving_avg_obs_fake_quant）一次做完，XNNPACKQuantizer 的 QAT 設定（is_qat=True）預設用它。代價是第 2 步由這顆 kernel 自己算，不經過 torchao observer 的 calculate_qparams()，問題就出在這裡。

A. 檢查：weight 用 XNNPACK 預設的 fused fake quant 時，int8／int4 下兩張圖差多少。
B. 原因：fused kernel 算 scale 的公式和 calculate_qparams()（convert 用的）不同。
C. 解法：weight 改用非 fused 的 FakeQuantize。

import copy

import torch

from common.data import calibration_batches, trained_tiny_net
from common.graph_utils import sqnr
from common.models import example_inputs
from common.pt2e import (
    FakeQuantize,
    FusedMovingAvgObsFakeQuantize,
    MovingAverageMinMaxObserver,
    MovingAveragePerChannelMinMaxObserver,
    convert_pt2e,
    disable_observer,
    export_module,
    prepare_qat_pt2e,
    quiet,
)
from common.qat import EPS, freeze_for_eval, qat_quantizer, train_qat
from common.viz import table

quiet()
_ = torch.manual_seed(0)

A. fake quant 圖 vs convert 後的圖

qat_quantizer(weight_fq=...)（common/qat.py）可以選 weight 的 fake quant："fused" 是 XNNPACK 預設的 FusedMovingAvgObsFakeQuantize，"plain" 是非 fused 的 FakeQuantize。activation 兩者都用 fused。每個設定訓練 2 個 epoch，freeze_for_eval（M06-2）之後 convert，比較兩張圖對同一批輸入的輸出 SQNR：

model = trained_tiny_net()
(x,) = next(calibration_batches(64, seed=9))


def fq_vs_converted(bits, weight_fq):
    prepared = prepare_qat_pt2e(export_module(copy.deepcopy(model), example_inputs(), train=True),
                                qat_quantizer(bits=bits, weight_fq=weight_fq))
    freeze_for_eval(train_qat(prepared, epochs=2))
    converted = convert_pt2e(copy.deepcopy(prepared))
    with torch.no_grad():
        return sqnr(prepared(x), converted(x))


a = {(wfq, bits): fq_vs_converted(bits, wfq) for wfq in ("fused", "plain") for bits in (8, 4)}
table([["FusedMovingAvgObsFakeQuantize（XNNPACK 預設）", f"{a['fused', 8]:.1f}", f"{a['fused', 4]:.1f}"],
       ["FakeQuantize（非 fused）", f"{a['plain', 8]:.1f}", f"{a['plain', 4]:.1f}"]],
      ["weight 的 fake quant", "int8 (dB)", "int4 (dB)"])

weight 的 fake quant	int8 (dB)	int4 (dB)
FusedMovingAvgObsFakeQuantize（XNNPACK 預設）	30.1	6.5
FakeQuantize（非 fused）	inf	inf

assert a["fused", 8] < 40 and a["fused", 4] < a["fused", 8]
assert a["plain", 8] == float("inf") and a["plain", 4] == float("inf")

預設設定下兩張圖不一樣：int8 只有 30 dB 左右，int4 掉到個位數 dB。訓練時模型適應的量化網格，和 convert 後 Q/DQ 圖的網格不是同一個，QAT 訓練的不是 convert 出來的那個模型。
兩個 row 的 activation 設定相同，所以問題出在 weight，也就是 per-channel symmetric 的 fake quant。

B. 原因：fused kernel 的 scale 公式

兩邊各自怎麼算 scale：

訓練時：fused kernel 在 forward 裡自己算。torch v2.14.0 的 ChooseQuantizationParams（QuantUtils.h 第 86～93 行）在 symmetric 時先把 \([q_{min}, q_{max}]\) 換成 \([-(n+1), n]\)，\(n = \lfloor (q_{max}-q_{min})/2 \rfloor\)，再算

\[ s_{train} = \max\!\Big(\frac{|min|}{n+1}, \frac{|max|}{n}\Big) \cdot \frac{2n+1}{q_{max}-q_{min}} \tag{1} \]

convert 時：convert_pt2e 呼叫 calculate_qparams()，委派給 torchao 的 observer：

\[ s_{convert} = \frac{\max(|min|, |max|)}{(q_{max}-q_{min})/2} \tag{2} \]

以 int8 weight 的 \([-127, 127]\) 為例，式 (1) 用的是 \([-128, 127]\)。下面用一組 weight 讓 fused fake quant 跑一次，比較它的 scale buffer（訓練用的）和 calculate_qparams()（convert 用的），並分別和式 (1)、(2) 比對：

def scale_eq1(mn, mx, qmin, qmax):
    n = (qmax - qmin) // 2
    return torch.maximum(mn.abs() / (n + 1), mx.abs() / n) * (2 * n + 1) / (qmax - qmin)


def scale_eq2(mn, mx, qmin, qmax):
    return torch.maximum(-mn, mx) / ((qmax - qmin) / 2)


def fused(qmin, qmax, per_channel=True):
    if per_channel:
        return FusedMovingAvgObsFakeQuantize(observer=MovingAveragePerChannelMinMaxObserver, ch_axis=0, dtype=torch.int8,
                                             qscheme=torch.per_channel_symmetric, quant_min=qmin, quant_max=qmax, eps=EPS)
    return FusedMovingAvgObsFakeQuantize(observer=MovingAverageMinMaxObserver, dtype=torch.int8,
                                         qscheme=torch.per_tensor_affine, quant_min=qmin, quant_max=qmax, eps=EPS)


w = torch.randn(6, 3, 3, 3)
w[0] *= 3            # 有的 channel 由 |min| 主導、有的由 max 主導，兩種情況公式差的方向不同
w[1, 0, 0, 0] = 5.0
rows, checks = [], []
for qmin, qmax in ((-127, 127), (-7, 7)):
    fq = fused(qmin, qmax)
    with torch.no_grad():
        fq(w)
    obs = fq.activation_post_process
    s_train, s_convert = fq.scale.flatten(), fq.calculate_qparams()[0].flatten()
    ok1 = torch.allclose(s_train, scale_eq1(obs.min_val, obs.max_val, qmin, qmax), rtol=1e-5)
    ok2 = torch.allclose(s_convert, scale_eq2(obs.min_val, obs.max_val, qmin, qmax), rtol=1e-5)
    ratio = s_train / s_convert
    rows.append([f"per-channel symmetric [{qmin}, {qmax}]", f"{float(ratio.min()):.4f} ~ {float(ratio.max()):.4f}", ok1, ok2])
    checks += [ok1, ok2]
fq = fused(-128, 127, per_channel=False)
with torch.no_grad():
    fq(w)
rows.append(["per-tensor affine [-128, 127]（activation 用的）", f"{float(fq.scale / fq.calculate_qparams()[0]):.4f}", "—", "—"])
table(rows, ["設定", "scale 比值：訓練 / convert", "訓練的 scale 符合式 (1)", "convert 的 scale 符合式 (2)"])

設定	scale 比值：訓練 / convert	訓練的 scale 符合式 (1)	convert 的 scale 符合式 (2)
per-channel symmetric [-127, 127]	0.9961 ~ 1.0039	True	True
per-channel symmetric [-7, 7]	0.9375 ~ 1.0714	True	True
per-tensor affine [-128, 127]（activation 用的）	1.0000	—	—

assert all(checks)

int8 差 ±0.4%，int4 差 −6%～+7%。per-tensor affine 不走 symmetric 這條分支，兩邊一致。

那 disable_observer 之後呢？同一版本的 fused_obs_fake_quant.cpp 寫得很清楚：observer_on 只決定要不要更新 min／max；只要 fake_quant_on，每次 forward 都會用式 (1) 重算 scale。驗證：關掉 observer，手動把 scale buffer 改成 convert 用的值，再 forward 一次：

fq = fused(-127, 127)
with torch.no_grad():
    fq(w)
fq.apply(disable_observer)
s_before = fq.scale.clone()
with torch.no_grad():
    fq.scale.copy_(fq.calculate_qparams()[0])  # 改成 convert 用的值
    fq(w)                                      # 再 forward 一次
print("forward 之後 scale 又變回式 (1) 的值：", torch.allclose(fq.scale, s_before))
assert torch.allclose(fq.scale, s_before)

forward 之後 scale 又變回式 (1) 的值： True

所以 disable_observer 救不了：只要 weight 的 fake quant 是 fused 的，訓練時用的就是式 (1)。

C. 解法：weight 改用非 fused 的 `FakeQuantize`

FakeQuantize.forward 三步分開做：observer → calculate_qparams() → fake_quantize_per_channel_affine。第二步和 convert 是同一條公式：

plain = FakeQuantize(observer=MovingAveragePerChannelMinMaxObserver, ch_axis=0, dtype=torch.int8,
                     qscheme=torch.per_channel_symmetric, quant_min=-7, quant_max=7, eps=EPS)
with torch.no_grad():
    plain(w)
ratio = plain.scale / plain.calculate_qparams()[0]
print(f"FakeQuantize：forward 用的 scale ÷ calculate_qparams() = {float(ratio.min()):.4f} ~ {float(ratio.max()):.4f}")
assert torch.equal(plain.scale, plain.calculate_qparams()[0])

FakeQuantize：forward 用的 scale ÷ calculate_qparams() = 1.0000 ~ 1.0000

A 段表格的第二個 row 因此是 inf。common/qat.py 的 qat_quantizer() 預設就這樣設定：activation 維持 fused（per-tensor affine，公式一致），weight 用 FakeQuantize。M06 之後的 QAT 都用它。

代價是少了 fused kernel 的速度。本章的小模型感覺不到，這裡沒有量測。
為什麼要對齊的是 convert 圖：XNNPACK 後端部署時讀的是 convert 圖 q/dq 上的 scale（M08），convert 算出什麼就部署什麼。

以後檢查的方法都一樣：fake quant 圖 vs convert 圖的輸出 SQNR 應該是 inf。有一個例外：activation 的 fused fake quant，scale buffer 偶爾和 calculate_qparams() 差 1 ULP（依 seed 而定），這時輸出最大差在 1 個輸出 LSB 內（輸出 dq 的 scale），SQNR 可能掉到 50 dB 左右，但這只是浮點捨入，也算通過。這裡的 30 dB、6.5 dB 則是公式不同，不是捨入。

這個檢查只保證「訓練圖 = Q/DQ 圖」；Q/DQ 圖到後端整數 kernel 的那一段，M08～M09 另外量。

小結

XNNPACK 預設的 QAT 設定，weight 用 fused fake quant，訓練圖和 convert 圖不同（int8 約 30 dB、int4 個位數 dB）。
原因：fused kernel 的 symmetric scale 公式（式 1）和 observer 的（式 2）不同，而且 disable_observer 後仍每次重算。
解法：weight 用非 fused 的 FakeQuantize；判準是 SQNR 為 inf，或最大差在 1 個輸出 LSB 內。

A. fake quant 圖 vs convert 後的圖

B. 原因：fused kernel 的 scale 公式

C. 解法：weight 改用非 fused 的 FakeQuantize

小結

C. 解法：weight 改用非 fused 的 `FakeQuantize`