最近、大規模言語モデル（LLM）の学習することが多いので、大規模なモデルの学習方法について記載します。

並列学習について

巨大なパラメータをもつモデルの学習は非常に時間がかかってしまいます。そのため、計算時に並列を行っていく必要があります。大きく分けて３パターンあります。

Data parallelism (DP)
Pipeline parallelism (PP)
Tensor parallelism (TP)

Data parallelism (DP)

モデルがGPU上に乗る場合、DPは一般的に利用されると思います。同一条件のモデルを各GPU上にコピーし、データを並列に流しながら計算させて、一気に勾配を更新します。大変なのは、モデルが1つのGPU上に乗り切らない場合です。

70億（7B）パラメータ級の小型LLMかつパラメータの浮動小数点がfloat16であれば、約15GBのGPUメモリが必要です。特に、学習する場合では、オプティマイザの状態を32bitとすると、トータルでGPUメモリは60GB弱が必要とで、なんとかA100（80GB）に乗り学習することができます。しかし、13B級になるとフルチューニングでは、何も工夫をしないと複数GPUが必要になってしまいます。しかしながら、DPに関しては後述するZeRO（DeepSpeed）を使ってメモリを削減すると、割と簡単にDPを実装することができます。

ZeRO

Zero Redundacy Opimizer (ZeRO)は、S. Rajbhandariらによって2019年に提案されています。こちらの実装は、DeepSpeedとしてGitHub上に公開されています。 DeepSpeedの詳細については以下のブログや論分を参照ください。また、ZeRO動作の仕組みについては、公式ブログの動画を見るのが一番わかりやすいと思います。

arxiv.org

github.com

www.microsoft.com

概要

ZeROは、stage 1-3の3段階（モデル分割の度合い）まであります。

ステージが上がるほどメモリ効率が上がりますが、各デバイスとのコミニュケーションコストがかかってしまうため、計算速度は落ちます。

Zero1: $P_{os}$ : 4倍メモリ削減、optimizer state partioning across GPUs
Zero2: $P_{os+g}$: 8倍メモリ削減、gradient partitioning across GPUs
Zero3: $P_{os+g+p}$: GPU数$N_d$によって線形にメモリ削減可能。$N_d=64$ならば64倍削減可能。

Figure 1: Comparing the per-device memory consumption of model states, with three stages of ZeRO-DP optimizations.

ベースライン（特に何もしていない状態）

モデルのパラメーター・勾配情報・オプティマイザの状態（勾配、バリアンス、モメンタム、パラメータなど）が、各GPUに配置されています。横軸はTransormerのブロックだとみなすことができます。メモリ消費の公式は、モデルサイズを$\Psi$、オプティマイザの状態でバイト数をKとすると、浮動小数点を16bit (2 byte)とすると、$(2 + 2 + K ) * \Psi$で表現できます。例えば、7.5BのモデルでK=12 (mixed-precison training with Adam optimizer)とすると、16bytes * 7.5B = 120GBのメモリが必要となります。

$P_{os}$の部分

図に着目するとoptimizer stateが一部分だけ表示されています。これは、DPにおけるoptimizer stateの冗長性を除去し、$N_d$個のGPUに振り分けることで、メモリを削減することができています。よって、削減量は　$\frac{K \times \Psi}{N_d}$となります。

実際にDPを行う際に、どのようにしてoptimizer stateを削減しているかは、公式ブログの動画を見ていただくのが一番わかりやすいです。各GPUで勾配操作のやり取りは、割と複雑のため、文章で説明するのは非常に難しいです。

実装レベルで解説されている方の記事も参考になります。

zenn.dev

実装例

DeepSpeedを使ってMNISTに対して学習する例を考えてみます。

installation

以下でインストールできます。CudaやOpen MPI等の設定が正しくされている必要があります。なお、これらについてはcondaで入れました。

pip install deepspeed

DeepSpeedの設定は、非常に簡単で、deepspeed.initialize()で初期化するだけです。

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"
os.environ["OMP_NUM_THREADS"] = "16"
os.environ["MKL_NUM_THREADS"] = "16"

import time

import deepspeed
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

torch.set_num_threads(16)


# モデル定義
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# データセットの準備
transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
)
train_dataset = datasets.MNIST(
    root="./data", train=True, download=True, transform=transform
)
train_loader = DataLoader(
    train_dataset, batch_size=2048, shuffle=True, num_workers=4, pin_memory=True
)

