.. _sec_multi_gpu:
複数GPUでの学習
===============
| これまで、CPUとGPU上でモデルを効率よく学習する方法について議論してきた。さらに、
:numref:`sec_auto_para`
では、深層学習フレームワークがそれらの間で計算と通信を自動的に並列化できることも示した。また、
:numref:`sec_use_gpu` では、\ ``nvidia-smi``
コマンドを使ってコンピュータ上で利用可能なGPUをすべて一覧表示する方法も示した。
| しかし、実際に深層学習の学習をどのように並列化するのかについては、まだ議論していなかった。
| その代わりに、データを何らかの形で複数デバイスに分割して動作させる、ということを暗黙に仄めかしていた。本節ではその詳細を補い、ゼロからネットワークを並列に学習する方法を示す。高レベルAPIの機能を活用する方法については、
:numref:`sec_multi_gpu_concise` に譲る。
| ここでは、 :numref:`sec_minibatch_sgd`
で説明したようなミニバッチ確率的勾配降下法に慣れていることを前提とする。
問題の分割
----------
| まずは、単純なコンピュータビジョン問題と、やや古典的なネットワーク、たとえば複数層の畳み込み、プーリング、そして最後にいくつかの全結合層を持つものから始めよう。
| つまり、LeNet :cite:`LeCun.Bottou.Bengio.ea.1998` や AlexNet
:cite:`Krizhevsky.Sutskever.Hinton.2012`
にかなり似たネットワークから始める。
| 複数のGPU（デスクトップサーバなら2枚、AWS g4dn.12xlarge
インスタンスなら4枚、p3.16xlarge なら8枚、p2.16xlarge
なら16枚）があるとき、単純で再現性のある設計方針の恩恵を受けつつ、良好な高速化を達成できるように学習を分割したいと考える。結局のところ、複数GPUは
*メモリ* と *計算能力*
の両方を増やする。要するに、分類したい学習データのミニバッチが与えられたとき、次の選択肢がある。
| まず、ネットワークを複数GPUにまたがって分割する方法がある。つまり、各GPUが特定の層へ流れ込むデータを入力として受け取り、いくつかの後続層にわたってデータを処理し、その後データを次のGPUへ送りる。
| これにより、単一GPUでは扱えないより大きなネットワークでデータを処理できる。
| さらに、GPUごとのメモリ使用量も十分に制御できる（全ネットワークのメモリ使用量の一部で済みます）。
| しかし、層間のインターフェース（したがってGPU間のインターフェース）には厳密な同期が必要である。特に、層ごとの計算負荷が適切に釣り合っていない場合、これは厄介である。GPUの数が増えると、この問題はさらに悪化する。
| 層間のインターフェースでは、活性化や勾配など、大量のデータ転送も必要になる。
| これはGPUバスの帯域幅を圧迫する可能性がある。
| さらに、計算集約的でありながら逐次的な操作を分割するのは容易ではない。この点についての最善の試みとしては、たとえば
:cite:t:`Mirhoseini.Pham.Le.ea.2017`
を参照しよ。これは依然として難しい問題であり、非自明な問題に対して良い（線形な）スケーリングを達成できるかどうかは明らかではない。複数GPUを連結するための優れたフレームワークやオペレーティングシステムの支援がない限り、推奨しない。
| 第二に、作業を層単位で分割する方法がある。たとえば、単一GPUで64チャネルを計算する代わりに、問題を4枚のGPUに分割し、それぞれが16チャネル分のデータを生成するようにできる。
| 同様に、全結合層では出力ユニット数を分割できる。
:numref:`fig_alexnet_original`\ （:cite:t:`Krizhevsky.Sutskever.Hinton.2012`
より引用）はこの設計を示しており、当時はGPUのメモリ容量が非常に小さかった（当時2GB）ため、この戦略が用いられていた。
| チャネル数（またはユニット数）が小さすぎない限り、計算の観点では良好なスケーリングが可能である。
| さらに、利用可能メモリは線形に増えるため、複数GPUを使えばより大きなネットワークを処理できる。
.. _fig_alexnet_original:
.. figure:: ../img/alexnet-original.svg
メモリ制約のために元のAlexNet設計で用いられたモデル並列。
| しかし、各層が他のすべての層の結果に依存するため、同期操作やバリア操作が
*非常に多く* 必要になる。
| さらに、転送が必要なデータ量は、GPU間で層を分散する場合よりもさらに大きくなる可能性がある。したがって、帯域幅コストと複雑さのため、この方法は推奨しない。
| 最後に、データを複数GPUに分割する方法がある。この方法では、すべてのGPUが異なる観測に対してではあるものの、同じ種類の作業を行う。各ミニバッチの学習後に、勾配がGPU間で集約される。
| これは最も単純な方法であり、どのような状況にも適用できる。
| 必要なのは、各ミニバッチの後に同期することだけである。とはいえ、他の勾配がまだ計算中のうちから、勾配パラメータの交換を始められると非常に望ましいである。
| さらに、GPU数が増えるとミニバッチサイズも大きくなるため、学習効率が向上する。
| ただし、GPUを増やしても、より大きなモデルを学習できるようにはならない。
.. _fig_splitting:
.. figure:: ../img/splitting.svg
複数GPU上での並列化。左から右へ：元の問題、ネットワーク分割、層単位分割、データ並列。
| 複数GPUでのさまざまな並列化方法の比較を :numref:`fig_splitting`
に示す。
| 概して、十分に大きなメモリを持つGPUにアクセスできるなら、データ並列が最も便利な方法である。分散学習のための分割についての詳細な説明は
:cite:`Li.Andersen.Park.ea.2014`
も参照しよ。深層学習初期にはGPUメモリが問題でしたが、現在では非常に特殊な場合を除いてこの問題は解決されている。以下ではデータ並列に焦点を当てる。
データ並列
----------
| 1台のマシンに :math:`k`
枚のGPUがあると仮定する。学習するモデルが与えられると、各GPUは独立にモデルパラメータの完全な集合を保持する。ただし、GPU間のパラメータ値は同一であり、同期されている。
| 例として、 :numref:`fig_data_parallel` は :math:`k=2`
のときのデータ並列による学習を示している。
.. _fig_data_parallel:
.. figure:: ../img/data-parallel.svg
2枚のGPU上でデータ並列を用いてミニバッチ確率的勾配降下法を計算する様子。
一般に、学習は次のように進みる。
- 学習の任意の反復で、ランダムなミニバッチが与えられたら、バッチ内の例を
:math:`k` 個の部分に分け、GPUに均等に分配する。
- 各GPUは、自分に割り当てられたミニバッチ部分に基づいて、モデルパラメータの損失と勾配を計算する。
- :math:`k`
枚のGPUの局所勾配を集約して、現在のミニバッチ確率的勾配を得る。
- 集約された勾配を各GPUに再配布する。
- 各GPUは、このミニバッチ確率的勾配を使って、自分が保持しているモデルパラメータの完全な集合を更新する。
| 実際には、\ :math:`k` 枚のGPUで学習するときにはミニバッチサイズを
:math:`k` 倍に *増やす*
のが普通である。そうすることで、各GPUが単一GPUで学習する場合と同じ量の仕事をすることになる。16枚GPUのサーバでは、ミニバッチサイズがかなり大きくなるため、それに応じて学習率も増やす必要があるかもしれない。
| また、 :numref:`sec_batch_norm`
のバッチ正規化は調整が必要で、たとえばGPUごとに別々のバッチ正規化係数を保持する方法がある。
| 以下では、複数GPU学習を説明するためにおもちゃのネットワークを使う。
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
%matplotlib inline
from d2l import torch as d2l
import torch
from torch import nn
from torch.nn import functional as F
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
%matplotlib inline
from d2l import mxnet as d2l
from mxnet import autograd, gluon, np, npx
npx.set_np()
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# 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')
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
# 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()
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
[07:25:19] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
train(num_gpus=1, batch_size=256, lr=0.2)
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
train(num_gpus=1, batch_size=256, lr=0.2)
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
train(num_gpus=2, batch_size=256, lr=0.2)
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
train(num_gpus=2, batch_size=256, lr=0.2)
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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)
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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])
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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])
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
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')
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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])
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
[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)
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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)
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
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'))
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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)
.. raw:: latex
\diilbookstyleoutputcell
.. parsed-literal::
:class: output
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))]
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
#@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))
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
#@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))
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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
.. raw:: html
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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)}')
.. raw:: html
.. raw:: html
.. raw:: latex
\diilbookstyleinputcell
.. code:: python
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)}')
.. raw:: html
.. raw:: html