SageMaker Training Compiler のベストプラクティスと考慮事項

一般的なベストプラクティス

• サポートされるインスタンスタイプ と テスト済みモデル のいずれかを使用していることを確認します。

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

• GPU メモリの使用率を測定して、GPU メモリに収まる最大バッチサイズを使用していることを確認します。Amazon SageMaker Training Compiler は、トレーニング中のモデルのメモリフットプリントを削減し、通常は GPU メモリにより大きな batch_size を収めることができます。より大きな batch_size を使用すると GPU の使用率が向上し、合計トレーニング時間が短縮されます。

  バッチサイズを調整する場合、learning_rate も適切に調整する必要があります。例えば、バッチサイズを k 倍に増やした場合は、learning_rate を直線的に (単純に k の乗算)、もしくは k の平方根を乗算して調整する必要があります。これは、トレーニング時間を短縮しながら、同じか同様の収束動作を実現するためです。人気モデルでテスト済みの batch_size のリファレンスについては、「テスト済みモデル」を参照してください。

• コンパイラで高速化されたトレーニングジョブをデバッグするには、compiler_config パラメータ の debug フラグを有効化します。これにより、SageMaker AI はデバッグログを SageMaker トレーニングジョブログに入れることができます。

  huggingface_estimator=HuggingFace(
      ...
      compiler_config=TrainingCompilerConfig(debug=True)
  )

  コンパイラでトレーニングジョブの完全なデバッグを有効化すると、オーバーヘッドが増す可能性があることに注意してください。

PyTorch のベストプラクティス

• PyTorch モデルを持ち込んでチェックポイントする場合は、Pytorch/XLA のモデル保存関数を使用してモデルを適切にチェックポイントしていることを確認します。この関数の詳細については、XLA デバイス上の PyTorch に関するドキュメントにある「torch_xla.core.xla_model.save」を参照してください。

  PyTorch スクリプトに変更を追加する方法については、「PyTorch を直接使用する大規模言語モデル (Hugging Face Transformers Trainer API なし)」を参照してください。

  モデル保存関数を使用する実際のアプリケーションの詳細については、PyTorch/XLA TPU での Hugging Face: より速くより安いトレーニングのブログにある「チェックポイントの書き込みと読み込み」を参照してください。

• 分散トレーニングで最適なトレーニング時間を達成するには、以下を考慮します。

  • 単一の GPU インスタンスではなく、複数の GPU を備えたインスタンスを使用します。例えば、単一の ml.p3dn.24xlarge インスタンスは 8 x ml.p3.2xlarge インスタンスと比較してトレーニング時間が短縮されます。

  • ml.p3dn.24xlarge や ml.p4d.24xlarge などの EFA サポートのあるインスタンスを使用します。これらのインスタンスタイプでは、ネットワーク速度が加速され、トレーニング時間が短縮されます。

  • データセットの preprocessing_num_workers パラメータを調整して、前処理が遅いためにモデルトレーニングが遅れることがないようにします。

• 明示的な評価につながるモデルテンソルのログ記録、チェックポイント、プロファイリングは避けます。明示的な評価とは何かを理解するために、次のコードコンパイル例を考慮します。

  a = b+c
  e = a+d

  コンパイラはコードを次のように解釈し、変数 a のメモリフットプリントを削減します。

  e = b+c+d

  次に、コードを変更して変数 a の print 関数を追加した例を考慮します。

  a = b+c
  e = a+d
  print(a)

  コンパイラは、変数 a を次のように明示的に評価します。

  e = b+c+d
  a = b+c    # Explicit evaluation
  print(a)

  例えば、PyTorch では、明示的な評価が導入される可能性がある torch.tensor.items() の使用を避けます。深層学習では、このような明示的な評価によってモデルのコンパイルグラフ内の融合オペレーションが中断され、テンソルの再計算につながるため、オーバーヘッドが発生する可能性があります。

  それでも SageMaker Training Compiler を使用しているとき、トレーニング中にモデルを定期的に評価したい場合は、明示的な評価によるオーバーヘッドを減らすために、ログ記録とチェックポイントの頻度を低くすることをお勧めします。例えば、エポックごとではなく 10 エポックごとにログを記録します。

• グラフのコンパイルはトレーニングの最初の数ステップで実行されます。そのため、最初の数ステップは非常に時間がかかると予想されます。ただし、これは 1 回限りのコンパイルコストであり、コンパイルによって今後のステップがはるかに速くなるため、より長い期間トレーニングすることで償却できます。初期コンパイルのオーバーヘッドは、モデルのサイズ、入力テンソルのサイズ、入力テンソルの形状の分布によって異なります。

• PyTorch モデルをコンパイルする際の最も一般的なエラーの 1 つは、演算子とテンソルのデバイスタイプが間違っていることによるものです。PyTorch モデルを適切にコンパイルするには、必ず XLA デバイス (xm.xla_device()) を使用し、CUDA を使用したり、CUDA デバイスと XLA デバイスを混在させたりしないでください。

• mark_step() は XLA 専用の障壁です。正しく設定しないと、トレーニングジョブが停止します。

• Pytorch/XLA は、追加の分散トレーニング API を提供しています。API を適切にプログラミングしないと、勾配が正しく収集されず、トレーニングが収束しなくなります。