# AI and Memory Wall
> [!abstract] 概要
> 前例のない教師なし学習データの利用可能性とニューラルスケーリング則により、大規模言語モデル(LLM)のサービングおよび学習に必要なモデル規模と計算量が前例のない勢いで増大している。
> しかし、主な性能ボトルネックは次第にメモリ帯域幅へ移りつつある。
> 過去20年間、サーバー・ハードウェアのピークFLOPSは2年ごとに3.0倍で伸びてきた一方、DRAMとインターコネクトの帯域幅はそれぞれ2年ごとに1.6倍と1.4倍しか伸びていない。
> この差により、AIアプリケーション、特にサービングでは、計算ではなくメモリが主要なボトルネックになっている。
> 本稿ではエンコーダー型とデコーダー型のTransformerモデルを分析し、デコーダー型モデルではメモリ帯域幅が支配的なボトルネックになりうることを示す。
> このメモリ制約を克服するため、モデル・アーキテクチャ、学習、デプロイ戦略の再設計を提案する。
## 論文情報
- タイトル: AI and Memory Wall
- 著者: Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Hooper, Michael W. Mahoney, Kurt Keutzer
- 所属: University of California, Berkeley、ICSI、LBNL
- 媒体: IEEE Micro, 44(3), 33–39, 2024(IEEE掲載版); arXiv拡張版は2024年3月21日
- arXiv: [2403.14123](https://arxiv.org/abs/2403.14123)
- DOI: [10.1109/MM.2024.3373763](https://doi.org/10.1109/MM.2024.3373763)
## 概要
LLMのパラメータ数と学習計算量が急増する一方で、ハードウェアの演算性能に対するメモリ・インターコネクト帯域幅の伸びは遅れている。著者らはTransformerのエンコーダーとデコーダーを比較し、特に小バッチの自己回帰デコーダー推論では演算量よりもメモリ操作と帯域幅がレイテンシを支配すると論じる。
## 問題設定
メモリウォールは、メモリの容量、転送帯域幅、レイテンシ、およびアクセラレーター間通信が演算性能の伸びに追いつかない問題である。キャッシュから取得できるデータとDRAMから取得するデータが混在する場合、DRAMからのデータ取得が十分に遅ければ、演算器がどれほど高速でも処理時間はDRAM帯域幅で決まる。
2018–2022年のSOTAモデルでは、Transformerのパラメータ数が2年ごとに410倍、学習計算量が2年ごとに750倍で増えた。一方、単一AIアクセラレーターのメモリ容量は2年ごとに2倍程度である。過去20年のピーク・ハードウェアFLOPSの増加は60,000倍だったのに対し、DRAM帯域幅は100倍、インターコネクト帯域幅は30倍にとどまる。
## 提案手法
本稿は新しい単一アルゴリズムを提案するのではなく、スケーリング傾向とTransformer推論のプロファイリングに基づき、学習・デプロイ・アクセラレーター設計を再考する方向を整理する。
### スケーリング分析
**Figure 1** は、1996–2023年のピークFLOPS、DRAM帯域幅、インターコネクト帯域幅を、R10000を基準に正規化して示す。ピークFLOPSの伸びが帯域幅を大きく上回り、メモリウォールが強まっている。
![[_attachments/Gholami-et-al.-2024---AI-and-Memory-Wall/fig01-scaling-bandwidth.png]]
(Figure 1. ピークFLOPSが過去20年で60,000倍に伸びる一方、DRAM帯域幅とインターコネクト帯域幅の伸びはそれぞれ100倍と30倍である。)
**Figure 2** は、SOTAモデルのパラメータ数・学習計算量とAIアクセラレーターのメモリ容量の推移を示す。モデル規模の伸びはアクセラレーター単体のメモリ容量を大きく上回る。
![[_attachments/Gholami-et-al.-2024---AI-and-Memory-Wall/fig02-scaling-models.png]]
(Figure 2. (a) Transformerモデルのパラメータ数は2年ごとに410倍で増えるが、単一GPUメモリは2倍にとどまる。(b) Transformerの学習計算量は2年ごとに750倍で増え、ムーアの法則の2倍を上回る。)
### 演算強度による分析
演算強度(arithmetic intensity)を、メモリから読み込む1バイトあたりに実行できるFLOPsとして定義する。論文では次式で計算する。
$
\mathrm{Arithmetic\ Intensity}=\frac{\#\mathrm{FLOPs}}{\#\mathrm{MOPs}}
$
ここでMOPsはメモリ操作数であり、対象の演算に必要な行列がローカルメモリに収まることを仮定する。実験ではモデルパラメータとアクティベーションを8ビット、バッチサイズを1とし、MACの乗算と加算を別々のFLOPとして数える。
### Transformer推論のプロファイリング
エンコーダー型(BERT)は全トークンを並列処理し、行列–行列演算を中心とする。デコーダー型(GPT)は自己回帰的に1トークンずつ生成し、行列–ベクトル演算を繰り返す。Intel Gold 6242 CPU上でBERT-Base、BERT-Large、GPT-2を系列長128–4096、バッチサイズ1で計測した。
**Figure 3** はFLOPs、MOPs、演算強度、正規化レイテンシを比較する。Figure 3aではエンコーダー型のFLOPsが同程度または大きい一方、Figure 3bではGPT-2のMOPsが桁違いに多い。Figure 3cではGPT-2の演算強度がエンコーダーより桁違いに低く、Figure 3dではFLOPsが少ないGPT-2の実行時間が最も長い。
![[_attachments/Gholami-et-al.-2024---AI-and-Memory-Wall/fig03-transformer-profiling.png]]
(Figure 3. BERT-Base、BERT-Large、GPT-2の系列長別プロファイリング。図3a–dはそれぞれFLOPs、MOPs、演算強度、CPU上の正規化レイテンシを示す。)
系列長128におけるBERT-Base基準の正規化レイテンシは、BERT-Largeが3、GPT-2が45である。系列長4096ではBERT-Baseが84、BERT-Largeが218、GPT-2が2,344となり、デコーダーのメモリ操作が長い系列で急速に不利になる。なお、長文書要約の入力処理や大バッチ推論では行列–行列演算が含まれるため、この結論がすべてのデコーダー用途に当てはまるわけではない。
## 新規性
過去の「メモリウォール」の予測を、近年のAIハードウェア、LLMのモデル規模・学習計算量、Transformerの実行特性に接続して再検討した点にある。FLOPS、MOPs、演算強度、エンドツーエンドレイテンシを同時に比較し、FLOPsだけではデコーダー推論のボトルネックを説明できないことを示す。
## 実験設定
- ハードウェア: Intel Gold 6242 CPU
- 対象モデル: BERT-Base、BERT-Large、GPT-2
- 系列長: 128、256、512、1024、2048、4096
- バッチサイズ: 1
- 数値精度: パラメータとアクティベーションを8ビットと仮定
- 指標: FLOPs、MOPs、演算強度、BERT-Baseの系列長128を基準にした正規化エンドツーエンドレイテンシ
- スケーリング分析: 1996–2023年のハードウェアデータ、2016–2022年のモデルパラメータ数、2012–2021年の学習計算量
## 実験結果
- Figure 1では、ピークFLOPSが2年ごとに3.0倍、DRAM帯域幅が1.6倍、インターコネクト帯域幅が1.4倍で伸びる。
- Figure 2では、Transformerのパラメータ数が2年ごとに410倍、学習計算量が750倍で増える一方、単一アクセラレーターのメモリは2倍程度である。
- Figure 3では、GPT-2はBERT-Baseとおおむね同じFLOPsでもMOPsが多く、演算強度が低いためレイテンシが大きい。系列長4096の正規化レイテンシはBERT-Base 84、BERT-Large 218、GPT-2 2,344である。
## 学習・デプロイ・ハードウェア設計への含意
### 効率的な学習アルゴリズム
一次SGD系最適化は実装しやすいが、学習率、減衰スケジュール、反復回数などのハイパーパラメータ調整に大きな試行錯誤を要する。二次確率的最適化は調整に頑健でSOTAに到達しうる一方、メモリ使用量が3–4倍になる。Zeroは冗長なオプティマイザー状態変数を除去・分割し、同じメモリ容量で8倍大きいモデルを学習できることを示した。
再マテリアライゼーションでは、順伝播時に全アクティベーションを保存せず一部だけをチェックポイントし、必要時に残りを再計算する。**Figure 4** はモデル世代に伴う学習メモリの増加とGPUメモリ上限の関係を示す。
![[_attachments/Gholami-et-al.-2024---AI-and-Memory-Wall/fig04-training-memory.png]]
(Figure 4. CNNおよびTransformer系モデルの学習に必要なメモリを、特徴マップ・オプティマイザー状態・パラメータに分けて示す。モデルの大型化がGPUメモリ上限を押し上げてきた。)
再計算により計算量は増えるが、メモリ使用量を最大5倍削減し、計算量の増加を20%に抑えられる場合がある。これは分散学習を使わずに大きなモデルを単一チップのメモリへ収める可能性を広げる。
低精度計算では、FP16がFP32に対してハードウェア計算能力を10倍以上高めた。INT8以下への学習精度削減は精度劣化が難しいが、FP8とFP16の混合、さらにFP4への展開が進められている。低精度最適化は計算資源の一部をメモリ改善へ回す余地も作る。
### 効率的なデプロイ
GPT-3や大規模推薦システムは分散メモリ上での推論を要し、データをアクセラレーター間で移動する通信がボトルネックになる。対策は量子化、枝刈り、小型言語モデルである。推論量子化はINT4までなら精度への影響を抑えつつモデルフットプリントとレイテンシを最大8倍削減できるが、INT4未満は難しい。
構造化スパース性ではニューロンを最大30%、非構造化スパース性では最大80%削減できる場合があるが、これを超えると精度劣化が大きい。小型言語モデルがチップ内に完全に収まれば、速度とエネルギー効率を桁違いに改善できる可能性がある。
### AIアクセラレーターの再設計
メモリ帯域幅とピーク計算能力を同時に高めることは難しいため、ピーク計算能力を一部犠牲にして計算量と帯域幅のバランスを取る設計が考えられる。CPUはキャッシュ階層を備え、帯域幅バウンドな処理ではGPUより高性能な場合があるが、ピークFLOPSはGPUやTPUよりおよそ一桁低い。
CPUとGPUの中間として、効率的なキャッシュ、高容量DRAM、異なる帯域幅を持つDRAM階層を備えるアーキテクチャが考えられる。高容量DRAMは、分散メモリ通信のボトルネック緩和にも役立つ。
## 考察
FLOPsだけを性能指標にすると、行列–行列演算中心のエンコーダーと行列–ベクトル演算中心のデコーダーの差を見落とす。演算強度とMOPsを併用することで、計算器の性能を上げてもデコーダー推論のレイテンシが短くならない理由を説明できる。
本稿のスケーリング値は各グラフのデータへ線形回帰して求めた傾きであり、学習計算量はハードウェア利用率に依存する近似を避けるためPFLOPsで報告している。実際の学習費用にはハイパーパラメータ探索などの追加オーバーヘッドが含まれるため、Figure 2(b)の値は下限である。
## 強み / 弱点・課題
- 強み: ハードウェアの長期スケーリングと、BERT/GPT-2のFLOPs・MOPs・演算強度・レイテンシを同じ議論へ接続している。
- 強み: 学習、デプロイ、アクセラレーター設計を個別の対策ではなくメモリ制約への連続した対応として整理している。
- 弱点・課題: Transformer以外のモデルや、デコーダーの大バッチ・長文書要約など行列–行列演算が支配的な条件では、プロファイリング結果の一般化範囲が限られる。
- 弱点・課題: スケーリング分析は公開ハードウェア・モデルデータと回帰に基づき、実運用の通信負荷、電力、コスト、利用率を直接評価していない。
- 未解決: メモリ帯域幅・容量・レイテンシ・通信を同時に最適化するAIアクセラレーターの具体的な設計点は定まっていない。
## 関連
- 概念: [[メモリウォール]] / [[Transformer]] / [[Rooflineモデル]] / [[AIアクセラレータ]] / [[LLM推論]] / [[量子化]] / [[再マテリアライゼーション]]
- エンティティ: [[Amir Gholami]] / [[Zhewei Yao]] / [[Sehoon Kim]] / [[Coleman Hooper]] / [[Michael W. Mahoney]] / [[Kurt Keutzer]] / [[BERT]] / [[GPT-2]] / [[GPT-3]] / [[University of California, Berkeley]] / [[Lawrence Berkeley National Laboratory]]
## 出典
- 原本PDF: [[.raw/papers/Gholami-et-al.-2024---AI-and-Memory-Wall.pdf]]
- 抽出テキスト: [[.raw/papers/Gholami-et-al.-2024---AI-and-Memory-Wall.txt]]