/w=3840,quality=90,fit=scale-down)
※本記事は、スタンフォード大学のオンラインコース「CS224N: NLP with Deep Learning」の講義12「効率的なトレーニング」の内容を基に作成されています。講義の詳細情報は https://web.stanford.edu/class/archive/ でご覧いただけます。本記事では、講義の内容を要約しております。なお、本記事の内容は原講義の内容を正確に反映するよう努めていますが、要約や解釈による誤りがある可能性もありますので、正確な情報や文脈については、オリジナルの講義動画（https://www.youtube.com/watch?v=UVX7SYGCKkA ）をご覧いただくことをお勧めいたします。
講師：Shikhar J. Murty（スタンフォード大学 コンピュータサイエンス PhD候補生） 指導教授：Christopher Manning（トーマス・M・シーベル機械学習教授、言語学・コンピュータサイエンス教授、スタンフォード人工知能研究所（SAIL）所長）
スタンフォード大学のオンライン人工知能プログラムについての詳細は https://stanford.io/ai をご覧ください。
1. 効率的なニューラルネットワークトレーニング
1.1. 講義の概要と目標
本日の講義12へようこそ。これまでの講義では、単語をベクトルに変換する方法、文をベクトルに変換する方法、そして実世界でアクションを起こすために、それらを使って文書を分類するといったことを学んできました。また、Transformerや事前学習についても学習しました。今日の講義は少し異なるアプローチを取ります。
今日は、GPUで大規模モデルをトレーニングする方法と、これらの機械学習システムがどのように機能するかについての基本的な内容を話します。これは自然言語処理とは直接関係ありませんが、皆さんの最終プロジェクトにとって役立つ内容になると思います。
具体的には、混合精度トレーニング（Mixed Precision Training）に時間を割き、DDP（Distributed Data Parallel）やFSDP（Fully Sharded Data Parallel）を使ったマルチGPUトレーニングについて解説し、パラメータ効率的な微調整（Parameter Efficient Fine-tuning）についても触れる予定です。この講義が終わる頃には、これらの用語が理解できるようになっていると思います。
講義に入る前に、いくつかお知らせがあります。プロポーザルの成績はまもなく、恐らく今日中に公開される予定です。課題4と最終プロジェクトのプロポーザルの締め切りが重なり、大変だったと思いますが、皆さんの努力に感謝します。また、プロジェクトのマイルストーンの詳細はすでにウェブサイトに掲載されているか、まもなく掲載される予定です。これは全体の成績の5%を占め、今から12日後が締め切りで、最大2ページです。このマイルストーンは、最終プロジェクトの作業を進めるための強制関数として考えてください。
それでは、本題に入りましょう。
1.2. 計算機における数値の表現方法
まず、コンピュータにおいてパラメータや勾配、一般的に数値がどのように表現されているかを考えることから始めましょう。これがディープラーニングにどう関係するかはすぐに明らかになります。
まず浮動小数点について話しましょう。この図式的な描写に馴染みのある方はどれくらいいますか？何人かいらっしゃいますね。FP32について振り返りましょう。FP32は32バイト、いや32ビットなので、メモリ要件は4バイトです。ニューラルネットワークの観点から考えると、各パラメータに対して4バイトのGPUメモリが必要になります。
この図で実際の数値に変換する方法ですが、最初のビットが符号を表し、緑色の部分が範囲（range）を表し、青色の部分が精度（precision）を表します。FP32では、かなり広範囲の数値を表現でき、かなり精度も高いと言えます。緑色の部分が大きいほど、より多くの数値（小さな数値から大きな数値まで）を表現できます。さらに、青色の部分が大きいほど、実際の数値表現の精度が高くなります。
FP32の半分のメモリしか使わないもう一つの人気のあるデータ型がFP16です。メモリを削減するために、緑色の部分（つまり範囲）と青色の部分（つまり精度）を減らします。これによりダイナミックレンジが狭くなり、精度も低下しますが、メモリ要件は半分になるという利点があります。
例えば、大きなニューラルネットワークをトレーニングしていて、モデルパラメータと勾配がFP32で表現されているとします。トレーニングを開始すると、突然「out of memory Cuda error」が発生したとします。ここまで見てきたことに基づけば、一つの解決策は全てをFP16にキャストすることです。これを行うと、メモリ使用量は半分になります。
しかし、このようなことを行うと、どのような問題が発生する可能性があるでしょうか。先ほど述べたように、緑色の部分が少ないため、範囲が狭くなります。これは、非常に小さな数値の多くがゼロに変換され、非常に大きな数値の多くがNaNに変換されることを意味します。また、青色の部分が少ないため精度も低くなり、丸め誤差が発生します。例えば、FP16では1.1が1に丸められてしまいます。
データ型の様々な特性をテストする方法のスクリーンショットを用意しました。注目すべき点はイプシロン（epsilon）です。イプシロンは、1に加えても精度が失われない最小の数値です。イプシロンより小さい数値を1に加えると、それは単に1に丸められます。また、「smallest normal」はFP16で表現できる最小の数値であり、それより小さい数値はゼロになります。
ニューラルネットワークのトレーニングでは、多くの小さな数値がゼロに丸められることは実際には良くありません。NVIDIAのブログ記事から取った図がありますが、これはトレーニング中の勾配を示しています。FP16では、これらの勾配の半分以上が文字通りゼロになってしまいます。これがFP16の範囲に関する問題です。二つ目の問題は精度に関するもので、FP16では精度が低いため、更新が正確ではなくなります。
1.3. 浮動小数点形式（FP32、FP16）の説明
FP32と呼ばれる32ビット浮動小数点形式は、ニューラルネットワークのパラメータを表現するための標準的な方法です。FP32は4バイトのメモリを必要とします。この形式では、最初のビットが符号を表し、その後に指数部（緑色で示した部分）があり、これが表現可能な数値の範囲を決定します。残りのビット（青色で示した部分）は仮数部で、これが精度を決定します。
FP32は広い範囲の数値を高い精度で表現できますが、大規模なニューラルネットワークではメモリ使用量が課題になります。各パラメータに4バイト必要であり、数十億のパラメータを持つモデルでは、このメモリ要件は非常に大きくなります。
一方、FP16は16ビット浮動小数点形式で、FP32の半分のメモリしか使いません。メモリ使用量を削減するために、指数部（範囲を表す部分）と仮数部（精度を表す部分）の両方が削減されています。これにより、表現できる数値の範囲と精度の両方が制限されますが、メモリ使用量が半分になるという大きな利点があります。
しかし、FP16には重要な制限があります。範囲が狭いため、非常に小さな数値の多くが単純にゼロにされてしまいます。また、精度が低いため、例えば1.1という数値が単に1に丸められるといった丸め誤差が発生します。
これらの制限を理解するため、各データ型の特性を確認できるテストを行うことができます。例えば、イプシロン（1に加えても精度が失われない最小の数値）や、表現可能な「最小の正規化数」（それより小さい数値はゼロになる）などがあります。FP16では、多くの小さな勾配値が単にゼロになるため、ニューラルネットワークのトレーニングには問題があります。
NVIDIAのブログ記事から引用した図によると、トレーニング中の勾配の半分以上がFP16では単にゼロに設定されてしまいます。これは大きな問題です。また、精度の問題により、パラメータ更新が正確でなくなるという二つ目の問題もあります。
このように、浮動小数点形式の選択は、メモリ効率と数値表現の正確さのバランスを取る重要な決断です。以降のセクションでは、これらの課題に対処するための混合精度トレーニングについて説明します。
2. 混合精度トレーニング
2.1. FP16の問題点と限界
FP16を使用する際の解決策として考えられる一つの方法は、FP32とFP16の両方を使用することです。これが基本的なアイデアです。具体的には、FP32のモデルのコピーを「マスターウェイト」として維持し、データのバッチを取得してフォワードパスを実行します。フォワードパスを実行するときは、FP32からFP16に変換して行います。
そして、勾配を計算し、バックワードパスを実行して、FP16で勾配を取得します。ここまでの処理はすべてFP16で行われています。次に、勾配をFP32にアップキャストし、マスターウェイトを更新します。マスターウェイトを更新したら、それをFP16バージョンのニューラルネットワークにコピーします。
これは合理的な方式のように思えます。GPUではFP16を使用していますが、より正確な更新のためにFP32の精度も別に保持しています。しかし、この方法にもまだ問題があります。
なぜこれが問題になるかというと、先ほどのNVIDIAの図に戻りましょう。これはバックワードパスにおける勾配を示しています。勾配をすべてFP16で計算すると言いましたが、その多くは単にゼロに変換されてしまいます。
そこで、考えられる解決策として、次のようなアプローチがあります。データのバッチを取得し、FP16でフォワードパスを実行し、損失を計算します。そして、損失を大きな値（例えば100や1000）でスケーリングします。これにより、勾配を大きな数値でスケーリングすることになり、この赤線の左側にあるすべてのものが右側にシフトすることになります。これにより、ゼロに丸められる値が少なくなることが期待できます。
その後、FP16で勾配を計算し、それをFP32にコピーして、スケーリング係数で割り、マスターウェイトを更新します。この方法で、先ほど述べた両方の問題を解決できます。
これが「混合精度トレーニング」と呼ばれるものであり、PyTorchでの実装は比較的簡単です。必要なのは、この「grad_scaler」オブジェクトをインスタンス化し、「AutoCast」のコンテキスト内でフォワードパスとバックワードパスを実行し、勾配をスケールダウンし、モデルパラメータを更新することです。
しかし、これはやや複雑に見えます。損失をスケーリングし、それを再度スケールダウンする必要があります。もし10,000倍してNaNになったらどうするか、次のイテレーションではスケーラーを更新し、1,000倍にする必要があるかもしれません。ネットワークのダイナミクスに合わせて調整する必要があります。
勾配スケーリングを行わない、より良い方法はないでしょうか？スケーリングが必要な理由は、FP16のダイナミックレンジがFP32と比較してはるかに小さいためです。そのため、FP16は非常に小さな数値を表現できません。
この問題をどう解決するかというと、精度を犠牲にすることです。これがBFloat16（Brain Float 16）のアイデアです。範囲を表すビット数を同じにし（8ビット）、つまりFP32と同じダイナミックレンジを持ちますが、精度が大幅に低くなっています。これはニューラルネットワークのトレーニングにとっては問題ないことが判明しています。
BFloat16を使用すると、勾配スケーラーを使用する必要がなくなり、モデルのフォワードパスとバックワードパスを適切なコンテキスト内でラップするだけで済みます。ただし、BFloat16はすべてのGPUで利用できるわけではなく、H100、A100、A6000などの最新のAmpereアーキテクチャが必要です。古いGPUでは、BFloat16を使用できない可能性があります。
これらの技術を使用した結果の例を見てみましょう。A100 GPUでの感情分類のファインチューニング実験では、64ビット浮動小数点（Float64）では約25分かかり、高い精度が得られますが、メモリも多く使用します。一方、BFloat16を使用した混合精度トレーニングでは、トレーニング時間が約3分の1に削減され、精度はほぼ同じ（実際には半精度表現の正則化効果により少し向上）、そしてメモリ使用量も大幅に削減されています。トレーニング時間が短縮される理由は、半精度での行列乗算が一般的に高速であるためです。
2.2. 勾配スケーリングによる解決法
FP16の問題に対する一つの解決策として、FP32とFP16の両方を使う方法があります。基本的なアイデアとしては、モデルのFP32コピーを「マスターウェイト」として維持します。データのバッチを取得し、FP32からFP16に変換してフォワードパスを実行します。そして損失を計算し、バックワードパスを実行してFP16で勾配を取得します。ここまでの処理はすべてFP16で行われています。
しかし、このままだとFP16の勾配の多くがゼロになってしまうという問題があります。そこで考えられる解決策が勾配スケーリングです。具体的には次のように行います：
データのバッチを取得し、FP16でフォワードパスを実行します。損失を計算したら、その損失を大きな値（例えば100や1000）でスケーリングします。このスケーリングされた損失を使って勾配を計算すると、勾配の値も大きくスケーリングされることになります。これにより、グラフの赤線の左側にあった小さな勾配値が右側に移動し、ゼロに丸められる勾配が減少します。
FP16で勾配を計算したら、それをFP32にコピーし、先ほどのスケーリング係数で割り戻します。そして、このスケールを戻した勾配を使ってマスターウェイトを更新します。この方法によって、精度の低下と小さな勾配値のゼロへの丸め込みという両方の問題を解決できます。
PyTorchでの実装は比較的簡単です。GradScalerオブジェクトをインスタンス化し、autocastコンテキスト内でフォワードパスとバックワードパスを実行します。そして勾配をスケールダウンし、モデルパラメータを更新します。
# PyTorch実装例（擬似コード）
scaler = GradScaler()
with autocast():
outputs = model(inputs)
loss = loss_fn(outputs, targets)
scaled_loss = scaler.scale(loss)
scaled_loss.backward()
scaler.unscale_(optimizer)
scaler.step(optimizer)
scaler.update()しかし、この方法はやや複雑です。損失を10,000倍にスケーリングしたらNaNになってしまう場合もあり、次のイテレーションではスケーラーを更新して1,000倍にするなど、ネットワークの動態に合わせて調整する必要があります。このような勾配スケーリングを行わずに済む、より良い方法が求められていました。
2.3. BFloat16の導入と利点
勾配スケーリングを行わない、より良い方法を考えてみましょう。スケーリングが必要になる根本的な理由を思い出してください。浮動小数点の構造において、緑色の部分（指数部）はデータ型のダイナミックレンジを表します。FP16はFP32に比べて範囲が大幅に小さいため、非常に小さな数値を表現できないという問題がありました。
この問題をどう解決するか？解決策は精度を犠牲にすることです。これがBFloat16（Brain Float 16）のアイデアです。BFloat16は、範囲（指数部）を表すビット数をFP32と同じ8ビットに保ち、代わりに精度（仮数部）を大幅に削減します。つまり、FP32と同じダイナミックレンジを持ちながら、精度は低くなっています。
興味深いことに、この精度の低下はニューラルネットワークのトレーニングにおいてはあまり問題にならないことが判明しています。BFloat16を使用すると、以前の混合精度トレーニングと比較して大きな利点があります。勾配スケーラーを使用する必要がなくなり、実装が非常にシンプルになるのです。
PyTorchでは、BFloat16を使用する場合、モデルのフォワードパスとバックワードパスを適切なautocastコンテキスト内でラップするだけで済みます：
# BFloat16の使用例（擬似コード）
with autocast(dtype=torch.bfloat16):
outputs = model(inputs)
loss = loss_fn(outputs, targets)
loss.backward()
optimizer.step()注意点として、BFloat16はすべてのGPUで利用できるわけではありません。最新のNVIDIA Ampereアーキテクチャ（H100、A100、A6000など）が必要です。古いGPUを使用している場合は、BFloat16を活用できない可能性があります。その場合は前述の勾配スケーリングを使用したFP16の混合精度トレーニングを使用する必要があります。
BFloat16の利点を示す具体的な結果として、A100 GPUを使用した感情分類のファインチューニングの例を見てみましょう。最も高精度の浮動小数点形式であるFloat64（64ビット）を使用した場合、トレーニングに約25分かかり、高い精度が得られましたが、メモリ使用量も多くなりました。
一方、BFloat16を使用した混合精度トレーニングでは、トレーニング時間が約3分の1に削減され、精度はほぼ同じ（実際には半精度表現による正則化効果で少し向上）、そしてメモリ使用量も大幅に削減されました。
トレーニング時間が短縮される理由は、半精度での行列乗算が一般的に高速であるためです。GPUのテンソルコアは低精度の計算に最適化されており、BFloat16はこれらのハードウェア加速を活用できます。
2.4. 混合精度トレーニングの性能比較
混合精度トレーニングの効果を示す実験結果について見てみましょう。単一のA100 GPUを使って感情分類のファインチューニングを行った結果があります。
一番上はFloat64（64ビット浮動小数点）で、トレーニングに約25分かかっています。非常に高い精度が得られますが、メモリ使用量も多くなっています。
一方、BFloat16を使用した混合精度トレーニングでは、トレーニング時間が約3分の1に削減されています。興味深いことに、精度はほぼ同じで、実際には半精度表現による正則化効果のためか、わずかに向上しています。さらに、メモリ使用量も大幅に削減されています。
トレーニング時間が短縮される主な理由は、半精度での行列乗算が高速であるためです。特に現代のGPUはこのような低精度計算に最適化されており、テンソルコアは半精度演算で大幅なパフォーマンス向上を提供します。
以下は、異なる浮動小数点形式による性能比較の概要です：
- Float64 (64ビット):
- トレーニング時間: 約25分
- メモリ使用量: 非常に大きい
- 精度: 高い
- Float32 (32ビット):
- トレーニング時間: Float64より短い
- メモリ使用量: Float64の半分
- 精度: 実用的に十分高い
- FP16 (勾配スケーリングあり):
- トレーニング時間: Float32より短い
- メモリ使用量: Float32の半分
- 実装の複雑さ: 勾配スケーリングが必要
- BFloat16:
- トレーニング時間: Float32の約1/3
- メモリ使用量: Float32の半分
- 実装の簡便さ: 勾配スケーリング不要
- ハードウェア要件: 最新のAmpereアーキテクチャが必要
この比較から明らかなように、BFloat16は精度、速度、メモリ効率、実装の簡便さのバランスが優れており、利用可能なハードウェアが対応している場合は優先して使用すべきです。古いハードウェアを使用している場合は、従来の勾配スケーリングによるFP16混合精度トレーニングが依然として有効な選択肢となります。
3. マルチGPUトレーニング
3.1. 分散データ並列（DDP）の基本
ここで話題を変えて、マルチGPUトレーニングについて考えていきましょう。複数のGPUを持っていて、それらすべてを使ってネットワークをトレーニングしたいという状況を想定します。
まずは基本から始めましょう。こちらに示しているのは、モデルとオプティマイザがデータセットからデータを受け取る様子を表した図です。これは多少簡略化した表現であり、後ほど修正点を説明します。
GPUのVRAM（ビデオメモリ）に何が保存されているかを考えてみましょう。まず、ニューラルネットワークのパラメータがあります。混合精度トレーニングを行っているとすると、これらはFP16で保存されています。次にオプティマイザがあります。
私がこれを初めて見たとき、オプティマイザもメモリを必要とすることに驚きました。しかし、Adamなどのオプティマイザを使用している場合、Adamのモーメント項と分散項を保存する必要があります。毎回勾配を取得するたびに、Adamのモーメントと分散を更新し、それを使ってパラメータを更新します。混合精度トレーニングを行っている場合、これらはFP32で表現される必要があります。
これが単一GPUの場合の図です。次に複数のGPUがある場合を考えてみましょう。この場合、まずデータセットを分割します。例えば4つのGPUがあるとすると、データセットを4つの部分に分割し、モデルの同期されたコピーを維持します。各モデルは自分のデータセットのスライスを受け取ります。
最初は全てのモデルが同期されていますが、フォワードパスを実行すると、各モデルは異なるデータポイントを受け取るため、異なる活性化値（アクティベーション）を持つことになります。その結果、各モデルは異なる勾配を持つことになります。
バックワードパスを実行した後、各モデルは異なるデータポイントに基づいて異なる勾配を持っています。ここで同期ステップを実行し、異なるワーカー間で勾配を通信します。
この講義で紹介する最初のMPI（Message Passing Interface）プリミティブは「all-reduce」演算です。all-reduceは、この例では4つの異なるGPU上の4つの情報を取り、それらを統合して、すべてのGPUに配布します。通信オーバーヘッドはパラメータあたり2バイトで、これはFP16勾配であるためです。
all-reduce操作によって勾配が通信された後、各オプティマイザは完全な勾配を持ち、モデルを更新して同期を維持できます。
これが基本的な「分散データ並列（Distributed Data Parallel、DDP）」と呼ばれる方法です。この方法は効果的ですが、メモリスケーリングが良くないという問題があります。
メモリ要件を計算してみましょう。モデルパラメータはFP16（混合精度トレーニングを行っているため）で、パラメータごとに2バイト必要です。勾配もFP16で、さらに2バイト必要です。そして緑色の部分、つまりオプティマイザの状態は、Adamを使用している場合、マスターウェイト（FP32）、モーメント（FP32）、分散（FP32）のために、パラメータごとに12バイト追加で必要になります。これらは全て各GPUに保存する必要があります。
この方法よりもメモリ効率の良い方法はないでしょうか？次のセクションでは、Zero Redundancy Optimizer（ZeRO）について説明し、さらにメモリ効率の良いアプローチを紹介します。
3.2. ZeRO (Zero Redundancy Optimizer)の段階的説明
メモリ効率を改善する方法として、ZeRO（Zero Redundancy Optimizer）と呼ばれる一連の技術を紹介します。これはMicrosoftがDeep Speedプロジェクトの一部としてリリースした技術です。基本的なアイデアは、すべてのGPUがすべての状態（青い部分、オレンジ色の部分、緑色の部分）を含む代わりに、それらをシャード（分割）するということです。つまり、すべてのGPUがすべてのパラメータやすべての勾配を持つのではなく、通信によって同期を取るというアプローチです。
ZeROには複数のステージがあります。ステージ1、2、3があり、それぞれのステージでシャードする対象が異なります。
ステージ1では、緑色の部分、つまりオプティマイザの状態をシャードします。同期を維持しながらシャードする方法は次のようになります。各GPUはFP16での完全なパラメータセットを持ち、各GPUは自分のデータに対する勾配を持ちます。しかし、完全なオプティマイザ状態のシャードされたコピーしか持たず、各GPUは自分のシャードに対応するパラメータの更新を担当します。
ステップバイステップで見ていくと、各GPUは自分のデータを持ち、そのデータのサブセットに対する勾配を取得します。次に「reduce-scatter」と呼ばれる操作を実行します。これが講義の2番目のMPI操作です。全てのGPUは自分のデータに対する完全な勾配を持っていて、例えばGPU 0からGPU 1が欲しい部分をGPU 1に通信したいとします。同様にGPU 2と3についても行います。つまり、完全な勾配から、他のワーカーが欲しい部分を該当するワーカーに通信するのです。これをreduce-scatterと呼びます。
全てのワーカーが自分のシャードに対応する勾配を受け取った後、それらのパラメータを更新し、all-gather操作を実行して同期を維持します。例えば、8つのパラメータを持つニューラルネットワークがあり、各GPUに2つのパラメータがあるとします。この過程の最後に、各GPUは自分のパラメータのサブセットを更新し、all-gatherで全てのGPUが更新されたパラメータの完全なセットを取得します。
この方法がなぜ従来のDDPよりも効率的なのかというと、オプティマイザの状態をシャードしているからです。例として、8つのパラメータを持つニューラルネットワークがあるとします。以前はすべてのGPUが8つのパラメータ全てに対するオプティマイザ状態を維持する必要がありましたが、今は各GPUは2つのパラメータに対するオプティマイザ状態だけを維持すればよいのです。
reduce-scatterが完了すると、シャードに対応する完全な勾配だけを持ち、部分的なオプティマイザ状態を使って、自分が担当するパラメータだけを更新します。その後、同期のためにパラメータを通信します。
重要なのは、これまで見てきた3つのMPI操作（all-gather、reduce-scatter、all-reduce）の関係です。実は、all-reduceはreduce-scatterとall-gatherを続けて実行するのと同等なのです。DDPでは全てall-reduce操作だけで済みました。ZeROでオプティマイザ状態をシャードする場合も、通信オーバーヘッドは全く同じです。なぜなら、all-reduceはreduce-scatterとall-gatherの組み合わせだからです。
つまり、基本的に無料でメモリを節約できるのです。常にこの手法を使うべきです。通信オーバーヘッドを増やすことなくメモリ節約ができます。
3.3. ZeRO Stage 1: オプティマイザ状態のシャーディング
ZeRO Stage 1では、オプティマイザの状態（緑色の部分）をシャーディングします。どのように機能するか詳細に見ていきましょう。
この段階では、各GPUは依然としてFP16形式での完全なモデルパラメータ（青色の部分）を保持しています。また、各GPUは自分のデータスライスに対する勾配（オレンジ色の部分）も完全に持っています。ただし、オプティマイザの状態（緑色の部分）については、シャーディングされた部分だけを持っています。
重要な点として、各GPUは自分のシャードに対応するパラメータの更新だけを担当します。これにより、メモリ使用量を削減しながら同期を維持できます。
実際の処理手順は次のようになります：
- 各GPUは自分のデータスライスを処理します。
- 各GPUが自分のデータサブセットに対する勾配を計算します。
- 「reduce-scatter」操作を実行します。これにより、各GPUの完全な勾配から、特定のGPUが担当するシャードに対応する部分だけを該当するGPUに送信します。
- 各ワーカーが自分のシャードに対応する勾配を受け取った後、自分が担当するパラメータのみを更新します。
- 最後に「all-gather」操作を実行し、更新されたパラメータを全てのGPUに配布して同期を維持します。
例として、8つのパラメータを持つニューラルネットワークと4つのGPUがある場合、従来のDDPでは各GPUが8つの全パラメータのオプティマイザ状態を保持する必要がありました。しかしZeRO Stage 1では、各GPUは2つのパラメータのオプティマイザ状態だけを保持すればよくなります。
重要な洞察は、DPPで使用されていた「all-reduce」操作が、実は「reduce-scatter」と「all-gather」の組み合わせと同等だということです。つまり：
all-reduce = reduce-scatter + all-gather
これにより、ZeRO Stage 1は追加の通信オーバーヘッドなしでメモリを節約できます。DDPと比較して通信量は同じですが、メモリ使用量が削減されるのです。
例えば、10億パラメータを持つモデルをトレーニングする場合、従来のDDPでは各GPUが全てのオプティマイザ状態（約12バイト/パラメータ）を保持する必要がありました。8つのGPUでZeRO Stage 1を使用する場合、各GPUは全パラメータの1/8だけのオプティマイザ状態を保持すればよいため、オプティマイザ状態に必要なメモリが8分の1になります。
このような利点があるため、マルチGPUトレーニングでは常にZeRO Stage 1を使用すべきです。追加コストなしでメモリ効率が向上するからです。
3.4. ZeRO Stage 2: 勾配のシャーディング
メモリを節約できたので、次はさらに多くの要素をシャードしていきましょう。ZeRO Stage 2では、緑色の部分（オプティマイザ状態）に加えて、勾配もシャードします。
これはより複雑になります。なぜなら、ワーカーは自分のデータスライスに対する完全な勾配をまだ必要としますが、各GPUは一度に小さなサブセットのパラメータに対する勾配を格納するメモリしか持っていません。この問題にどう対処するか？
解決策は、完全な勾配ベクトルを一度に全て確保することはせず、バックワードパスで勾配を受け取るたびに、受け取った勾配のパラメータに対して一時的にベクトルを確保し、勾配を計算し、適切なワーカーに送信した後、作成したメモリを解放するというアプローチです。
ステップバイステップで見ていきましょう。4つのワーカー（GPU）があり、各ワーカーがバックワードパスを実行します。バックワードパスはレイヤーごとに行われます（自動微分の講義を思い出してください）。損失から始まり、レイヤーごとに勾配を計算していきます。
例えば、レイヤーJにいるとします。上流からの勾配を受け取り、そのレイヤーのパラメータに対する勾配を計算します。勾配を計算したら、すぐにそのレイヤーを担当する適切なワーカーに送信します。つまり、レイヤーJを担当するワーカーが存在します。
レイヤーJに対する勾配を計算した各GPUは、その勾配を適切なワーカーに送信します。それが完了したら、作成したメモリを解放します。
これは4つ目のMPI操作ですが、reduce-scatterとそれほど違いはありません。これは単なる「reduce」です。4つのGPUが勾配を持ち、その勾配をそのレイヤーのメンテナンスを担当するワーカーに通信するだけです。
特定のレイヤーを担当するワーカーは、通信によって受け取った完全な勾配とオプティマイザ状態を使用して、そのパラメータシャードを更新します。最後に、すべてを同期するために、以前と同様にall-gather操作を実行する必要があります。
ZeRO Stage 1では、all-reduceがreduce-scatterとall-gatherの組み合わせと同等であるため、実質的に無料でメモリを節約できることを見ました。ここでも同様に、reduceとall-gatherを使用しています。これは実質的にもDDPと比較して追加の通信オーバーヘッドなしでメモリを節約できます。
これでZeRO Stage 1と2の両方について、DDPと比較して通信オーバーヘッドを増やすことなくメモリを節約できることがわかりました。次はさらにシャーディング範囲を広げたStage 3について見ていきます。
3.5. ZeRO Stage 3 (FSDP): モデルパラメータのシャーディング
さらにメモリを節約できるか検討してみましょう。今度はモデルパラメータ自体もシャードするとどうなるでしょうか。例えば、オプティマイザの状態を気にする前に、モデル自体が1台のGPUに収まらないような状況を考えてみましょう。その場合、モデルを異なるGPU間で分割する必要があります。つまり、青色で示したモデルパラメータをシャードするということです。
しかし、ここで重要な注意点があります。これまでのステージとは異なり、この方法ではメモリ節約を無料で得ることはできません。通信オーバーヘッドが発生します。これがZeRO Stage 3、別名FSDP（Fully Sharded Data Parallel）です。この用語を聞いたことがある方もいるかもしれません。
ここで高レベルな概念を説明しましょう。実は、Stage 1や2と比べて理解しやすいかもしれません。なぜなら、各ステップで通信が必要になるためです。まず最初に行うのは、モデル全体をFSDPユニットに変換することです。
こちらに例があります。シンプルなディープニューラルネットワークがあり、それを複数のFSDPユニット（ここでは3つ）に変換します。これはまだデータ構造に過ぎず、何も実質的な処理は行っていません。次にこのFSDPユニットを「flat parameter」と呼ばれる別のデータ構造に変換し、これらのパラメータのサブセットを各GPUに割り当てます。
この例では16台のGPUがあり、14個のパラメータ（プラス適切に分割するための追加パディング）からなるフラットパラメータがあります。各パラメータを異なるGPUに割り当てます。これは要するに、いくつかのデータ構造を作成し、モデルパラメータを各GPUに分割したということです。各GPUはモデルパラメータのサブセットしか持ちません。
ここでフォワードパスがどのようになるか考えてみましょう。どのGPUも完全なパラメータセットを持っていないため、通信が必要です。例えば、レイヤー4にいるとします。どのGPUもレイヤー4の全てを持っていないので、all-gather操作を実行してレイヤー4の全ての部分を集め、フォワードパスを実行します。フォワードパスが終わったら、レイヤー4は不要になるので、パラメータシャードを破棄します。
次にバックワードパスを実行する必要があります。損失を計算した後、バックワードパスを実行します。再びレイヤー4に戻ってきたとき、上流からの勾配はありますが、レイヤー4がないため、もう一度all-gatherを実行してレイヤー4の全てのパラメータを取得します。
各GPUは異なるデータポイントを持っているため、レイヤー4に対する勾配もGPUごとに異なります。all-gatherを実行してすべてのパラメータを取得し、勾配を計算します。各GPUは異なる勾配を持つため、reduce-scatterを実行して、完全な勾配をレイヤー4の担当部分を持つGPUに送信します。
こうして、各GPUはバックワードパスとフォワードパスを実行した後、受け取った完全な勾配を使って自分のパラメータシャードを更新します。その後、同期が行われます。
これまで見てきたすべてを簡単に振り返りましょう。DDPでは、シャーディングは行わず、すべてのGPUが完全なモデル、完全な勾配、完全なオプティマイザ状態を持ちます。分割されるのはデータセットだけです。大きなデータセット（例えば1000サンプル）があれば、各GPUは250サンプルを受け取ります。
フォワードパスとバックワードパスを計算すると、各GPUは異なる勾配を持ちます。その勾配を通信し、同期する必要があります。これがMPI用語でのall-reduce操作です。
次にZeROを見ました。ここではメモリを節約するため、すべてのGPUに完全なモデル、勾配、オプティマイザ状態のメモリ要件を持たせたくありません。ZeRO Stage 1では、オプティマイザ状態をシャードしました。これにより、各GPUが持つオプティマイザ状態が減少します。同期を維持するための通信オーバーヘッドはall-reduce（実質的にはreduce-scatterとall-gatherの組み合わせ）と同等になります。つまり、Stage 1と2では実質的に無料でメモリを節約できます。
ZeRO Stage 3（FSDP）では、複雑さが増します。モデルパラメータ、オプティマイザ状態、勾配をすべて分割する必要があります。フォワードパス実行中、任意のレイヤー（例えばレイヤー4）の完全なパラメータを取得するために通信が必要です。バックワードパスでも同様に、完全な勾配を取得するためのall-gatherと、適切なGPUに完全な勾配を送信するためのreduce-scatterが必要です。
全体として、これは2回のall-gatherとreduce-scatterという、Stage 1と2よりも多くのオーバーヘッドがあります。しかし、GPUのVRAMが不足していてモデルを1台のGPUに読み込めない場合は、これが唯一の選択肢になります。
3.6. MPI通信プリミティブ（All-Reduce、Reduce-Scatter、All-Gather）
マルチGPUトレーニングについて説明する中で、いくつかのMPI（Message Passing Interface）通信プリミティブについて触れてきました。ここでは、これらの重要な操作についてより詳細に説明します。
まず最初に紹介したのはAll-Reduce操作です。これは分散データ並列（DDP）の基本です。All-Reduce操作では、各GPUが自分のデータに基づいた勾配を計算した後、これらの勾配が集められ、平均化されるか合計され、結果が全てのGPUに配布されます。つまり、各GPUはグローバルな勾配情報を手に入れることができます。通信オーバーヘッドとしては、パラメータあたり2バイト（FP16勾配の場合）が必要です。
二つ目に紹介したのはReduce-Scatter操作です。これはZeRO Stage 1と2で使用されています。この操作では、各GPUが完全な勾配ベクトルを持っていて、その一部を担当するGPUに送信します。例えば、GPUが計算した勾配のうち、GPU 1が担当する部分をGPU 1に送り、GPU 2が担当する部分をGPU 2に送ります。これにより、各GPUは自分が担当するパラメータ部分に対する完全な（全データポイントからの）勾配を得ることができます。
三つ目に紹介したのはAll-Gather操作です。これは各GPUが保持している部分的な情報（例えばパラメータやその更新値）を収集し、全てのGPUが完全な情報を得られるようにします。ZeRO Stage 3（FSDP）では、フォワードパスとバックワードパスの両方でこの操作が必要です。各GPUはモデルの一部しか持っていないため、特定のレイヤーの計算を行う前に、そのレイヤーの全パラメータを集める必要があります。
これらの操作の間には重要な関係があります。All-Reduce操作は、Reduce-Scatterに続いてAll-Gatherを実行するのと同等です。この関係は、ZeRO Stage 1と2がなぜ追加の通信オーバーヘッドなしでメモリを節約できるのかを説明しています。DDPでは単一のAll-Reduce操作を使用していましたが、ZeROではReduce-ScatterとAll-Gatherに分解しています。しかし、通信量自体は変わらないのです。
FSDP（ZeRO Stage 3）では、モデルパラメータのシャーディングにより、フォワードパスとバックワードパスの両方で追加のAll-Gather操作が必要になります。これにより、Stage 1と2と比較して通信オーバーヘッドが増加しますが、モデルサイズが非常に大きい場合には必要な犠牲となります。
これらのMPI通信プリミティブを理解することは、分散トレーニングのパフォーマンスを最適化するために重要です。適切な操作を選択することで、メモリ使用量と通信オーバーヘッドのバランスを取ることができます。
3.7. GPUメモリ使用量の計算と最適化
ここで、以前に説明したGPUメモリ使用量の計算について、いくつか修正を加えたいと思います。実は、GPUのVRAM計算に関して、少し単純化した説明をしていました。
これまでは、モデルパラメータと勾配、そしてオプティマイザの状態についてのみ話してきましたが、もう一つ重要な要素があります。それは「モデルのアクティベーション（活性化値）」です。バッチサイズを増やしたいときに、あるポイントでGPUが「もう入らない」と言い始めるのは、このアクティベーションのためです。
バックワードパスでは、これらのモデルアクティベーションを保存する必要があり、それはバッチサイズに比例してメモリ使用量が増加します。バッチサイズが大きいほど、保存すべきモデルアクティベーションの数も多くなります。混合精度トレーニングを行っている場合でも、これはFP16かBFloat16で保存されますが、重要なのはバッチサイズに比例するという点です。
これまで見てきた技術（ZeRO Stage 1、2、3）はモデルアクティベーションのシャーディングには対応していないことに注意してください。
GPUメモリの計算をもう少し詳細に見てみましょう。混合精度トレーニングを行っている場合、モデルパラメータはFP16形式で保存され、パラメータあたり2バイト必要です。同様に、勾配もFP16で保存され、さらに2バイト必要です。
緑色の部分、つまりオプティマイザ状態については、Adamを使用している場合、マスターウェイト（FP32、4バイト）、モーメント（FP32、4バイト）、分散（FP32、4バイト）の合計でパラメータあたり12バイト追加で必要になります。
そして先ほど説明したように、モデルアクティベーションも考慮する必要があります。これはFP16/BFloat16で保存され、バッチサイズに比例します。具体的には、バッチサイズが大きいほど、より多くのモデルアクティベーションを保存する必要があります。
これらを合計すると、GPUメモリの使用量は：
- モデルパラメータ: 2バイト/パラメータ（FP16）
- 勾配: 2バイト/パラメータ（FP16）
- オプティマイザ状態: 12バイト/パラメータ（FP32）
- モデルアクティベーション: バッチサイズに比例
これらの要素を考慮することで、より正確にGPUメモリ要件を予測し、適切なバッチサイズや分散戦略を選択できます。例えば、バッチサイズを半分にすることで、アクティベーションに必要なメモリも半分になります。同様に、ZeRO Stage 1を使用することで、オプティマイザ状態のメモリ要件をGPU数で割ることができます（例：8GPUの場合は8分の1に削減）。
さらに、アクティベーションのメモリ使用量を削減するための手法として、「勾配チェックポインティング」や「アクティベーションチェックポインティング」があります。これらの技術では、全てのアクティベーションを保存する代わりに、一部だけを保存し、必要に応じて再計算します。計算時間とメモリ使用量のトレードオフになりますが、非常に大きなモデルをトレーニングする場合には有効な選択肢です。
最終的に、モデルサイズ、バッチサイズ、使用可能なGPU数に応じて、適切なシャーディング戦略を選択することが重要です。小さなモデルでは通常のDDPで十分かもしれませんが、モデルサイズが大きくなるにつれて、ZeRO Stage 1、2、そして最終的にはStage 3（FSDP）を検討する必要があります。
4. パラメータ効率的な微調整
4.1. なぜパラメータ効率的な微調整が必要か
マルチGPUトレーニングの基礎について学んだところで、今度はパラメータ効率的な微調整について話しましょう。パラメータ効率的な微調整とは何でしょうか？
完全な微調整（フルファインチューニング）では、フォワードパスとバックワードパスを実行し、モデルのすべてのパラメータを更新します。一方、パラメータ効率的な微調整では、パラメータの全体集合のうち、小さなサブセットだけを更新します。
なぜこのようなアプローチを取りたいのでしょうか。まず、たとえバッチサイズが1であっても、すべての可能なトリックを試しても、完全な微調整ができない状況があるかもしれません。その場合、パラメータ効率的な微調整が必要になるでしょう。
また、少し科学的な理由として、現代のモデルは大幅に過剰パラメータ化されており、小さなデータセットを持っている場合、パラメータ効率的な微調整を行うことで、より良い一般化性能が得られると考えられます。あるいは、完全な微調整に匹敵する性能が得られるかもしれません。
パラメータ効率的な適応を行いたい理由としては、他にも理由があります。右側のグラフは、赤線が最大のAIモデルをトレーニングするための計算量の推定増加を示し、青線がグローバルな計算能力を示しています。近い将来、グローバルな計算能力を超える計算量が必要になると予測されており、これは持続可能ではありません。
この道を進み続けると、AIの開発は少数の資金力のある組織にのみ集中することになり、学生である私たちはそれを行うことができなくなります。これは問題です。また、モデルのトレーニングと微調整を行うプレーヤーの数が少なければ、それらの組織は特定の方法でモデルにバイアスをかけ、それが彼らの価値観を反映する可能性があり、より広い一般大衆の価値観を反映しない可能性があります。
これはパラメータ効率的な適応を考慮する別の理由です。一般的に機械学習、特にNLPでは、効率性よりも精度に焦点を当てる傾向があります。右側のグラフは、主な貢献が単により正確なモデルを生成する方法である論文の割合と、同じ精度でより効率的なモデルを生成する方法の割合を示しています。ほとんどの会議では、論文の大多数が精度に関するもので、効率性に関する論文はわずかです。
これは一種の単一文化につながっている可能性があり、それが効率性に焦点を当てたい理由かもしれません。もう一つの、おそらくより大きな懸念は、大規模言語モデルのトレーニングと微調整には、膨大な環境コストが隠されているということです。
最近読んだレポートによれば、GPT-3のトレーニングコストは110万トンの炭素排出量に相当すると推定されており、これは石炭火力発電所を10時間連続で稼働させるのと同等とのことです。
より身近な例としては、強化学習のクラスでの出来事があります。宿題の一つで、多くの学生が一般的なアルゴリズムを実装しました。他のアルゴリズムより優れたアルゴリズムがいくつかありましたが、それらはより多くの電力を使用しました。ある計算によれば、もし全員がより効率的なアルゴリズムを使用していたら、クラス全体の電力消費が約880キロワット時削減できたとのことです。これはアメリカの一般家庭が1ヶ月で使用する電力量に相当します。
これらはすべて、効率性と、モデルを少ないリソースで微調整する方法について考える理由です。パラメータ効率的な微調整の概念に戻りましょう。
4.2. 計算資源の制約と環境への影響
大規模言語モデルのトレーニングと微調整には、膨大な計算資源が必要であり、それに伴う環境コストも無視できません。右側のグラフを見ると、赤線は最大のAIモデルをトレーニングするための計算量の推定増加を示し、青線はグローバルな計算能力を示しています。この図から明らかなように、非常に近い将来、AIモデルのトレーニングに必要な計算量がグローバルな計算能力を超えると予測されています。
これは明らかに持続不可能な状況です。計算リソースがこれほど莫大になると、AIの開発は少数の資金力のある組織にのみ集中することになり、学術機関や個人研究者、学生などはこの分野から実質的に排除されてしまいます。これは多様性の観点から大きな問題です。
また、モデルのトレーニングと微調整を行うプレーヤーの数が少ないと、彼らの特定の価値観や優先事項がモデルに反映され、より広い一般大衆の多様な価値観や需要を反映しない可能性があります。これはAI技術の公平性と包括性に関わる重要な懸念事項です。
さらに深刻なのは、大規模言語モデルのトレーニングによる環境への影響です。最近読んだレポートによれば、GPT-3のトレーニングコストは110万トンの炭素排出量に相当すると推定されています。これは石炭火力発電所を10時間連続で稼働させるのと同等のインパクトです。
より身近な例として、強化学習のクラスでの経験を共有します。あるホームワーク課題で、多くの学生が一般的なアルゴリズムを実装しました。その中で特に性能の良いアルゴリズムがいくつかありましたが、それらはより多くの電力を消費するものでした。誰かが計算したところによると、もし全員がより効率的なアルゴリズムを使用していたら、クラス全体の電力消費が約880キロワット時削減できたそうです。これはアメリカの一般家庭が1ヶ月で使用する電力量に相当します。
このように、AIモデルのトレーニングと微調整による環境への影響は無視できない規模に達しています。私たちが今、効率性に焦点を当て、同じ性能を少ないリソースで達成する方法を真剣に考えなければ、AIの進歩は環境的に持続不可能になるリスクがあります。
パラメータ効率的な微調整技術は、このような課題に対する一つの重要な解決策です。完全なモデルパラメータをすべて更新するのではなく、効率的に選ばれた少数のパラメータのみを更新することで、計算要件とそれに伴う環境コストを大幅に削減できます。
4.3. AIモデル開発の集中化問題
AIモデル開発が少数の組織に集中することの問題についてもう少し話しましょう。計算要件が増大し続ける現在の傾向は、AIの開発や微調整が、必要な計算リソースを持つ少数の企業や組織にのみ可能になるという状況を生み出しています。
この状況は、AIモデル開発の「モノカルチャー」を促進している可能性があります。右側のグラフが示すように、ほとんどの学術会議では、より正確なモデルを生成する方法に関する論文が圧倒的多数を占め、同じ精度でより効率的なモデルを生成する方法に関する論文はごくわずかです。ほとんどの会議において、論文の大多数が精度向上に焦点を当てており、効率性に関する研究は非常に少ないのです。
この「精度至上主義」は、計算資源を大量に必要とするアプローチを奨励し、結果としてAI開発の参入障壁を高めることになります。研究者がより効率的なアルゴリズムやパラメータ効率的な手法を模索するインセンティブが少なければ、イノベーションはますます少数の資金力のある大組織に集中してしまいます。
大規模モデルのトレーニングと微調整が少数のプレーヤーに限られると、これらの組織の価値観や優先事項が技術に埋め込まれることになります。例えば、特定の文化的文脈や倫理的視点、あるいは商業的利益が優先され、より多様な視点や価値観が反映されなくなる可能性があります。これは、AIが社会全体の多様なニーズと価値観を代表すべきだという理想に反します。
また、少数の組織がAI開発を独占すると、オープンな研究コミュニティの集合知や創造性が活用できなくなります。多様な研究者や開発者が参加できるエコシステムの方が、長期的には革新的で堅牢なソリューションを生み出す可能性が高いのです。
パラメータ効率的な微調整技術は、このような集中化の問題に対する一つの解決策です。少ないリソースでも高品質なモデルを開発・微調整できるようにすることで、より多くの研究者や組織がAI開発に参加できるようになります。これにより、多様な視点や価値観を反映したAIの発展が促進され、技術の進歩がより広範な社会的利益につながる可能性が高まります。
効率性に焦点を当てることは、単に環境的な持続可能性だけでなく、AIの発展における多様性と民主化にも貢献する重要な取り組みなのです。
4.4. LoRA（Low-Rank Adaptation）の仕組み
パラメータ効率的な微調整の方法はさまざまありますが、今日はLoRA（Low-Rank Adaptation）について説明します。LoRAは「低ランク適応」という意味で、大きな言語モデルを微調整する際に勾配が低い本質的なランクを持つという観察に基づいています。
ランクやSVD（特異値分解）を覚えていますか？基本的に、大規模言語モデルを微調整すると、そのパラメータの勾配は低い本質的なランクを持つ傾向があります。LoRAの開発者たちはこの特性を活かし、モデル内の各完全ランク行列に対して、はるかに小さいランクRの行列を微調整するという方法を考案しました。
具体的には、事前学習された重み行列W∈ℝ^(d×k)があるとします。通常の微調整では任意の更新を適用しますが、LoRAでは更新が特定の形を持つようにします。更新は2つの低ランク行列BとAの積になります。ここでAは(r×k)行列、Bは(d×r)行列で、rはランクであり、入力次元kや出力次元dよりもはるかに小さい値です。
さらに、αという項があり、これは事前学習モデルに既に格納されている知識と、モデルに追加したい新しいタスク固有の知識とのトレードオフを調整するパラメータと考えることができます。αが0の場合、何も変更しません。αが非常に小さい値の場合、モデルパラメータをあまり変更せず、非常に小さなタスク固有の知識を追加することを意味します。
ここで訓練可能なパラメータはAとBのみです。この方法の利点として、更新をU=B×Aという積として表現することで、rを増やすにつれて完全微調整に近づくという特性があります。つまり、どの程度の微調整を行いたいかをコントロールするスライダーのような役割をrが果たすのです。
また、推論レイテンシの面でも重要なポイントがあります。学習した行列を各タスクに対して保存しておき、別のタスクに切り替える際には、そのタスクに対して追加した項を取り除き、新しいタスク用のタスク固有の項を追加するだけでよいのです。これらのはるかに小さな行列を保存するコストも、完全なデルタを保存するよりもはるかに低くなります。
LoRAはどこに適用すべきかというと、一般的に自己注意機構（self-attention）の重み行列に適用するのが良いでしょう。
コード面では、実際にはかなりシンプルです。通常のフォワードパスでは、隠れ状態hを重み行列と入力特徴ベクトルの積として計算します。LoRAを使用する場合は、モデルパラメータを凍結し、以前と同様にhを計算し、そこに追加のオフセット項を加えるだけです。このオフセット項だけが訓練可能になります。これを各レイヤーの各重み行列に対して行う必要があります。
4.5. 低ランク行列による更新とアルファパラメータ
LoRAの具体的な数学的定式化を詳しく見ていきましょう。LoRAの基本的なアイデアは、元の重み行列Wに対する更新を低ランク行列として表現することです。
事前学習された重み行列W∈ℝ^(d×k)があるとします。LoRAでは、この行列を直接更新する代わりに、以下の形式で表現します：
W' = W + α(BA)
ここで：
- W'は更新された重み行列
- Wは元の事前学習された重み行列
- Bは(d×r)の行列
- Aは(r×k)の行列
- αはスケーリング係数（ハイパーパラメータ）
- rはランクであり、dやkよりも非常に小さい値（例：r = 8）
この式の中で訓練可能なパラメータはAとBのみです。元の重み行列Wはフリーズされ、変更されません。AとBが作る低ランク行列の積(BA)が、元の行列に対する「更新」として機能します。
αパラメータは非常に重要な役割を果たします。これは事前学習された知識と新しいタスク固有の知識のバランスを調整するスケーリングファクターです。具体的には：
- αが1の場合：事前学習された知識と新しいタスク知識の間で等しいトレードオフを意味します。これが一般的なデフォルト値です。
- αが1より大きい場合：モデルが事前学習で「知らない」タスクに対して使用します。タスク固有の知識をより強調します。
- αが1より小さい場合：モデルをあまり変更したくない場合に使用します。事前学習された知識をより保持します。
コード内では、αは2行目から1行目を引く際のスケーリング係数として現れます。通常は1に設定されますが、タスクの性質に基づいて調整することもできます。
ランクrの選択も重要なハイパーパラメータです。rを大きくすると、モデルの表現力が増しますが、パラメータ数も増加します。rが無限大に近づくと、理論的には完全な微調整と同等になります。実践では、r=8から始めて、必要に応じて調整するのが一般的です。
LoRAの重要な特性として、更新がBA製品として表現されるため、必要なパラメータ数が大幅に削減されることが挙げられます。例えば、元の行列が1000×1000（100万パラメータ）で、r=10の場合、LoRAでは(1000×10)+(10×1000)=20,000パラメータしか必要としません。これは元のパラメータ数の2%に過ぎません。
このように、LoRAは少数の訓練可能パラメータで効果的な適応を可能にし、メモリ効率が大幅に向上します。複数のタスクに対して微調整を行う場合、各タスクに対して小さなA、B行列のセットだけを保存すればよいため、ストレージ要件も大幅に削減されます。
4.6. 実装と適用方法
LoRAの実装は実際にはかなりシンプルです。以下にコードレベルでどのように実装されるかを説明します。
通常のフォワードパスでは、隠れ状態hを重み行列と入力特徴ベクトルの積として計算します：
h = W @ x # 通常の行列乗算LoRAを使用する場合は、モデルのパラメータを凍結し、通常通りにhを計算した後、そこに追加のオフセット項を加えます：
h = W @ x # 凍結されたモデルパラメータによる計算
h = h + alpha * (B @ (A @ x)) # LoRAによる追加項このオフセット項だけが訓練可能になります。ここでAlphaはハイパーパラメータで、通常は1に設定されます。これを各レイヤーの各重み行列に対して行う必要があります。
LoRAをモデルのどこに適用すべきかという問題については、一般的なルールとして、自己注意機構（self-attention）内の特定の重み行列に適用するのが最も効果的です。具体的には：
- 隠れ状態を「クエリ（Query）」に変換する行列
- 隠れ状態を「バリュー（Value）」に変換する行列
これらの行列にLoRAを適用することで、全体的に最高のパフォーマンスが得られることが実験的に示されています。
例えば、Transformerモデルでは、以下の行列にLoRAを適用します：
- W_q（クエリ変換行列）
- W_v（バリュー変換行列）
キー（Key）変換行列やFFNネットワークの行列にLoRAを適用することも可能ですが、実験結果によると、クエリとバリューの行列に絞って適用するだけで十分な性能が得られることが多いです。
LoRAのもう一つの重要なハイパーパラメータは最適なランクです。先ほど述べた2つの行列AとBは両方とも低ランクであり、実験結果では非常に小さいランク（r=8程度）でも高いパフォーマンスが得られることがわかっています。これは現代のモデルの隠れ層の次元（数百から数千）と比較するとはるかに小さい値です。
実装面での注意点として、LoRAを適用する際は通常、元のモデルパラメータを凍結（学習しないように設定）し、LoRAの行列AとBのみを訓練可能にします。これにより、メモリ効率が向上し、過剰適合のリスクも減少します。
また、推論時には、事前計算としてα(BA)を計算し、元の重み行列Wに加えることで、追加の計算オーバーヘッドなしで効率的に実行できます。異なるタスクに対しては、それぞれのタスク固有のAとB行列を保存しておき、必要に応じて切り替えることができます。
4.7. 他のパラメータ効率的手法との比較
パラメータ効率的な微調整には、LoRA以外にもさまざまな方法があります。全ての手法を挙げることはしませんが、アダプター、BitFit、Pチューニングなど多くの手法が存在します。
これらの異なる手法と比較すると、LoRAは多くの異なるタスクで高いパフォーマンスを示しています。比較的小さなモデルでのさまざまなタスクにおいても、LoRAは効果的であることが研究で示されています。
より大きなモデル、例えばGPT-3のような大規模モデルを微調整する場合の比較を見てみましょう。完全な微調整（フルファインチューニング）が最も上位にあり、次にBitFit（バイアス項のみを微調整する手法）があります。LoRAはアダプターと比較して、いくつかの重要な利点があります：
- 必要な追加パラメータがはるかに少ない
- 完全微調整と比較して精度のトレードオフが良好
- 場合によっては、モデルパラメータの小さなサブセットだけを微調整することによる正則化効果から、精度が向上することもある
LoRAをモデル内のどの部分に適用すべきかについての実験結果も興味深いものがあります。一般的には、隠れ状態をクエリに変換する行列と、隠れ状態をバリューに変換する行列にLoRAを適用するのが最も効果的です。これらの2つに適用するだけで、全体的に最高のパフォーマンスが得られることが多いです。
LoRAのもう一つの重要なハイパーパラメータは最適なランクです。実験によると、非常に小さいランク（r=8程度）でも高いパフォーマンスが達成できることがわかっています。これは現代のモデルの隠れ層の次元（数百から数千）と比較するとはるかに小さい値であり、パラメータ効率の高さを示しています。
LoRAとその他のパラメータ効率的な手法の主な違いは：
- パラメータ効率：LoRAは同じ性能を達成するために必要なパラメータ数が少ない
- 計算効率：推論時にオーバーヘッドを最小限に抑えられる
- 実装の簡便さ：既存のモデルアーキテクチャに対する変更が最小限
- 適応性：異なるタスク間での切り替えが容易
BitFitはさらにパラメータ効率が高いかもしれませんが（バイアス項のみを更新）、タスクによってはパフォーマンスが低下する可能性があります。アダプターはより多くの追加パラメータを必要とし、アーキテクチャの変更も必要です。Pチューニングは特定のタスクには有効ですが、汎用性に欠ける場合があります。
これらの比較から、LoRAはパラメータ効率と性能のバランスが優れた選択肢であり、特に計算リソースが限られている状況や、複数のタスクに対して効率的に微調整を行いたい場合に適しています。
5. 実践的な適用ガイドライン
5.1. フローチャートによる実践的なアプローチ
これまで、浮動小数点形式、マルチGPUトレーニング、そしてLoRAについて多くの基本概念を説明してきました。この全ての情報は、最終的には皆さんのプロジェクトで使えるシンプルなフローチャートにまとめることができます。もし講義中ずっと寝ていたとしても、今こそ目を覚まして、このフローチャートを見てください。
まず最初に、常に混合精度トレーニングを使用してください。性能（一般化能力やF1スコア、精度など）にほとんど影響を与えることなく、メモリと速度の両方で大きな利点があります。最新のAmperアーキテクチャ（H100、A100、A6000など）を使用している場合は、必ずBFloat16を使用してください。これは単純に優れています。torch.cuda.is_bf16_supportedコマンドで対応を確認できます。
次に、バッチサイズ1がGPU1台に収まるかどうかを自問してください。収まる場合は、より大きなバッチサイズを試してください。バッチサイズ1は小さすぎます。より大きなバッチサイズを試して、ZeRO Stage 2を使用してください。ZeRO Stage 2はほぼ無料で使えるので、ぜひ使ってバッチサイズを増やしてください。
バッチサイズ1でさえも収まらない場合は、ZeRO Stage 3があなたのメモリ不足問題を解決するかどうかを確認する必要があります。この段階ではモデルパラメータをシャードします。
これらすべては完全微調整（フルファインチューニング）、つまりモデルのすべてのパラメータを更新する文脈での話です。時にはこの質問の答えもノーになることがあります。例えば、A100を4台使っても、ZeRO Stage 3を試し、混合精度トレーニングを使用し、バッチサイズを1にし、おそらく勾配チェックポインティング（アクティベーションチェックポインティング）も試したけれど、何も上手くいかない場合があります。
そのような場合、完全微調整はできないので、パラメータ効率的な微調整を試す必要があります。これによりメモリ使用量を大幅に削減できます。
このフローチャートは、限られたリソースで大規模モデルを効率的に微調整するための実践的なガイドとなります。各ステップで問題が解決しない場合は、次のより高度な技術に進むことができます。最終的には、ほとんどの状況に対応できる解決策を見つけることができるでしょう。
5.2. 混合精度トレーニングの常時使用
最終プロジェクトで利用できる実践的なガイドラインの第一のポイントは、常に混合精度トレーニングを使用することです。これはもはや選択肢ではなく、基本的な前提と考えるべきです。
混合精度トレーニングを使用すると、性能（一般化能力やF1スコア、精度など）にほとんど影響を与えることなく、メモリ使用量と計算速度の両方で大きな利点が得られます。実際、半精度表現による正則化効果のため、わずかに性能が向上することさえあります。
特に、最新のAmpereアーキテクチャ（H100、A100、A6000など）のGPUを使用している場合は、必ずBFloat16を使用してください。BFloat16はFP16よりも優れており、勾配スケーリングなどの複雑な処理が不要です。BFloat16はFP32と同じダイナミックレンジを維持しながら、メモリ要件を半分に削減します。
あなたのGPUがBFloat16をサポートしているかどうかを確認するには、PyTorchで以下のコマンドを実行できます：
torch.cuda.is_bf16_supported()もし古いGPUを使用していてBFloat16がサポートされていない場合は、勾配スケーリングを使用したFP16での混合精度トレーニングを代わりに使用してください。これは少し複雑ですが、それでもFP32のみを使用するよりも大幅に効率的です。
混合精度トレーニングの利点は単にメモリ使用量を削減するだけではありません。半精度での行列乗算は一般的に高速であるため、トレーニング時間も短縮されます。A100 GPUでの感情分類実験では、混合精度トレーニングによりトレーニング時間が約3分の1に短縮されました。
最新のGPUはテンソルコアを通じて低精度計算に最適化されているため、混合精度トレーニングを使用しないことは、利用可能な計算能力の大部分を無駄にしているのと同じです。
このように、混合精度トレーニングは、ディープラーニングワークフローの標準的な部分と考えるべきであり、特に大規模モデルの微調整においては必須の技術です。常に使用することで、より大きなモデルをトレーニングし、より速くイテレーションを行い、より多くの実験を実行できるようになります。
5.3. ZeROステージの選択基準
混合精度トレーニングを使用することを決めた後、次の重要な決断はZeROのどのステージを使用するかです。これはバッチサイズとメモリ制約に基づいて決定します。
まず、バッチサイズ1が単一のGPUに収まるかどうかを確認します。これが最初のチェックポイントです。もしバッチサイズ1が収まるなら、より大きなバッチサイズを試してみてください。バッチサイズ1は通常小さすぎるため、より大きなバッチサイズでトレーニングすることで、勾配の推定がより安定し、収束が速くなる可能性があります。
バッチサイズを増やしながら、ZeRO Stage 2を常に使用することをお勧めします。ZeRO Stage 2は「無料で」使えるという重要な特性があります。これは、通信オーバーヘッドを増やすことなく、メモリ使用量を削減できることを意味します。これが可能な理由は、Stage 2で使用される通信パターン（reduce-scatterとall-gather）が、従来のDDPで使用されるall-reduce操作と等価だからです。
具体的には、ZeRO Stage 2では：
- オプティマイザの状態をシャード（分割）します（Stage 1の機能）
- 勾配もシャードします
これにより、特に大規模モデルでは大幅なメモリ節約が可能になります。オプティマイザの状態は、Adamなどを使用する場合、パラメータあたり約12バイトものメモリを使用します。これをシャードすることで、使用可能なGPU数に応じてこのメモリ要件を分散できます。
バッチサイズ1でさえGPUに収まらない場合は、ZeRO Stage 3（FSDP: Fully Sharded Data Parallel）を検討する時です。Stage 3では、前の2つのステージに加えて、モデルパラメータ自体もシャードします。これにより、非常に大きなモデルでも複数のGPUに分散できますが、通信オーバーヘッドが増加するというトレードオフがあります。
Stage 3でも十分なメモリが確保できない場合は、「勾配チェックポインティング」または「アクティベーションチェックポインティング」も試してみてください。これらの技術では、すべてのアクティベーションを保存する代わりに、一部のみを保存し、必要に応じて再計算します。これは計算時間とメモリ使用量のトレードオフになりますが、非常に大きなモデルでは必要なアプローチかもしれません。
ZeROステージの選択に関する簡潔なガイドライン：
- 常に混合精度トレーニングを使用する
- バッチサイズ1がGPUに収まる場合は、バッチサイズを増やしてZeRO Stage 2を使用する
- バッチサイズ1が収まらない場合は、ZeRO Stage 3を試す
- それでも不十分な場合は、勾配チェックポインティングを追加する
- すべての方法を試しても不十分な場合は、パラメータ効率的な微調整（例：LoRA）に移行する
5.4. LoRAのハイパーパラメータ設定
もしZeRO Stage 3を試してもバッチサイズ1でさえ収まらない状況に直面した場合、次の選択肢はLoRAです。LoRAを使用する際の主なハイパーパラメータは以下の3つです：
- アルファ値（α）：事前学習された知識と新しいタスク知識のバランスを調整するスケーリング係数です。アルファをどのように設定すべきか？ 1に設定することをお勧めします。これが良い出発点になります。モデルが特定のタスクに全く馴染みがないと思われる場合は、1より大きい値を試すこともできますが、通常は1から始めるのが最適です。
- ランク（r）：低ランク行列の次元を決定します。ランクをどのように設定すべきか？ ランクを8に設定することをお勧めします。実験によれば、非常に小さなランク（r=8程度）でも高いパフォーマンスが得られることがわかっています。これは現代のモデルの隠れ層の次元（数百から数千）と比較するとはるかに小さい値です。
- 適用する重み行列：モデルのどの部分にLoRAを適用すべきか？ クエリ行列とバリュー行列に適用することをお勧めします。具体的には、隠れ状態をクエリに変換する行列と、隠れ状態をバリューに変換する行列にLoRAを適用すると、全体的に最高のパフォーマンスが得られることが多いです。
これらのハイパーパラメータを使用して、以下のようなシンプルな初期設定で始めることができます：
# LoRAの基本設定例
config = {
'r': 8, # ランク
'alpha': 1, # スケーリング係数
'target_modules': ['q_proj', 'v_proj'] # クエリとバリュー行列に適用
}これが良い出発点となり、モデルが合理的に良い性能を発揮するはずです。必要に応じて、これらの値を調整することができます。例えば、パフォーマンスが不十分な場合はランクを16に増やすことを検討するか、より多くの層や行列（例えば、キー行列やFFN層も）にLoRAを適用することを試みることができます。
しかし、多くの場合、この基本設定（r=8、α=1、クエリとバリュー行列に適用）から始めて、そこから調整するのが最も効率的なアプローチです。これらの設定で、完全微調整の性能に近い結果を得ながら、メモリ要件を大幅に削減することができます。
これらのガイドラインに従うことで、限られた計算リソースでも大規模モデルを効果的に微調整することが可能になります。モデルを適切に微調整し、良好な結果を得るために、まず混合精度トレーニングを使用し、必要に応じてZeROの適切なステージを選択し、それでも不十分な場合はLoRAに移行するという段階的なアプローチを取ることをお勧めします。
