EasyCkptを使用して基礎モデルのトレーニングを保存および再開する - Platform For AI

PyTorch基礎モデルのトレーニングが失敗した場合、EasyCkptによって保存された最新のチェックポイントを使用して、計算を繰り返すことなくトレーニングを再開できます。これにより、時間とコストを節約できます。 EasyCkptは、Platform for AI (PAI) が提供するPyTorch基礎モデルのトレーニングシナリオに使用される高性能チェックポイントフレームワークです。 EasyCkptを使用すると、モデルトレーニングの進行状況をほぼ無料で保存できます。 EasyCkptは、基礎モデルトレーニングのプロセス全体を損なうことなく、モデルトレーニングを保存および再開する機能を提供します。 EasyCkptは、MegatronとDeepSpeedの基礎モデルトレーニングフレームワークをサポートしています。このトピックでは、EasyCkptの技術原則とEasyCkptの使用方法について説明します。

背景情報

基礎モデルのトレーニングの典型的な課題は、プロセスが中断されないようにすることです。トレーニングプロセスでは、ハードウェア障害、システム問題、接続エラー、およびその他の未知の問題が発生する可能性があります。頻繁な中断は、基礎モデルのトレーニングの進捗に悪影響を与える可能性があり、完了するまでにかなりの時間とリソースが必要になります。特定のケースでは、チェックポイント操作を実行して、トレーニングの進行状況を保存および再開できます。チェックポイント操作を完了するのに必要な時間は、モデルのサイズによって異なります。ほとんどの場合、数千億から数千億のパラメータを含む基礎モデルは、チェックポイント操作を完了するために数分から数十分を必要とする。トレーニングタスクは期間中中断され、ユーザーが頻繁にチェックポイント操作を実行できないようにします。トレーニングが中断されると、システムが回復したときに失われた反復が再計算されます。これには数時間かかる場合があります。例えば、千個のGPUを使用する基礎モデルは、モデルを実行するために必要とされる多数のGPUに起因して、重大な損失をもたらし得る。

したがって、エラーが発生したときに最新のチェックポイントを保存するために使用できるコスト削減方法が必要です。トレーニング回復中に計算を繰り返す必要はありません。これは時間とコストを節約するのに役立ちます。

原則

次のセクションでは、以前の発生に基づくGPUとディープラーニングの失敗の特性について説明します。

特徴1: 失敗は特定の労働者に影響を与える
ほとんどの場合、障害の根本原因は1台または2台のマシンにまでさかのぼることができますが、これは複数のワーカーにのみ影響します。すべてのワーカーが大規模な分散トレーニングジョブに失敗するわけではありません。
特徴2: 障害はサーバの特定のコンポーネントに影響を与えます
ほとんどの場合、クラスターでは次のシナリオが発生します。
- GPUでエラーが発生すると、サーバーのCPUとメモリは期待どおりに動作します。
- ノードのアイドルメモリ空間は大きく、通常は予想されるモデル状態よりもはるかに大きい。
- エラーは、各ノードの特定のネットワークインターフェイスでのみ発生します。したがって、ノードは、ノードが期待通りに動作していなくても、依然として通信することができる。
特徴3: 故障はモデルの特定の部分に影響を与えます
ほとんどの場合、基礎モデルのトレーニングは、3D並列処理やZero Redundancy Optimizer (ZeRO) などの最適化手法を使用します。ほとんどのタスクには、複数の並列データレプリカがあります。これにより、モデルトレーニングパラメータが複数のレプリカでバックアップされます。マシン上のGPUに障害が発生した場合、他のマシンのGPUに保持されているレプリカを使用してトレーニングを回復できます。

