Nebula を使用してチェックポイント処理の速度を向上させ、コストを削減する

[アーティクル]
09/01/2024

Nebula を使用して、大規模な Azure Machine Learning トレーニングモデルのチェックポイント処理の速度を向上させ、チェックポイントコストを削減する方法について説明します。

概要

PyTorch 用 Azure コンテナー (ACPT) の Nebula は、高速でシンプル、ディスクレスのモデル対応チェックポイントツールです。 Nebula では、PyTorch を使用した分散型の大規模モデルトレーニングジョブに対して、シンプルで高速なチェックポイント処理ソリューションを提供します。 Nebula は、最新の分散コンピューティングテクノロジを利用して、チェックポイント時間を数時間から数秒に短縮でき、95% から 99.9% の時間を節約できる可能性があります。大規模なトレーニングジョブは、Nebula のパフォーマンスから大きな恩恵を受けることができます。

Nebula をトレーニングジョブで使用できるようにするには、スクリプトに nebulaml Python パッケージをインポートします。 Nebula は、PyTorch Lightning、DeepSpeed など、PyTorch のさまざまな分散トレーニング戦略に完全に互換しています。 Nebula API では、チェックポイントのライフサイクルを監視および表示する簡単な方法を提供します。 API ではさまざまなモデルの種類をサポートし、チェックポイントの一貫性と信頼性を確保します。

重要

nebulaml パッケージは、パブリック PyPI Python パッケージインデックスでは使用できません。これは、Azure Machine Learning 上の PyTorch 用 Azure コンテナー (ACPT) キュレーション環境でのみ使用できます。問題を回避するため、PyPI から nebulaml をインストールしたり、pip コマンドを使用したりしないでください。

このドキュメントでは、Azure Machine Learning 上の ACPT で Nebula を使用して、モデルトレーニングジョブをすばやくチェックポイント処理する方法について説明します。さらに、Nebula チェックポイントデータを表示および管理する方法についても説明します。また、Azure Machine Learning の中断、失敗、または終了が発生した場合に、利用可能な最後のチェックポイントからモデルトレーニングジョブを再開する方法についても説明します。

大規模なモデルトレーニングのチェックポイント最適化が重要な理由

データ量が増加し、データ形式がより複雑になるにつれて、機械学習モデルもより洗練されています。このような複雑なモデルのトレーニングは、GPU メモリ容量の制限と長いトレーニング時間のため、困難な場合があります。その結果、分散トレーニングは多くの場合、大規模なデータセットや複雑なモデルを操作するときに使用されます。ただし、分散アーキテクチャでは、予期しない障害やノード障害が発生する可能性があり、機械学習モデルのノード数が増えるにつれて、それがますます問題になる可能性があります。

チェックポイントは、特定の時点で完全なモデル状態のスナップショットを定期的に保存するため、これらの問題を軽減するのに役立ちます。障害が発生した場合、このスナップショットを使用して、その時点からトレーニングを再開できるように、スナップショット時の状態にモデルを再構築できます。

大規模なモデルトレーニング操作で障害または終了が発生した場合、データサイエンティストや研究者は、以前に保存したチェックポイントからトレーニングプロセスを復元できます。ただし、チェックポイントと終了の間に行われた進行状況は失われるため、保存されていない中間結果を復旧するには、評価を再実行する必要があります。チェックポイント間隔を短くすると、この損失を減らすことができます。この図は、チェックポイントから終了までのトレーニングプロセス間に無駄になった時間を示しています。

ただし、チェックポイント自体を保存するプロセスで、大幅なオーバーヘッドが発生する可能性があります。 TB サイズのチェックポイントを保存すると、トレーニングプロセスのボトルネックになることが多く、同期されたチェックポイントプロセスによってトレーニングが何時間もブロックされます。チェックポイント関連のオーバーヘッドは、平均して、トレーニング時間全体の 12% を占め、最大 43% に達する場合があります (Maeng 他 2021 年)。

要約すると、大規模なモデルチェックポイント管理には、大きなストレージが必要となるうえ、ジョブ回復時間のオーバーヘッドも発生します。チェックポイントを頻繁に保存しつつ、利用可能な最新のチェックポイントからトレーニングジョブを再開することは、多大な困難が伴います。

Nebula によるサポート

大規模な分散モデルを効果的にトレーニングするには、データの損失とリソースの浪費を最小限に抑えるために、トレーニングの進行状況を保存して再開する、信頼できる効率的な方法を用意することが重要です。 Nebula は、より高速で簡単なチェックポイント管理を提供して、大規模なモデルの Azure Machine Learning トレーニングジョブに対するチェックポイントの節約時間と GPU 時間の需要を減らすのに役立ちます。

