マネージド階層チェックポイントをセットアップする

このセクションでは、Amazon SageMaker HyperPod のマネージド階層チェックポイントのセットアッププロセスについて説明します。クラスターで機能を有効にし、トレーニングコードにチェックポイントを実装する方法を説明します。

前提条件

マネージド階層チェックポイントを設定する前に、以下の要件を満たしていることを確認してください。

ステップ 1: クラスターのマネージド階層チェックポイントを有効にする

クラスターを作成または更新する際に、HyperPod API を介してマネージド階層チェックポイントを有効にします。TieredStorageConfig パラメータを指定すると、サービスはメモリ管理システムを自動的にインストールします。新しいクラスターでは、以下が作成されます。

aws sagemaker update-cluster \
    --cluster-name my-training-cluster \
    --tiered-storage-config {
            "Mode": "Enable"
            "InstanceMemoryAllocationPercentage": percentage
            }

InstanceMemoryAllocationPercentage パラメータは、チェックポイントに割り当てるクラスターメモリの percentage (int) を指定します。範囲は 20～100 です。

ステップ 2: トレーニングイメージに Python ライブラリをインストールする

Amazon SageMaker チェックポイントライブラリとその依存関係を Dockerfile に追加して、トレーニングイメージにインストールします。

# Add this line to your training image Dockerfile
RUN pip install amzn-sagemaker-checkpointing

ステップ 3: チェックポイント設定を作成する

チェックポイントの動作を指定する CheckpointConfig オブジェクトを作成します。これには、以下が含まれます。

次の例は、チェックポイント設定を説明しています。

from amzn_sagemaker_checkpointing.config.sagemaker_checkpoint_config import SageMakerCheckpointConfig
from amzn_sagemaker_checkpointing.checkpointing.filesystem import SageMakerTieredStorageWriter, SageMakerTieredStorageReader

checkpoint_config = sm_ckpt.CheckpointConfig(
    world_size = 100,
    in_memory_namespace: my-ml-workload, # Logical grouping for checkpoints
    s3_base_path: "s3://bucket-name/checkpointing-path-prefix/",
    s3_every_n_steps: 100,  # Every 100 steps, save to S3
)

ステップ 4: SageMaker ファイルシステムライターを定義する

チェックポイントファイルシステムライターを定義します。オプションで、初期化中にステップ番号を指定できます。

ベーシックライター (セーブコールで指定されたステップ):

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(checkpoint_config)

ステップパラメータを持つライター (初期化時に指定されたステップ):

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(
    checkpoint_config, 
    step=step_number
)

ステップ 5: トレーニングループにチェックポイントを保存する

トレーニングループで、FileSystemWriter を使って PyTorch DCP を使用してチェックポイントを保存します。

FileSystemWriter で PyTorch DCP を使用する

FileSystemWriter を入力として dist_cp.save() メソッドを呼び出します。

オプション 1: ステップ形式で checkpoint_id を使用する (ライターでステップが指定されていない場合)

# Construct checkpoint directory with step number
checkpoint_dir = f"step_number"

dist_cp.save_state_dict(
    state_dict=state_dict,  # state_dict is a dictionary containing model parameters, optimizer state, etc.
    checkpoint_id=checkpoint_dir,  # Should contain step number
    storage_writer=smWriter
)

オプション 2: ステップパラメータでライターを使用する (checkpoint_id はオプションとなる)

dist_cp.save_state_dict(
    state_dict=state_dict,
    storage_writer=smWriter  # Step already specified in writer initialization
)

ステップ 6: リカバリ用のチェックポイントをロードする

チェックポイントをロードする必要がある場合は、FileSystemReader で PyTorch DCP を使用します。

FileSystemReader で PyTorch DCP を使用する

FileSystemReader を入力として DCP load メソッドを呼び出します。

# Define FileSystemReader
smReader = sagemaker_checkpointing.SageMakerTieredStorageReader(
    config=checkpoint_config
)

# Load checkpoint
dist_cp.load_state_dict(
    state_dict=state_dict,
    checkpoint_id=checkpoint_dir,
    storage_reader=smReader
)

モニタリングと検証

メトリクスとログを使用して、マネージド階層チェックポイントオペレーションをモニタリングして検証できます。

カスタムログ記録 (オプション)

カスタムロガーをライブラリに渡すことで、チェックポイントログを他のログと統合できます。例えば、カスタムロガーをトレーニングコードに追加して、ライブラリのすべてのログがトレーニングロガーにも収集されるようにできます。

拡張サービスログ記録 (オプション)

デバッグとサービスの可視性を強化するために、ポッド内のチェックポイントログパス /var/log/sagemaker_checkpointing をホスト上のパス /var/logs/sagemaker_checkpointing にマウントできます。これにより、ライブラリ固有のログのみが個別に収集されることになり、サービスチームはデバッグとサポートの可視性を向上させることができます。