💻

AVX VNNI での INT8 演算サポートのメモ

2023/02/05に公開

avx

vnni

int8

tech

ChatGPT とかで LLM(large language model)がはやっている...
GPU で INT8 でサイズ小さくするとかはやってる...

メモリ強つよな GPU もクラウドの CPU もお高いので,
ローカルの x86 CPU で INT8 つかって LLM うごかしたいな
(昨今の民生 PC だとメモリ 128 GB は搭載できる. 幸い(?)メモリ不況のため, DDR4 で速度にこだわらなければ 32GB x 1 が 1 万円くらいで調達できる https://akiba-pc.watch.impress.co.jp/docs/price/monthly_repo/1473254.html)

x86 の場合, AVX VNNI で INT8 演算ができました
(VNNI は DL boost とも呼ばれているようです)

ただ, uint8 x int8 の mul のみのようで, フレキシブルに int8 演算するにはうまく量子化を考えないといけなさそうです!

情報

ありがとうございます.

AVX-VNNI と AVX512-VNNI

ややこしいですが,

AVX 版(not AVX512)の VNNI(128bit, 256bit)
AVX512 版の VNNI(512bit)

の 2 つがあります(たぶん).

Intel の場合, AVX-VNNI は民生 CPU の Alder Lake(i9 12900K とか)以降でサポートされています.

11 世代では AVX512 が一部サポートされており, AVX-512-VNNI で 8bit 演算が使えます!

ただ, AVX-VNNI での 8bit 演算とはちょっと異なります(ややこしいね)

とりあえずは Linux では cat /proc/cpuinfo して vnni があれば int8 命令使えます.

i9 12900K で cat /prc/cpuinfo すると vnni サポートのコア数は 24 でした. スレッド数と同じとなり, E コアでもサポートされています.

$ cat /proc/cpuinfo | grep vnni | wc
     24    3336   22512

ちなみに feaure flags は以下となりました.

flags           : fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb intel_pt sha_ni xsaveopt xsavec xgetbv1 xsaves split_lock_detect avx_vnni dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req umip pku ospke waitpkg gfni vaes vpclmulqdq tme rdpid movdiri movdir64b fsrm md_clear serialize pconfig arch_lbr flush_l1d arch_capabilities
vmx flags       : vnmi preemption_timer posted_intr invvpid ept_x_only ept_ad ept_1gb flexpriority apicv tsc_offset vtpr mtf vapic ept vpid unrestricted_guest vapic_reg vid ple shadow_vmcs ept_mode_based_exec tsc_scaling usr_wait_pause
bugs            : spectre_v1 spectre_v2 spec_store_bypass swapgs eibrs_pbrsb

~~AVX512 は Intel では民生 CPU ではサポートされておらず(2023/02 時点の 13 世代含む), AVX512-VNNI はサーバ用 CPU でのサポートになります.~~
AVX512-VNNI は 11 世代ではサポートされていました. 12 世代から廃止.

(ちなみに Alder Lake の場合, E コア off にすると AVX512 が使えるリビジョンもあるようである(初期ロット?) https://northwood.blog.fc2.com/blog-entry-11204.html )

AMD では, Ryzen 7000(ZEN4) から AVX512 と AVX512-VNNI がサポートされています.

VPDPBUSD で int8 x int8 する

uint8 x int8 の計算に落とし込めれば, vpdpsubd でいけますが,

int8 x int8 が必要な場合は, 実質演算量は二倍になりますが,

(x_ui - 127) * w_i 
= x_i([0, 127]) * w_i + (127 * (-w_i))

と, int8 量子化した w の符号反転を用意して処理することになるでしょうか.
64 要素ごとなど, x と w の取りうる最大値(= x = 256, w = 256) で結果の和をとっても int(32bit) で飽和しなければ, 計算後に最後に補正も行けるでしょう.

AVX-VNNI-INT8

さらにややこしいですが, AVX-VNNI-INT8 があります. 低消費電力 xeon サーバ向け CPU 用でしょうか.

組み合わせできる int8 のタイプが増えています.

専用の cpuid が追加されています.
こちらも民生品に降りてきてくれるといいのですけどね.

性能でるの?

不明です. 理論上は fp32 に比べて 4 倍くらい速くなりそうですが, 実効はよくて +50 ~ 100%(up to 2 倍)でしょうか.

Transformer みたいなのだと, メモリ速度律速の場合は量子化により消費メモリが減るだけのメリットになるかもしれません.

対応コンパイラは?

AVX-VNNI は clang-12, gcc-11 から対応しています.

Linux kernel 対応は?

少なくとも Ubuntu 22.04 LTS 5.15 では AVX-VNNI 使えました.
Ryzen Zen4(AVX512-VNNI) は不明です. 現状対応してなかったとしても 1 年以内には対応されるかなとは思います.

INT4(4bit int) は...?

RDNA3, Ada/Hopper では INT4 命令があります.
VNNI では今のところ INT4 命令はなさそうです.

AMX(Advanced Matrix eXtension)

別途 AMX(Advanced Matrix eXtension)で INT8 行列の matmul があります! より GPU 的なかんじでしょうか. ただ Xeon などのサーバ向けなので, 民生 CPU にはおりてこないかもしれません(降りてきても 5,6 年後くらい?)

ライブラリでのサポートは?

XNNPACK では AVX VNNI(での int8 演算) サポートはありませんでした.
oneDNN(mkl-dnn) でのサポートはあります.
xbyak も一応ありますが, なんか AVX-VNNI モードでコンパイルしてもなんかうまくいきませんでした

Alder Lake で動かすのを考えるのであれば,

このあたりを参考にして自前で intrinsic 書くのが手っ取り早そうな気がします!

BF16, FP16

BF16 は現状 VNNI(DL boost) の枠組みでは提供されておらず, AVX512 の枠組みでの提供でした.

FP16 も AVX-VNNI ではまだありません. AVX-512 かつ FP16 対応があれば, AVX512 の枠組みで使えます.
(Alder Lake では FP16 対応とありましたが, AVX-512 disable になったので使えなくなったっぽい?)

Ryzen 7000 series(ZEN4)でも AVX-512 の FP16 対応はされていません(BF16 は対応).

したがって現状(2023/02)の民生 CPU では,

Intel は AVX の BF16, FP16 共に未サポート
Ryzen 7000(ZEN) は AVX の BF16 が使える

となります.

別途, F16C fp32 <-> fp16 変換命令(SIMD(AVX) では無い?)は Alder Lake や Ryzen(少なくとも ZEN2)では使えます.

おまけ

ray tracing での BVH build/traversal あたりになんかつかえそうです!

quantized BVH での bbox 判定とか.