テンソル並列性 - Amazon SageMaker

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

テンソル並列性

テンソル並列処理は、特定のモデルの重み、勾配、オプティマイザ状態がデバイス間で分割されるモデル並列処理の一種です。個々の重みはそのままにし、重み、勾配、オプティマイザのセットをデバイス間で分割するパイプライン並列処理とは対照的に、テンソル並列処理は個々の重みをシャードします。これには、通常、モデルの特定のオペレーション、モジュール、またはレイヤーの分散計算が含まれます。

1 つのパラメータがGPUメモリの大部分を消費する場合 (語彙サイズが大きい大きな埋め込みテーブルや、クラス数が多い大きなソフトマックスレイヤーなど)、テンソル並列処理が必要です。この場合、この大きなテンソルまたはオペレーションを原子単位として扱うことは非効率的であり、メモリ負荷のバランスを妨げます。

SMP v2 は、テンソル並列処理の実装のために Transformer Engine と統合され、 PyTorch FSDP 上で実行されますAPIs。とSMPテンソル並列処理を同時に有効に PyTorch FSDPしてテンソル並列処理を決定し、最適なパフォーマンスを得るための最適なモデル並列処理を決定できます。

実際には、テンソル並列処理は、次のシナリオで特に役立ちます。

  • コンテキスト長が長いトレーニングを として行うと、FSDP単独でアクティベーションメモリが高くなります。

  • グローバルバッチサイズが目的の制限を超える、本当に大きなクラスターでトレーニングする場合。

SMP テンソル並列処理と互換性のあるHugging Face Transformer モデル

SMP v2 は現在、以下の Hugging Face トランスフォーマーモデルに対してテンソル並列処理をサポートしています。

  • GPT-NeoX

  • ラマ 2

これらのモデルにテンソル並列処理を適用するためのリファレンス設定については、「」を参照してください設定のヒント

テンソル並列処理を設定する

ではtensor_parallel_degree、テンソル並列処理の程度の値を選択します。値は、クラスターGPUs内の の数を均等に分割する必要があります。例えば、8 のインスタンスを使用してモデルをシャードするにはGPUs、2、4、または 8 を選択します。少数から始めて、モデルがGPUメモリに収まるまで徐々に増やすことをお勧めします。

次のコードスニペットは、 で導入された 2 ステッププロセスに従って、SMP初期化モジュールtorch.sagemaker.init()をトレーニングスクリプトに追加し、トレーニングジョブランチャー用の JSON 形式でSMP設定ディクショナリを設定する方法を示しています SageMaker モデル並列処理ライブラリ v2 を使用する。 PyTorch モデルやPyTorch FSDP設定を変更する必要はありません。tensor_parallel_degree、および random_seed、パラメータの詳細については、SMP v2 コア機能設定パラメータ を参照してください。

SMP 設定

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

トレーニングスクリプトで

で初期化torch.sagemaker.init()して SMP v2 をアクティブ化し、モデルを torch.sagemaker.transform でラップしますAPI。

import torch.sagemaker as tsm tsm.init() from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_config(..) model = tsm.transform(model)

Hugging Face Transformer チェックポイントの保存とロード

SMP ライブラリがモデルを変換すると、モデルの状態ディクショナリ (state_dict) が変更されます。これは、モデルが元の Hugging Face Transformer チェックポイント機能と互換性がないことを意味します。これを処理するために、SMPライブラリは、変換されたモデルからのチェックポイントAPIsを Hugging Face Transformer 表現に保存し、Hugging Face Transformer モデルチェックポイントをロードtorch.sagemaker.transformAPIして微調整するための を提供します。

SMP v2 のテンソル並列処理機能の使用中にチェックポイントを保存する方法の詳細については、「」を参照してくださいを使用したチェックポイント SMP

SMP v2 のテンソル並列処理機能を適用するモデルの微調整の詳細については、「」を参照してくださいファインチューニング