Nebula を使用すると、次のことができます。

トレーニングプロセスと非同期で動作するシンプルな API を使用して、チェックポイント処理の速度を最大 1000 倍に向上できます。 Nebula では、チェックポイント処理にかかる時間を数時間から数秒に短縮でき、95% から 99% の削減を実現します。

この例では、Hugging Face GPT2、GPT2-Large、GPT-XL トレーニングジョブの 4 つのチェックポイントを節約するためのチェックポイント処理とエンドツーエンドのトレーニング時間の削減を示しています。中規模の Hugging Face GPT2-XL チェックポイント保存の場合 (20.6 GB)、Nebula では 1 つのチェックポイントに対して 96.9% の時間短縮を達成しました。

チェックポイント処理の速度は、モデルサイズと GPU 数によってさらに向上する可能性があります。たとえば、128 個の A100 Nvidia GPU で 97 GB のトレーニングポイントチェックポイント保存をテストすると、20 分から 1 秒に短縮できます。
チェックポイントのオーバーヘッドを最小限に抑え、ジョブの回復に費やされる GPU 時間の数を減らして、大規模なモデルのエンドツーエンドのトレーニング時間と評価コストを削減できます。 Nebula ではチェックポイントを非同期的に保存し、トレーニングプロセスのブロックを解除して、エンドツーエンドのトレーニング時間を短縮します。また、チェックポイントの保存頻度を増やすこともできます。これにより、中断後に最新のチェックポイントからトレーニングを再開し、ジョブ回復と GPU のトレーニング時間に費やされる時間とコストを節約できます。
PyTorch を使用して完全な互換性を提供できます。 Nebula は PyTorch と完全な互換性があり、DeepSpeed (>=0.7.3) や PyTorch Lightning (>=1.5.0) などの分散トレーニングフレームワークと完全に統合できます。これは、Azure Machine Learning や AKS など、さまざまな Azure Machine Learning コンピューティング先で使用することもできます。
チェックポイントの一覧表示、取得、保存、読み込みに役立つ Python パッケージを使用して、チェックポイントを簡単に管理できます。チェックポイントのライフサイクルを示すために、Nebula ではAzure Machine Learning スタジオに関する包括的なログも提供します。チェックポイントをローカルまたはリモートのストレージ場所に保存するよう選択し、
- Azure Blob Storage
- Azure Data Lake Storage
- NFS
わずか数行のコードでいつでもアクセスできます。

前提条件

Azure サブスクリプションと Azure Machine Learning ワークスペース。ワークスペースリソースの作成の詳細については、ワークスペースリソースの作成に関するページを参照してください。
Azure Machine Learning コンピューティングターゲット。コンピューティングターゲットの作成の詳細については、 Manage トレーニングとコンピューティングのデプロイを参照してください
PyTorch を使用するトレーニングスクリプト。
ACPT キュレーション (PyTorch 用 Azure コンテナー) 環境。 ACPT イメージを取得する方法については、「キュレーションされた環境」を参照してください。キュレーションされた環境の使用方法について説明します

Nebula の使用方法

Nebula により、既存のトレーニングスクリプトで、迅速で簡単なチェックポイントエクスペリエンスが提供されます。素早く Nebula を起動する手順は次のとおりです。

ACPT 環境の使用
Nebula の初期化
API を呼び出してチェックポイントを保存して読み込む

ACPT 環境の使用

PyTorch モデルトレーニング用のキュレーション環境である PyTorch 用 Azure コンテナー (ACPT) には、プレインストールされた依存 Python パッケージとして Nebula が含まれています。キュレーション環境を表示するには PyTorch 用 Azure コンテナー (ACPT) に関するページを参照してください。ACPT イメージの詳細については、「Azure Machine Learning で PyTorch 用 Azure コンテナーを使用してディープラーニングを有効にする」を参照してください。

Nebula の初期化

ACPT 環境で Nebula を有効にするには、トレーニングスクリプトを変更して nebulaml パッケージをインポートし、適切な場所で Nebula API を呼び出すだけで済みます。 Azure Machine Learning SDK または CLI の変更を回避できます。また、Azure Machine Learning Platform で大規模なモデルをトレーニングするために、他の手順に変更を加える必要もありません。

Nebula をトレーニングスクリプトで実行するには、初期化が必要です。次のコードスニペットに示すように、初期化フェーズでは、チェックポイントの保存場所と頻度を決定する変数を指定します。

  import nebulaml as nm
  nm.init(persistent_storage_path=<YOUR STORAGE PATH>) # initialize Nebula

