跳至主要内容

M06 概覽 | QAT 的訓練流程：讓訓練時的圖等於 convert 之後的圖

Open In Colab

M05 用 PTQ 量化一個訓練好的模型，並學會怎麼查精度問題。位元數再壓低時，PTQ 可能真的不夠，這時可以讓模型在訓練中就看到量化誤差、自己去適應它，這就是 QAT（quantization-aware training）。 QAT 的圖長什麼樣在 M04-2 看過；這一章只管訓練。後面六本 notebook（M06-1～M06-6，M06-5、M06-6 是選讀）逐一動手驗證。

這一章回答的問題

問題 一句話的答案 動手驗證
prepared 的 QAT 圖能直接訓練嗎？ 可以，它是普通的 GraphModule；round 的梯度靠 STE M06-1
評估前要做什麼？ BN 切 eval、observer 關掉，否則結果不可重現 M06-2
訓練時的圖等於 convert 之後的圖嗎？ 不一定：XNNPACK 預設的 weight fake quant 就不相等 M06-3
低位元時，QAT 真的比 PTQ 好嗎？ int4 下明顯較好；同樣配方再訓練 fp32 後做 PTQ 則沒有改善，所以效益來自量化適應 M06-4

選讀的兩本：M06-5 把 scale 也交給梯度學（range learning）；M06-6 處理 weight 由別的參數算出來（weight_norm）時的 QAT。見後面的〈選讀〉兩節。

先看一個例子：int4 下 QAT 追回 PTQ 掉的精度

模型和 M05 相同（訓練好的 TinyConvNet）。int8 下 PTQ 已經貼著 fp32，沒有東西可救（M06-1），所以這裡用 int4：activation [-8, 7]、weight [-7, 7]。 比較 PTQ int4 與 QAT int4（訓練 5 個 epoch）的 test loss，並檢查 QAT 訓練完的圖和它 convert 之後的圖輸出是否一致：

模型 test loss
fp32 0.826
PTQ int4 1.071
QAT int4 0.920 
QAT 追回了 PTQ int4 比 fp32 多出的 test loss 的 62%
QAT 訓練完的 fake quant 圖 vs convert 之後的圖：SQNR 141.1837615966797 dB，最大差 0.00 個輸出 LSB
  • PTQ int4 的 test loss 比 fp32 高了一截；QAT 只訓練 5 個 epoch 就追回一部分（比例印在上面）。這只是一個 seed，5 個 seed 與對照組的比較在 M06-4。
  • 最後一行是這一章最重要的檢查：訓練時模型看到的圖（fake quant 圖），和 convert 之後的 Q/DQ 圖，輸出一致：SQNR 是 inf，或最大差在 1 個輸出 LSB 內（判準見後面〈全章的原則〉）。這件事不是自動成立的，M06-3 與 M06-5 各抓到一個讓它不成立的地方。

這裡用的 qat_quantizer、train_qat、freeze_for_eval 都是本專案的工具（common/qat.py），不是 torchao 的 API。M06-1 會把訓練迴圈逐行寫出來，之後才換成這些工具。

定義

fake quant

QAT 在圖上插的不是 observer，而是 fake quant：forward 時先量化再立刻反量化，輸出仍是 fp32，但值已經落在量化的格點上。scale \(s\)、zero point \(z\)、整數範圍 \([q_{min}, q_{max}]\)：

\[ \hat{x} = \big(\operatorname{clamp}(\operatorname{round}(x/s) + z,\ q_{min},\ q_{max}) - z\big)\cdot s \tag{1} \]

STE（straight-through estimator）

round 幾乎處處梯度為 0，直接微分就訓練不動。STE 在 backward 時把 round 當成恆等函數，只保留 clamp 的效果：

\[ \frac{\partial \hat{x}}{\partial x} = \begin{cases} 1 & q_{min} \le \operatorname{round}(x/s) + z \le q_{max} \\ 0 & \text{otherwise} \end{cases} \tag{2} \]

用一個 int4 的 FakeQuantize 和式 (1)、(2) 比對：

fq = FakeQuantize(observer=MovingAverageMinMaxObserver, dtype=torch.int8, qscheme=torch.per_tensor_affine,
                  quant_min=-8, quant_max=7, eps=2**-12)
