複数のGPUにわたるLLMの分割：テクニック、ツール、およびベストプラクティス

紹介

LLM（大規模言語モデル）は、会話型AIシステム、テキスト生成アプリケーション、および機械翻訳プロセスを含む高度なNLPタスクにおいて基本的な役割を果たします。数百億のパラメータを持つモデルは卓越した性能を発揮しますが、大量のGPUリソースを必要とします。従来の単一GPUシステムでは、大規模な言語モデルをロードまたはトレーニングする際にメモリがすぐに不足することがあります。LLMを複数のGPUに分割することは、運用効率を維持するために不可欠になっています。このガイドでは、複数のGPUにわたるLLMの分割とロードに関する包括的な知見を提供し、GPUメモリ制約に対処し、モデル推論速度を改善する方法を説明します。モデルとデータの並列処理を含む並列処理手法を取り上げ、分散LLMトレーニングへの適用方法を解説しました。

前提条件

• GPT、BERT、T5などのNLPモデルを理解し、それらが自然言語処理タスクにどのように応用されるかを知ること。
• 人気のあるディープラーニングフレームワークを使用したPythonスクリプトの開発と実行に関する熟練度。
• CUDAテクノロジーによるGPUアクセラレーションの理解と、深層学習がGPU計算から得る利点。
• トレーニングプロセスとモデル評価方法、その予測タスクへの応用に関する知識。
• データ並列処理、モデル並列処理、および分散トレーニングの原則に関する基礎知識。

なぜ大規模言語モデルを複数のGPUに分割するのか？

PaLMのような現代の大規模言語モデルやMegatronに類似したネットワークは数十億のパラメータを含んでいます。単一のGPUでは、12GB、24GB、または80GBのような利用可能なVRAMオプションでは、すべてのモデルパラメータ、アクティベーション、およびオプティマイザ状態を読み込むには不十分なことが多いため、巨大なモデルをホストできない場合があります。
LLMを複数のGPUに分割することで、ボトルネックを2つの明確な方法で解決します:

• メモリのスケーラビリティ: モデルパラメータを複数のGPUに分散させることで、アウトオブメモリー（OOM）エラーが発生する可能性が減ります。巨大なモデルは、トレーニングおよび推論の段階でこのアプローチを必要とします。
• パフォーマンス向上: 計算が並行して実行されると、トレーニングと推論の速度が大幅に向上します。適切に実装されれば、このプロセスはマルチGPUによる推論とトレーニングを加速します。

LLMを分割することは、高度なAIタスクにとって基本的な習慣であり、単一のマシンでマルチGPUディープラーニングを行う場合でも、複数のサーバーでLLMの分散トレーニングを行う場合でも同様です。

モデルの並列処理とデータの並列処理

LLMを用いて複数のGPUを使用するための主要な方法は2つあり、さまざまなアプリケーションに特有の利点を提供します:

データ並列処理 この方法では、各GPUにフルモデルのレプリカを実行し、各GPUにユニークなデータセグメントを割り当てて処理を行います。各GPUは、トレーニング中にデータサブセットに基づいて勾配を計算し、その後、すべてのGPU間で同期を取ります。

モデルパラレリズム モデルパラレリズムのアプローチは、モデルを複数のGPUに分割し、各GPUがモデルの特定のレイヤーまたはパラメータを処理します。モデルのパラメータは、テンソル、レイヤー、パイプラインステージなど、さまざまな粒度でGPUに分配できます。

モデル並列性の種類

モデル並列性はさらに特定の技術に分けることができます。

テンソル並列処理
テンソル並列処理は、各層の重みをテンソルレベルで複数のGPUに分割することを含みます。大規模な行列乗算操作では、GPU間の分配により、各GPUが行列の異なる部分を独立して処理できるようになります。

パイプライン並列処理
パイプライン並列処理は、異なるレイヤーを複数のGPUに分散させることで、各GPUがモデルの特定のセグメントを処理します。例えば、GPU 0がフォワードパス中に最初のバッチを完了すると、すぐに出力をGPU 1に送ります。これにより、GPU 1は次のステージのデータを処理でき、このサイクルが続きます。これらのミニバッチを巧みにずらすことで、すべてのGPUが同時に動作できます。

