大規模言語モデル (LLM) トレーニングは、Llama 2、Falcon、Mistral などのいくつかの人気モデルのリリースにより、ここ 1 年で人気が急上昇しました。顧客は現在、175 億から XNUMX 億を超えるパラメータにわたる LLM を事前トレーニングおよび微調整し、ヘルスケアから金融、マーケティングに至るまで、さまざまな業界のアプリケーションのモデルのパフォーマンスを最適化しています。

この規模でパフォーマンスの高いモデルをトレーニングするのは困難な場合があります。高精度の LLM では、目標の精度を達成するために、テラバイトのトレーニング データと数千時間、さらには数百万時間のアクセラレータの計算時間が必要になる場合があります。トレーニングを完了してタイムリーに製品を発売するために、お客様は並列処理技術を利用して、この膨大なワークロードを最大数千のアクセラレータ デバイスに分散します。ただし、これらの並列処理手法は使用が難しい場合があります。さまざまな手法やライブラリは特定のワークロードとのみ互換性があるか、特定のモデル アーキテクチャに制限されており、トレーニングのパフォーマンスは不明瞭な構成に非常に敏感であり、最先端技術は急速に進化しています。その結果、機械学習の実践者は、LLM ワークロードを GPU の大規模クラスターに拡張するための準備に数週間を費やす必要があります。

この投稿では、の新機能を紹介します。 アマゾンセージメーカー モデル並列 (SMP) ライブラリは、大規模なモデルのトレーニング プロセスを簡素化し、LLM をより速くトレーニングするのに役立ちます。特に、オープンソースの PyTorch Fully Sharded Data Parallel (FSDP) API に基づいて構築された SMP ライブラリの新しい簡素化されたユーザー エクスペリエンス、数千億のパラメーターを使用したモデルのトレーニングを可能にする拡張されたテンソル並列機能、モデルのトレーニング時間を短縮するパフォーマンスの最適化について説明します。コストは最大 20% 削減されます。

SageMaker モデル並列ライブラリの詳細については、以下を参照してください。 SageMaker モデル並列処理ライブラリ v2 ドキュメンテーション。もご参照ください。 ノートブックの例 始めるために。

大規模なモデルのトレーニングを簡素化および高速化する新機能

この投稿では、SageMaker モデル並列ライブラリの v2.0 リリースに含まれる最新の機能について説明します。これらの機能により、ライブラリの使いやすさが向上し、機能が拡張され、トレーニングが高速化されます。次のセクションでは、新機能を要約し、ライブラリを使用して大規模モデルのトレーニングを高速化する方法について説明します。

SMP とオープンソース PyTorch の連携

2020 年の発売以来、SMP は SageMaker コンピューティング インスタンス上での高パフォーマンスで大規模なトレーニングを可能にしてきました。 SMP のこの最新メジャー バージョン リリースでは、ライブラリは API をオープン ソースの PyTorch と連携させることでユーザー エクスペリエンスを簡素化します。

PyTorch のオファー 完全にシャード化されたデータ並列処理 (FSDP) は、多くのコンピューティング デバイスにわたる大規模なトレーニング ワークロードをサポートするための主な方法です。次のコード スニペットに示されているように、シャード データ並列処理などの手法のための SMP の更新された API は、PyTorch の API を反映しています。単純に実行できます import torch.sagemaker の代わりにそれを使用します torch.

## training_script.py
import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
...
)

optimizer = ...
...

SMP の API に対するこれらの更新により、既存の PyTorch FSDP トレーニング スクリプトを全面的に見直すことなく、SageMaker と SMP ライブラリのパフォーマンス上の利点を実現できるようになりました。このパラダイムにより、オンプレミスでトレーニングする場合も SageMaker でトレーニングする場合と同じコード ベースを使用できるようになり、複数の環境でトレーニングする顧客のユーザー エクスペリエンスが簡素化されます。

既存の PyTorch FSDP トレーニング スクリプトを使用して SMP を有効にする方法の詳細については、次を参照してください。 SMP を使ってみる.

テンソル並列処理を統合して大規模なクラスターでのトレーニングを可能にする

SMP のこのリリースでは、PyTorch FSDP の機能も拡張され、テンソル並列処理技術が組み込まれています。シャード データ並列処理を単独で使用する場合の問題の XNUMX つは、クラスター サイズをスケールアップすると収束の問題が発生する可能性があることです。これは、データ並列ランク全体でパラメーター、勾配、およびオプティマイザーの状態をシャーディングすると、グローバル バッチ サイズも増加するためです。大規模なクラスターでは、このグローバル バッチ サイズが、モデルが収束するしきい値を超える可能性があります。クラスターを拡張するときにグローバル バッチ サイズを増やす必要のない追加の並列処理手法を組み込む必要があります。

