PyTorch

独自の PyTorch モデルを SageMaker AI に持ち込み、SageMaker Training Compiler でトレーニングジョブを実行します。

Hugging Face Transformer を使用する PyTorch モデル

Hugging Face Transformers を使用する PyTorch モデルは、PyTorch の torch.nn.Module API に基づいています。Hugging Face Transformers には、自然言語処理 (NLP) モデルを構成する労力を軽減するのに役立つ PyTorch 用のトレーナークラスと事前トレーニング済みのモデルのクラスも用意されています。トレーニングスクリプトを準備したら、 の次のトピックに進むときに、SageMaker Training Compiler HuggingFace 設定で SageMaker AI PyTorchまたは推定器を使用してトレーニングジョブを起動できますSageMaker Training Compiler を有効化する。

ヒント

トレーニングスクリプトで Transformers を使用して NLP モデルのトークナイザを作成する場合は、必ず padding='max_length' を指定して静的な入力テンソル形状を使用してください。バッチ内の最長シーケンスにパディングすると、各トレーニングバッチのテンソル形状が変化する可能性があるため、padding='longest' は使用しないでください。動的入力形状は、モデルの再コンパイルをトリガーし、合計トレーニング時間が長くなる可能性があります。Transformer トークナイザのパディングオプションの詳細については、「Hugging Face Transformer ドキュメント」の「Padding and truncation」を参照してください。

トピック

• Hugging Face Transformers Trainer クラスを使用する大規模言語モデル

• PyTorch を直接使用する大規模言語モデル (Hugging Face Transformers Trainer API なし)

Hugging Face Transformers Trainer クラスを使用する大規模言語モデル

Transformer ライブラリの Trainer クラスを使用する場合、トレーニングスクリプトに変更を加える必要はありません。SageMaker Training Compiler は、推定器クラスを通じて有効化すると、Trainer モデルを自動的にコンパイルします。次のコードは、Hugging Face Trainer API を使用する PyTorch トレーニングスクリプトの基本形式を示しています。

from transformers import Trainer, TrainingArguments

training_args=TrainingArguments(**kwargs)
trainer=Trainer(args=training_args, **kwargs)

トピック

• シングル GPU のトレーニングの場合

• 分散トレーニングの場合

• Trainer を用いた SageMaker Training Compiler の使用のベストプラクティス

シングル GPU のトレーニングの場合

transformers.Trainer クラスを使用する場合、コードを変更する必要はありません。

分散トレーニングの場合

PyTorch v1.1.0 以降

SageMaker Training Compiler で分散トレーニングを実行するには、トレーニングスクリプトに次の _mp_fn() 関数を追加し、main() 関数をラップする必要があります。PyTorch の SageMaker AI 分散ランタイム (pytorchxla) からの_mp_fn(index)関数呼び出しをトレーニングスクリプトの main() 関数にリダイレクトします。 PyTorch

def _mp_fn(index):
    main()

この関数は、分散トレーニング用のクラスター内にある現在の GPU のランクを示す index 引数を受け入れます。サンプルスクリプトをさらに見つけるには、「Hugging Face Transformers language modeling example scripts」を参照してください。

PyTorch v1.10.2 以前を使用する v4.17 以前の Transformers の場合

SageMaker Training Compiler は、分散トレーニングジョブを起動する代替メカニズムを使用するため、トレーニングスクリプトを変更する必要はありません。代わりに、SageMaker Training Compiler では、SageMaker AI 分散トレーニングランチャースクリプトを entry_point引数に渡し、トレーニングスクリプトを SageMaker AI Hugging Face 推定器の hyperparameters引数に渡す必要があります。

Trainer を用いた SageMaker Training Compiler の使用のベストプラクティス

• transformers.TrainingArgument の設定時に optim 引数を adamw_torch_xla に設定して、SyncFree オプティマイザを使用していることを確認してください。「Hugging Face Transformers のドキュメント」の「Optimizer」も参照してください。

• データ処理パイプラインのスループットがトレーニングスループットよりも高いことを確認してください。このために、transformers.TrainingArgument クラスの dataloader_num_workers 引数と preprocessing_num_workers 引数を微調整できます。通常、これらは GPU の数以上で、CPU の数より少ない必要があります。

トレーニングスクリプトの調整が完了したら、「SageMaker Training Compiler を使用して PyTorch トレーニングジョブを実行する」に進みます。

