LatentSpaceがAmazonSageMakerモデルの並列処理ライブラリを使用して大規模なトランスフォーマーのフロンティアを推進した方法

ソースノード: 1204406

このブログは、Sarah Jane Hong CSO、Darryl Barnhart CTO、およびLatentSpaceのCEOであるIanThompsonとAWSのPremRangaによって共同執筆されています。

潜在空間は、機械学習(ML)モデルが学習する抽象的なアイデアの隠された表現です。 たとえば、「犬」、「花」、「ドア」は、潜在空間の概念または場所です。 で 潜在空間、私たちは、言語と視覚の両方のプロンプトでこの空間を操作および探索できるエンジンに取り組んでいます。 潜在空間チームは、グラフィックスと自然言語処理(NLP)という、長い間ほとんど重複していなかったXNUMXつの分野から来ています。 従来、画像とテキストのモダリティは別々に処理され、それぞれが複雑で、高価で、壊れやすい機能エンジニアリングの独自の歴史を持っていました。 ドキュメントの理解や質問応答などのNLPタスクは、通常、シーンの理解やレンダリングなどのビジョンタスクとほとんど共通点がなく、通常、タスクごとに非常に異なるアプローチとモデルを使用します。 しかし、これは急速に変化しています。

この単一の共有潜在空間でのモダリティのマージにより、ゲームからドキュメントの理解まで、新世代のクリエイティブおよび商用アプリケーションのロックが解除されます。 しかし、これらの新しいアプリケーションをXNUMXつのモデルでロック解除すると、RichardSuttonによる「TheBitterLesson」で強調されているように、新しいスケーリングの課題と、スケーリング法に関する過去数年間のエキサイティングな作業が開かれます。 これを可能にするために、Latent Spaceは、これらのモダリティを単一のモデルに融合するだけでなく、スケーリングして効率的に行うための最先端の研究に取り組んでいます。 これがモデルの並列処理の出番です。

アマゾンセージメーカーの独自の自動モデル分割と効率的なパイプラインアプローチにより、わずかなエンジニアリング作業でモデルの並列処理を採用でき、モデルのトレーニングを1億個のパラメーターを超えて拡張しました( p4d.24xlargeA100インスタンス)、これは私たちにとって重要な要件です。 さらに、SageMakerモデル並列処理ライブラリを使用して16ノード38 GPUトレーニングセットアップでトレーニングした場合、以前のトレーニング実行と比較して効率がXNUMX%向上したことを確認しました。

大規模な変圧器の訓練に関する課題

Latent Spaceでは、トランスフォーマーモデルの言語とビジョンを数十億のパラメーターと融合して、ユーザーの想像力からの「分布外」のユースケースをサポートしています。これは、現実の世界では発生しますが、トレーニングデータでは発生しません。 私たちは、XNUMXつの異なる方法で、数十億以上のパラメーターへのスケーリングに固有の課題に対処しています。

情報検索技術は、長い間、検索エンジンとQAタスクの重要なコンポーネントでした。 最近、古典的なIR技術と最新のトランスフォーマーを組み合わせてエキサイティングな進歩が見られました。特に、質問に答えるために関連ドキュメントを取得することを学習するニューラルレトリーバーとモデルを共同でトレーニングする質問応答タスク向けです。 概要については、FAIRの最近の作品をご覧ください。 検索拡張世代:インテリジェントな自然言語処理モデルの作成を合理化します & デコーダ内の融合、グーグルブレインの レルム、そしてNvidia's ニューラルレトリーバー 質問応答用。

検索拡張技術はコストと効率に役立ちますが、最大のモデルの単一のGPUにモデルを適合させることはできません。 これは、モデルの並列処理を使用してトレーニングする必要があることを意味します。 ただし、取得アーキテクチャの性質上、トレーニング入力全体で取得されたコンテキスト間の相互依存性のため、モデル分割の設計は困難でした。 さらに、モデルをどのように分割するかを決定したとしても、モデルの並列処理を導入することは、研究開発ライフサイクル全体で手動で行う重要なエンジニアリングタスクでした。

SageMakerモデルの並列処理ライブラリ