with torch.no_grad():
    fq(torch.linspace(-1, 1, 100))  # 校準 range
fq.apply(disable_observer)
s, z = (float(v) for v in fq.calculate_qparams())

x = torch.linspace(-2, 2, 41, requires_grad=True)
y = fq(x)
y.sum().backward()
q = torch.round(x.detach() / s) + z
by_eq1 = (q.clamp(-8, 7) - z) * s                       # 式 (1)
by_eq2 = ((q >= -8) & (q <= 7)).float()                 # 式 (2)
assert x.grad is not None
print(f"s = {s:.4f}, z = {z:.0f}；式 (1) 相符：{torch.allclose(y, by_eq1)}；式 (2) 相符：{torch.equal(x.grad, by_eq2)}")
assert torch.allclose(y, by_eq1) and torch.equal(x.grad, by_eq2)
s = 0.1333, z = -1；式 (1) 相符：True；式 (2) 相符：True

range 外的值梯度是 0，這一步得不到更新。STE 的推導見部落格〈搞懂 Quantization Aware Training 中的 Fake Quantization〉；STE 的改良（EWGS、MAD 等）整理在部落格〈量化技術路線〉的 QAT 一節。

指標：test loss，比較對象是「同樣配方再訓練的 fp32」

M05 用「對 fp32 輸出的 SQNR」。QAT 會改 weight，這個 SQNR 量到的是「離原模型多遠」，不是「量化後好不好」，所以這一章改用 test loss（cross-entropy，連續、而且是 QAT 直接最佳化的量）。 QAT 多做了「訓練」這件事，要分開訓練本身與量化的影響，基準是用同樣配方再訓練、不量化的 fp32（M06-1 C）；int4 時再加上「再訓練後做 PTQ」這組對照（M06-4）。seed 只改變資料的順序，test loss 就會不同，所以 M06-4 每組跑 5 個 seed。

流程

export(train=True)        BN 要是 train 版的 op
  → prepare_qat_pt2e      插 fake quant、建 conv + BN 的 QAT 子圖（M04-2）
  → 訓練                  每個 epoch 切 train／eval（M06-1）
  → freeze_for_eval       BN 用 running stats、observer 關掉（M06-2）
  → convert_pt2e          得到 Q/DQ 圖

全章的原則：訓練時的圖 = convert 之後的圖

QAT 的前提是「訓練時模擬的量化誤差 = 部署時的量化誤差」。檢查方法固定：同一批資料餵 fake quant 圖（freeze 之後）和 convert 圖，比較輸出。 這一章抓到兩個讓等號不成立的地方：

出處 等號為什麼不成立 解法
M06-3 weight 的 fused fake quant 用自己的公式算 scale weight 改用非 fused 的 FakeQuantize
M06-5 learnable 的 scale 被梯度推成負數，convert 時被夾成 eps 每步把 scale 夾回正數

判準是：輸出 SQNR 為 inf，或最大差在 1 個輸出 LSB 內（輸出 dequantize 的 scale）。 後者是因為 activation 的 fused fake quant，scale 偶爾和 convert 用的值差 1 ULP，只是浮點捨入（M06-3 C）。

這個檢查只保證「訓練圖 = Q/DQ 圖」；Q/DQ 圖到後端整數 kernel 的那一段，M07 起再看。

選讀一：把 scale 也交給梯度學（M06-5）

到 M06-4 為止，scale 都由 observer 的 min／max 統計決定，訓練只改 weight。LearnableFakeQuantize 把 scale 與 zero point 做成 nn.Parameter，讓 loss 的梯度去調整量化範圍（LSQ 的做法）。 接進 PT2E 只要把 spec 的 fake quant 類別換掉，但有兩件事要自己處理：

  • scale 是第一次 forward 才建立的參數，optimizer 要在那之後才建。
  • torchao 只在 convert 時把 scale 夾成正數；訓練中梯度可能把它推成負數，訓練圖就不等於 convert 圖了，所以要每步夾回正數。

這個模型在 int4 上，range learning 3 個 seed 都優於 observer 版；代價是多了 grad scaling、qparams 的 lr 這組超參數，調錯會發散。 梯度的推導見部落格〈Learning Zero-Point and Scale in Quantization Parameters〉。

