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から変更なし

```python

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利用法に焦点を当てて無意味なデータを学習

```python
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 （上のネットワーク定義も入れる）

```python
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

```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メモリに乗り切らないモデルを作ってみた

```python
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部分を

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

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

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

## 現時点での感想

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

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

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