Amazon SageMaker トレーニングで PyTorch Lightning とネイティブ PyTorch DDP を実行し、Amazon Search を利用する

非常に多くのデータ、非常に短い時間。 機械学習 (ML) の専門家、データ サイエンティスト、エンジニア、愛好家は、世界中でこの問題に遭遇しています。 自然言語処理からコンピューター ビジョン、表形式から時系列、およびその間のすべてに至るまで、可能な限り多くの GPU に対してデータを実行する際の速度を最適化するという古くからの問題は、無数のソリューションに影響を与えてきました。 本日、ネイティブのオープンソース フレームワークを使用する PyTorch 開発者向けの機能を発表できることをうれしく思います。 PyTorchライトニング & PyTorch DDP、これにより、クラウドへの道が合理化されます。

アマゾンセージメーカー は ML 向けのフルマネージド サービスであり、SageMaker モデル トレーニングは大規模な高性能トレーニング用に最適化されたコンピューティング環境です。 SageMaker モデルトレーニングは、シームレスなコントロールプレーンを使用したリモートトレーニング体験を提供し、ML モデルを簡単にトレーニングして再現し、高性能かつ低コストにします。 SageMaker トレーニング ポートフォリオの新機能を発表できることを嬉しく思います。これにより、大規模な PyTorch の実行がさらに簡単になり、アクセスしやすくなります。

1. PyTorch Lightning を統合できるようになりました SageMaker の分散データ並列ライブラリ たったXNUMX行のコード変更で。
2. SageMaker モデルのトレーニング は、NCCL バックエンドを使用したネイティブ PyTorch Distributed Data Parallel をサポートするようになり、開発者はこれまで以上に簡単に SageMaker に移行できるようになりました。

この投稿では、これらの新機能について説明し、Amazon Search が SageMaker の最適化された分散トレーニング バックエンドで PyTorch Lightning を実行して、モデルのトレーニング時間を短縮した方法についても学びます。

Amazon Search のケース スタディに入る前に、なじみのない人のために、いくつかの背景を説明したいと思います。 SageMaker の分散データ並列ライブラリ. 2020 年に、クラスター全体の効率を向上させる大規模な分散勾配降下のカスタム クラスター構成を開発し、開始しました。 ニシン. SageMaker Distributed Data Parallel (SMDDP) は、パラメーターサーバーとリングベースのトポロジーの両方の長所を利用して最適化されています。 アマゾン エラスティック コンピューティング クラウド (Amazon EC2) ネットワーク トポロジー (以下を含む) EFA. より大きなクラスター サイズの場合、SMDDP は、Horovod (TensorFlow) および PyTorch Distributed Data Parallel と比較して、20 ～ 40% のスループット向上を実現できます。 クラスター サイズが小さく、サポートされているモデルの場合は、 SageMakerトレーニングコンパイラ、全体のジョブ時間を最大 50% 削減できます。

お客様のスポットライト: Amazon Search を使用した SageMaker の最適化されたバックエンドでの PyTorch Lightning

Amazon Search は、 Amazon.com. これは、Amazon で購入する商品を探している顧客の検索エクスペリエンスを強化します。 大まかに言うと、Amazon 検索は、Amazon.com で販売されているすべての製品のインデックスを作成します。 顧客がクエリを入力すると、Amazon Search は深層学習モデルを含むさまざまな ML 手法を使用して、関連性のある興味深い製品を顧客のクエリに一致させます。 次に、顧客に結果を表示する前に製品をランク付けします。

Amazon Search の科学者は、PyTorch に追加された使いやすさの機能により、検索ランキングを強化する深層学習モデルをトレーニングするための主要なフレームワークの XNUMX つとして PyTorch Lightning を使用しています。 この新しい SageMaker のリリース前は、PyTorch Lightning で記述された深層学習モデルでは SMDDP はサポートされていませんでした。 これにより、PyTorch Lightning を使用することを好む Amazon Search の科学者は、データ並列技術を使用してモデル トレーニングをスケーリングできなくなり、トレーニング時間が大幅に遅くなり、よりスケーラブルなトレーニングを必要とする新しい実験をテストできなくなりました。

チームの初期のベンチマーク結果では、単一ノードのトレーニング ベースラインと比較して、7.3 つのノードでトレーニングした場合、サンプル モデルのトレーニング時間が 200 倍高速であることが示されています。 これらのベンチマークで使用されるベースライン モデルは、3.16 つの密な全結合層と XNUMX 以上のパラメーターを持つ多層パーセプトロン ニューラル ネットワークです。 次の表は、ml.pXNUMXxlarge SageMaker トレーニング インスタンスのベンチマーク結果をまとめたものです。

次に、新しいローンチの詳細について説明します。 必要に応じて、対応する手順を実行できます 例のノート。

SageMaker 分散トレーニング ライブラリで PyTorch Lightning を実行する

SageMaker Data Parallel が、SageMaker トレーニング内で PyTorch Lightning とシームレスに統合されるようになったことをお知らせします。