シャーデッドデータ並列処理 この技術はデータ並列処理とパラメータシャーディング（各GPUはパラメータの一部のみを保存）を利用して、メモリの要件を減らし、効率的なモデルのトレーニングを維持します。

GPUメモリ管理：隠れた課題

マルチGPUシステムは、GPUメモリ管理を主なパフォーマンスボトルネックとして頻繁に経験します。単純な方法は、複数のGPU間でモデルスライスをプッシュすることですが、GPU間の通信コストと潜在的なメモリフラグメンテーションの問題を無視しています。LLMタスクによるマルチGPU推論中には、レイヤー、テンソル、およびパイプラインステージの慎重な割り当てが不可欠です。

重要な考慮事項:

• バッチサイズ: バッチサイズを大きくするとGPUの利用効率が向上しますが、注意深い管理をしないとOOM（Out of Memory）エラーが発生するリスクがあります。PyTorchの内蔵プロファイラーや外部ライブラリなどのプロファイリングツールを使用することで、ユーザーはメモリのホットスポットを特定することができます。
• アクティベーションチェックポイント：アクティベーションチェックポイント技術は、逆伝播中に選択した層のフォワードパスを再計算することによってメモリ消費を削減します。
• オフロード：いくつかのフレームワークは、未使用のGPUメモリをCPUやNVMe（不揮発性メモリ拡張）ストレージデバイスに転送することを可能にします。このアプローチは巨大なモデルをサポートしますが、追加の処理遅延を引き起こす可能性もあります。

CUDAの入門チュートリアルを読むことで、GPUメモリの基本知識を包括的に得ることができます。

LLMを複数のGPUに分割するためのツールとライブラリ

多くのオープンソースフレームワークは、大規模モデルのためにマルチGPUディープラーニングを実現します。次のセクションでは、いくつかの主要なソリューションと、それらが大規模言語モデルの機械学習におけるGPU並列処理にどのように貢献しているかを概説します。

PyTorchのDistributedDataParallel

PyTorchのDistributedDataParallel (DDP)は、大規模言語モデルの分散トレーニングに最も広く使用されているアプローチの1つです。これは、複数のGPUとノード間での勾配の同期を簡単に行えるようにします。それぞれのプロセスは、データのサブセットに対して作業しながら、各トレーニング反復後に勾配を平均化することで、同じモデルを実行します。主な利点

• 使いやすさ: DDPはモデルをラップし、自動的に勾配を同期します。
• スケーラビリティ: 大規模クラスターに対して非常にスケーラブルです。
• 汎用性: シングルノードのマルチGPUおよびマルチノードのセットアップでうまく機能します。

HuggingFace アクセラレート

HuggingFaceのAccelerateライブラリは、最小限のコード変更でマルチGPU推論を簡単に行える方法を提供します。このライブラリは、device_map=“auto ”パラメータを介して自動モデルシャーディングを提供し、分散推論のためのツールを用意しています。次の例を通じてこれを説明します：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/LLaMA-2-7B-32K",
    torch_dtype=torch.float16,
    device_map="auto" # Automatically distributes across available GPUs
)

Accelerateは最初にGPU 0を満たし、その後GPU 1に移動し、残りのGPUに進むことでモデル全体をロードします。この方法は、手動でのパーティショニングなしに、複数のGPUにLLMをロードするための自動的で簡単な解決策を提供します。詳細な説明については、Hugging FaceのAccelerateライブラリを使用したRaw PyTorchのマルチGPUをチェックしてください。

Ollama 複数の GPU

Ollamaフレームワークは、効率的なCPUおよびGPU推論機能を備えたLLMの実行ソリューションを提供します。複数のGPUがOllamaを実行することができ、環境変数を設定したり、設定ファイルを編集してモデルのパーティショニングを指定したりできます。公式のOllamaドキュメントには、モデルの重みを分割するための詳細な手順が記載されています。

GPUパーティショニングのための環境変数。GPU設定の構成には、適切な環境変数のエクスポートが必要です。例えば:

export OLLAMA_GPU_COUNT=5
export OLLAMA_GPU_MEMORY_LIMIT=16GB

ここです：

• OLLAMA_GPU_COUNTは使用するGPUの数を指定します。
• OLLAMA_GPU_MEMORY_LIMIT は GPU メモリ割り当ての上限を定義します。