モデルの並列処理は、モデルを複数のデバイスまたはノード(GPUを備えたインスタンスなど)に分割し、効率的なパイプラインを作成してこれらのデバイス間でモデルをトレーニングし、GPUの使用率を最大化するプロセスです。 ザ・ モデル並列処理ライブラリ SageMakerでは、自動化されたモデル分割を提供することにより、モデルの並列処理をより利用しやすくします。 自動モデル分割 洗練されたパイプライン実行スケジューリング。 モデル分割アルゴリズムは、速度またはメモリ消費を最適化できます。 ライブラリは、メモリのバランスを取り、デバイス間の通信を最小限に抑え、パフォーマンスを最適化するパーティショニングアルゴリズムを使用します。

自動モデル分割

PyTorchのユースケースでは、モデル並列ライブラリは、モデルグラフを作成し、テンソルとパラメーターの形状を決定するトレースステップ(最初のトレーニングステップ)を内部的に実行します。 次に、ネストされたツリーで構成されるツリーを構築します nn.Module モデル内のオブジェクト、および保存された量など、トレースから収集された追加データ nn.Parameters、およびそれぞれのランタイム nn.Module.

次に、ライブラリはルートからこのツリーをトラバースし、計算負荷とメモリ使用量のバランスを取り、インスタンス間の通信を最小限に抑えるパーティショニングアルゴリズムを実行します。 複数のnn.Moduleが同じnn.Parameterを共有する場合、同じパラメーターの複数のバージョンを維持しないように、これらのモジュールは同じデバイスに配置されます。 パーティショニングの決定が行われた後、割り当てられたモジュールとウェイトがデバイスにロードされます。

パイプライン実行スケジューリング

SageMaker分散モデル並列ライブラリのもうXNUMXつのコア機能は次のとおりです。 パイプライン実行、モデルのトレーニング中にデバイス間で計算が行われ、データが処理される順序を決定します。 パイプライン化は、ミニバッチをマイクロバッチに分割することに基づいています。マイクロバッチは、トレーニングパイプラインにXNUMXつずつ供給され、ライブラリランタイムで定義された実行スケジュールに従います。

マイクロバッチパイプラインは、すべてのGPUが完全に利用されることを保証します。これは、自分で構築する必要があるものですが、モデルの並列処理ライブラリを使用すると、これは舞台裏でうまく行われます。 最後に、私たちは使用することができます アマゾンFSx、これは、取得を伴うマルチモーダルモデルのトレーニング中に読み取られるファイルの数を考慮して、読み取り速度を高速にするために重要です。

トレーニングアーキテクチャ

次の図は、トレーニングアーキテクチャの設定方法を表しています。 私たちの主な目的は、トレーニング速度を向上させ、コストを削減することでした。 私たちがトレーニングしている画像と言語のトランスフォーマーは非常に複雑で、内部に非常に多くのレイヤーと重みがあり、数十億のパラメーターが実行されているため、単一ノードのメモリに収まりません。 各ノードはモデルのサブセットを運び、それを介してデータフローと変換が共有およびコンパイルされます。 16を設定します p4d.24xlarge 次のアーキテクチャ表現を使用するXNUMXつのGPUを備えたインスタンス:

モデルをスケールアップするとき、一般的な傾向は、すべてをネットワークの重みに格納することです。 ただし、実用的な目的のために、モデルを拡張して、レンダリングのタスクに役立つ関連コンテキストを探す方法を学習したいと思います。 これにより、画質を犠牲にすることなく、サービスコストを抑えることができます。 大規模なトランスベースのNLPモデルを使用し、前述のように、次のように、SageMakerモデルの並列処理ライブラリを使用してトレーニング効率が38%向上することを確認しました。

  • テンソルレベルの並列処理の場合、すべての計算にallreduceが必要です。 これには、O(log_2 n)の並列ステップが必要です。 つまり、O(n log_2 n)の合計操作に対して、O(n)ステップを実行するn台のマシンです。
  • パイプラインの並列処理では、データをパイプラインに渡すためのO(1)並列ステップが必要です。
  • 16つのGPUを備えた1台のマシンを考えると、パイプライン並列のO(2)コスト、および深度ごとのモデル並列のO(log_8(3))= O(XNUMX)コストがあります。
  • この場合、SageMakerモデルの並列処理で使用するパイプライン並列に切り替えることでネットワークコストが1/3に削減され、全体的なトレーニングコストが1/2 + 1/2 * 1 / log_2(16)に削減されることがわかります。 )=元のコストの0.625で、対応する効率の向上につながります。