PyTorch を直接使用する大規模言語モデル (Hugging Face Transformers Trainer API なし)

PyTorch を直接使用するトレーニングスクリプトがある場合は、PyTorch トレーニングスクリプトに変更を加えて、PyTorch/XLA を実装する必要があります。手順に従ってスクリプトを変更し、PyTorch /XLA プリミティブを適切にセットアップします。

トピック

• シングル GPU のトレーニングの場合

• 分散トレーニングの場合

• PyTorch/XLA を用いた SageMaker Training Compiler の使用のベストプラクティス

シングル GPU のトレーニングの場合

1. 最適化ライブラリをインポートします。

   import torch_xla
   import torch_xla.core.xla_model as xm

2. ターゲットデバイスを torch.device("cuda") ではなく XLA に変更します。

   device=xm.xla_device()

3. PyTorch の自動混合精度 (AMP) を使用する場合は、次のようにします。

   1. torch.cuda.amp を次のように置き換えます。

      import torch_xla.amp

   2. torch.optim.SGD および torch.optim.Adam を次のように置き換えます。

      import torch_xla.amp.syncfree.Adam as adam
      import torch_xla.amp.syncfree.SGD as SGD

   3. torch.cuda.amp.GradScaler を次のように置き換えます。

      import torch_xla.amp.GradScaler as grad_scaler

4. AMP を使用しない場合は、optimizer.step() を次のように置き換えます。

   xm.optimizer_step(optimizer)

5. 分散データローダーを使用する場合は、データローダーを PyTorch/XLA の ParallelLoader クラスにラップします。

   import torch_xla.distributed.parallel_loader as pl
   parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)

6. parallel_loader を使用しない場合は、トレーニングループの最後に mark_step を追加します。

   xm.mark_step()

7. トレーニングをチェックポイントするには、Pytorch/XLA のモデルチェックポイントメソッドを使用します。

   xm.save(model.state_dict(), path_to_save)

トレーニングスクリプトの調整が完了したら、「SageMaker Training Compiler を使用して PyTorch トレーニングジョブを実行する」に進みます。

分散トレーニングの場合

前記 シングル GPU のトレーニングの場合 のセクションで記載した変更点に加え、GPU 間にワークロードを適切に配分するには、次の変更を追加します。

1. AMP を使用する場合は、all_reduce を scaler.scale(loss).backward() の後に追加します。

   gradients=xm._fetch_gradients(optimizer)
   xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())

2. local_ranks および world_size に対して変数を設定する必要がある場合、次と同様のコードを使用します。

   local_rank=xm.get_local_ordinal()
   world_size=xm.xrt_world_size()

3. 1 より大きな world_size (num_gpus_per_node*num_nodes) については、トレーニングサンプラーを定義する必要があります。これは次のようなものです。

   import torch_xla.core.xla_model as xm
   
   if xm.xrt_world_size() > 1:
       train_sampler=torch.utils.data.distributed.DistributedSampler(
           train_dataset,
           num_replicas=xm.xrt_world_size(),
           rank=xm.get_ordinal(),
           shuffle=True
       )
   
   train_loader=torch.utils.data.DataLoader(
       train_dataset, 
       batch_size=args.batch_size,
       sampler=train_sampler,
       drop_last=args.drop_last,
       shuffle=False if train_sampler else True,
       num_workers=args.num_workers
   )

4. torch_xla distributed モジュールによって提供される parallel_loader を確実に使用するには、次の変更を加えます。

   import torch_xla.distributed.parallel_loader as pl
   train_device_loader=pl.MpDeviceLoader(train_loader, device)

   train_device_loader は、次に示すように、通常の PyTorch ローダーのように機能します。

   for step, (data, target) in enumerate(train_device_loader):
       optimizer.zero_grad()
       output=model(data)
       loss=torch.nn.NLLLoss(output, target)
       loss.backward()

   これらのすべての変更により、Transformer Trainer API を使用せずに、どの PyTorch モデルでも分散トレーニングを開始できます。これらの手順は、シングルノードマルチ GPU とマルチノードマルチ GPU の両方に使用できます。

