13.5. 複数GPUでの学習

これまで、CPUとGPU上でモデルを効率よく学習する方法について議論してきた。さらに、 13.3 章 では、深層学習フレームワークがそれらの間で計算と通信を自動的に並列化できることも示した。また、 6.7 章 では、nvidia-smi コマンドを使ってコンピュータ上で利用可能なGPUをすべて一覧表示する方法も示した。

しかし、実際に深層学習の学習をどのように並列化するのかについては、まだ議論していなかった。

その代わりに、データを何らかの形で複数デバイスに分割して動作させる、ということを暗黙に仄めかしていた。本節ではその詳細を補い、ゼロからネットワークを並列に学習する方法を示す。高レベルAPIの機能を活用する方法については、 13.6 章 に譲る。

ここでは、 12.5 章 で説明したようなミニバッチ確率的勾配降下法に慣れていることを前提とする。

13.5.1. 問題の分割

まずは、単純なコンピュータビジョン問題と、やや古典的なネットワーク、たとえば複数層の畳み込み、プーリング、そして最後にいくつかの全結合層を持つものから始めよう。

つまり、LeNet (LeCun et al., 1998) や AlexNet (Krizhevsky et al., 2012) にかなり似たネットワークから始める。

複数のGPU（デスクトップサーバなら2枚、AWS g4dn.12xlarge インスタンスなら4枚、p3.16xlarge なら8枚、p2.16xlarge なら16枚）があるとき、単純で再現性のある設計方針の恩恵を受けつつ、良好な高速化を達成できるように学習を分割したいと考える。結局のところ、複数GPUは メモリ と 計算能力 の両方を増やする。要するに、分類したい学習データのミニバッチが与えられたとき、次の選択肢がある。

まず、ネットワークを複数GPUにまたがって分割する方法がある。つまり、各GPUが特定の層へ流れ込むデータを入力として受け取り、いくつかの後続層にわたってデータを処理し、その後データを次のGPUへ送りる。

これにより、単一GPUでは扱えないより大きなネットワークでデータを処理できる。

さらに、GPUごとのメモリ使用量も十分に制御できる（全ネットワークのメモリ使用量の一部で済みます）。

しかし、層間のインターフェース（したがってGPU間のインターフェース）には厳密な同期が必要である。特に、層ごとの計算負荷が適切に釣り合っていない場合、これは厄介である。GPUの数が増えると、この問題はさらに悪化する。

層間のインターフェースでは、活性化や勾配など、大量のデータ転送も必要になる。

これはGPUバスの帯域幅を圧迫する可能性がある。

さらに、計算集約的でありながら逐次的な操作を分割するのは容易ではない。この点についての最善の試みとしては、たとえば Mirhoseini et al. (2017) を参照しよ。これは依然として難しい問題であり、非自明な問題に対して良い（線形な）スケーリングを達成できるかどうかは明らかではない。複数GPUを連結するための優れたフレームワークやオペレーティングシステムの支援がない限り、推奨しない。

第二に、作業を層単位で分割する方法がある。たとえば、単一GPUで64チャネルを計算する代わりに、問題を4枚のGPUに分割し、それぞれが16チャネル分のデータを生成するようにできる。

同様に、全結合層では出力ユニット数を分割できる。 図 13.5.1（Krizhevsky et al. (2012) より引用）はこの設計を示しており、当時はGPUのメモリ容量が非常に小さかった（当時2GB）ため、この戦略が用いられていた。

チャネル数（またはユニット数）が小さすぎない限り、計算の観点では良好なスケーリングが可能である。

さらに、利用可能メモリは線形に増えるため、複数GPUを使えばより大きなネットワークを処理できる。

図 13.5.1 メモリ制約のために元のAlexNet設計で用いられたモデル並列。

しかし、各層が他のすべての層の結果に依存するため、同期操作やバリア操作が 非常に多く 必要になる。

さらに、転送が必要なデータ量は、GPU間で層を分散する場合よりもさらに大きくなる可能性がある。したがって、帯域幅コストと複雑さのため、この方法は推奨しない。

最後に、データを複数GPUに分割する方法がある。この方法では、すべてのGPUが異なる観測に対してではあるものの、同じ種類の作業を行う。各ミニバッチの学習後に、勾配がGPU間で集約される。

これは最も単純な方法であり、どのような状況にも適用できる。

必要なのは、各ミニバッチの後に同期することだけである。とはいえ、他の勾配がまだ計算中のうちから、勾配パラメータの交換を始められると非常に望ましいである。

さらに、GPU数が増えるとミニバッチサイズも大きくなるため、学習効率が向上する。

ただし、GPUを増やしても、より大きなモデルを学習できるようにはならない。