一般に、必要に応じて分散トレーニング(モデルサイズまたはトレーニングデータのスケーリングに関する問題)が必要な場合は、一連のベストプラクティスに従って、どのアプローチが最適かを判断できます。

分散トレーニングのベストプラクティス

私たちの経験に基づいて、分散データ並列アプローチから始めることをお勧めします。 次のような分散データの並列処理 SageMaker分散データ並列ライブラリ モデルレプリカに関するネットワークの問題のほとんどを解決するため、モデルを最小数のノードに適合させてから、必要に応じてバッチサイズをスケーリングするように複製する必要があります。

このシナリオで行ったように、トレーニング中にメモリが不足した場合は、モデル並列アプローチに切り替えることをお勧めします。 ただし、モデル並列トレーニングを試す前に、次の代替案を検討してください。

  • NVIDIA Tensor Coreを搭載したハードウェアでは、 混合精度トレーニング スピードアップを作成し、メモリ消費を削減します。
  • バッチサイズを減らします(または、可能であれば、画像の解像度またはNLPシーケンスの長さを減らします)。

さらに、で説明されているように、バッチ正規化がないモデル設計をお勧めします。 正規化なしの高性能大規模画像認識。 回避できない場合は、バッチの正規化がデバイス間で同期されていることを確認してください。 分散トレーニングを使用する場合、バッチはGPU間で分割されるため、正確なバッチ統計ではすべてのデバイス間で同期する必要があります。 これがないと、正規化によってエラーが増加し、収束が損なわれます。

次の制約がある場合は、モデル並列トレーニングから始めます。

  • モデルが単一のデバイスに適合しない
  • モデルのサイズが原因で、モデルの重みがGPUメモリの大部分を占め、より小さく、最適ではないバッチサイズを選択せざるを得ない場合など、より大きなバッチサイズを選択する際に制限に直面しています。

パフォーマンスを最適化するときは、次のようにします。

  • ノード間通信にパイプラインを使用して、遅延を最小限に抑え、スループットを向上させます
  • 気泡を最小限に抑えるために、パイプラインはできるだけ短くしてください。 マイクロバッチの数は、計算効率とバブルサイズのバランスを取り、少なくともパイプラインの長さになるように調整する必要があります。 必要に応じて、で説明されているように、トークンレベルでマイクロバッチを形成できます。 TeraPipe:大規模な言語モデルをトレーニングするためのトークンレベルのパイプライン並列処理

コストを最適化するときは、トレーニングにSageMakerマネージドスポットインスタンスを使用してください。 これにより、オンデマンドインスタンスよりもモデルのトレーニングコストを最大90%最適化できます。 SageMakerは、ユーザーに代わってスポットの中断を管理します。

考慮すべきその他の要因:

  • 高速相互接続がある場合のノード内では、より微妙な違いがあります。 十分なノード内ネットワーク容量がある場合は、より最適な計算のためにデータを再シャッフルすることでメリットが得られる場合があります。
  • アクティベーションがウェイトテンソルよりもはるかに大きい場合は、シャードオプティマイザーも役立つ場合があります。 参照してください ゼロ のガイドをご参照ください。

次の表は、いくつかの一般的なトレーニングスケールアップシナリオと、AWSでそれらを設定する方法を示しています。