Ollamaのより高度な使用法については、最大の効率を得るためにH100 GPUでLLMを実行する方法をご覧ください。

vLLM

vLLM（バースタイル大規模言語モデル）ライブラリは、大規模言語モデルの推論タスクを効率的に実行するための最近の開発を表しています。このライブラリは、高度に最適化されたトランスフォーマーインタープリターを提供し、長いプロンプト処理におけるKV（キー・バリュー）キャッシュのメモリ要件を扱うためにPagedAttentionを導入しています。vLLMは分散推論とサービングをサポートしています。

単一のGPUの容量を超える場合、vLLMを介してモデルを提供するために複数のGPUまたはマシンを使用できます。サービングインスタンスを初期化する際に、ユーザーはテンソルパラレルサイズを設定できます。以下の例で確認できます。

from vllm import LLM
Initialize model with tensor parallelism across 4 GPUs
llm = LLM(model="meta-llama/Llama-2-70b-hf", tensor_parallel_size=4)
Generate text for multiple prompts in parallel
outputs = llm.generate(["Write a book", "Explain artificial intelligence"])

DeepSpeed

マイクロソフトは、大規模モデルのトレーニングプロセスを最適化するためのソフトウェアライブラリとしてDeepSpeedを開発しました。DeepSpeedのZero Redundancy Optimizerアルゴリズムは、GPU間でモデルの状態を分割してメモリの冗長性を排除します。
DeepSpeedのステージには以下が含まれます:

• ZeRO-1: シャードオプティマイザーステート,
• ZeRO-2: シャード最適化器 + 勾配,
• ZeRO-3: シャードオプティマイザー + 勾配 + パラメータ (モデルの重み自体).

Deepspeedは、ZeRO-OffloadおよびZeRO-Infinityとして知られる機能を通じて、CPUおよびNVMeオフロードのサポートを含んでいます。

ZeROステージ3を有効にし、必要に応じてCPUにoffload_paramを直接オフロードすることができます。以下の例は、DeepSpeed設定ファイルで通常見られる部分を表しています。

{
 "train_batch_size": 8,
 "fp16": { "enabled": true },
 "zero_optimization": {
 "stage": 3,
"offload_param": { "device": "cpu" }
 }
}

DeepSpeedシステムは、指定された構成に従ってモデルと勾配の配布を管理します。

メガトロン-LM

NVIDIAは、このフレームワークをGitHubリポジトリを通じて提供しており、GPT-2、GPT-3、T5などの大規模なトランスフォーマーモデルのトレーニングを可能にしています。Megatron-LMはテンソル並列処理とパイプライン並列処理を統合して、膨大な並列処理能力を実現しています。

これはユーザーがテンソル並列サイズを定義できることを可能にし、各層のGPU分配とモデルのステージセグメンテーションを制御するパイプライン並列サイズを指定します。Megatron-LMには、数十億のパラメータを持つモデルを最初からトレーニングするのに便利な高度な技術が含まれています。

複数のマシンにわたる分散トレーニング

LLM用の複数のマシンでは、各ノードで複数のGPUが動作する分散環境を設定する必要があります。PyTorchの分散通信バックエンド（NCCL）を使用すると、すべてのプロセス間でシームレスな接続を作成できます。

• マスターノードのセットアップ: すべてのノードを管理するマスターノードのIPアドレスとポートを特定します。
• ランクと世界のサイズ: 各プロセスまたはノードにはランクが割り当てられ、プロセスの総数が世界のサイズを表します。
• 起動: torch.distributed.launch または torchrun を使用して、異なるノード間で複数のプロセスを開始します。
• ネットワーク最適化: 同期オーバーヘッドを減らすために、InfiniBandなどの高帯域幅ネットワークを使用するべきです。

PyTorch DDPを使用した分散LLM訓練: 最小限の例

PyTorch の DistributedDataParallel (DDP) システムは、複数の GPU 上でモデルを訓練するために、各 GPU 上にモデルを複製し、勾配を同期させて単一デバイスでの訓練をシミュレートします。このガイドは、DDP を通じて分散訓練を実装するための重要なステップを要約しています。

1. プロセスグループの初期化

プロセス間の通信を処理するために、分散バックエンドを構成する必要があります。以下の例では、GPUオペレーションのためにNCCL（NVIDIA Collective Communications Library）バックエンドを使用する方法を示します。