5. PyTorch v1.11.0 以降の場合

   SageMaker Training Compiler で分散トレーニングを実行するには、トレーニングスクリプトに次の _mp_fn() 関数を追加し、main() 関数をラップする必要があります。PyTorch の SageMaker AI 分散ランタイム (pytorchxla) からの_mp_fn(index)関数呼び出しをトレーニングスクリプトの main() 関数にリダイレクトします。 PyTorch

   def _mp_fn(index):
       main()

   この関数は、分散トレーニング用のクラスター内にある現在の GPU のランクを示す index 引数を受け入れます。サンプルスクリプトをさらに見つけるには、「Hugging Face Transformers language modeling example scripts」を参照してください。

   PyTorch v1.10.2 以前を使用する v4.17 以前の Transformers の場合

   SageMaker Training Compiler は分散トレーニングジョブを起動するための代替メカニズムを使用し、SageMaker AI 分散トレーニングランチャースクリプトを entry_point引数に渡し、トレーニングスクリプトを SageMaker AI Hugging Face 推定器の hyperparameters引数に渡す必要があります。

トレーニングスクリプトの調整が完了したら、「SageMaker Training Compiler を使用して PyTorch トレーニングジョブを実行する」に進みます。

PyTorch/XLA を用いた SageMaker Training Compiler の使用のベストプラクティス

ネイティブの PyTorch トレーニングスクリプトで SageMaker Training Compiler を利用するには、まず、XLAデバイス上のPyTorch に習熟することが望まれます。以下のセクションでは、PyTorch 用に XLA を有効化するためのベストプラクティスをいくつか示します。

注記

ベストプラクティスに関するこのセクションでは、次の PyTorch/XLA モジュールを使用していることを前提としています。

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl

Pytorch/XLA の lazy モードを理解する

Pytorch/XLA とネイティブ PyTorch の大きな違いは、Pytorch/XLA システムが lazy モードで実行されるのに対し、ネイティブ PyTorch は eager モードで実行される点です。lazy モードのテンソルは、コンパイルと評価が完了した後それらがマテリアライズされるまで計算グラフを構築するためのプレースホルダーです。PyTorch/XLA システムは、テンソルと演算子を使用する計算の構築のために PyTorch API が呼び出されたときに、その場で計算グラフを構築します。計算グラフは、xm.mark_step() が pl.MpDeviceLoader/pl.ParallelLoader によって明示的または暗黙的に呼び出されたとき、または、loss.item() や print(loss) を呼び出すなどテンソルの値が明示的にリクエストされたときに、コンパイルされ、実行されます。

コンパイル・実行の数を pl.MpDeviceLoader/pl.ParallelLoader および xm.step_closure を使用して最小限に抑える

最高のパフォーマンスを得るには、Pytorch/XLA の lazy モードを理解するの説明のように、コンパイル・実行を開始するために可能な方法を念頭に置き、コンパイル・実行の数を最小限に抑えるようにしてください。理想的には、トレーニングの反復ごとに必要なコンパイル・実行は 1 つだけであり、pl.MpDeviceLoader/pl.ParallelLoader によって自動的に開始されます。MpDeviceLoader は XLA 用に最適化されており、最高のパフォーマンスを得るために可能な限り常に使用してください。トレーニング中に、損失値などの中間結果をいくつか調べることもできます。そのような場合、lazy テンソルの印刷は、不要なコンパイル・実行を避けるため、xm.add_step_closure() を使用してラップしてください。

AMP および syncfree オプティマイザを使用する

自動混合精度 (AMP) モードでのトレーニングは、NVIDIA GPU のテンソルコアを活用することで、トレーニング速度を大幅に高めます。SageMaker Training Compiler は、XLA 用に最適化された syncfree オプティマイザを備えており、AMP のパフォーマンスを向上させます。現在、以下の 3 つの syncfreeオプティマイザが利用可能で、最高のパフォーマンスを得るために可能な限り使用してください。

torch_xla.amp.syncfree.SGD
torch_xla.amp.syncfree.Adam
torch_xla.amp.syncfree.AdamW

これら syncfree オプティマイザは、グラデーションのスケーリング/スケーリング解除のために torch_xla.amp.GradScaler と組み合わせます。

ヒント

PyTorch 1.13.1 以降、SageMaker Training Compiler により、PyTorch/XLA が torch.optim または torch_xla.amp.syncfree (torch_xla.amp.syncfree.SGD、torch_xla.amp.syncfree.Adam、torch_xla.amp.syncfree.AdamW など) の syncfree バージョンを使用する transformers.optimization オプティマイザ (SGD、Adam、AdamW など) を自動的にオーバーライドすることにより、パフォーマンスが向上しています。トレーニングスクリプト内でオプティマイザを定義しているコード行を変更する必要はありません。