DeepSpeedのインストール&ZeRO-Offload利用法

巨大なニューラルネットを学習するZeRO-Offloadと、それを含むDeepSpeedライブラリが話題だったので使ってみました。(2021/1/24時点)
本家ドキュメント https://www.deepspeed.ai/getting-started/

インストール

apt install openmpi-bin
pip3 install mpi4py
pip3 install torch
pip3 install deepspeed

注意点

• GPUが無いと動かない
• インストールされているCUDAのバージョンとPyTorchのCUDAのバージョンが違うと動かない

参考

• CUDA10.1で動作確認（PyTorchはCUDA10.1のものを指定してインストール https://pytorch.org/get-started/locally/）
• CUDA11.0では今の所、deepspeedインストール時にtritonのインストール失敗で動かせていない ※サポート外？ 情報求む

サンプル動作

ニューラルネット定義

• 約1億パラメータのモデル
• 記法は通常のPyTorchから変更なし

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(200, 10000)
        self.fc2 = nn.Linear(10000, 10000)
        self.fc3 = nn.Linear(10000, 1)

    def forward(self, x):
        h = self.fc1(x)
        h = F.relu_(h)
        h = self.fc2(h)
        h = F.relu_(h)
        h = self.fc3(h)
        return h

学習（DeepSpeedなし）

※DeepSpeed利用法に焦点を当てて無意味なデータを学習

import numpy as np
import torch.optim as optim

model = Model()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
batch_size = 1000
model.cuda()

for e in range(1000):
    x = torch.FloatTensor(np.random.rand(batch_size, 200)).cuda()
    y_ = torch.FloatTensor(np.random.randn(batch_size, 1)).cuda()

    y = model(x)
    loss = (y - y_).pow(2).mean()
    loss.backward()
    optimizer.step()
    print(e, loss.item())

学習(DeepSpeed)

sample.py （上のネットワーク定義も入れる）

import argparse
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('--deepspeed', action='store_true')
parser.add_argument('--deepspeed_config', type=str)
parser.add_argument('--local_rank', type=int)

cmd_args = parser.parse_args()

model = Model()
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=model.parameters())

batch_size = 1000  # ここは設定ファイル（下）と合わせた

for e in range(1000):
    x = torch.HalfTensor(np.random.rand(batch_size, 200)).cuda()  # 半精度
    y_ = torch.HalfTensor(np.random.randn(batch_size, 1)).cuda()  # 半精度

    y = model_engine(x)  # 変更
    loss = (y - y_).pow(2).mean()
    model_engine.backward(loss)  # 変更
    model_engine.step()  # 変更
    print(e, loss.item())

ds_config.json

{
    "train_batch_size": 1000,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 1e-4
        }
    },
    "fp16": {
        "enabled": true
    },
    "zero_optimization": true
}

コマンド

deepspeed sample.py --deepspeed --deepspeed_config=ds_config.json

動作結果

データ型はFloat, Halfのどちらでも実験

DeepSpeed	浮動小数点数	学習結果	時間
なし	Float	OK	3分49秒
なし	Half	nan	45秒
あり	Float	ZeRO非対応	---
あり	Half	OK	1分15秒

このサンプルでは、

• DeepSpeedにより同じ計算における速度向上はなかった
• DeepSpeedは半精度でもオーバーフローに対応して学習できていた（そういったメッセージが出ていたことを確認）
• 結果として、（もし学習精度が同程度なら）半精度学習として速く軽くなる利用価値はある

巨大なモデルの学習

15GBのGPUメモリに乗り切らないモデルを作ってみた

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(200, 20000)
        self.fc2 = nn.Linear(20000, 20000)
        self.fc3 = nn.Linear(20000, 20000)
        self.fc4 = nn.Linear(20000, 1)

    def forward(self, x):
        h = self.fc1(x)
        h = F.relu_(h)
        h = self.fc2(h)
        h = F.relu_(h)
        h = self.fc3(h)
        h = F.relu_(h)
        h = self.fc4(h)
        return h

結果

RuntimeError: CUDA out of memory.

GPUに乗らないモデルはこのままでは扱えなかった。
対応するドキュメント https://www.deepspeed.ai/tutorials/zero-offload/ に従ってds_config.jsonのzero_optimization部分を

{
    "zero_optimization": {
        "stage": 2,
        "cpu_offload": true,
        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

のように書き換えてZeRO-Offloadを有効したがやはりダメ、どんな容量オーバーモデルに対してもすぐ使える方法という訳ではなさそう。
引き続き要調査。

※追記：あくまでTransfomerのような大量のTensorを使うモデルの最適化が目的で、巨大なTensorを分解してまでGPUに載せようとするわけではない、とのこと。

現時点での感想

長所

• 呼び出しを変えるだけで半精度モデル学習をオーバーフロー対処付きで使える

今後に注目

• コード側で引数を準備して呼び出し方を変えるので、コードの一部だけ変えるだけ、とまでは言えない（並列システムだから仕方ない気も）
• 巨大なモデルの学習はできるとは限らない

情報はぜひコメントお願いします。