FlashAttention のスケーリング、マルチ GPU 学習、重みの量子化

現代の AI インフラは、ハードウェアを意識したアテンションカーネル、分散マルチ GPU 並列処理、低ビット量子化を活用し、巨大なモデルを効率的に学習・提供しています。この 3 つの技術は、理論上のモデル規模と GPU の物理的限界の隔たりを埋めます。

1. Hopper (H100) GPU における FlashAttention-2 と FlashAttention-3

初代 FlashAttention はメモリ帯域幅のボトルネックを解消しましたが、最新 GPU の素の演算性能には、まだ大きな未使用部分が残っていました。FlashAttention-2 (2023) と FlashAttention-3 (2024) は、GPU スレッドとハードウェアエンジンの連携方法を作り直しました。

• FlashAttention-2（作業分割の改善）。

• 非行列積演算の削減。GPU の Tensor Cores は、標準 CUDA コアが Softmax の再スケーリングなどのスカラー演算を行う場合と比べ、行列乗算 (matmul) を約 16 倍の速度で実行します。FlashAttention-2 はオンライン Softmax の計算を書き直し、再スケーリングを内側ループの最後まで遅らせることで、低速な非行列積命令を減らしました。
• 系列長に沿った並列化。初代は主にバッチサイズとアテンションヘッドに沿って処理を並列化していたため、長い系列と小さなバッチで学習すると GPU コアが遊んでいました。FlashAttention-2 は系列長の次元にも並列性を加え、ループ順序を逆転させて（外側で Q ブロック、内側で K と V を処理）、低速な共有メモリへの書き込みを最小限に抑えます。

• FlashAttention-3（Hopper H100 ハードウェアの活用）。

• 非同期データ移動 (TMA)。NVIDIA の Hopper アーキテクチャは、演算スレッドを占有せずにグローバル HBM とオンチップ SRAM の間でデータをコピーする専用ハードウェア、Tensor Memory Accelerator (TMA) を導入しました。
• ワープの役割分担（生産者・消費者モデル）。FlashAttention-3 は GPU のワープ（32 スレッドのグループ）に専用の役割を割り当てます。生産者ワープは TMA を使って後続の K と V のタイルを SRAM に先読みし、消費者ワープは現在のタイルを Tensor Cores に渡します（WGMMA 命令）。これにより、メモリ転送、行列演算、Softmax 計算を同時に重ねて進めます。
• FP8 低精度のサポート。FlashAttention-3 は Hopper の FP8 Tensor Cores を活用し、FP16 の 2 倍の素のスループットを実現します。FP8 の外れ値となる活性値による精度低下を防ぐため、ブロック量子化（小さなタイルを個別にスケーリング）と非コヒーレント処理（量子化前に Q と K にランダムな直交行列を掛けて極端な値をならす処理）を使います。

2. マルチ GPU のスケーリング戦略（3D 並列処理）

モデルのパラメータ数が数百億から数千億になると、重み、勾配、オプティマイザの状態は 1 基の GPU の VRAM に収まりません。大規模学習では、通常 3D 並列処理と呼ばれる、異なる 3 つの並列化戦略を組み合わせます。

• データ並列処理と ZeRO / FSDP。標準のデータ並列処理では、各 GPU がモデル全体のコピーを保持するため、メモリが無駄になります。最新のフレームワークは ZeRO (Zero Redundancy Optimizer) または PyTorch FSDP (Fully Sharded Data Parallel) を使い、パラメータ、勾配、オプティマイザの状態を全 GPU に分割し、各層の重みを計算の直前に集めます。

• テンソル並列処理（層内）。巨大なフィードフォワード行列積 Y = XW を重複して計算する代わりに、テンソル並列処理は W を列方向に分割して GPU に配置します (W = [W_1, W_2])。各 GPU が出力の半分を同時に計算し、高速 NVLink（H100 では 900 GB/s）で部分的な結果を交換します。

• パイプライン並列処理（層間）。第 1–8 層を GPU 1、第 9–16 層を GPU 2 というように配置します。下流の GPU が前段の層の終了を待って遊ばないように（「パイプラインバブル」）、フレームワークは各学習バッチをさらに小さなマイクロバッチに分け、組み立てラインのように全 GPU を稼働させ続けます。

3. 量子化 (GPTQ, AWQ) とオンザフライ逆量子化カーネル

推論中（トークンを 1 つずつ生成する際）、大規模言語モデルは演算能力ではなく、メモリ帯域幅に制約されます。たった 1 トークンを生成するために、700 億個の 16 ビットパラメータ (140 GB) を VRAM から読み込むと、Tensor Cores はメモリ転送を待つことになります。量子化は重みを 16 ビット浮動小数点数から 4 ビット整数 (INT4) に圧縮し、メモリ容量と転送時間をほぼ 4 分の 1 にします。

• GPTQ（ヘッセ行列による較正を使う学習後量子化）。

• 16 ビットの重みを単純に 4 ビット整数へ丸めると、モデルの精度が損なわれます。GPTQ は小さな較正用データセットをモデルに通し、二次の数学的手法（ヘッセ行列の逆行列）で各重みの感度を測ります。
• 重みを列ごとに量子化し、ある重みの丸めで誤差が生じるたびに、その行に残る未量子化の重みを数学的に調整して誤差を打ち消します。

• AWQ（活性値を考慮した重み量子化）。

• 研究者は、重みチャネルの約 0.1% から 1% が、大きな活性値と乗算されるため、残りよりもはるかに重要であることを発見しました。
• AWQ は、そうした重要な重みを 16 ビットのまま保持する（複雑で低速な混合精度メモリ配置を招く）代わりに、量子化前に重要な重みチャネルを拡大し、実行時に対応する入力活性値を縮小します。これにより、重みの 100% を一様な 4 ビット形式に保ちながら精度を維持します。

• カスタム CUDA カーネルによるオンザフライ逆量子化の仕組み (W4A16)。

• GPU は 16 ビットの活性値 (FP16) と 4 ビット整数の重み (INT4) を直接掛けることができません。代わりに、Marlin や ExLlama などの専用推論カーネルが W4A16（4 ビット重み、16 ビット活性値）の実行を行います。
• 1. 圧縮された HBM の読み出し：カーネルは低速なグローバル VRAM からパックされた 4 ビット重み（1 個の 32 ビット整数に 8 個の INT4 値を格納）を読み込み、メモリ帯域幅の使用量を 4 分の 1 にします。
• 2. レジスタ内の展開：GPU の超高速レジスタ内で、CUDA カーネルは高速ビットシフト (>>) とビットマスク (&) により INT4 整数を展開し、保存済みのスケール係数とゼロ点係数を適用して、その場で FP16 値に戻します。
• 3. 即座の Tensor Core 演算：展開直後の FP16 重みを FP16 活性値とともに直接 Tensor Cores に送り、展開した 16 ビット重みを VRAM に書き戻すことはありません。