# モデルとオプティマイザーの設定
model = SimpleNet()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# DeepSpeedの設定
deepspeed_config = {
    "train_batch_size": 2048,
    "gradient_accumulation_steps": 2,
    "fp16": {"enabled": True},
}

# DeepSpeedの初期化
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, optimizer=optimizer, config_params=deepspeed_config
)

rank = dist.get_rank()
start_time = time.time()
for epoch in range(100):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(model_engine.local_rank).half(), target.to(
            model_engine.local_rank
        )
        model_engine.train()
        optimizer.zero_grad()
        output = model_engine(data)
        loss = nn.functional.cross_entropy(output, target)
        model_engine.backward(loss)
        model_engine.step()
        if batch_idx % 100 == 0 and rank == 0:
            print(
                f"Train Epoch: {epoch+1} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}"
            )
    if rank == 0:
        epoch_end_time = time.time()
        print(f"Epoch {epoch+1}, Time: {epoch_end_time - start_time:.2f} seconds")

if rank == 0:
    total_time = time.time() - start_time
    print(f"Total training time with DeepSpeed: {total_time:.2f} seconds")

実際に実行する場合は、以下のように実行できます。GPUは自動的にすべて割り当てられるので以下のように設定もできます。

deepspeed --num_gpus 2 train.py

RunningAvgSamplesPerSec=21802.901847271874, CurrSamplesPerSec=16403.02205927284, MemAllocated=0.02GB, MaxMemAllocated=0.02GB
Epoch 10, Time: 100.81 seconds
Total training time with DeepSpeed: 100.81 seconds

モデルが小すぎるので、完全コピーでも良いと思いますが、動くことは確認できました。Huggingfaceでも実行する場合も簡単です。

huggingface.co

Pipeline parallelism (PP)

Pipeline parallelismもしくはModel parallelismでは、モデルが複数GPU間で垂直（レイヤーレベル）に分割されます。結果として、特定のレイヤーは単一のGPUで計算することができます。各GPUは並列に異なるパイプラインのステージを処理し、小さいチャンクのバッチとして作用します。これは、GPipe (Y. Huang et al. 2019)で提案されています。Fは順方向の計算、Bをバックプロパゲーションの計算を表すと、PPの概略図を次のような図で表すことができます。

上段：ナイーブなモデル並列（MP）戦略では、ネットワークのシークエンシャルな特性のために、フルに活用されていません。下段：パイプライン並列では、入力のミニバッチをより小さなチャンクに分割することで、異なるデバイスが同時に分割されたチャンクを処理することができています。

この図からもわかるようにナイーブなMPでは、モデルレイヤーを複数のGPUに配置します。PyTorchで考えると、特定のレイヤーを.to()で移動させることに対応します。これは、垂直な並列（vertical model parallel）とみなすことができます。例えば、8層のモデルを考えた時に、[0, 1, 2, 3]をGPU0に、[4, 5, 6, 7]をGPU1に振り分けることを考えると、計算を行う時はGPU0にて0→1→2→3と計算された後に、layer3の結果がGPU1に投げ込まれ、4→5→6→7と計算されます。

ここでの問題は、特にマルチノードで計算しているときの物理的・通信的なオーバーヘッドです（この例だと、layer 3 → layer 4に移動させるとき）。また、バックプロパゲーションをする場合は、その逆の操作が行われるため、ここでも問題になります。これは「ナイーブな」MPの問題といえます。

一方で、PPでは、ほとんどMPと同一なのですが、入力バッチをマイクロバッチ（MBS）に分割し、人工的にパイプラインを作ることによって、このGPUのアイドリング問題を解決しています。GPU0, 1, 2, 3の場合を考えます。F0, F1, F2, F3と計算する場合に、さらにチャンクを分割して考ます。図ではchunks=4としており、これがPPのハイパーパラメータです。まずGPU0がF00, F01, F02, F03を計算します。次にF00の計算が完了した直後にF10を計算することによって、待ちが減っています。この無駄が生じない部分はBubbleと呼ばれています。

概念的にはgradient accumulation steps (GAS)と同じです。PyTorchではchunksを使っており、DeepSpeedではGASとしています。

research.google

次はTPとMegatronについて記載しようと思います。