Nebula は DeepSpeed と PyTorch Lightning に統合されています。そのため、初期化をシンプルかつ簡単に行えます。これらの例では、Nebula をトレーニングスクリプトに統合する方法を示します。

重要

Nebula を使用してチェックポイントを保存するには、チェックポイントを保存するためのメモリが必要です。メモリは、チェックポイントの少なくとも 3 つのコピーより大きくなるようにしてください。

メモリがチェックポイントを保持するのに十分でない場合は、チェックポイントを保存するときに各ノードごとのメモリの使用を制限するように、コマンドで環境変数 NEBULA_MEMORY_BUFFER_SIZE を設定することをお勧めします。この変数を設定すると、Nebula はこのメモリをバッファーとして使用してチェックポイントを保存します。メモリ使用量に制限がない場合、Nebula はできるだけメモリを使用してチェックポイントを保存します。

同じノードで複数のプロセスが実行されている場合、チェックポイントを保存するための最大メモリは、制限の半分をプロセス数で割った値になります。残りの半分を Nebula がマルチプロセスの調整のために使用します。たとえば、各ノードあたりのメモリ使用量を 200 MB に制限する場合は、コマンドで環境変数を export NEBULA_MEMORY_BUFFER_SIZE=200000000 (バイト単位で約 200 MB) として設定できます。この場合、Nebula は各ノードにチェックポイントを格納するために 200 MB のメモリのみを使用します。同じノードで 4 つのプロセスが実行されている場合、Nebula は各プロセスごとに 25 MB のメモリを使用してチェックポイントを保存します。

API を呼び出してチェックポイントを保存して読み込む

Nebula には、チェックポイントの保存を処理するための API が用意されています。これらの API は、PyTorch torch.save() API と同様に、トレーニングスクリプトで使用できます。これらの例では、トレーニングスクリプトで Nebula を使用する方法を示します。

チェックポイント履歴を表示する

トレーニングジョブが完了したら、[ジョブ Name> Outputs + logs] ウィンドウに移動します。左側のパネルで [Nebula] フォルダーを展開し、checkpointHistories.csv を選択して、Nebula チェックポイントの保存に関する詳細情報 (期間、スループット、チェックポイントサイズ) を表示します。

例

これらの例では、さまざまな種類のフレームワークで Nebula を使用する方法を示します。トレーニングスクリプトに最適な例を選択できます。

Nebula と PyTorch ベースのトレーニングスクリプトとの完全な互換性を可能にするには、必要に応じてトレーニングスクリプトを変更します。

まず、必須の nebulaml パッケージをインポートします。

  # Import the Nebula package for fast-checkpointing 
  import nebulaml as nm

Nebula を初期化するには、次に示すように、main() で nm.init() 関数を呼び出します。

  # Initialize Nebula with variables that helps Nebula to know where and how often to save your checkpoints
  persistent_storage_path="/tmp/test",
  nm.init(persistent_storage_path, persistent_time_interval=2)

チェックポイントを保存するには、元の torch.save() ステートメントを置き換えて、チェックポイントを Nebula に保存します。チェックポイントインスタンスが "global_step" で始まる ("global_step500" や "global_step1000" など) ことを確認してください。
```
checkpoint = nm.Checkpoint('global_step500')  
checkpoint.save('<CKPT_NAME>', model)  
```
Note

<'CKPT_TAG_NAME'> はチェックポイントの一意の ID です。通常、タグはステップ数、エポック番号、またはユーザー定義名です。省略可能な <'NUM_OF_FILES'> パラメーターは、このタグに対して保存する状態番号を指定します。
次に示すように、最新の有効なチェックポイントを読み込みます。
```
latest_ckpt = nm.get_latest_checkpoint()
p0 = latest_ckpt.load(<'CKPT_NAME'>)
```
チェックポイントまたはスナップショットには多数のファイルが含まれている可能性があるため、1 つ以上のファイルを名前で読み込むことができます。最新のチェックポイントを使用すると、トレーニング状態を最後のチェックポイントで保存した状態に復元できます。

チェックポイント管理は他の API で処理できる
- すべてのチェックポイントを一覧表示する
- 最新のチェックポイントを取得する
```
# Managing checkpoints
## List all checkpoints
ckpts = nm.list_checkpoints()
## Get Latest checkpoint path
latest_ckpt_path = nm.get_latest_checkpoint_path("checkpoint", persisted_storage_path)
```

次に示すように、ds_config.json 構成ファイルで Nebula を有効にした場合、DeepSpeed (>=0.7.3) に基づくトレーニングスクリプトで Nebula を使用できます。

