複数GPUを用いたAI分散学習におけるモデル並列とメモリ分割

GPU増設は最終手段。モデル並列とメモリ分割で挑む、LLM分散学習の最適解とアーキテクチャ設計論

この記事は急速に進化する技術について解説しています。最新情報は公式ドキュメントをご確認ください。

約19分で読めます
文字サイズ:
GPU増設は最終手段。モデル並列とメモリ分割で挑む、LLM分散学習の最適解とアーキテクチャ設計論
目次

この記事の要点

  • 大規模AIモデル(LLM等)のGPUメモリ要件を最適化
  • モデル並列とメモリ分割(FSDP, ZeRO)による分散学習手法
  • ハードウェア増設に依存しない効率的な学習環境構築

はじめに:その「メモリ不足」、GPUを買わずに解決できるかもしれません

「CUDA out of memory」

この冷酷なエラーメッセージを前に、何度ため息をついたことでしょうか。限られた予算とリソースの中で巨大なモデルを動かそうと必死になり、深夜までログと格闘するという課題は、多くの開発現場で決して珍しくありません。多くのエンジニアやプロジェクトマネージャーは、このエラーを見ると反射的に「GPUのVRAMが足りない、もっと上位のGPUが必要だ」あるいは「台数を増やして分散させよう」と考えがちです。

近年では、VRAM容量が16GBから32GBクラスに達する最新のハイエンドGPUの導入を検討したり、最新のCUDA 13系環境やNGCコンテナを活用して環境を刷新することで解決を試みるケースも増えています。また、最新のGPUアーキテクチャではFP4量子化などによりVRAM消費を大幅に抑制する技術も発表されています。一方で、古い世代のGPUは最新のCUDA環境で非サポートとなりつつあるため、ハードウェアのライフサイクルに依存した解決策には限界があります。

あえて言わせてください。GPUの増設や最新ハードウェアへの依存は、最初の解決策ではなく「最終手段」であるべきです。

なぜなら、現代のLLM(大規模言語モデル)学習におけるメモリ消費のメカニズムは複雑怪奇であり、単にハードウェアをスケールアウトさせるだけでは、通信ボトルネックという別の壁に激突するだけだからです。適切な分散戦略を持たずにGPUを並べても、期待したほどの速度向上は得られず、投資対効果(ROI)は悪化する一方です。モデルの巨大化は、ハードウェアの進化を上回るスピードで進み続けています。経営者視点で見れば、無尽蔵なハードウェア投資は避けるべきであり、エンジニア視点で見れば、ソフトウェアアーキテクチャの工夫で乗り越えるべき壁です。

LLM開発の最前線では、「どの分散学習ライブラリを使えばいいのかわからない」「DeepSpeedとFSDP、結局どちらの環境に合っているのか」という悩みが頻繁に報告されています。Tensor Parallelism(TP)、Pipeline Parallelism(PP)、ZeRO(Zero Redundancy Optimizer)など、選択肢はあまりに多く、それぞれの特性を深く理解していなければ最適な組み合わせを見つけることは困難です。

この記事では、既存の表面的な手順書や難解な数式解説とは一線を画し、「リソース制約のある現場で、どの技術スタックを採用すべきか」という意思決定のプロセスに焦点を当てます。物理的なメモリ消費の内訳から、通信オーバーヘッドとのトレードオフ、そして具体的なシナリオ別の構成案まで、システム思考のアプローチで紐解いていきます。

開発現場が直面している「メモリの壁」を、技術と戦略で乗り越えるための羅針盤として、本記事を活用してください。

1. 「メモリの壁」の正体:GPUメモリはなぜ枯渇するのか

分散学習の技術論に入る前に、まず敵を知る必要があります。GPUメモリ(VRAM)は何に食いつぶされているのでしょうか? 「モデルが大きいから」というのは半分正解で、半分間違いです。実は、学習時にはモデル本体の何倍ものメモリが、「学習を成立させるための付帯情報」によって消費されています。

モデルパラメータだけではないメモリ消費の内訳

一般的に、ディープラーニングモデルのパラメータ(重み)は、FP16(16ビット浮動小数点)や、現在のLLM開発において標準となっているBF16(BFloat16)で保存される場合、1パラメータあたり2バイトを消費します。つまり、70億パラメータ(7B)のモデルであれば、重みだけで約14GBのVRAMが必要です。「お、24GBや48GBのGPUなら余裕で入るじゃないか」と思いますよね。ここに罠があります。

