ハードウェアとアクセラレーション
機械学習の性能の大半を決めるのが演算能力ではなくメモリ帯域幅である理由、GPU・TPU・Apple Silicon・NPU・FPGAの違い、そして混合精度とカスタムカーネルで実際に何が得られるかを解説します。
概要
機械学習の性能に関する問いはほぼすべて、2つの数値のどちらかに帰着します。アクセラレータが1秒間に何バイト動かせるか、そして1秒間に何回の浮動小数点演算ができるかです。この2つの比と、ワークロードにおける演算数とバイト数の比を合わせれば、コードを1行も書く前に、計算律速なのかメモリ律速なのかがわかります。ルーフラインモデル（Williams, Waterman and Patterson, 2009）はこの議論をグラフに描いたもので、このハブで最も役に立つ唯一の考え方です。
このページは、ハードウェアを選ぶ、ジョブの規模を見積もる、あるいは学習が遅いのにGPUの使用率が10%と表示されるのはなぜかと悩んでいるエンジニアのためのものです。主なアクセラレータのファミリー、メモリ階層、混合精度、カスタムカーネルを扱います。購入のアドバイスはすぐに古くなるので、ここに挙げる数値は本稿執筆時点でのベンダーの公表仕様であり、問題の形を示すために掲げるもので、特定の部品を推奨するものではありません。
重要な考え方
帯域幅が予算である
NVIDIA H100 SXMを例に取ります。NVIDIAのデータシートには、HBM3が80 GBで3.35 TB/s、密なBF16テンソルコアのスループットが989 TFLOPSとあります。この2つを割ると、1バイトあたり約295演算になります。大きなバッチサイズの行列乗算は読み込んだ1バイトあたりそれよりはるかに多くの演算を行うので計算律速であり、学習がチップを飽和させるのはそのためです。7 Bパラメータのモデルからbf16でトークンを1つデコードするには、14 GBの重みを読んでおよそ14 GFLOPの仕事をします。強度は1バイトあたり約1演算で、したがって完全にメモリ律速です。演算がどれだけ速くても、最大で3.35 TB/sを14 GBで割った、毎秒約240トークンにしかなりません。リクエストのバッチ処理、重みの量子化、KVキャッシュを小さく保つことは、すべてこの同じ限界への対処であり、/topics/mlops-and-deployment/と/topics/efficient-ml/がこれらに多くの紙幅を割いているのはそのためです。
同じ計算で民生品も順位付けできます。NVIDIAの公表ではRTX 4090はGDDR6Xが24 GBで1,008 GB/sです。Appleの公表ではM4 Maxはユニファイドメモリが最大128 GBで546 GB/s、M3 Ultraは最大512 GBで819 GB/sです。128 GBのMacは24 GBのGPUには入らないモデルを保持でき、トークンあたりおよそ半分の速度でデコードします。
アクセラレータのファミリー
GPUは汎用的な選択肢です。数千の単純なコア、行列乗算のためのテンソルコア、成熟したコンパイラ、そして最大のライブラリエコシステムをもちます。NVIDIAはCUDAを通じて支配的です。AMDのInstinct製品はROCmを通じてPyTorchを動かし、スタックは薄いものの改善しつつあります。GoogleのTPUは密な行列乗算のために作られたシストリックアレイ型のマシンで、JAXとXLAを通じてプログラムし、クラウドインスタンスとしてのみ提供されます。Googleのドキュメントでは、TPU v6eはチップあたりHBMが32 GBで1,640 GB/sとされ、大規模な学習のためにチップをポッドに接続します。Apple SiliconはCPU、GPU、ニューラルエンジンを共有メモリプールとともに1つのダイに載せているため、モデルはどの単体の民生用GPUのメモリよりも大きくできますが、その代償として帯域幅は低く、CUDAもありません。スマートフォンやノートPCのNPUは数ワットを狙った小さな固定機能の行列エンジンで、すでにint8以下に量子化されたモデルのオンデバイス推論に向いています。FPGAはその対極にあります。再構成可能なロジックであり、レイテンシ律速、特殊な精度、センサー直結の推論では勝ち、大きな密行列乗算に見えるものではすべて負けます。姉妹サイトにそのケースを扱う専用ハブがあり、下にリンクしています。
混合精度とカーネル
完全なfp32での学習はテンソルコアの大半を無駄にします。混合精度（Micikevicius et al., 2017）は重みのマスターコピーをfp32で保持し、順伝播と逆伝播を16ビットで実行します。bf16はfp32の指数範囲を保ち、損失スケーリングを必要としないため、対応ハードウェアではfp16に取って代わりました。Hopper世代で導入されたFP8は行列乗算のバイト数をさらに半減させ、テンソルごとのスケーリングで範囲内に収めます。
残りの2倍から3倍はカーネルに隠れています。融合カーネルはメモリへの1往復で1つではなく複数の演算を行います。FlashAttention（Dao et al., 2022）が典型例で、共有メモリに留まるタイル単位で厳密な注意機構を計算し、注意行列全体を一度も実体化しません。torch.compileとTritonによって、CUDAを書かなくても融合カーネルに手が届きます。
どこから始めるか
まず/learn/beginner/04-training-loop-and-loss/と/learn/beginner/07-your-first-neural-network/を読み、次に/reference/pytorch-training-loop/の学習ループをCPU、MPSまたはCUDA、そしてbf16のautocastで計時して比較してください。/reference/tensor-shapes-and-broadcasting/は、カーネルが速いかどうかを決めるレイアウトの問題を扱っています。/projects/lora-fine-tune-apple-silicon/プロジェクトは現実的なユニファイドメモリのワークロードです。/topics/deep-learning/ハブは、これらの数値が適用されるモデルを扱っています。
ツールとプラットフォーム
| ツール | 用途 | ライセンス | プラットフォーム |
|---|---|---|---|
| CUDA Toolkit、cuDNN、cuBLAS | NVIDIA GPUのプログラミングとPyTorchが呼び出すライブラリ | プロプライエタリ、無償利用可 | /platforms/nvidia/ |
| Triton（OpenAI） | Pythonで融合GPUカーネルを書く | MIT | /platforms/openai/ |
| torch.compileとPyTorchプロファイラ | グラフのキャプチャ、カーネル融合、遅い演算子の特定 | BSD-3 | /platforms/pytorch/ |
| JAXとXLA | コンパイラを介したTPUとGPUのプログラミング | Apache-2.0 | /platforms/jax/ |
| Cloud TPU | Googleのアクセラレータ。チップ単位またはポッド単位でレンタル | プロプライエタリ、従量課金 | /platforms/google/ |
| MLXとCore ML | Apple Siliconでの学習とオンデバイス推論 | MIT（MLX）、プロプライエタリ（Core ML） | /platforms/apple/ |
| ROCm | PyTorch向けのAMD GPUスタック | 大部分はMITおよびオープンソース | /platforms/opensource/ |
| FPGA Guide | FPGAでの推論フローを扱う姉妹サイト | 該当なし | fpgaguide.com/topics/ai-ml/ |
よくある間違い
- ピークTFLOPSを比較する。ピーク値はすべてのユニットが稼働する密な行列乗算を前提にしています。自分のワークロードでの実測スループットを比較し、引用された数値に構造化スパース性が含まれていないか確認してください。
- ジョブがメモリ律速なのに演算能力で購入する。推論では帯域幅と容量が、大規模な学習では相互接続が決め手になります。
- 損失スケーリングなしでfp16で学習する。あるいはbf16テンソルコアをもつハードウェアでfp32で学習する。
- GPU使用率を効率の尺度として読む。メモリを待っているカーネルも使用率100%と表示されます。
- データローダーをCPUワーカー2つで動かして、GPUのせいにする。
- メモリに収まるモデルなら動くと思い込む。KVキャッシュ、オプティマイザの状態、活性化にも場所が必要です。
さらに深く
- Williams, Waterman and Patterson, Roofline: an insightful visual performance model for multicore architectures, Communications of the ACM, 2009。
- Micikevicius et al., Mixed Precision Training, 2017。
- Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022。
- NVIDIA H100 Tensor Core GPUデータシート。上記の容量と帯域幅の数値の出典です。
- Google Cloud TPUシステムアーキテクチャのドキュメント。TPUの数値の出典です。
- 姉妹サイト: FPGA GuideのFPGAにおけるAI・ML高速化のハブ。データフロー型とオーバーレイ型のアクセラレータ、FPGA向けの量子化、FPGAがGPUやNPUに勝つのはどんな場合かを扱っています。