import torch.distributed as dist
dist.init_process_group(backend="nccl")

すべてのプロセスを接続する通信グループを初期化します。各計算プロセスは単一のGPUに対応します。

2. デバイスを構成し、モデルをDDPでラップする

LOCAL_RANK環境変数を確認して現在のプロセスに対するGPUを特定し、次にそのGPUにモデルを割り当てます。それから、モデルをDistributedDataParallelでラップします:

import os
import torch
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
model = MyModel().to(local_rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

loss.backward() が実行されると、DDPはすべてのプロセス間で勾配を同期させて平均化します。

3. DataLoaderでDistributedSamplerを使用する

データをランダムにシャッフルするのではなく、shuffle=Trueを使用する代わりに、DistributedSamplerを使用して、各プロセスがデータセットのユニークなサブセットを受け取ることを保証します。例えば:

from torch.utils.data import DataLoader, DistributedSampler
sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=sampler)

この方法は、複数のプロセスが同じトレーニングサンプルにアクセスするのを防ぎます。

4. トレーニングループの実装

各プロセス内で、通常のトレーニングループを実行します：train_loaderからデータを取得し、lossを計算し、loss.backward()を実行し、次にoptimizer.step()を実行する前に、それをGPUに送ります。
DDPは、追加のコーディング要件なしで、バックワードパスにおける勾配の同期を自動的に処理します。トレーニングが終了したら、次のことでクリーンアップします：
dist.destroy_process_group()

5. torchrunを使用してトレーニングを開始する

torchrun ユーティリティを使用してトレーニングスクリプトを起動します: torchrun --nproc_per_node=4 train_ddp.py -nproc_per_node=4:

• これにより、各ノード（マシン）で起動するプロセスの数が確定されます。
• 通常、この値はノード上で利用可能なGPUの数と等しく設定します。
• システムはローカルマシン上で4つのプロセスを実行します（つまり、4つのGPUを持っていると仮定します）。

-train_ddp.py:

• DDPベースの分散トレーニングのためのトレーニングスクリプト。

• このスクリプト内では、あなたは次のことをするべきです:

  • プロセスグループを初期化する（torch.distributed.init_process_group）
  • torch.nn.parallel.DistributedDataParallelでモデルをラップします
  • データセットにはDistributedSamplerを使用してください。

初期設定のおかげで、それぞれのプロセスは適切なLOCAL_RANKを取得し、あらかじめ定義されたプロセスグループを通じて他のプロセスと同期します。

一般的なエラーとデバッグ

この要約は、複数のGPUにおける大型言語モデルの頻繁な分配エラーの概要と、それらの原因および推奨される解決策について詳述しています。

メモリオーバーフローの問題、遅い同期、非効率的な並列処理に対処できれば、マルチGPU LLM環境の安定性とパフォーマンスが大幅に向上します。

マルチモーダルLLMの考慮事項

テキストと画像、さらには音声や動画を組み合わせたマルチモーダルLLMモデルが一般的になりつつあります。そのサイズはテキストのみのモデルを超えることが多く、マルチGPU戦略が不可欠です。
重要なポイント:

• 追加のモダリティ: 各モダリティは、専門化されたエンコーダまたはデコーダブロックを導入し、より高いパラメータ数に寄与します。
• カスタムレイヤー: 画像エンコーダー（ViTなど）と音声エンコーダー（Wav2Vec2のような）には、異なる分配戦略が必要になる場合があります。
• インターモーダルフュージョン: テキスト要素と視覚的または聴覚的コンポーネントが新しいパイプライン段階を導入する可能性があります。
• ツールサポート: 選択したフレームワークは、多モーダル入出力機能および大規模モデルの並行設定との統合能力をサポートしている必要があります。

FAQセクション

Q1: LLMは複数のGPUで動作できますか？
はい。PyTorchやTensorFlowなどの最新の深層学習フレームワークでは、複数のGPUや分散ノードを使用してモデルをトレーニングすることができます。データ並列性とモデル並列性の技術により、LLMは複数のGPUで効率的に動作することができます。

Q2: 複数のGPUでモデルを実行するにはどうすればよいですか？
DistributedDataParallel（PyTorch）またはtf.distribute（TensorFlow）を使用して、分散トレーニングのためにモデルをラップすることができます。また、Hugging Face AccelerateやDeepSpeedは、自動的にモデルのパラメータとデータを複数のGPUに分散します。