この問題を軽減するために、SMP v2.0 では、シャード データ並列処理とテンソル並列処理を合成する機能を導入しています。テンソル並列処理により、グローバル バッチ サイズを変更したり、モデルの収束に影響を与えたりすることなく、クラスター サイズを増やすことができます。この機能を使用すると、256 ノード以上のクラスターをプロビジョニングすることで、トレーニングのスループットを安全に向上させることができます。

現在、PyTorch FSDP によるテンソル並列処理は SMP v2 でのみ利用可能です。 SMP v2 を使用すると、数行のコード変更でこの手法を有効にし、大規模なクラスターでも安定したトレーニングを実現できます。 SMP v2 との統合 変圧器エンジン テンソル並列処理の実装に優れており、PyTorch FSDP API との互換性を備えています。 PyTorch モデルまたは PyTorch FSDP 構成を変更せずに、PyTorch FSDP と SMP テンソル並列処理を同時に有効にすることができます。次のコード スニペットは、JSON 形式で SMP 構成辞書を設定し、SMP 初期化モジュールを追加する方法を示しています。 torch.sagemaker.init()、トレーニング ジョブの開始時にバックエンドで構成ディクショナリをトレーニング スクリプトに受け入れます。

SMP 構成は次のとおりです。

{
"tensor_parallel_degree": 8,
"tensor_parallel_seed": 0
}

トレーニング スクリプトでは、次のコードを使用します。

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

SMP でのテンソル並列処理の使用の詳細については、「 テンソル並列処理 ドキュメントのセクション。

高度な機能を使用してモデルのトレーニングを最大 20% 高速化します

SMP は、数百のインスタンスを含むクラスターでの分散トレーニングを可能にするだけでなく、モデルのトレーニングを最大 20% 高速化できる最適化テクニックも提供します。このセクションでは、これらの最適化のいくつかを取り上げます。詳細については、を参照してください。 コア機能 ドキュメントのセクション。

ハイブリッドシャーディング

シャード データ並列処理は、モデルの状態 (モデル パラメーター、勾配、オプティマイザーの状態) をデバイス間で分割する、メモリを節約する分散トレーニング手法です。このメモリ フットプリントの小ささにより、より大きなモデルをクラスターに適合させたり、バッチ サイズを増やすことができます。ただし、トレーニング中にシャード モデルのアーティファクトがさまざまなデバイスから頻繁に収集されるため、シャード データの並列処理によりトレーニング ジョブの通信要件も増加します。このように、シャーディングの度合いは、メモリ消費量と通信オーバーヘッドをトレードオフする重要な構成です。

デフォルトでは、PyTorch FSDP はクラスター内のすべてのアクセラレータ デバイスにわたってモデル アーティファクトをシャーディングします。トレーニング ジョブによっては、このシャーディング方法により通信のオーバーヘッドが増加し、ボトルネックが発生する可能性があります。これを支援するために、SMP ライブラリは、PyTorch FSDP 上で構成可能なハイブリッド シャード データ並列処理を提供します。この機能を使用すると、トレーニング ワークロードに最適なシャーディングの程度を設定できます。構成 JSON オブジェクトでシャーディングの程度を指定し、それを SMP トレーニング スクリプトに含めるだけです。

SMP 構成は次のとおりです。

{ "hybrid_shard_degree": 16 }

ハイブリッド シャード データ並列処理の利点について詳しくは、次を参照してください。 AWS での巨大なモデル トレーニングのほぼ線形のスケーリング。既存の FSDP トレーニング スクリプトを使用したハイブリッド シャーディングの実装の詳細については、を参照してください。 ハイブリッド共有データ並列処理 私たちのドキュメントにあります。

AWS インフラストラクチャに最適化された SMDDP 集合通信操作を使用する

SMP ライブラリは、 SageMaker 分散データ並列処理 (SMDDP) ライブラリ 分散トレーニングのワークロードを加速します。 SMDDP には最適化されたものが含まれています AllGather SageMaker p4d および p4de で高速化されたインスタンスで最高のパフォーマンスが得られるように設計された集合通信操作。分散トレーニングでは、GPU ワーカー間で情報を同期するために集合的な通信操作が使用されます。 AllGather これは、前方および後方の計算ステップの前にレイヤー パラメーターを具体化するためにシャード データ並列処理で通常使用されるコア集合通信操作の XNUMX つです。通信がボトルネックとなっているトレーニング ジョブの場合、集合操作を高速化することで、コンバージェンスに副作用を与えることなく、トレーニング時間とコストを削減できます。

SMDDP ライブラリを使用するには、トレーニング スクリプトに XNUMX 行のコードを追加するだけです。

import torch.distributed as dist

# Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp") # Replacing "nccl"

# Initialize with SMP
import torch.sagemaker as tsm
tsm.init()

