Transformers documentation
GPU メモリ使用量
GPU メモリ使用量
バッチサイズ16で混合精度により4Bパラメータのモデルを学習するには、約85GBのGPUメモリが必要です。GPUメモリを占有している要素や、学習中に実行される計算操作を理解することで、メモリ使用量を削減できる箇所を特定できます。
┌─────────────────────────── TENSORS ──────────────────────────────────┐ │ │ │ MODEL WEIGHTS ████████████████████████ 6 bytes/param │ │ (mixed precision) ├── fp16 copy (2B) — forward/backward pass │ │ └── fp32 copy (4B) — stable weight updates │ │ │ │ OPTIMIZER STATES ████████████████████████████████ 8 bytes/param │ │ (Adam) ├── fp32 momentum (4B) │ │ └── fp32 variance (4B) │ │ ↳ quantized Adam (bitsandbytes) → 2 bytes/param │ │ │ │ GRADIENTS ████████████████ 4 bytes/param (fp32) │ │ └── computed in backward pass, per parameter │ │ │ │ ACTIVATIONS ████ varies — batch × seq_len × depth × hidden │ │ (forward cache) └── cached for backward; can OOM even if model │ │ fits (long seqs / large batches) │ │ bf16/fp16 if using mixed precision │ │ │ │ TEMPORARY TENSORS ▓ short-lived (softmax, matmul scratch) │ │ └── peak spikes can cause OOM │ │ │ │ OTHER OVERHEAD ▒ beam search caches, large embedding tables │ └──────────────────────────────────────────────────────────────────────┘
GPU メモリには、保存されたテンソルとそれらを処理するオペレーションという2つのカテゴリのものが保持されます。
テンソル
学習には、保存が必要な多くのテンソル型が必要です。
モデルの重みは GPU に保存されます。混合精度学習では、重みのコピーが 2 つ必要になります。1 つは forward/backward pass 用の fp16、もう 1 つは安定した重み更新のための「メインコピー」としての fp32 です。これにより、パラメータあたり 6 バイトのメモリが必要になります。
Adam などの Optimizer state は、パラメータごとにモーメンタムと分散という 2 つの追加テンソルを保存します。これらはどちらも fp32 です。つまり、パラメータごとに 8 バイトの追加メモリが必要になります。
bitsandbytes の Adam の量子化バージョンなど、別のオプティマイザーを使用して、パラメータあたり2バイトに圧縮することもできます。
バックワードパスでは、各パラメータについて勾配テンソルが計算されます。これは fp32 で保持されるため、パラメータあたり 4 バイトになります。
Forward activations は forward pass で計算され、backward pass で勾配を計算するためにキャッシュされます。これらの activations のサイズは、batch size、sequence length、model depth、hidden size に応じて変化します。そのため、モデル自体が収まっていなくても、batch size や sequence length によって GPU メモリが枯渇することがあります。
softmaxや行列積などのオペレーションによって一時テンソルが作成され、各オペレーションの後に解放されます。単一のオペレーションのピークが非常に激しい場合、メモリ不足を引き起こす一時的なスパイクが発生する可能性があります。
一部の機能には独自のオーバーヘッドが伴います。例えば、beam searchでは複数の出力を保持する必要があり、大規模な語彙の埋め込みテーブルは非常に大きくなることがあります。
運用
トレーニングオペレーションには、主に3つの種類があります。
行列積（matmul）が主な演算タイプであり、線形層、QKV 射影、アテンション出力射影、FFN 層をカバーします。主なメモリ消費要因は、シーケンス長の二乗に比例して増加するアテンションスコア行列です。これが、長いシーケンスが非常にコストがかかる理由であり、メモリ上に完全な行列をマテリアライズすることを避けるために、異なるアテンションバックエンドが存在します。
softmax や layer norm などの reduction 演算では、テンソル全体を読み取り、ある次元にわたる統計量を計算し、その後、テンソルを再度読み取ってその統計量を適用します。これにより、1 つの演算ごとにメモリへ複数回アクセスする必要があります。
活性化やdropoutなどの要素ごとの演算は、各要素に対して独立して関数を適用し、そのメモリ使用量はテンソルのサイズに比例します。
次のステップ
- GPU メモリが不足することなく、より大きな実効バッチサイズでの学習をシミュレートする方法を学ぶには、Gradient accumulation ガイドを参照してください。
- 一部の中間アクティベーションのみを保存してメモリ使用量を削減する方法を学ぶには、Gradient checkpointingガイドを参照してください。
- メモリを削減し、トレーニングを高速化するために低精度のデータ型を使用する方法については、Mixed precision trainingガイドを参照してください。
- カスタム融合カーネルによる学習の高速化方法については、Kernelsガイドを参照してください。