図 13.5.2 複数GPU上での並列化。左から右へ：元の問題、ネットワーク分割、層単位分割、データ並列。

複数GPUでのさまざまな並列化方法の比較を 図 13.5.2 に示す。

概して、十分に大きなメモリを持つGPUにアクセスできるなら、データ並列が最も便利な方法である。分散学習のための分割についての詳細な説明は (Li et al., 2014) も参照しよ。深層学習初期にはGPUメモリが問題でしたが、現在では非常に特殊な場合を除いてこの問題は解決されている。以下ではデータ並列に焦点を当てる。

13.5.2. データ並列

1台のマシンに \(k\) 枚のGPUがあると仮定する。学習するモデルが与えられると、各GPUは独立にモデルパラメータの完全な集合を保持する。ただし、GPU間のパラメータ値は同一であり、同期されている。

例として、 図 13.5.3 は \(k=2\) のときのデータ並列による学習を示している。

図 13.5.3 2枚のGPU上でデータ並列を用いてミニバッチ確率的勾配降下法を計算する様子。

一般に、学習は次のように進みる。

• 学習の任意の反復で、ランダムなミニバッチが与えられたら、バッチ内の例を \(k\) 個の部分に分け、GPUに均等に分配する。

• 各GPUは、自分に割り当てられたミニバッチ部分に基づいて、モデルパラメータの損失と勾配を計算する。

• \(k\) 枚のGPUの局所勾配を集約して、現在のミニバッチ確率的勾配を得る。

• 集約された勾配を各GPUに再配布する。

• 各GPUは、このミニバッチ確率的勾配を使って、自分が保持しているモデルパラメータの完全な集合を更新する。

実際には、\(k\) 枚のGPUで学習するときにはミニバッチサイズを \(k\) 倍に 増やす のが普通である。そうすることで、各GPUが単一GPUで学習する場合と同じ量の仕事をすることになる。16枚GPUのサーバでは、ミニバッチサイズがかなり大きくなるため、それに応じて学習率も増やす必要があるかもしれない。

また、 8.5 章 のバッチ正規化は調整が必要で、たとえばGPUごとに別々のバッチ正規化係数を保持する方法がある。

以下では、複数GPU学習を説明するためにおもちゃのネットワークを使う。

pytorchmxnetjaxtensorflow

%matplotlib inline
from d2l import torch as d2l
import torch
from torch import nn
from torch.nn import functional as F

%matplotlib inline
from d2l import mxnet as d2l
from mxnet import autograd, gluon, np, npx
npx.set_np()

new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)

new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)

13.5.3. おもちゃのネットワーク

7.6 章 で導入した LeNet を少し変更して使う。パラメータ交換と同期を詳細に示すため、これをゼロから定義する。

pytorchmxnetjaxtensorflow

# Initialize model parameters
scale = 0.01
W1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
W2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
W3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
W4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# Define the model
def lenet(X, params):
    h1_conv = F.conv2d(input=X, weight=params[0], bias=params[1])
    h1_activation = F.relu(h1_conv)
    h1 = F.avg_pool2d(input=h1_activation, kernel_size=(2, 2), stride=(2, 2))
    h2_conv = F.conv2d(input=h1, weight=params[2], bias=params[3])
    h2_activation = F.relu(h2_conv)
    h2 = F.avg_pool2d(input=h2_activation, kernel_size=(2, 2), stride=(2, 2))
    h2 = h2.reshape(h2.shape[0], -1)
    h3_linear = torch.mm(h2, params[4]) + params[5]
    h3 = F.relu(h3_linear)
    y_hat = torch.mm(h3, params[6]) + params[7]
    return y_hat

# Cross-entropy loss function
loss = nn.CrossEntropyLoss(reduction='none')

# Initialize model parameters
scale = 0.01
W1 = np.random.normal(scale=scale, size=(20, 1, 3, 3))
b1 = np.zeros(20)
W2 = np.random.normal(scale=scale, size=(50, 20, 5, 5))
b2 = np.zeros(50)
W3 = np.random.normal(scale=scale, size=(800, 128))
b3 = np.zeros(128)
W4 = np.random.normal(scale=scale, size=(128, 10))
b4 = np.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# Define the model
def lenet(X, params):
    h1_conv = npx.convolution(data=X, weight=params[0], bias=params[1],
                              kernel=(3, 3), num_filter=20)
    h1_activation = npx.relu(h1_conv)
    h1 = npx.pooling(data=h1_activation, pool_type='avg', kernel=(2, 2),
                     stride=(2, 2))
    h2_conv = npx.convolution(data=h1, weight=params[2], bias=params[3],
                              kernel=(5, 5), num_filter=50)
    h2_activation = npx.relu(h2_conv)
    h2 = npx.pooling(data=h2_activation, pool_type='avg', kernel=(2, 2),
                     stride=(2, 2))
    h2 = h2.reshape(h2.shape[0], -1)
    h3_linear = np.dot(h2, params[4]) + params[5]
    h3 = npx.relu(h3_linear)
    y_hat = np.dot(h3, params[6]) + params[7]
    return y_hat