学習プロセス、特にAdamのような適応的最適化アルゴリズムを使用する場合、以下の要素がメモリを圧迫します。

  • モデルパラメータ(Weights): モデルそのものの重み。
  • 勾配(Gradients): 逆伝播(Backpropagation)時に計算される、各パラメータの更新量。通常、パラメータと同じサイズが必要です。
  • オプティマイザ状態(Optimizer States): これが最大の伏兵です。標準的なAdamオプティマイザは、精度維持のためにパラメータごとに「モーメンタム(勢い)」と「分散(二乗平均)」をFP32(32ビット単精度)で保持することが一般的です。これはパラメータ数の数倍のメモリを消費します。
  • アクティベーション(Activations): 順伝播(Forward pass)時に計算された各層の出力結果。逆伝播で勾配を計算するために一時的に保持しておく必要があります。

具体的に、標準的な混合精度学習(Mixed Precision Training)を行う場合のメモリ消費を見てみましょう。7Bモデルを例にします。

  1. モデルパラメータ(FP16/BF16): 14GB
  2. 勾配(FP16/BF16): 14GB
  3. オプティマイザ状態(FP32): パラメータのマスターコピー(FP32)、モーメンタム(FP32)、分散(FP32)で、合計12バイト/パラメータ。つまり、7B × 12bytes = 84GB

これだけで合計112GBに達します。さらにアクティベーションや一時バッファが加わります。つまり、7Bクラスのモデル学習には、80GBのメモリを持つNVIDIA A100や、その後継であるH100/B100といったハイエンドGPUであっても、単体ではメモリ容量が不足するというのが現実なのです。推論だけなら動作するモデルが、学習となると途端に牙をむく理由はここにあります。

データ並列(DDP)の限界点

従来の分散学習の主流であった「データ並列(Data Parallelism: DDP)」は、モデルの複製(レプリカ)を全GPUに持たせ、学習データを分割して投入する手法です。計算が終わるたびに、各GPUで計算された勾配を平均化(All-Reduce)して同期します。

この手法はシンプルで実装も容易ですが、「各GPUがモデル全体とオプティマイザ状態のコピーを保持しなければならない」という致命的な制約があります。先ほどの計算通り、モデルが巨大化すると1枚のGPUに収まりきらなくなり、そもそも学習を開始することすらできません。

「メモリ不足だからGPUを増やしてDDPしよう」としても、個々のGPUに必要なメモリ量は変わらないため、OOM(Out Of Memory)は解消しないのです。ここで必要になるのが、モデルや状態そのものを分割するアプローチです。

OOM(Out Of Memory)が発生するメカニズム

OOMは、メモリ割り当て(Allocation)が限界に達した時だけでなく、メモリの断片化(Fragmentation)によっても発生します。PyTorchなどのフレームワークは、高速化のためにメモリをキャッシュしますが、可変長のシーケンスを扱うLLMでは、確保と解放が頻繁に行われ、大きな連続領域が確保できなくなることがあります。

また、バッチサイズコンテキスト長(Sequence Length)も重要な変数です。バッチサイズを上げれば学習は安定し高速化しますが、アクティベーションメモリは比例して増大します。特に近年は、RAG(検索拡張生成)や長文要約の需要から、数万〜数十万トークンを扱う長文脈(Long Context)対応が進んでおり、アクティベーションがメモリ消費の主役になるケースも増えています。

このように、メモリ消費の構造を定量的に理解することで、初めて「どこを削り、どこを分割すべきか」という戦略が見えてきます。最新のBlackwellアーキテクチャなどではFP4などの低精度演算によるメモリ削減も提案されていますが、まずは基本となるこのメモリ消費構造を理解することが、アーキテクチャ設計の第一歩です。

2. アプローチの比較:モデル並列 vs メモリ分割

2. アプローチの比較:モデル並列 vs メモリ分割 - Section Image

DDPの限界を突破するために、エンジニアたちは様々な「分割」の手法を編み出してきました。大きく分けると、「計算を分割する(モデル並列)」アプローチと、「データを分割しつつメモリの冗長性を排除する(ZeRO/FSDP)」アプローチがあります。これらは排他的なものではなく、組み合わせることも可能ですが、まずはそれぞれの特性を理解しましょう。

モデル並列(Tensor Parallelism / Pipeline Parallelism)のアーキテクチャ

モデル並列は、巨大なモデルそのものを複数のGPUに切り分ける技術です。