Q3: 機械学習に複数のGPUを使用できますか？
はい、もちろんです。ディープラーニングの専門家は、トレーニングや推論プロセスを加速するために複数のGPUをよく使用します。データ分配やモデルレイヤーの分割を通じて、パフォーマンスのスケーリングが可能です。

Q4: 一度に2つのGPUを使用しても大丈夫ですか？
2つ以上のGPUを使用することは許可されており、パフォーマンスを向上させることができます。正しく設定すれば、1台のマシンにインストールされた2つのGPUは、クラスタを持っていなくてもトレーニングや推論の時間をほぼ半分にすることができます。

Q5: LLMをどのように並列化しますか？
LLMを並列化するには、データ並列処理、モデル並列処理、そして時々それらのハイブリッドを使用できます。DeepSpeed、Megatron-LM、Hugging Face Accelerateなどのライブラリは、このプロセスをより簡単にします。

Q6: 複数のLLMを持つことはできますか？
十分な計算能力があれば、複数のLLMインスタンスを作成することができます。同じGPUを使用して複数の大規模モデルを実行すると、リソースの競合が発生し、アウトオブメモリーエラーが引き起こされる可能性があります。

Q7: 同じGPUで複数のディープラーニングモデルを実行できますか？
同じGPU上で実行することは可能ですが、複数のモデルを複数のGPUに分配すると、より良いパフォーマンスを得られます。単一のGPU上で複数のモデルを実行すると、メモリの制限が生じ、データ処理速度が低下する可能性があります。

Q8: ダブルLLMとは何ですか？
「ダブルLLM」という用語は、タスクのパフォーマンス、精度、または効率を向上させるために、協力して動作する2つの大規模言語モデルを組み合わせたシステムアーキテクチャを指します。このアプローチでは、2つのモデルの補完的な強みを活用し、各モデルに異なる責任が割り当てられます。

Q9: 自分のLLMを訓練できますか？

LLMのトレーニングは、十分なデータ、計算リソース（GPUなど）、信頼できるトレーニングパイプラインがあれば可能です。開発者は、Megatron-LMやDeepSpeedなどのオープンソースプログラムを使用して、大規模言語モデルをゼロからトレーニングできます。

Q10: シングルモーダルLLMとマルチモーダルLLMの違いは何ですか？
シングルモーダルLLMは1つの入力タイプのみを処理します。マルチモーダルLLMは、テキストと画像や音声など、さまざまなデータタイプを管理できます。

Q11: モデル並列処理とは何ですか？LLMにどのように適用されるのですか？

モデルの並列処理は、単一のモデルのパラメータを複数のGPUに分割し、各GPUがネットワークの一部を管理することを含みます。これは、非常に大きなLLM（大規模言語モデル）にとって不可欠です。なぜなら、モデルが限られたGPUメモリに適合し、同時処理をサポートするからです。

結論

最先端のNLPおよびマルチモーダルAIシステムに取り組む研究者や開発者にとって、大規模言語モデルを複数のGPUに分割することは不可欠となっています。
この記事では、現代のLLMの計算およびメモリの制限を扱うための実用的な方法を提示しました。これには、モデル並列配置およびデータ並列配置、DeepSpeed、Hugging Face Accelerate、Megatron-LMなどのツールが含まれます。

適切な戦略を適用することで、マルチGPUシステムはスケーラブルなトレーニングと高速な推論時間をサポートできるようになります。今後のAI開発は、効率的なGPUメモリ管理と、さまざまなデータタイプを統合するモデルが進化する中で最適化された分散トレーニング手法に依存するでしょう。並列性のアーキテクチャを理解し、オープンソースツールを使用しながら一般的なパフォーマンスの問題を予測できる実務者は、大規模でマルチモーダルなLLMの潜在能力を最大限に引き出すことができます。

リソース

• 分散データ並列処理の始め方
• https://gist.github.com/sgraaf/5b0caa3a320f28c27c12b5efeb35aa4c
• ZeRO & DeepSpeed
• メガトロン-LM: モデル並列性を使用した数十億パラメータの言語モデルの訓練
• メガトロンLM — モデル並列性が言語モデルを新たな高みに押し上げる方法