PyTorch Lightning は、PyTorch でのカスタム モデルの記述を簡素化するオープン ソース フレームワークです。 Keras が TensorFlow や Hugging Face に対して行ったことと同様の方法で、PyTorch Lightning は、PyTorch 自体の低レベル機能の多くを抽象化した高レベル API を提供します。 これには、モデルの定義、プロファイリング、評価、プルーニング、モデルの並列処理、ハイパーパラメーター構成、転移学習などが含まれます。

以前は、PyTorch Lightning 開発者は、トレーニング コードを高性能 SageMaker GPU クラスターにシームレスに移行する方法について確信が持てませんでした。 さらに、SageMaker Data Parallel によって導入された効率の向上を利用する方法がありませんでした。

PyTorch Lightning の場合、一般的に言えば、SageMaker Training でこれらの API を実行するだけのコード変更はほとんどまたはまったく必要ありません。 ノートブックの例では、 DDP頭脳系 & DDPPプラグイン 方法。

最適化されたバックエンドとして SageMaker Data Parallel で PyTorch Lightning を使用するには、XNUMX つのステップがあります。

1. サポートされている AWSディープラーニングコンテナ (DLC) を基本イメージとして使用するか、オプションで独自のコンテナーを作成して、 SageMaker Data Parallel バックエンドを自分でインストールする. 必要なパッケージに PyTorch Lightning が含まれていることを確認します。 requirements.txt ファイルにソフトウェアを指定する必要があります。
2. 最適化されたバックエンドを有効にするために、トレーニング スクリプトにいくつかの小さなコード変更を加えます。 これらには以下が含まれます：
   1. SM DDP ライブラリをインポートします。

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. SageMaker 用に PyTorch Lightning 環境をセットアップします。

      from pytorch_lightning.plugins.environments.lightning_environment 
        import LightningEnvironment
      
      env = LightningEnvironment()
      env.world_size = lambda: int(os.environ["WORLD_SIZE"])
      env.global_rank = lambda: int(os.environ["RANK"])

   3. 1.5.10 より古いバージョンの PyTorch Lightning を使用している場合は、さらにいくつかの手順を追加する必要があります。
      1. まず、環境変数を追加します。

         os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

      2. 第二に、必ず使用してください DDPPluginよりむしろ DDPStrategy. より新しいバージョンを使用している場合は、 requirements.txt セクションに source_dir あなたの仕事では、これは必要ありません。 次のコードを参照してください。

         ddp = DDPPlugin(parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)], cluster_environment=env)

   4. オプションで、プロセス グループ バックエンドを次のように定義します。 "smddp" セクションに DDPSTrategy 物体。 ただし、サポートされている PyTorch DDP バックエンドで PyTorch Lightning を使用している場合は、単にこの ` を削除してください。process_group_backend` パラメータ。 次のコードを参照してください。

      ddp = DDPStrategy(
        cluster_environment=env, 
        process_group_backend="smddp", 
        accelerator="gpu")

