Amazon Web Services ブログ

数十億ものパラメータを持つ深層学習モデルのトレーニングをシンプル化する Amazon SageMaker

今日は、ハードウェアの制限が原因で、これまでトレーニングすることが難しかった超大型深層学習モデルのトレーニングを Amazon SageMaker がシンプル化することをご紹介したいと思います。

過去 10 年の間、深層学習 (DL) と呼ばれる機械学習のサブセットが一世を風靡してきました。ニューラルネットワークを基盤とする DL アルゴリズムは、膨大な量の非構造化データ (画像、動画、スピーチ、またはテキストなど) に隠された情報パターンを抽出する、類いまれな能力を備えています。DL は、さまざまな複雑かつ人間的なタスク、特にコンピュータビジョンと自然言語処理において、瞬く間に目覚ましい成果を達成しました。現に、DL は ImageNet Large Scale Visual Recognition Challenge (ILSVRC)、the General Language Understanding Evaluation (GLUE)、または Stanford Question Answering Dataset (SQUAD) といったリファレンスタスクにおける結果を向上させ続けているため、イノベーションがかつてない速さで進んでいます。

これまで以上に複雑なタスクに挑戦するために、DL 研究者はますます高度なモデルを設計し、さらなるニューロン層と結合を追加してパターン抽出と予測精度を向上させており、モデルサイズに直接的な影響を及ぼしています。たとえば、画像分類では 100 メガバイトの ResNet-50 モデルで極めて良好な結果を得ることができますが、オブジェクト検出やインスタンスセグメンテーションなどのより困難なタスクには、約 250 メガバイトの Mask R-CNN または YOLO v4 などのより大きなモデルを使用しなければならなくなります。

想像がつくと思いますが、モデルの増大もモデルのトレーニングに必要な時間とハードウェアリソースに影響します。Graphical Processing Units (GPU) が以前から大型 DL モデルのトレーニングと微調整に好まれるオプションであるのはこのためです。GPU の超並列的なアーキテクチャと大型のオンボードメモリのおかげで、ミニバッチトレーニングと呼ばれる手法の使用が可能になります。複数のデータサンプルを、ひとつずつではなく、一度に GPU に送信することによって、通信オーバーヘッドが削減され、トレーニングジョブが大幅に高速化されます。たとえば、Amazon Elastic Compute Cloud (EC2) p4 ファミリーで利用できる NVIDIA A100 は、7,000 個を超えるコンピューティングコアと、40 ギガバイトの高速オンボードメモリを備えています。超大型モデルで大量のデータをトレーニングするには、これで十分なはずではないでしょうか?

ところが、そうではありません。OpenAI GPT-2 (15 億パラメータ)、T5-3B (30 億パラメータ)、および GPT-3 (1,750 億パラメータ) などの巨大な自然言語処理では、数十ギガバイト、あるいは数百ギガバイトもの GPU メモリが消費されます。高解像度 3D 画像を処理する最先端のモデルも同様に、バッチサイズを 1 にした場合でさえも、大きすぎて GPU メモリに収まらないことがあります。

DL 研究者は、不可能なことを可能するために、以下のような手法の組み合わせを使用します。

  • より強力な GPU を購入する (モデルによってはこれが選択肢にさえならないことを確認したところではありますが)。
  • 性能が低いモデルで作業して、精度を犠牲にする。
  • 勾配チェックポイントを実装する。これは、メモリにすべてを保存しておくのではなく、トレーニングの中間結果をディスクに保存することに頼る手法で、トレーニング時間が 20～30% 遅くなります。
  • モデル並列化を実装する。つまり、モデルを手動で分割して、モデルの (小さな) 断片を異なる GPU でトレーニングするということです。言うまでもなく、これはエキスパートプラクティショナーにとっても極めて困難で、時間がかかる不確かなタスクです。

お客様からは、超大型モデルでの作業に対しては上記のいずれも十分満足のいくソリューションではないという声が聞かれました。よりシンプルで、コスト効率性に優れたソリューションを求められた AWS は、これのソリューションの実現に乗り出しました。

Amazon SageMaker のモデル並列化のご紹介
SageMaker のモデル並列化は、自動的に複数の GPU 全体でモデルを効率的にパーティション化し、精度面での妥協や複雑な手動作業を行う必要性を排除します。さらに、モデルトレーニングに対するこのスケールアウトアプローチにより、メモリのボトルネックを発生させずに超大型モデルで作業することを可能にするだけでなく、小型でコスト効率性に優れた GPU を多数活用することも可能になります。

ローンチ時点で、この機能は TensorFlow と PyTorch でサポートされ、コードには最小限の変更しか必要ありません。トレーニングジョブを開始するときは、モデルを速度またはメモリ使用率のどちらに最適化するかを指定できます。その後、Amazon SageMaker がユーザーに代わって最初のプロファイリングジョブを実行し、モデルのコンピューティングおよびメモリ要件を分析します。この情報は、通信を最小限に抑えながら、モデルを分割する方法とモデルパーティションを GPU にマッピングする方法を決定するパーティショニングアルゴリズムに提供されます。パーティショニング判定の結果はファイルに保存され、実際のトレーニングジョブに入力として渡されます。

