新機能 – Amazon SageMaker の管理されたデータ並列化による大規模なデータセットを使用したトレーニングのシンプル化

今日は、数百から数千ギガバイトにおよぶデータセットでのモデルのトレーニングを容易にする、新しいデータ並列化ライブラリの Amazon SageMaker によるサポートが開始されたことをお知らせしたいと思います。

データセットとモデルがますます大きくなり、高度化するにつれて、大規模な分散型トレーニングジョブを扱う機械学習 (ML) プラクティショナーは、Amazon Elastic Compute Cloud (EC2) p3 および p4 インスタンスなどの強力なインスタンスを使用している場合でさえも、長くなる一方のトレーニング時間に対応しなければなりません。たとえば、8 個の NVIDIA V100 GPU を搭載した ml.p3dn.24xlarge インスタンスを使用しても、一般公開されている COCO データセットでの Mask RCNN および Faster RCNN などの高度なオブジェクト検出モデルのトレーニングには 6 時間以上かかります。これと同じく、最先端の自然言語処理モデルである BERT のトレーニングにも、同一のインスタンスで 100 時間以上かかります。自律走行車企業などのお客様には、大規模な GPU クラスターで何日もかけて実行される、さらに大きなトレーニングジョブを定期的に処理するお客様もおられます。

ご想像どおり、これらの長いトレーニング時間は ML プロジェクトの深刻なボトルネックであり、生産性を損なうと共に、イノベーションを遅らせています。お客様から助けを求められた AWS は、この問題の解決に乗り出しました。

Amazon SageMaker のデータ並列化のご紹介
SageMaker Data Parallelism (SDP) ライブラリのおかげで、Amazon SageMaker を使って ML チームによる分散型トレーニングの時間とコストの削減を実現することが可能になりました。TensorFlow および PyTorch で利用できる SDP は、コンピューティングのより効率的な分散を実装し、ネットワーク通信を最適化して、AWS で最も高速な p3 と p4 GPU インスタンスをフル活用します。

このため、GPU リソースの最大 90% を、データ転送ではなくトレーニングに使用できるようになります。分散型トレーニングジョブは、使用される GPU の数にかかわらず、ほぼリニアなスケーリング効率性を実現することができます。言い換えると、単一のインスタンスで 8 時間かかるトレーニングジョブの場合、8 個のインスタンスでは約 1 時間しかかからず、コストの増加も最小限ですみます。SageMaker は、トレーニングコストとトレーニング時間におけるトレードオフを効果的に排除することで、ML チームがより早く結果を取得し、より迅速に反復して、イノベーションを加速化することを可能にします。

AWS re:Invent 2020 の基調講演で、Swami Sivasubramanian は T5-3B および Mask-RCNN で史上最速のトレーニング時間を実証しました。

• T5-3B モデルには 30 億個のパラメータがあり、自然言語処理ベンチマークで最新鋭の精度を達成しますが、トレーニングとパフォーマンスの調整には通常数週間におよぶ取り組みが必要です。AWS は、このモデルのトレーニングを、256 個の ml.p4d.24xlarge インスタンスを使って 6 日で完了しました。
• Mask-RCNN は、お客様によって今もなお広く使用されるインスタンスセグメンテーションモデルです。去年の re:Invent では、Mask-RCNN を PyTorch を用いて 26 分、TensorFlow を用いて 27 分でトレーニングしました。今年は、Mask-RCNN のトレーニングにおいて、TensorFlow では 6 分 12 秒、PyTorch では 6 分 45 秒という史上最速のトレーニング時間を記録しました。

Amazon SageMaker がこのような高速化をどのように実現するかを説明する前に、まずデータ並列化の仕組みと、スケーリングが難しい理由について説明したいと思います。

データ並列性の基礎
単一の GPU でモデルをトレーニングしている場合は、モデルパラメータ、オプティマイザパラメータ、および勾配 (パラメータ更新はバックプロパゲーションで計算) などのすべての内部状態をローカルで利用できますが、トレーニングジョブを GPU のクラスターに分散するとなると事情が違ってきます。

トレーニングセットは、「データ並列化」という名前の手法を使用して、GPU 全体に均等に分散されるミニバッチに分割されます。このため、各 GPU はデータセット全体のごく一部のみでモデルをトレーニングします。GPU がそれぞれ異なるバッチを処理するため、各 GPU でのモデル状態にわずかな違いが生じることになるのは明らかです。トレーニングのコンバージェンスを確実にするためにも、モデル状態はすべてのノードで定期的に更新される必要があります。これは、同期または非同期で実行できます。

• 同期トレーニング: すべての GPU が、その勾配更新をその他すべての GPU に報告 (多対多コミュニケーション)、またはそれらを再配布する中央パラメータサーバーに報告 (多対 1、その後 1 対多) します。すべての更新が同時に適用されるため、モデル状態はすべての GPU で同期化されており、次のミニバッチを処理できます。
• 非同期トレーニング: 勾配更新は、その他すべてのノード、または中央サーバーに送信されますが、それらは直ちに適用されるので、モデル状態は GPU ごとに異なることになります。