3. 次のように、見積もりツールに記載されている配布方法があることを確認してください。 distribution={"smdistributed":{"dataparallel":{"enabled":True} Herring バックエンドを使用している場合、または distribution={"pytorchddp":{"enabled":True}.

• 適切なパラメータの完全なリストについては、 distribution パラメータについては、ドキュメントを参照してください こちら.

これで、SageMaker トレーニングジョブを開始できます! Python SDK、Boto3、SageMaker コンソール、 AWSコマンドラインインターフェイス (AWS CLI) など、数え切れないほどの方法があります。 AWS の観点からは、これは 単一の API コマンド: create-training-job. このコマンドをローカル ターミナルから起動するかどうかに関係なく、 AWSラムダ 関数、 Amazon SageMakerスタジオ ノートブック、KubeFlow パイプライン、またはその他のコンピューティング環境は完全にあなた次第です。

PyTorch Lightning と SageMaker Data Parallel の統合は現在、1.11 以降の新しいバージョンの PyTorch でのみサポートされていることに注意してください。 さらに、このリリースは、PyTorch 1.12 以降の SageMaker の AWS DLC でのみ利用できます。 この画像をベースとして指定してください。 の us-east-1、このアドレスは次のとおりです。

ecr_image = '763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker'

次に、これをベース イメージとして使用して Docker コンテナを拡張するか、これを変数として image_uri の引数 SageMaker トレーニング推定器。

その結果、AWS で利用可能な最高のパフォーマンスで、SageMaker Training の最適化された GPU で PyTorch Lightning コードを実行できるようになります。

SageMaker で PyTorch 分散データを並列に実行する

最大の問題 PyTorch 分散データ並列 (DDP) solves は一見シンプルです: 速度です。 優れた分散トレーニング フレームワークは、安定性、信頼性、そして最も重要なこととして、優れたパフォーマンスを大規模に提供する必要があります。 PyTorch DDP は、シングルノード設定とマルチノード設定の両方で、複数の GPU デバイスを介してモデルをレプリケートするための API をトーチ開発者に提供することで、これを実現します。 次に、フレームワークは、トレーニング データセットから各モデル コピーへのさまざまなオブジェクトのシャーディングを管理し、各モデル コピーの勾配を平均化して、各ステップでそれらを同期します。 これにより、完全なトレーニング実行が完全に完了すると、XNUMX つのモデルが生成されます。 次の図は、このプロセスを示しています。

PyTorch DDP は、大規模なデータセットを使用するプロジェクトで一般的です。 各データセットの正確なサイズは大きく異なりますが、一般的なガイドラインは、データセットをスケーリングし、サイズを計算し、サイズをモデル化することです。 とも呼ばれている スケーリング則、これら XNUMX つの最適な組み合わせは非常に議論の余地があり、アプリケーションによって異なります。 AWS では、複数の顧客との協力に基づいて、データセット全体のサイズが少なくとも数十 GB の場合に、データ並列戦略の利点を明確に見ることができます。 データセットがさらに大きくなった場合、ある種のデータ並列戦略を実装することは、実験全体を高速化し、価値実現までの時間を短縮するための重要な手法です。

以前は、オンプレミスまたは他のコンピューティング環境での分散トレーニングに PyTorch DDP を使用していたお客様は、プロジェクトを SageMaker Training に簡単に移行して、シームレスなコントロールプレーンで高性能 GPU を活用するためのフレームワークがありませんでした。 具体的には、データ並列フレームワークを SMDDP に移行するか、SageMaker Training で PyTorch DDP の機能を手動で開発およびテストする必要がありました。 現在、SageMaker トレーニングは、PyTorch DDP コードをオンボーディングするお客様にシームレスなエクスペリエンスを提供できることをうれしく思います。

これを効果的に使用するために、トレーニング スクリプトを変更する必要はありません。

この新しいパラメーターは、次のコードで確認できます。 の中に distribution パラメータ、単純に追加 pytorchddp として有効に設定 true.

estimator = PyTorch(
    base_job_name="pytorch-dataparallel-mnist",
    source_dir="code",
    entry_point = "my_model.py",
    ... 
    # Training using SMDataParallel Distributed Training Framework
    distribution = {"pytorchddp": {"enabled": "true"}}
)

この新しい構成は、SageMaker Python SDK バージョン 2.102.0 および PyTorch DLC の 1.11 で開始されます。

人気のある PyTorch DDP 開発者向け たいまつ これは、堅牢な耐障害性をすでに提供している SageMaker トレーニング環境では必要ないことを知っておくと役立ちます。 ただし、コードの書き換えを最小限に抑えるために、このコマンドを実行する別のランチャー スクリプトをエントリ ポイントとして使用できます。

現在、PyTorch 開発者はスクリプトを SageMaker に簡単に移動できるため、スクリプトとコンテナを複数のコンピューティング環境でシームレスに実行できます。

これにより、将来的に利用する準備が整います。 SageMaker の分散トレーニング ライブラリ AWS に最適化されたトレーニング トポロジを提供し、最大 40% の高速化を実現します。 PyTorch 開発者にとって、これは XNUMX 行のコードです! PyTorch DDP コードの場合、単純にバックエンドを smddp 初期化で（参照 PyTorchトレーニングスクリプトを変更する)、次のコードに示すように:

import smdistributed.dataparallel.torch.torch_smddp
import torch.distributed as dist
dist.init_process_group(backend='smddp')

上で見たように、バックエンドを設定することもできます DDPStrategy 〜へ smddp ライトニング使用時。 これにより、最大で 全体の速度が 40% 向上 大規模クラスター向け！ SageMaker での分散トレーニングの詳細については、次を参照してください。 オンデマンドのウェビナー、サポート ノートPC、関連 ドキュメント, 論文.

まとめ

この記事では、SageMaker Training ファミリーに XNUMX つの新機能を紹介しました。 これらにより、PyTorch 開発者は、PyTorch DDP と PyTorch Lightning の両方で、SageMaker で既存のコードを使用することがはるかに簡単になります。

また、Amazon Search が深層学習モデルのトレーニングに SageMaker Training を使用する方法、特に SageMaker Data Parallel 最適化集合ライブラリをバックエンドとして使用する PyTorch Lightning を使用する方法も示しました。 分散トレーニング全体に移行することで、Amazon Search はトレーニング時間を 7.3 倍高速化しました。

著者について

エミリー・ウェバー SageMakerがリリースされた直後にAWSに参加し、それ以来、世界中にそのことを伝えようとしています。 エミリーは、顧客のために新しいML体験を構築する以外に、チベット仏教の瞑想と勉強を楽しんでいます。

カラン・ディマン カナダのトロントを拠点とする AWS のソフトウェア開発エンジニアです。 彼は、機械学習の分野と、分散コンピューティング ワークロードを加速するためのソリューションの構築に非常に情熱を注いでいます。

ヴィシュワ・カリア AWS Deep Engine のソフトウェア開発エンジニアです。 彼女の関心は、機械学習と分散システムの交差点にあり、テクノロジーと AI における女性のエンパワーメントにも情熱を注いでいます。

エイマン・エルナハラヴィ Amazon Search のプリンシパル ソフトウェア エンジニアであり、機械学習の加速、スケーリング、自動化の取り組みをリードしています。 彼女の専門知識は、機械学習、分散システム、パーソナライゼーションなど、複数の分野に及びます。