<span class="notranslate">シナリオ</span> いつ適用されますか? ソリューション
単一のGPUから多数のGPUへのスケーリング トレーニングデータの量またはモデルのサイズが大きすぎる場合 3.16つのGPUを備えたpXNUMXxlargeなどのマルチGPUインスタンスに変更し、データと処理をXNUMXつのGPUに分割し、モデルのトレーニングにかかる​​時間をほぼ直線的に高速化します。
単一インスタンスから複数インスタンスへのスケーリング スケーリングの必要性がインスタンスサイズの変更を超えて拡張される場合 instance_typeをp3.16xlargeに設定し、instance_countを2に設定して、SageMaker Python SDKの推定関数を使用してインスタンスの数をスケーリングします。単一のp3.16xlargeに16つのGPUを配置する代わりに、XNUMXつの同一のインスタンスにXNUMXのGPUを配置します。 の使用を検討してください SageMaker分散データ並列ライブラリ.
トレーニングのためのモデル並列アプローチの選択 トレーニング中にメモリ不足エラーが発生した場合 を使用してモデル並列アプローチに切り替えます SageMaker分散モデル並列ライブラリ.
ノード間通信のネットワークパフォーマンス 複数のインスタンスを使用する分散トレーニングの場合(たとえば、AllReduce操作を実行するときのクラスター内のノード間の通信) インスタンスは、同じリージョンと同じアベイラビリティーゾーンにある必要があります。 SageMaker Python SDKを使用すると、これが処理されます。 トレーニングデータも同じアベイラビリティーゾーンにある必要があります。 の使用を検討してください SageMaker分散データ並列ライブラリ.
最適化されたGPU、ネットワーク、およびストレージ 大規模な分散トレーニングのニーズに対応 p4d.24xlargeインスタンスタイプは、高速ローカルストレージと最大400ギガビットの高速ネットワークバックプレーン用に設計されており、分散トレーニングの最もパフォーマンスの高いオプションとして強くお勧めします。

まとめ

SageMakerのモデル並列ライブラリを使用すると、自動化されたモデルのパーティション分割や効率的なパイプライン処理など、すぐに使用できる多くのメリットが得られます。 この投稿では、MLのユースケース、さまざまなトレーニングアプローチに関する考慮事項、AmazonSageMakerモデルの並列処理ライブラリを使用してトレーニングを高速化する方法について課題を共有しました。 何よりも、ここで説明するモデルの並列処理とパフォーマンスの向上のためのベストプラクティスを採用するのに数時間しかかからないようになりました。 この投稿が問題の解決に役立つ、または刺激を与える場合は、ぜひお知らせください。 コメントやフィードバックを共有してください。

参考文献

詳細については、以下を参照してください。


著者について

プレムランガ ジョージア州アトランタを拠点とするエンタープライズソリューションアーキテクトです。 彼は機械学習技術分野コミュニティの一員であり、MLとAIの旅で顧客と協力するのが大好きです。 Premはロボット工学に情熱を注いでおり、自動運転車の研究者であり、ヒューストンやその他の場所にAlexaが制御するBeerPoursを構築しました。

サラジェーンホン LatentSpaceの共同創設者兼最高科学責任者です。 彼女の経歴は、人間とコンピューターの相互作用と機械学習の交差点にあります。 彼女は以前、会話型AIスペースのビジネスにサービスを提供するSonar(Marchexが買収)でNLP研究を主導していました。 彼女はまた、尊敬されているAR / VR開発者であり、Oculus、Mozilla Mixed Reality、およびMicrosoftHololensから賞やフェローシップを受け取っています。

ダリル・バーンハート LatentSpaceの共同創設者兼最高技術責任者です。 彼は、GPUアクセラレーション、コンピューターグラフィックス、大規模データ、機械学習の経験を持つ経験豊富な開発者です。 他の情熱には、数学、ゲーム開発、および情報の研究が含まれます。

イアン・トンプソン LatentSpaceの創設者兼CEOです。 イアンは、私たちの生活に大きな影響を与えようとしているテクノロジーである「隣接する可能性」に触発されたエンジニア兼研究者です。 現在、安全で創造的なAIの構築を支援するために、マルチモーダル表現学習の簡素化とスケーリングに重点を置いています。 彼は以前、グラフィックス/バーチャルリアリティ(AltspaceVR、マイクロソフトが買収)および教育/ NLP(HSE)で企業を構築するのを支援しました。

ソース:https://aws.amazon.com/blogs/machine-learning/how-latent-space-used-the-amazon-sagemaker-model-parallelism-library-to-push-the-frontiers-of-large-スケールトランスフォーマー/

タイムスタンプ:

より多くの AWS機械学習ブログ