Tensor Parallelism (TP)
これは、1つの層(レイヤー)の中にある巨大な行列演算を分割します。例えば、行列Aと行列Bの積を計算する際、行列を行や列で分割し、複数のGPUで分担して計算します。Megatron-LMなどが採用している手法です。

  • メリット: 巨大な層を分割できるため、単体GPUに乗らない巨大モデルも扱える。
  • デメリット: 計算のたびにGPU間で頻繁な通信(All-Reduce)が発生します。そのため、GPU同士がNVLinkのような超高速インターコネクトで接続されていることが必須条件となります。Ethernet接続の別ノード間でTPを行うと、通信待ちで計算が止まり、極端に遅くなります。

Pipeline Parallelism (PP)
こちらは、層(レイヤー)単位でモデルを分割します。GPU-1が1〜10層、GPU-2が11〜20層を担当するイメージです。データはバケツリレーのようにGPU間を流れていきます。

  • メリット: GPU間の通信量はTPに比べて少ない(境界のデータのみ送れば良い)。
  • デメリット: 「パイプラインバブル」と呼ばれるアイドル時間が発生します。前のGPUが計算している間、次のGPUはデータ待ちで暇をしてしまうのです。これを防ぐためにマイクロバッチなどの工夫が必要ですが、実装やチューニングの難易度は高くなります。

メモリ分割(ZeRO / FSDP)のアーキテクチャ

一方、Microsoftが開発したZeRO (Zero Redundancy Optimizer) や、それをPyTorchに統合したFSDP (Fully Sharded Data Parallel) は、データ並列の進化形と言えます。

DDPでは全GPUが同じモデルとオプティマイザ状態を持っていましたが、ZeRO/FSDPではこれらを「シャード(断片)」として各GPUに分散保持します。計算が必要な瞬間だけ、他のGPUから必要なパラメータを取り寄せ(All-Gather)、計算が終わったら破棄します。

  • ZeRO Stage 1: オプティマイザ状態のみを分割(メモリ削減効果:4倍)
  • ZeRO Stage 2: オプティマイザ状態 + 勾配を分割(メモリ削減効果:8倍)
  • ZeRO Stage 3: オプティマイザ状態 + 勾配 + パラメータ全てを分割(メモリ削減効果:GPU数に比例してリニアに削減)

このアプローチの革新的な点は、「GPUを増やせば増やすほど、1台あたりのメモリ負担が減る」というスケーラビリティを実現したことです。これにより、数十億〜数千億パラメータのモデル学習への道が開かれました。

通信オーバーヘッドと計算効率のトレードオフ

ここで重要なトレードオフが発生します。「メモリ効率」と「通信コスト」はトレードオフの関係にあります。

ZeRO Stage 3やFSDPのFull Shardingは、メモリ効率は最強ですが、計算のたびにパラメータを全GPU間で通信するため、通信量が莫大になります。ネットワーク帯域が狭い環境でこれを行うと、GPUの使用率(Utilization)が低下し、学習が進まないという事態に陥ります。

一方、TPは通信頻度は高いものの、データ量は比較的小さい傾向にあります。PPは通信頻度は低いですが、計算リソースの遊び(バブル)が生じます。

「どの手法が優れているか」ではなく、「自社のインフラ環境(特にネットワーク帯域)において、どのボトルネックが許容できるか」で選ぶ必要があるのです。

3. 統合アーキテクチャとデータフローの可視化

3. 統合アーキテクチャとデータフローの可視化 - Section Image

複数のGPUが連携して動くとき、裏側ではどのようなデータフローが発生しているのでしょうか。ここをイメージできるかどうかが、トラブルシューティングの早さを決めます。

GPU間通信(NCCL)の役割とボトルネック

NVIDIAのGPU分散学習では、NCCL(NVIDIA Collective Communications Library)という通信ライブラリが黒子として働いています。ここで頻出する通信パターンを理解しておきましょう。

  • All-Reduce: 全GPUのデータを集計し、その結果を全GPUに配る(DDPの勾配同期で使用)。
  • All-Gather: 各GPUが持っている断片データを集めて、全GPUで完全なデータを作る(FSDPでパラメータを復元する際に使用)。
  • Reduce-Scatter: 全GPUのデータを集計し、結果を分割して各GPUに配る(FSDPで勾配を同期・分割する際に使用)。

順伝播・逆伝播時のデータ移動