SMP に加えて、SMDDP はオープンソースの PyTorch FSDP と DeepSpeed をサポートしています。 SMDDP ライブラリの詳細については、次を参照してください。 SageMaker 分散データ並列処理ライブラリを使用して分散トレーニングを実行する.

アクティベーションオフロード

通常、モデル トレーニングの前方パスは各レイヤーでアクティベーションを計算し、対応するレイヤーの後方パスが終了するまでそれらを GPU メモリに保持します。これらの保存されたアクティベーションは、トレーニング中に大量の GPU メモリを消費する可能性があります。アクティベーション オフロードは、フォワード パスの後にこれらのテンソルを CPU メモリに移動し、後で必要になったときに GPU にフェッチして戻す手法です。このアプローチにより、トレーニング中の GPU メモリの使用量を大幅に削減できます。

PyTorch はアクティベーション オフロードをサポートしていますが、その実装は非効率的であり、バックワード パス中にアクティベーションが CPU からフェッチバックされている間、GPU がアイドル状態になる可能性があります。これにより、アクティベーション オフロードを使用するときにパフォーマンスが大幅に低下する可能性があります。

SMP v2 は、トレーニング パフォーマンスを向上させる、最適化されたアクティベーション オフロード アルゴリズムを提供します。 SMP の実装は、GPU で必要になる前にアクティベーションをプリフェッチし、アイドル時間を短縮します。

SMP は PyTorch の API 上に構築されているため、最適化されたアクティベーション オフロードを有効にするには、わずか数行のコード変更が必要です。関連する構成を追加するだけです (sm_activation_offloading & activation_loading_horizon パラメータ) を作成し、トレーニング スクリプトに含めます。

SMP 構成は次のとおりです。

{
"activation_loading_horizon": 2,
"sm_activation_offloading": True
}

トレーニング スクリプトでは、次のコードを使用します。

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
)

model = offload_wrapper(model)

アクティベーションオフロード用のオープンソース PyTorch チェックポイント ツールの詳細については、 チェックポイント_wrapper.py PyTorch GitHub リポジトリ内のスクリプトと アクティベーションチェックポイント PyTorch ブログ投稿内 Pytorch 分散を使用した TorchMultimodal でのマルチモーダル基盤モデルのスケーリング。 SMP のアクティベーション オフロードの最適化された実装の詳細については、 アクティベーションオフロード ドキュメントのセクション。

SMP は、ハイブリッド シャーディング、SMDDP、アクティベーション オフロードに加えて、大規模なモデルのトレーニング ワークロードを高速化できる追加の最適化を提供します。これには、最適化されたアクティブ化チェックポイント設定、遅延パラメーター初期化などが含まれます。詳細については、を参照してください。 コア機能 ドキュメントのセクション。

まとめ

データセット、モデルのサイズ、トレーニング クラスターが拡大し続けるにつれて、タイムリーかつ手頃な価格でモデルと製品を提供するには、効率的な分散トレーニングがますます重要になっています。 SageMaker モデル並列ライブラリの最新リリースは、コードの変更を減らし、PyTorch FSDP API と連携することでこれを達成するのに役立ちます。これにより、テンソル並列処理と最適化によって大規模なクラスターでのトレーニングが可能になり、トレーニング時間を最大 20% 削減できます。

SMP v2 の使用を開始するには、次のドキュメントを参照してください。 ドキュメント と サンプルノートブック.

著者について

ロバート・ヴァン・デューセン Amazon Sagemakerのシニアプロダクトマネージャーです。 彼は、ディープラーニングトレーニングのためのフレームワーク、コンパイラ、および最適化技術をリードしています。

ルイ・キンテラ AWS SageMaker モデル並列ライブラリのソフトウェア開発者マネージャーです。暇なときは、サンフランシスコのベイエリアでハーレーに乗っている姿が見られます。

ゴータム・クマール は、AW​​S AI Deep Learning のソフトウェア エンジニアです。彼は AI 用のツールとシステムの構築に情熱を注いでいます。余暇には、サイクリングと読書を楽しんでいます。

ラフル・ウィゴル アマゾン ウェブ サービスの分散ディープラーニングのシニア ソフトウェア開発エンジニアです。

• SEO を活用したコンテンツと PR 配信。 今日増幅されます。
• PlatoData.Network 垂直生成 Ai。 自分自身に力を与えましょう。 こちらからアクセスしてください。
• プラトアイストリーム。 Web3 インテリジェンス。 知識増幅。 こちらからアクセスしてください。
• プラトンESG。 カーボン、 クリーンテック、 エネルギー、 環境、 太陽、 廃棄物管理。 こちらからアクセスしてください。
• プラトンヘルス。 バイオテクノロジーと臨床試験のインテリジェンス。 こちらからアクセスしてください。
• 情報源： https://aws.amazon.com/blogs/machine-learning/amazon-sagemaker-model-parallel-library-now-accelerates-pytorch-fsdp-workloads-by-up-to-20/