基礎モデルのトレーニングシナリオでのチェックポイントの特性に基づいて、PAIは高性能チェックポイントを提供するEasyCkptフレームワークを提供します。 EasyCkptは、非同期階層チェックポイント、重複するモデルチェックポイントと計算、ネットワーク対応の非同期チェックポイントなどの戦略を採用することにより、ほとんどコストをかけずにモデル保存メカニズムと、基礎モデルトレーニングのプロセス全体を損なうことなくモデルトレーニングを保存および再開する機能を提供します。 EasyCkptは、MegatronとDeepSpeedの基礎モデルトレーニングフレームワークをサポートしています。最小限のコード変更でEasyCkptを使用できます。

手順

SDK for AIMasterのインストール

EasyCkptを使用するには、AIMaster用のSDKをインストールする必要があります。次の操作を実行して、AIMaster用のSDKをインストールします。

# py36
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp36-cp36m-linux_x86_64.whl

# py38
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp38-cp38-linux_x86_64.whl

# py310
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp310-cp310-linux_x86_64.whl

メガトロン

サンプルコード

Megatronフレームワークに基づいて、training.pyファイルのコードを変更します。次の図に例を示します。

また、実際のトレーニングコードファイルにコード行をインポートする必要があります。次の図の例では、pretrain_gpt.pyファイルを使用しています。

次のコードは、変更されたtraining.pyファイルの例を示しています。

from megatron.core.utils import get_model_config
from megatron import print_rank_0
from megatron import print_rank_last
# from megatron.checkpointing import load_checkpoint
from megatron.checkpointing import save_checkpoint
from megatron.model import Float16Module
from megatron.model import GPTModel
from megatron.utils import report_memory
from megatron.model.vision.knn_monitor import compute_feature_bank
from aimaster.python.torch.easyckpt.megatron import (load_checkpoint,
                                                    initialize_easyckpt,
                                                    save_checkpoint_if_needed)

def print_datetime(string):
    """Note that this call will sync across all ranks."""
    timers('interval-time', log_level=0).start(barrier=True)
    print_datetime('before the start of training step')
    report_memory_flag = True

    initialize_easyckpt(save_mem_interval=1, save_storage_interval=5, max_ckpt_num=5, log_file_path='./test.log')

    while iteration < args.train_iters:
        if args.profile and \
           iteration == args.profile_step_start and \
                                       args.micro_batch_size * \
                                       get_num_microbatches()

        save_checkpoint_if_needed(iteration, model, optimizer, opt_param_scheduler)

        # Logging.
        loss_scale = optimizer.get_loss_scale().item()
        params_norm = None

次のコードは、変更された実際のトレーニングファイルの例を示します。この例では、pretrain_gpt.pyファイルが使用されています。

from megatron.utils import average_losses_across_data_parallel_group
from megatron.arguments import core_transformer_config_from_args

import aimaster.python.torch.easyckpt.megatron.hook

def model_provider(pre_process=True, post_process=True):
    """Build the model."""

説明

EasyCkptフレームワークは、メガトロンに次のインターフェイスを提供します。

load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True, concat=False): インターフェイスは、Megatronフレームワークのネイティブのload_checkpoint() 関数のシグネチャに基づいて、concatパラメーターを追加します。メガトロン2304を使用する場合は、メガトロンのload_checkpointを交換するだけです。メガトロン2305または2306を使用する場合は、次の注意事項に注意してください。
initialize_easyckpt(save_mem_interval, save_storage_interval, max_ckpt_num, log_file_path=None): インターフェイスは、Easyckptフレームワークを初期化するために使用されます。 save_mem_intervalを使用してメモリコピーの頻度を指定します。 save_storage_intervalを使用して、非同期ストレージ頻度を指定します。 max_ckpt_numを使用して、ストレージデバイスに保存できるチェックポイントの最大数を指定します。ログの詳細情報を保存する必要がある場合は, log_file_pathでログパスを指定してください。
save_checkpoint_if_needed(iteration、model、optimizer、opt_param_scheduler): インターフェイスは、メモリ内のチェックポイントを実行するためにEasyCkptフレームワークを呼び出すために使用されます。パラメータは、メガトロン符号における既存の変数である。パラメーターを指定する必要はありません。

