前幾天我們已經把幾個零件拆開看過了。
Day 6 先把問題拆成 pipeline:
選模數 / 建 CRT 空間
-> 把輸入轉成餘數矩陣
-> INT8 GEMM
-> CRT 重建
-> scaling 回 double
Day 7 看 CRT 表和重建。
Day 8 把 INT8 GEMM 先寫成 naive kernel,確認每個餘數矩陣真的能乘起來。
Day 9 再把 INT8 GEMM 搬上 Tensor Core,開始碰 cp.async、ldmatrix、mma.sync。
到這裡為止,我們其實已經有一條可以跑的路了:
double
-> 整數切片
-> modulo
-> INT8 Tensor Core
-> CRT
-> double
但「可以跑」和「精度真的補齊」是兩件事。
今天要補的是這個洞:
前面幾天的設計可以展示 CRT / Tensor Core 怎麼接起來,但如果要把精度往 FP64 的 53-bit mantissa 靠近,模數數量、模數大小、切片數量、pair schedule 都要重新補齊。
這篇不講怎麼使用 repository,也不講哪個 mode 怎麼切。
今天只講一件事:
要怎麼把前幾天那個簡化版,補成比較完整的高精度 GEMM 設計?
前面為了好懂,我們一直用「小於 128 的 prime」來講。
例如 Day 7 用的是:
127, 113, 109, 107, 103, 101, 97
這樣做很直覺,因為 signed int8 的範圍是:
-128 .. 127
所以如果 residue 都是 0..126,看起來就一定塞得進 int8_t。
但這個說法只是一個教學上的簡化。
實際上,INT8 能放的是 8 個 bit 的位元圖樣。如果我們用的是 signed int8,硬體看到的是 -128..127,但在 modular arithmetic 裡,我們可以把它解讀成某個模數下的代表元。
例如對模數 256:
0 -> 0
1 -> 1
...
127 -> 127
128 -> -128
129 -> -127
...
255 -> -1
這些 signed 值雖然看起來有負數,但它們在 modulo 256 的世界裡是等價的:
-1 ≡ 255 (mod 256)
-128 ≡ 128 (mod 256)
所以重點不是「數字表面上不能超過 127」。
重點是:
每個 residue 能不能用一個 signed 8-bit representative 表示,而且後面的乘法累加再取 mod 之後,是否仍然和原本的 residue 乘法同餘。
如果 a' ≡ a (mod p),b' ≡ b (mod p),那麼:
a' * b' ≡ a * b (mod p)
所以我們可以用 signed int8 的 -128..127 去表示更大的模數,例如 256、255、253。儲存時看起來可能是負數,但只要 representation 和後面的 modular reduction 一致,數學上仍然是同一個 residue。
Day 5 到 Day 7 都用 prime 講 CRT,因為那樣最好理解。
但 CRT 真正需要的是:
模數彼此互質。
不是每個模數都一定要是 prime。
例如下面這組:
256, 255, 253, 251, 247, 241, 239, 233,
229, 227, 223, 217, 211, 199, 197, 193,
191, 181, 179, 173
它裡面 256 不是 prime,255 也不是 prime。
但只要整組模數彼此互質,CRT 還是可以把每個模數下的 residue 拼回同一個整數空間。
為什麼這組數字有用?
因為它們都可以用一個 signed int8 representative 來承載 residue,而且每個模數都盡量靠近 256。模數越大,乘起來的 CRT 空間就越快變大。
前面 7 個 prime 的乘積大概提供:
log2(127 * 113 * 109 * 107 * 103 * 101 * 97)
≈ 47.26 bits
也就是說,它能重建的整數空間大概是 47 bits。
這對前面簡化過的 34-bit-ish representation 來說可以展示整個流程,但對 full FP64 mantissa 來說不夠。
如果換成剛剛那組更大的 modulus budget:
前 15 個 modulus:約 117.78 bits
全部 20 個 modulus:約 155.37 bits
這個差距非常大。
用表格看會比較清楚:
| 模數集合 | 大約 CRT bits |
|---|---|
| 7 個小 prime:127 到 97 | 47.26 bits |
| 前 15 個大 modulus:256 到 197 | 117.78 bits |
| 20 個大 modulus:256 到 173 | 155.37 bits |
這就是為什麼「把模數從 7 個增加到 15 或 20 個」不是小修小補。
它直接把可以重建的整數空間放大很多。
GEMM 的一個輸出元素是 dot product:
C[i,j] = sum_k A[i,k] * B[k,j]
如果我們希望 A 的整數化表示有 t bits,B 也有 t bits,那單次乘法大概需要:
t + t = 2t bits
但 GEMM 不是只乘一次。
它還要沿著 K 維度累加:
sum over K terms
最壞情況下,K 個值相加還會多需要:
ceil(log2(K)) bits
所以一個簡化但很好用的估算是:
required_crt_bits = 2 * target_bits + ceil(log2(K))
如果目標只是 FP32-class 的 24-bit mantissa,K=4096:
2 * 24 + 12 = 60 bits
舊的 7 個 prime 只有約 47 bits,所以其實也偏緊。
如果目標是 FP64 的 53-bit mantissa,K=4096:
2 * 53 + 12 = 118 bits
這時候 7 個 prime 完全不夠。
但前 15 個大 modulus 大約是:
117.78 bits
它已經非常接近這個 118-bit 需求。
如果想留更多 buffer,20 個 modulus 則有:
155.37 bits
這就是高精度版本為什麼會往 15 到 20 個 modulus 走。
這裡還有第三個誤會:
CRT 空間夠大,就代表 FP64 精度完成了嗎?
還不夠。
CRT 只負責重建「已經被送進去的整數」。
如果一開始 double 只被切成 high / low 兩片,那沒有被切出來的 mantissa bits 就不存在於後面的 GEMM 裡。
這就像你拍照時只存低解析度圖片。
後面就算用再大的硬碟、再好的壓縮格式,也不會把當初沒拍到的細節變回來。
所以要補齊精度,需要同時補兩個方向:
1. CRT 空間要夠大
2. mantissa slice 要切得夠完整
前面幾天的 high / low 設計比較像:
x ≈ x_hi * 2^-17 + x_lo * 2^-34
也就是兩片。
如果要往 FP64 的 53-bit mantissa 靠近,兩片 17-bit 大概只覆蓋到 34 bits,後面還有一截沒有被表示進來。
所以高精度版本要做的是把 double 拆成更多片。
概念上可以想成:
x ≈ x0 * 2^-b
+ x1 * 2^-2b
+ x2 * 2^-3b
+ ...
每一片都用 INT8 能承載的 bit budget 表示。
b 不一定永遠是 8,也不一定永遠是 17。它要看 K、accumulator range、模數設計和誤差目標。
前面 high / low 只有兩片,所以 A 和 B 相乘只有四組:
A_hi * B_hi
A_hi * B_lo
A_lo * B_hi
A_lo * B_lo
但如果 A 被切成很多片:
A0, A1, A2, A3, ...
B 也被切成很多片:
B0, B1, B2, B3, ...
那理論上的 pair 會變成:
A0 * B0
A0 * B1
A1 * B0
A0 * B2
A1 * B1
A2 * B0
...
這些 pair 的重要性不一樣。
越高位的 slice pair 權重越大,越低位的 slice pair 權重越小。
所以正式的高精度實作通常不只是暴力做所有 pair,而是會有一個 pair schedule:
哪些 pair 必須做
哪些 pair 可以省
每個 pair 的 shift 是多少
累加時要放到哪個位權
這就是為什麼 reference 裡會看到類似:
num_split
bits_per_int8
gemm_pair_config_list
accumulate_in_f64
這些東西不是裝飾。
它們就是在回答:
double 要切幾片?每片幾 bits?哪些 slice pair 要相乘?乘完之後要怎麼用正確位權加回去?
把今天的補齊版整理起來,流程會變成:
double A, double B
|
v
選 scaling / exponent
|
v
把 A/B 拆成多個 INT8 slice
|
v
產生 pair schedule
|
v
對每個重要 slice pair 做 INT8 GEMM
|
v
對每個 modulus 做 residue / modular accumulation
|
v
用更多 modulus 做 CRT reconstruction
|
v
依照 pair shift 和 scaling 累回 double
|
v
得到高精度 GEMM 結果
前幾天做的是其中幾個關鍵零件:
CRT 怎麼拼
residue matrix 怎麼來
INT8 GEMM 怎麼算
Tensor Core 怎麼接
今天補的是:
模數範圍要放大
模數數量要增加
mantissa slice 要補足
pair schedule 要完整
precision budget 要能算
精度補齊不是免費的。
如果從 7 個 modulus 變成 15 或 20 個 modulus,代表 modular GEMM 的工作量也會增加。
如果從 high / low 兩片變成更多 slice,代表 slice pair 數量也會增加。
非常粗略地看,成本會跟下面幾個東西有關:
modulus 數量
slice pair 數量
每個 pair 的 GEMM 成本
CRT reconstruction 成本
global memory 暫存成本
所以高精度版一開始一定比較難快。
前面幾天把 Tensor Core 接起來,只是讓我們有機會把 INT8 GEMM 的部分加速。
但如果 slice pair 和 modulus 數量變多,後面的問題就會變成:
怎麼少寫 global memory?
怎麼把多個 modulus 合批?
怎麼把 CRT 往 accumulator 附近推?
怎麼避免 register pressure 爆掉?
怎麼讓 Tensor Core 不要等 scalar epilogue?
這些才是後面要優化的地方。
今天其實是在修正前幾天簡化版留下的精度缺口。
要把那條路補成更完整的高精度 GEMM,需要四件事一起成立:
模數可以用到 256 附近,不是只能小於 128。
signed int8 可以存 balanced residue。-1 可以代表 255 mod 256,-128 可以代表 128 mod 256。
CRT 需要的是 pairwise-coprime modulus,不是每個都一定要 prime。
所以 256, 255, 253, ... 這種設計是合理的,只要整組模數彼此互質。
模數數量要依 precision budget 增加。
7 個小 prime 大約 47 bits;15 個大 modulus 大約 118 bits;20 個大 modulus 大約 155 bits。
high / low 兩片不夠補滿 FP64 mantissa。
要往 53-bit mantissa 走,需要更多 slice 和更完整的 pair schedule。
所以 Day 10 的結論不是「換一組模數就結束」。
比較正確的結論是:
前幾天我們把 CRT、重建、INT8 GEMM、Tensor Core 都接起來了。今天開始,要把它從簡化版補成高精度版,就必須同時補 modulus budget 和 mantissa slice budget。
明天 Day 11 再來談另一個問題:
精度補齊之後,工作量變多了。那要怎麼把它重新加速回來? (速度真的好難@@ 這幾天一直只能用AIslob就是因為再補跟cublass的速度 太難了...)