"nebula": {
      "enabled": true,
      "persistent_storage_path": "<YOUR STORAGE PATH>",
      "persistent_time_interval": 100,
      "num_of_version_in_retention": 2,
      "enable_nebula_load": true
  }

この JSON スニペット関数は、nebulaml.init() 関数と同様に機能します。

ds_config.json ファイル構成で初期化すると Nebula が有効になり、チェックポイントの保存が可能になります。モデルチェックポイント API model_engine.save_checkpoint() を使用する元の DeepSpeed 保存メソッドでは、Nebula が自動的に使用されます。この保存メソッドを使用すると、コードを変更する必要がなくなります。

PyTorch Lightning (Nebula ではバージョン >=1.5.0 をサポート) では、トレーナーの使用時に自動的にチェックポイント処理が行われます。きめ細かく制御するために、カスタマイズされた動作を使用してチェックポイントを保存することがよくありますが、PyTorch Lightning では、ModelCheckpoint() による条件付き保存と trainer.save_checkpoint() による手動保存の 2 通りのチェックポイント保存方法があります。これらの手法は、PyTorch (>=0.15.0) トレーニングスクリプトに適用されます。

ModelCheckpoint を使用してチェックポイントを条件付きで保存する場合は、ModelCheckpoint ではなく NebulaCallback を初期化に使用できます。

# import Nebula package
import nebulaml as nm

# define NebulaCallback
config_params = dict()
config_params["persistent_storage_path"] = "<YOUR STORAGE PATH>"
config_params["persistent_time_interval"] = 10

nebula_checkpoint_callback = nm.NebulaCallback(
   ****, # Original ModelCheckpoint params
   config_params=config_params, # customize the config of init nebula
)

次に、Trainer へのプラグインとして nm.NebulaCheckpointIO() を追加し、次のように trainer.save_checkpoint() ストレージパラメーターを変更します。

trainer = Trainer(plugins=[nm.NebulaCheckpointIO()],   # add NebulaCheckpointIO as a plugin
                  callbacks=[nebula_checkpoint_callback]) # use NebulaCallback as a plugin

trainer.save_checkpoint() を使用してチェックポイントを手動で保存する場合は、Trainer で NebulaCheckpointIO プラグインを使用し、trainer.save_checkpoint() のストレージパラメーターを次のように変更できます。

# import Nebula package
import nebulaml as nm

# initialize Nebula
nm.init(persistent_storage_path=<YOUR STORAGE PATH>) 

trainer = Trainer(plugins=[nm.NebulaCheckpointIO()])  # add NebulaCheckpointIO as a plugin

# Saving checkpoints
storage_options = {}
storage_options['is_best'] = True
storage_options['persist_path'] = "/tmp/tier3/checkpoint"

trainer.save_checkpoint("example.ckpt",
  storage_options=storage_options, # customize the config of Nebula saving checkpoint
)

チェックポイントの読み込み

PyTorch および PyTorch Lightning と一致するチェックポイントを読み込みます。唯一の修正点は、ストレージパスを指定することです。

最新のチェックポイントを読み込む場合は、次に示すように、MyLightningModule.load_from_checkpoint() が引き続き動作します。

  persistent_path = "/tmp/tier3/checkpoint"
  latest_ckpt_path = nebulaml.get_latest_checkpoint_path("checkpoint", persist_path)
  model = MyLightningModule.load_from_checkpoint(latest_ckpt_path)

Trainer() で nebula_checkpoint_callback を使用した場合、次に示すように、元の変更されていないスクリプトが引き続き動作します。

  trainer = Trainer(
      default_root_dir=tmpdir, 
      max_steps=100,
      plugins=[NebulaCheckpointIO()],
      callbacks=[nebula_checkpoint_callback],
  )
  trainer.fit(model, ckpt_path="/path/example.ckpt")

次の方法で共有

Nebula を使用してチェックポイント処理の速度を向上させ、コストを削減する

概要

大規模なモデルトレーニングのチェックポイント最適化が重要な理由

Nebula によるサポート

前提条件

Nebula の使用方法

ACPT 環境の使用

Nebula の初期化

API を呼び出してチェックポイントを保存して読み込む

チェックポイント履歴を表示する

例

次の手順

フィードバック

その他のリソース

次の方法で共有

Nebula を使用してチェックポイント処理の速度を向上させ、コストを削減する

概要

大規模なモデル トレーニングのチェックポイント最適化が重要な理由

Nebula によるサポート

前提条件

Nebula の使用方法

ACPT 環境の使用

Nebula の初期化

API を呼び出してチェックポイントを保存して読み込む

チェックポイント履歴を表示する

例

フィードバック

その他のリソース

大規模なモデルトレーニングのチェックポイント最適化が重要な理由