FSDP(ZeRO-3)を使用している時の動きを追ってみましょう。

  1. 順伝播(Forward)開始: 各GPUは自分の担当する層のパラメータしか持っていません(例:GPU-0は一部の層の1/Nの断片しか持っていない)。
  2. All-Gather: 計算に必要な層のパラメータを、他の全GPUからかき集めます。一瞬だけ、完全な層のパラメータがメモリ上に再構築されます。
  3. 計算: 計算を実行し、アクティベーションを保持します。
  4. 破棄: 次の層の計算に移るため、集めたパラメータは即座にメモリから削除されます。
  5. 逆伝播(Backward): 同様にパラメータを再集結(All-Gather)させて勾配を計算します。
  6. Reduce-Scatter: 計算された勾配を平均化しつつ、各GPUの担当分だけを残して分割保存します。

このように、計算の裏で常にデータが飛び交っています。まるでジャグリングのように、必要なボール(パラメータ)を空中に投げ合いながら処理を進めているのです。

インターコネクト帯域(NVLink vs Ethernet)の影響

この「ジャグリング」をスムーズに行うためには、手の動き(通信速度)が重要です。

  • NVLink / NVSwitch: サーバー内部のGPU間をつなぐ高速道路。帯域幅は数百GB/s〜900GB/sに達します。TPやFSDPのような頻繁な通信を行う場合、この帯域が必須です。
  • Ethernet / InfiniBand: サーバー間(ノード間)をつなぐ道路。100Gbps〜400Gbps(12.5GB/s〜50GB/s)程度が一般的です。NVLinkに比べると桁違いに遅いため、ここをまたぐ通信は極力減らす必要があります。

したがって、「TPはノード内(NVLink)に留め、ノード間はDDPやFSDPでつなぐ」というハイブリッド構成が、大規模クラスタにおける定石となります。

4. シナリオ別・最適構成選定ガイド

4. シナリオ別・最適構成選定ガイド - Section Image 3

理論はわかりましたが、現場では「で、ウチの場合はどうすればいいの?」という問いに答えなければなりません。ここでは、よくある3つのシナリオに基づいて、推奨される構成を提示します。

【ケースA】7B-13Bモデル × コンシューマーGPU(24GB VRAM)

状況: 小規模な開発環境や研究開発部門。RTX 3090/4090(24GB)を1〜2枚、あるいは4枚搭載したワークステーションを使用。

  • 課題: 24GBでは7Bモデルのフルパラメータ学習は不可能です(前述の通り100GB以上必要)。
  • 推奨構成: QLoRA + FSDP (ZeRO-2)
    • QLoRA (Quantized LoRA): モデルを4-bitで量子化し、学習可能なパラメータをLoRAアダプタのみに限定します。これにより、メモリ消費を劇的に(数分の一に)削減できます。
    • FSDP / ZeRO-2: 複数枚あるなら、ZeRO-2を使ってオプティマイザ状態を分散させます。ZeRO-3は通信オーバーヘッド(PCIe経由)が大きいため、コンシューマー機では逆に遅くなる可能性があります。
    • 判断基準: フルファインチューニングにこだわらず、LoRAで十分な精度が出るならこれが最強のコスパです。まずは動くプロトタイプを作り、仮説検証を回すには最適な構成と言えます。

【ケースB】70Bクラス × A100/H100(80GB)複数台

状況: 大規模なAIプロジェクト。クラウド上のA100 80GBインスタンスを4枚〜8枚(1ノード)確保できる。

  • 課題: 70Bモデルは巨大ですが、A100 80GBがあれば選択肢が広がります。
  • 推奨構成: FSDP (ZeRO-3) + CPU Offloading (Optional)
    • FSDP (ZeRO-3): 8枚のGPUがあれば、70Bモデルをパラメータ分割して載せることが可能です。NVLinkがあるため、通信コストも許容範囲内です。
    • CPU Offloading: もし80GBでも足りない(バッチサイズを大きくしたい)場合、オプティマイザ状態をCPUメモリ(メインメモリ)に退避させる「オフローディング」を使います。計算速度は落ちますが、OOMを回避する最後の砦として機能します。
    • Flash Attention 2: 必須です。これを使うことで、学習速度とメモリ効率の両方を向上させられます。

【ケースC】超巨大モデル(100B+) × マルチノードクラスタ

