大規模言語モデルパラメータとメモリ要件の理解：深い理解

パラメータ数の計算

事前トレーニング済みモデルによる効率的な Transformer トレーニング

Transformer ベースの LLM のパラメータ計算

Transformer ベースの LLM の各コンポーネントのパラメータ計算を分解してみましょう。元の論文の表記を使用します。ここで、d_model はモデルの隠れ状態の次元を表します。

1. 埋め込みレイヤー:
   • パラメータ = vocab_size * d_model
2. マルチヘッド注意:
   • h ヘッドの場合、d_k = d_v = d_model / h:
   • パラメータ = 4 * d_model^2 (Q、K、V、出力投影の場合)
3. フィードフォワードネットワーク:
   • パラメータ = 2 * d_model * d_ff + d_model + d_ff
   • d_ff は通常 d_model の 4 倍です
4. レイヤー正規化:
   • パラメータ = 2 * d_model (スケールとバイアス用)

1 つの Transformer レイヤーのパラメータの合計:

• Parameters_layer = Parameters_attention + Parameters_ffn + 2 * Parameters_layernorm

N レイヤを持つモデルの場合:

• 合計パラメータ = N * Parameters_layer + Parameters_embedding + Parameters_output

計算例

次の仕様を持つモデルを考えてみましょう:

• d_model = 768
• h (注意ヘッドの数) = 12
• N (レイヤーの数) = 12
• vocab_size = 50,000

1. 埋め込みレイヤー:
   • 50,000 * 768 = 38,400,000
2. マルチヘッド注意:
   • 4 * 768^2 = 2,359,296
3. フィードフォワードネットワーク:
   • 2 * 768 * (4 * 768) + 768 + (4 * 768) = 4,719,616
4. レイヤー正規化:
   • 2 * 768 = 1,536

1 レイヤーのパラメータの合計:

• 2,359,296 + 4,719,616 + (2 * 1,536) = 7,081,984

12 レイヤーのパラメータの合計:

• 12 * 7,081,984 = 84,983,808

モデルの合計パラメータ数:

• 84,983,808 + 38,400,000 = 123,383,808

このモデルには約 123 百万のパラメータがあります。

メモリ使用の種類

LLM を使用する場合、2 つの主要なメモリ使用の種類を考慮する必要があります:

1. モデルメモリ: モデルのパラメータを格納するために必要なメモリです。
2. ワーキングメモリ: 中間アクティベーション、勾配、オプティマイザ状態を格納するために、推論またはトレーニング中に必要なメモリです。

モデルメモリの計算

モデルメモリは、パラメータ数に直接関係しています。各パラメータは、通常、32 ビット浮動小数点数として格納されますが、一部のモデルでは、16 ビット浮動小数点数を使用した混合精度トレーニングを使用します。

モデルメモリ (バイト) = パラメータ数 * パラメータあたりのバイト数

例モデル (123 百万パラメータ) の場合:

• モデルメモリ (32 ビット) = 123,383,808 * 4 バイト = 493,535,232 バイト ≈ 494 MB
• モデルメモリ (16 ビット) = 123,383,808 * 2 バイト = 246,767,616 バイト ≈ 247 MB

ワーキングメモリの推定

ワーキングメモリ要件は、特定のタスク、バッチサイズ、シーケンス長によって大きく異なります。推論時のワーキングメモリの粗い見積もりは次のとおりです:

ワーキングメモリ ≈ 2 * モデルメモリ

これは、モデルのパラメータと中間アクティベーションの両方を格納することを考慮しています。トレーニング中のメモリ要件は、勾配とオプティマイザ状態を格納する必要があるため、さらに高くなることがあります:

トレーニングメモリ ≈ 4 * モデルメモリ

例モデルでは:

• 推論ワーキングメモリ ≈ 2 * 494 MB = 988 MB ≈ 1 GB
• トレーニングメモリ ≈ 4 * 494 MB = 1,976 MB ≈ 2 GB

定常メモリ使用とピークメモリ使用

Transformer アーキテクチャに基づく大規模言語モデルのトレーニングでは、メモリ使用を理解することが、効率的なリソース割り当てとパフォーマンスの向上に不可欠です。メモリ要件を 2 つの主要なカテゴリに分解してみましょう: 定常メモリ使用とピークメモリ使用。

定常メモリ使用

定常メモリ使用には、次のコンポーネントが含まれます:

1. モデルの重み: FP32 形式のモデルのパラメータ、4N バイトが必要です。ここで、N はパラメータ数です。
2. オプティマイザ状態: Adam オプティマイザの場合、8N バイト (パラメータごとに 2 つの状態) が必要です。
3. 勾配: FP32 形式の勾配、4N バイトが必要です。
4. 入力データ: int64 形式の入力データ、8BD バイトが必要です。ここで、B はバッチサイズ、D は入力次元です。

定常メモリ使用の合計は、次のように推定できます:

• M_steady = 16N + 8BD バイト

ピークメモリ使用

ピークメモリ使用は、バックワードパス中にアクティベーションが勾配計算のために格納されるときに発生します。ピークメモリの主な寄与者は次のとおりです:

1. レイヤー正規化: 各レイヤー正規化に 4E バイトが必要です。ここで、E = BSH (B: バッチサイズ、S: シーケンス長、H: 隠れ層のサイズ) です。
2. 注意ブロック:
   • QKV 計算: 2E バイト
   • 注意行列: 4BSS バイト (S: シーケンス長)
   • 注意出力: 2E バイト
3. フィードフォワードブロック:
   • 最初の線形層: 2E バイト
   • GELU アクティベーション: 8E バイト
   • 2 番目の線形層: 2E バイト
4. クロスエントロピーロス:
   • ログits: 6BSV バイト (V: 語彙サイズ)

アクティベーションメモリの合計は、次のように推定できます:

• M_act = L * (14E + 4BSS) + 6BSV バイト

ここで、L は Transformer レイヤーの数です。

合計ピークメモリ使用

トレーニング時のピークメモリ使用は、定常メモリとアクティベーションメモリを組み合わせることで推定できます:

• M_peak = M_steady + M_act + 4BSV バイト

追加の 4BSV 項は、バックワードパスの開始時に追加の割り当てを考慮しています。

これらのコンポーネントを理解することで、トレーニングと推論時のメモリ使用を最適化し、大規模言語モデルのパフォーマンスを向上させることができます。