残念ながら、これらの手法ではスケーリングが困難です。GPU の数が増加するとともに、パラメータサーバーがどうしてもボトルネック化してしまいます。パラメータサーバーがない場合でも、ネットワーク輻輳がすぐに問題化します。これは、n 個の GPU がイテレーションごとに n*(n-1) 個のメッセージを交換する必要があり、合計で n*(n-1)*model size バイトになるからです。たとえば、ResNet-50 はコンピュータビジョンアプリケーションで使用される一般的なモデルです。2,600 万個のパラメータがあり、32 ビットの勾配更新にはそれぞれ約 100 メガバイトが必要です。8 個の GPU を使用すると、各イテレーションで 56 個の更新の送信と受信が必要となり、合わせて 5.6 ギガバイトになります。これは、高速ネットワークを使用したとしてもオーバーヘッドの原因となり、トレーニング時間を長引かせます。

2017年、Holovod プロジェクトのおかげで大幅な前進が遂げられました。Horovod は、分散型トレーニング用に最適化された「ring-allreduce」という名前の通信アルゴリズムを実装し、このアルゴリズムは瞬く間に人気の深層学習ライブラリに統合されました。

一言で言うと、ring-allreduce は分散型の非同期アルゴリズムです。パラメータサーバーはなく、ノードは有向閉路グラフ (つまり、一方向のリング) に編成されます。ノードは、イテレーションごとにその前にあるノードから勾配更新を受け取ります。ノードが独自のバッチを処理すると、両方の更新 (独自の更新と受け取った更新) を適用し、結果をネイバーに送信します。n 個の GPU では、各 GPU が 2*(n-1) 個のメッセージを処理してから、すべての GPU が更新されます。このため、GPU ごとに交換されるデータの総量は 2*(n-1)*model size になり、n*(n-1)*model size よりもはるかに優れています。

それでも、データセットが拡大し続けるにつれて、ネットワークのボトルネック問題も再発しやすくなります。そこで登場するのが SageMaker と新しい AllReduce アルゴリズムです。

Amazon SageMaker の新しいデータ並列化アルゴリズム
AllReduce アルゴリズムを使用すると、GPU は相互に通信しなくなります。各 GPU はその勾配更新を GPU メモリに保存します。特定のしきい値を超えると、これらの更新がシャーディングされ、GPU インスタンスの CPU で実行されているパラメータサーバーに送信されます。これにより、専用のパラメータサーバーが不要になります。

各 CPU がモデルパラメータのサブセットに対する責任を担い、すべての GPU からの更新を受信します。たとえば、単一の GPU を備えた 3 つのトレーニングインスタンスを使用する場合、トレーニングクラスター内の各 GPU が、3 つの CPU それぞれに勾配更新の 3 分の 1 を送信します。

次に、各 CPU が受信したすべての勾配更新を適用し、統合された結果をすべての GPU に配布します。

アルゴリズムの仕組みを理解したところで、インフラストラクチャを管理することなく、このアルゴリズムを独自のコードで使用する方法を見ていきましょう。

Amazon SageMaker でのデータ並列化を使ったトレーニング
SageMaker の Data Parallelism API は、使いやすさを考慮して設計されており、既存の分散型トレーニングのツールキットとシームレスに統合されます。ほとんどの場合、トレーニングコードで変更しなければならないのは、Horovod (TensorFlow)、または Distributed Data Parallel (PyTorch) の import 文のみです。

PyTorch の場合はこのようになります。

import smdistributed.dataparallel.torch.parallel.distributed as dist
dist.init_process_group()

次に、各 GPU を単一の SDP プロセスに固定する必要があります。

torch.cuda.set_device(dist.get_local_rank())

その後、通常通りにモデルを定義します。以下はその例です。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
...

最後に、モデルをインスタンス化し、それを使用して DistributedDataParallel オブジェクトを次のように作成します。

import torch
from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
device = torch.device("cuda")
model = DDP(Net().to(device))

残りのコードは vanilla PyTorch なので、SageMaker SDK で利用できる PyTorch 推定器を使用してトレーニングできます。ここでは、8 個の NVIDIA V100 GPU が搭載された ml.p3.16xlarge インスタンスを使用しています。

from sagemaker.pytorch import PyTorch
estimator = PyTorch(
    entry_point='train_pytorch.py',
    role=sagemaker.get_execution_role(),
    framework_version='1.6.0',
    instance_count=1,
    instance_type='ml.p3.16xlarge',
    distribution={'smdistributed':{'dataparallel':{enabled': True}}}
)
estimator.fit()

この先は SageMaker が引き継いで、必要なインフラストラクチャのすべてをプロビジョニングします。トレーニングジョブが実行されている間は、他のタスクに集中できます。

使用の開始
トレーニングジョブが複数の GPU で数時間または数日間継続されるという場合は、SageMaker Data Parallelism ライブラリが時間と費用を節約し、実験とイノベーションをより迅速に行うことを可能にしてくれるはずです。この機能は、本日から SageMaker が利用できるすべてのリージョンでご利用いただけます。

使用をすばやく開始するための例をご用意しましたので、ぜひお試しいただき、感想をお聞かせください。AWS では、お客様からのフィードバックをお待ちしております。フィードバックは、通常の AWS サポート担当者、または SageMaker の AWS フォーラム経由でお送りください。

– Julien