状況: 基盤モデル構築や大規模な継続事前学習。数十台〜数百台のGPUを使用。

  • 課題: 単一ノードには収まらず、ノード間通信がボトルネックになりやすい。
  • 推奨構成: 3D Parallelism (TP + PP + DP)
    • TP (Tensor Parallelism): ノード内のGPU(例:4〜8枚)でTPを行い、巨大な層を分割します。
    • PP (Pipeline Parallelism): モデルがさらに巨大な場合、ノードをまたいでパイプラインを組みます。
    • DP (Data Parallelism): 残りの次元でデータ並列を行います。
    • Megatron-DeepSpeed: このレベルになると、PyTorch標準機能だけでは管理が難しく、NVIDIAのMegatron-LMやMicrosoftのDeepSpeedをフル活用した複雑なスタックが必要になります。

3D Parallelismの適用条件

3D Parallelismは強力ですが、構築と運用のコストは跳ね上がります。「モデルが1つのノード(8枚のGPU)に収まるかどうか」が境界線です。収まるならFSDP単体の方が圧倒的に楽です。収まらない規模になって初めて、TPやPPの導入を検討してください。

5. 実装に向けたエコシステムとツールチェーン

最後に、これらの理論をコードに落とし込むためのツールについて触れておきます。2025年現在、主要なプレイヤーは絞られてきています。

PyTorch Native (FSDP) vs DeepSpeed vs Accelerate

  1. PyTorch FSDP: PyTorchの標準機能として組み込まれています。外部ライブラリへの依存がなく、安定性が高いのが特徴です。最近のアップデートで使い勝手が向上しており、第一選択肢として推奨します。
  2. DeepSpeed: Microsoft製。非常に多機能で、ZeROの生みの親だけあって最先端の機能(ZeRO-Offload, ZeRO-Infinity)が使えます。ただし、設定項目が膨大で(ds_config.json)、エラー時のデバッグが難しいのが難点です。極限のパフォーマンスを求めるならこちら。
  3. Hugging Face Accelerate: 上記2つをラップして抽象化したライブラリです。コードをほとんど書き換えずにFSDPやDeepSpeedを切り替えられるため、開発効率は最強です。まずは動くものを作って検証するプロトタイプ思考のアプローチには、Accelerate経由での実装が非常に適しています。

既存コードへの統合難易度と移植性

「とりあえず動くコード」から「分散学習対応コード」への書き換えは、意外と骨が折れます。

  • Accelerateを使う場合: 数行の変更で済みます。accelerator = Accelerator() を宣言し、model, optimizer = accelerator.prepare(...) でラップするだけです。
  • FSDPを直接使う場合: モデルのラップ方法(Sharding Strategy)を明示的に記述する必要があります。特に、カスタムモデルを使っている場合は、どの層をラップするか(Wrapping Policy)の指定にノウハウが必要です。

デバッグとパフォーマンスチューニングの勘所

分散学習のデバッグは地獄です。エラーログは複数のプロセスから非同期に出力され、真の原因が埋もれてしまいます。

  • コツ1: まずは単一GPU、小規模モデルでロジックが正しいか確認する。
  • コツ2: TORCH_DISTRIBUTED_DEBUG=DETAIL などの環境変数を活用する。
  • コツ3: プロファイラ(PyTorch Profiler)を使って、通信待ち(Communication Wait)が発生していないか確認する。もし通信待ちが長いなら、バッチサイズを大きくして計算密度(Arithmetic Intensity)を上げるか、通信圧縮技術を検討します。

まとめ:アーキテクチャ設計がAIプロジェクトの成否を分ける

「GPUメモリ不足」は、単なるハードウェアの問題ではなく、ソフトウェアアーキテクチャの問題です。モデル並列、ZeRO/FSDP、そして量子化技術。これらを適切に組み合わせることで、手持ちのリソースでも驚くほど大規模なモデルを学習させることが可能になります。

しかし、最適な組み合わせは、モデルの構造、データ量、ネットワーク環境、そしてビジネス上の制約(予算・納期)によって千差万別です。誤った構成を選べば、高価なGPUクラスタがただの暖房器具と化してしまいます。

大規模な学習基盤の設計で迷っている場合や、現在の学習速度に不満がある場合は、専門家に相談して現状のボトルネックを診断し、コストパフォーマンスを最大化するアーキテクチャを検討することをおすすめします。GPUへの投資を決断する前に、まずは「知恵」への投資を検討してみてはいかがでしょうか。

AIプロジェクトが、メモリの壁を越え、次のステージへと進むことを応援しています。

GPU増設は最終手段。モデル並列とメモリ分割で挑む、LLM分散学習の最適解とアーキテクチャ設計論 - Conclusion Image

コメント

コメントは1週間で消えます
コメントを読み込み中...