注: Megatron 2305または2306を使用し、分散オプティマイザを有効にした場合、training.pyファイルのload_checkpoint() 関数のconcatパラメーターをTrueに設定する必要があります。ロードプロセス中にインスタンス数を変更するか、分散オプティマイザのパラメーターをマージします。

DeepSpeed

ほとんどの場合、TransformerモデルのトレーナーはDeepSpeedタスクを開始するために使用されます。 EasyCkptは、必要な変更を最小限に抑えるメソッドをサポートしています。

サンプルコード

スタートアップパラメーター: DeepSpeedのEasyCkptフレームワークは、Transformerのチェックポイントパラメーターを再利用します。パラメータの意味は、Transformerで定義されている意味と同じです。詳細については、以下の説明セクションをご参照ください。サンプルコードでは、チェックポイントは2つのミニバッチごとに保存されます。 2つまでの最近のチェックポイントコピーが同時に永続的に保持されます。

変更後のコード (右上の図)

--max_steps=10 \
--block_size=2048 \
--num_train_examples=100000 \
--gradient_checkpointing=false \
--save_strategy="steps" \
--save_steps="2" \
--save_total_limit="2"

コードの変更: TransformerのTrainerは、EasyCkptが提供するTrainerWrapperでラップし、resume_from_Checkpointパラメーターを有効にする必要があります。

変更後のコード (右上の図)

import datasets
import transformers

from aimaster.python.torch.easyckpt.transformers import TrainerWrapper

logger = logging.getLogger(__name__)

    tokenizer=tokenizer,
    data_collator=transformers.default_data_collator,
  )
  trainer = TrainerWrapper(trainer)
  trainer.train(resume_from_checkpoint=True)

if __name__ = ""__main__":
  main()

説明

EasyCkptフレームワークは、DeepSpeedに次のインターフェイスを提供します。

save_strategy: トレーニング中のチェックポイントの保存モード。有効な値：
- no: トレーニング中にチェックポイントを保存しません。
- epoch: 各エポックの最後にチェックポイントを保存します。
- steps: save_stepsの指定値に基づいてチェックポイントを保存します。
save_steps: トレーニング中にチェックポイントが保存されるステップの数。インターフェイスは、save_strategyをstepsに設定した場合にのみ有効です。
save_total_limit: 保持できるチェックポイントの最大数。

注意: save_total_limitを有効にすると、古いチェックポイントフォルダが削除されます。フォルダを削除する前に、必要なデータが保存されていることを確認してください。詳細については、Transformerの公式ドキュメントを参照してください。

データセキュリティノート

EasyCkptは、指定したストレージ内のデータを読み書きする必要があります。 EasyCkptは、保持されるチェックポイントの最大数を制御するためにデータを削除する必要があります。データのセキュリティを確保するために、PAIはEasyCkptによって実行されるすべての読み取りおよび書き込み操作を定義し、EasyCkptのデータセキュリティを確保し、推奨される使用方法を提供します。

EasyCkptは、次の読み取りおよび書き込み操作を実行します。 EasyCkptを使用する場合、デフォルトの権限が付与されます。

ロードディレクトリからチェックポイントデータを読み取り、そのデータを新しいチェックポイントデータに統合します。
チェックポイントデータを保存ディレクトリに保存し、特定の設定に基づいて保存ディレクトリ内のメガトロンまたはトランスフォーマー形式のチェックポイントフォルダを削除します。

EasyCkptは、次の要件を満たすことを保証します。

保存ディレクトリとロードディレクトリ以外のデータに対する操作は実行されません。
EasyCkptによって実行されるすべての保存および削除操作がログに記録されます。

モデルのsaveまたはloadディレクトリに他のデータを格納しないことを推奨します。そうしないと、EasyCkptが期待どおりに機能しない場合があります。発生する可能性のあるデータリスクとデータ損失については、お客様が責任を負います。