選讀二：weight 由別的參數算出來時（M06-6）

M05 概覽的〈選讀〉介紹過 torch.nn.utils.parametrize：註冊之後，參數變成「由別的參數算出來」，每次存取都重新計算（PyTorch 教學 Parametrizations）。 例子是 weight_norm，每個輸出 channel 各自算

\[ W = g \cdot \frac{v}{\lVert v \rVert} \tag{3} \]

export 之後，式 (3) 變成圖上的 aten._weight_norm，conv 的 weight 是它的輸出、不是 get_attr。PTQ 不訓練，在 prepare 之前把它換成常數就好；QAT 不一樣，訓練的正是 \(g\)、\(v\)。

fake quant 插在哪裡、梯度怎麼走

和 PTQ 的 observer 一樣，weight 的 fake quant 接在 _weight_norm 的輸出上，量化的是算出來的 \(W\)。 backward 時，STE 的梯度穿過 fake quant，再經 _weight_norm 傳回 \(g\)、\(v\)。所以 QAT 期間 weight_norm 照常有效，被訓練的是 \(g\) 與 \(v\)。

有 BN 時，哪一步報錯

讀常數的 torchao 內部函式只接受 get_attr，所以「要把 conv 的 weight 當常數讀出來」的那一步會報錯：

ValueError: Expected get_attr node, got call_function

PTQ 和 QAT 在不同的步驟碰到它：

流程 什麼時候折 BN 結果
PTQ prepare_pt2e 一開始（M04-1） prepare 就報錯
QAT prepare_qat_pt2e 只建 QAT 子圖、不折（M04-2）；convert_pt2e 才折 prepare 能過、訓練照常，訓練全部做完、convert 時才報錯

怎麼解

把 \(W\) 算好、存成一般的參數，BN 就能照常折。關鍵是什麼時候做：

時機 做法 結果
export 之前 remove_parametrizations(..., leave_parametrized=True) 等於放棄訓練 \(g\)、\(v\)
訓練之後、convert 之前 strip_weight_norm(gm)：把圖上的 _weight_norm 換成常數 訓練了 \(g\)、\(v\)，convert 也能折 BN

strip_weight_norm 是本專案的工具（common/weight_norm.py），不是 torch 或 torchao 的 API，M05-5 已經介紹過。讀者只要知道它做的事；換完之後輸出逐位元不變、convert 圖和 fake quant 圖一致，都在 M06-6 C 段驗證。

有 BN 又要保留 weight_norm 的 QAT 流程因此是：export(train=True) → prepare_qat_pt2e → 訓練 → freeze_for_eval → strip_weight_norm → convert_pt2e。

小結與本章地圖

  • QAT 的圖可以直接訓練，round 的梯度靠 STE；評估與 convert 前要 BN 切 eval、observer 關掉。
  • 全章的原則是「訓練時的圖 = convert 之後的圖」，檢查方法是比較兩張圖的輸出。
  • int4 下 QAT 明顯優於 PTQ；同樣配方再訓練 fp32 後做 PTQ 沒有改善，所以效益來自訓練時模擬量化，不是多訓練。
  • （選讀）range learning 有效但多了要調的超參數；weight_norm 要保留到訓練結束，convert 前才拔。
notebook 回答的問題
M06-1 最小 QAT STE 的梯度；最小訓練迴圈；convert 後和 PTQ 同形
M06-2 評估前的兩個開關 BN 與 observer 不切會怎樣；allow_exported_model_train_eval
M06-3 fused fake quant 的陷阱 為什麼 XNNPACK 預設設定的訓練圖 ≠ convert 圖；scale 公式
M06-4 int4 的對照實驗 低位元時 QAT 是否比 PTQ 好；效益來自量化適應還是多訓練
M06-5（選讀）range learning LearnableFakeQuantize 怎麼接、要注意什麼、效果
M06-6（選讀）weight_norm 保留 \(g\)、\(v\) 訓練；有 BN 時 convert 報錯與解法

這一章刻意跳過的：Q/DQ 圖之後會被改寫成什麼、和後端的整數運算差多少（M07 起）。

參考資料