# Cross-entropy loss function
loss = gluon.loss.SoftmaxCrossEntropyLoss()

[07:25:19] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU

train(num_gpus=1, batch_size=256, lr=0.2)

train(num_gpus=1, batch_size=256, lr=0.2)

13.5.4. データ同期

効率的な複数GPU学習には、2つの基本操作が必要である。

まず、パラメータのリストを複数デバイスに配布する 能力と、勾配を付加する能力（get_params）が必要である。パラメータがなければ、GPU上でネットワークを評価することはできない。

第二に、複数デバイス間でパラメータを合計する能力、つまり allreduce 関数が必要である。

pytorchmxnetjaxtensorflow

def get_params(params, device):
    new_params = [p.to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

def get_params(params, device):
    new_params = [p.copyto(device) for p in params]
    for p in new_params:
        p.attach_grad()
    return new_params

train(num_gpus=2, batch_size=256, lr=0.2)

train(num_gpus=2, batch_size=256, lr=0.2)

モデルパラメータを1枚のGPUにコピーして試してみよう。

new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)

b1 weight: tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       device='cuda:0', requires_grad=True)
b1 grad: None

まだ何の計算もしていないので、バイアスパラメータに関する勾配はまだゼロである。

次に、複数GPUに分散されたベクトルがあると仮定しよう。次の allreduce 関数は、すべてのベクトルを加算し、その結果をすべてのGPUにブロードキャストする。これを動かすには、結果を集約するデバイスへデータをコピーする必要があることに注意しよう。

pytorchmxnet

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i][:] = data[0].to(data[i].device)

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].copyto(data[0].ctx)
    for i in range(1, len(data)):
        data[0].copyto(data[i])

異なるデバイス上で異なる値を持つベクトルを作成し、それらを集約してみよう。

pytorchmxnet

data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('before allreduce:\n', data[0], '\n', data[1])
allreduce(data)
print('after allreduce:\n', data[0], '\n', data[1])

before allreduce:
 tensor([[1., 1.]], device='cuda:0')
 tensor([[2., 2.]], device='cuda:1')
after allreduce:
 tensor([[3., 3.]], device='cuda:0')
 tensor([[3., 3.]], device='cuda:1')

data = [np.ones((1, 2), ctx=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('before allreduce:\n', data[0], '\n', data[1])
allreduce(data)
print('after allreduce:\n', data[0], '\n', data[1])

[07:25:20] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for GPU
before allreduce:
 [[1. 1.]] @gpu(0)
 [[2. 2.]] @gpu(1)
after allreduce:
 [[3. 3.]] @gpu(0)
 [[3. 3.]] @gpu(1)

13.5.5. データの分配

ミニバッチを複数GPUに均等に 分配する ための簡単なユーティリティ関数が必要である。たとえば、2枚のGPUでは、データの半分をそれぞれのGPUにコピーしたいとする。

より便利で簡潔なので、まずは深層学習フレームワークの組み込み関数を使って、\(4 \times 5\) 行列で試してみよう。

pytorchmxnet

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

input : tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]])
load into [device(type='cuda', index=0), device(type='cuda', index=1)]
output: (tensor([[0, 1, 2, 3, 4],
        [5, 6, 7, 8, 9]], device='cuda:0'), tensor([[10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]], device='cuda:1'))

data = np.arange(20).reshape(4, 5)
devices = [npx.gpu(0), npx.gpu(1)]
split = gluon.utils.split_and_load(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

input : [[ 0.  1.  2.  3.  4.]
 [ 5.  6.  7.  8.  9.]
 [10. 11. 12. 13. 14.]
 [15. 16. 17. 18. 19.]]
load into [gpu(0), gpu(1)]
output: [array([[0., 1., 2., 3., 4.],
       [5., 6., 7., 8., 9.]], ctx=gpu(0)), array([[10., 11., 12., 13., 14.],
       [15., 16., 17., 18., 19.]], ctx=gpu(1))]

後で再利用できるように、データとラベルの両方を分割する split_batch 関数を定義する。

pytorchmxnet

#@save
def split_batch(X, y, devices):
    """Split `X` and `y` into multiple devices."""
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X, devices),
            nn.parallel.scatter(y, devices))