このように、SageMaker はすべての作業を実行します。希望に応じて、モデルのプロファイリングとパーティション化を手動で実行してから、SageMaker でトレーニングすることも可能です。

コードを検証する前に、内部の構造を簡単に説明したいと思います。

モデルパーティションとマイクロバッチを使用したトレーニング
異なる GPU で実行されているモデルパーティションは、互いにフォワードパス入力 (活性値) を期待するため、一連のパーティション全体でのトレーニングミニバッチの処理では、常に 1 つのパーティションのみがビジー状態になり、残りのパーティションは停止されます。

この非効率な動作を避けるため、ミニバッチは、異なる GPU で並列的に処理されるマイクロバッチに分割されます。たとえば、GPU #1 がマイクロバッチ n のフォワードプロパゲーションを実行する一方で、GPU #2 がマイクロバッチ n+1 に同じことを実行できます。活性値は保存して、次のパーティションの受け入れ態勢が整ったときに渡すことができます。

バックプロパゲーションについても、パーティションは互いに入力値を期待します (勾配)。パーティションはフォワードプロパゲーションとバックプロパゲーションを同時に実行できないため、すべての GPU がそれぞれのマイクロバッチでのフォワードパスを完了するまで待ってから、対応するバックワードパスを実行させることができます。このシンプルモードは、Amazon SageMaker で使用することができます。

インターリーブモードと呼ばれる、さらに効率的なオプションもあります。このモードでは、SageMaker がマイクロバッチの数に従ってパーティションをレプリケートします。たとえば、2 つのマイクロバッチの処理では、各 GPU が受け取ったパーティションの 2 つのコピーを実行します。各コピーは、フォワードプロパゲーションまたはバックプロパゲーションのために、他の GPU で実行されているパーティションと連携します。

これは、4 つの異なるマイクロバッチが 2 つの重複するパーティションによって処理される状況を図解したものです。

説明図

要約すると、SageMaker が GPU の使用率を最大限に活かす方法は、異なるマイクロバッチのフォワードパスとバックワードパスをインターリブするということになります。

それでは、TensorFlow でこれを動作させる方法を見ていきましょう。

Amazon SageMaker のモデル並列化の実装
SageMaker Model Parallelism (SMP) ライブラリのおかげで、独自の TensorFlow コードへのモデル並列化の実装は簡単に実行できます (PyTorch でも同じようなプロセスになります)。実行する必要がある事柄は以下のとおりです。

  • パーティショニング設定を定義して初期化する。
  • 標準の Keras subclassing を使用して、モデルを DistributedModel クラスのサブクラスにする。
  • モデルのフォワードステップとバックワードステップを表すトレーニング関数 @smp.step を作成し、修飾する。この関数は、前のセクションで説明したアーキテクチャに従ってパイプライン処理されます。
  • オプションで、同様にパイプライン処理される evaluation 関数にも同じ手順を実行する。

4 個の NVIDIA V100 GPU が搭載された ml.p3.8xlarge インスタンスを使用して、これを MNIST データセットでのシンプルな畳み込みネットワークトレーニングに適用しましょう。

まず、SMP API を初期化します。

import smdistributed.modelparallel.tensorflow as smp
smp.init()

次に、DistributedModel をサブクラス化して、モデルを構築します。

class MyModel(smp.DistributedModel):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv = Conv2D(32, 3, activation="relu")
        self.flatten = Flatten()
        self.dense1 = Dense(128)
        self.dense2 = Dense(10)
. . .

トレーニング関数はこのようになります。

@smp.step
def forward_backward(images, labels):
    predictions = model(images, training=True)
    loss = loss_obj(labels, predictions)
    grads = optimizer.get_gradients(loss, model.trainable_variables)
    return grads, loss

これで、いつものように、SageMaker SDK で利用できる TensorFlow 推定器を使用してトレーニングできるようになります。あとは、モデル並列化設定、つまり、2 つのパーティション (2 つの GPU でトレーニングが行われる) と 2 つのマイクロバッチ (パーティションごとに 2 つのコピーができる) をインターリービングと共に追加するだけです。

smd_mp_estimator = TensorFlow(
    entry_point="tf2.py",
    role=role,
    framework_version='2.3.1',
    pv_version='py3',
    instance_count=1,
    instance_type='ml.p3.16xlarge',
    distribution={
        "smdistributed": {
            "modelparallel": {
                "enabled":True,
                "parameters": {
                    "microbatches": 2,
                    "partitions": 2,
                    "pipeline": "interleaved",
                    "optimize": "memory",
                    "horovod": True,
                }
             }
         },
        "mpi": {
            "enabled": True,
            "processes_per_host": 2, # Pick your processes_per_host
            "custom_mpi_options": mpioptions
        },
    }
)

使用の開始
ご覧いただいたように、モデルの並列化は最先端の超大型深層学習モデルのトレーニングを容易にします。この機能は、本日から Amazon SageMaker が利用できるすべてのリージョンでご利用いただけます。

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

– Julien