#@save
def split_batch(X, y, devices):
    """Split `X` and `y` into multiple devices."""
    assert X.shape[0] == y.shape[0]
    return (gluon.utils.split_and_load(X, devices),
            gluon.utils.split_and_load(y, devices))

13.5.6. 学習

これで、単一ミニバッチに対する複数GPU学習 を実装できる。その実装は主として、本節で説明したデータ並列のアプローチに基づいている。複数GPU間でデータを同期するために、先ほど説明した補助関数 allreduce と split_and_load を使う。並列性を実現するために特別なコードを書く必要がないことに注意しよう。ミニバッチ内では計算グラフにデバイス間の依存関係がないため、計算は 自動的に 並列実行される。

pytorchmxnet

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    # Loss is calculated separately on each GPU
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:  # Backpropagation is performed separately on each GPU
        l.backward()
    # Sum all gradients from each GPU and broadcast them to all GPUs
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce([device_params[c][i].grad for c in range(len(devices))])
    # The model parameters are updated separately on each GPU
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0]) # Here, we use a full-size batch

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    with autograd.record():  # Loss is calculated separately on each GPU
        ls = [loss(lenet(X_shard, device_W), y_shard)
              for X_shard, y_shard, device_W in zip(
                  X_shards, y_shards, device_params)]
    for l in ls:  # Backpropagation is performed separately on each GPU
        l.backward()
    # Sum all gradients from each GPU and broadcast them to all GPUs
    for i in range(len(device_params[0])):
        allreduce([device_params[c][i].grad for c in range(len(devices))])
    # The model parameters are updated separately on each GPU
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0])  # Here, we use a full-size batch

次に、学習関数 を定義できる。これは前の章で使ったものとは少し異なる。GPUを割り当て、すべてのモデルパラメータをすべてのデバイスにコピーする必要があるからである。

もちろん、各バッチは複数GPUを扱うために train_batch 関数で処理される。便宜上（そしてコードを簡潔にするため）、精度の計算は1枚のGPU上で行うが、他のGPUはアイドル状態になるため、これは 非効率的 である。

pytorchmxnet

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # Copy model parameters to `num_gpus` GPUs
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            # Perform multi-GPU training for a single minibatch
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        # Evaluate the model on GPU 0
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
          f'on {str(devices)}')

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # Copy model parameters to `num_gpus` GPUs
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            # Perform multi-GPU training for a single minibatch
            train_batch(X, y, device_params, devices, lr)
            npx.waitall()
        timer.stop()
        # Evaluate the model on GPU 0
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
          f'on {str(devices)}')

これが 単一GPU上 でどの程度うまく動くか見てみよう。

まず、バッチサイズを256、学習率を0.2にする。

train(num_gpus=1, batch_size=256, lr=0.2)

test acc: 0.82, 3.7 sec/epoch on [device(type='cuda', index=0)]

バッチサイズと学習率を変えずに GPU数を2に増やす と、テスト精度は前の実験と比べておおむね同じままであることがわかる。

最適化アルゴリズムの観点では、両者は同一である。残念ながら、ここでは意味のある高速化は得られない。モデルが単純に小さすぎるうえ、データセットも小さいためである。また、複数GPU学習の実装にやや洗練されていない方法を使っているため、Pythonのオーバーヘッドも大きくなっている。今後は、より複雑なモデルと、より洗練された並列化の方法に出会うことになる。

とはいえ、Fashion-MNIST で何が起こるか見てみよう。

train(num_gpus=2, batch_size=256, lr=0.2)

test acc: 0.82, 3.8 sec/epoch on [device(type='cuda', index=0), device(type='cuda', index=1)]

13.5.7. まとめ

• 深層ネットワークの学習を複数GPUに分割する方法はいくつかある。層間で分割する方法、層ごとに分割する方法、データに分割する方法である。前の2つは、綿密に調整されたデータ転送を必要とする。データ並列が最も単純な戦略である。

• データ並列学習は簡単である。ただし、効率を高めるために実効ミニバッチサイズを増やする。

• データ並列では、データを複数GPUに分割し、各GPUが独自に順伝播と逆伝播を実行し、その後で勾配を集約して結果をGPUにブロードキャストする。

• より大きなミニバッチに対しては、学習率を少し大きくしてもよいだろう。

13.5.8. 演習

1. \(k\) 枚のGPUで学習するとき、ミニバッチサイズを \(b\) から \(k \cdot b\) に変更し、つまりGPU数に応じて拡大しなさい。

2. 異なる学習率で精度を比較しなさい。GPU数に応じてどのようにスケールするか。

3. 異なるGPU上の異なるパラメータを集約する、より効率的な allreduce 関数を実装しなさい。なぜそれがより効率的なのか。

4. 複数GPUでのテスト精度計算を実装しなさい。