より高速な SageMaker AI ジョブのための SOCI インデックス付きイメージの発表

ディープラーニングのコンテナイメージのダウンロードに 1 分かかるたびに、その 1 分分の費用を支払い、チームはその 1 分間待たされます。コンテナイメージのダウンロードは純粋なオーバーヘッドであり、モデルの進捗に何ら寄与しない課金対象のジョブ時間です。最新のディープラーニングイメージでは、このオーバーヘッドはすべてのジョブで数分に及びます。本日、私たちは次の発表を嬉しく思います Amazon SageMaker AI Training および Amazon SageMaker AI 処理ジョブ向けの Seekable OCI (SOCI) インデックスこれにより、すべてのジョブのクリティカルパスからコンテナイメージ全体のダウンロードが取り除かれます。

SOCI インデックスを使用すると、SageMaker AI はコンテナイメージを遅延ロードします。コンテナの起動に必要なファイルのみを取得し、コードの実行中に残りをバックグラウンドでストリーミングするため、ジョブは数分ではなく数秒で開始されます。導入には、事前にインデックスが付与された AWS Deep Learning Container (DLC) を指定するか、独自のイメージにインデックスを追加します。いずれの方法でも、ジョブ定義、エントリポイント、データチャネルの変更は不要です。その結果、ジョブの起動が速くなり、ジョブコストが下がり、インフラを待つ時間ではなくモデル構築に費やす時間が増えます。

SOCI インデックス付きイメージは、トレーニングパフォーマンスに影響を与えることなく、Amazon SageMaker AI のトレーニングコンテナを数秒で起動します。 測定した内容 のセクションで示すように、ベンチマークではイメージのプル中の GPU アイドル時間を最大 90 パーセント削減できました。

この記事では、SOCI インデックスの導入方法として、その概要、測定した結果、および自分のジョブで有効にする方法を説明します。

課題

SageMaker AI のトレーニングジョブと処理ジョブは、Amazon Elastic Container Registry (Amazon ECR) に保存されたコンテナイメージ内で実行されます。ジョブがコードを実行する前に、サービスはイメージ全体をインスタンスにダウンロードし、すべてのレイヤーを展開してディスクに書き込みます。

小さなイメージであれば数秒で済みます。しかし、ディープラーニングのイメージは 10 GB を超えることも珍しくありません。このダウンロードによりすべてのジョブに数分が追加され、その分はジョブ時間として課金されます。影響が最も大きいのは、まさに反復速度が最も重要となる場面、すなわち数百のジョブを伴うハイパーパラメータチューニングのスイープ、迅速な実験、そして短い処理ジョブであり、イメージのダウンロードが総実行時間の大半を占めることもあります。

SOCI インデックスの仕組み

Seekable OCI (SOCI) は、コンテナイメージを遅延ロードするオープンソーステクノロジーです。起動前にイメージ全体をダウンロードするのではなく、ランタイムは SOCI インデックスを使用して、コンテナを即座に起動するために必要なファイルのみを取得します。その後、コードがアクセスするのに合わせて残りをバックグラウンドでストリーミングします。

SOCI インデックスは、イメージ内のすべてのファイルをマッピングするメタデータです。Amazon ECR に、元のイメージにリンクされた OCI 準拠のアーティファクトとして保存します。インデックスは独立したアーティファクトであるため、イメージのダイジェストや既存の署名は有効なままです。

SageMaker AI はインデックスを自動的に確認します

トレーニングジョブまたは処理ジョブを作成すると、SageMaker AI は指定したイメージに SOCI インデックスがあるかどうかを確認し、それに応じて動作します:

必要に応じてジョブは常に標準のダウンロードにフォールバックするため、SOCI インデックスを追加してもジョブの実行が妨げられることはありません。得られるのは起動の高速化だけです。既存のワークフローが壊れるリスクはありません。

測定した内容

実際の SageMaker AI トレーニングジョブで、同じコンテナイメージを標準ダウンロードと SOCI インデックス付き遅延ロードの 2 つの方法でプルしました。次に、両方の方法で完全なトレーニングジョブを実行し、トレーニング速度への影響を測定しました。

セットアップ

• 画像は8点圧縮時のサイズは1.8～35.4 GB。
• GPUおよびCPUインスタンス ml.g4dn、ml.g5、ml.g6、ml.g6e、ml.g7、ml.g7e、ml.p4d、ml.p5、ml.m7i の各インスタンスを、xlargeから48xlargeまで、米国西部 (オレゴン) リージョン (us-west-2) および米国東部 (バージニア北部) リージョン (us-east-1) で提供。
• ダウンロード時間 画像のプル開始からコンテナが実行可能になるまでを計測。すべての数値は繰り返し実行の中央値です。

結果1: すべての画像サイズでダウンロード時間が数分から数秒に短縮

次のチャートは、1.8 GB から 35.4 GB までの 8 つのイメージサイズについて、標準的な Docker pull のコンテナイメージダウンロード時間と SOCI レイジーローディングを比較したものです。

各イメージサイズには 2 本のバーが表示されます。Docker のダウンロード時間のバー(青)はイメージサイズに比例して増加し、1.8 GB で 48 秒から 35.4 GB で 7 分 17 秒まで伸びています。

各ペアの上の高速化ラベルは、Docker 時間に対する SOCI 時間の比率を示しています。高速化率は、最小のイメージで約 10 倍から、最大のイメージで 35 倍超までの範囲です。

図 1: 8 つのイメージサイズにおける標準的な Docker pull と SOCI レイジーローディングのコンテナイメージダウンロード時間の比較

標準ダウンロード時間がイメージサイズとともに増加するのは、すべてのファイルの展開と書き込みが行われるためです。レイジーローディングはサイズに関係なく 3〜11 秒で推移し、8〜46 倍の改善となりました。イメージが大きいほど、節約できる時間は長くなります。

結果 2: この節約効果はすべてのインスタンスタイプとサイズで得られます

13.2 GB のイメージを 13 種類のインスタンスタイプ/サイズの組み合わせで実行したところ、標準ダウンロードは 3 分 16 秒から 4 分 20 秒の範囲だったのに対し、レイジーローディングは 6〜8 秒で推移しました。すべての構成で約 34 倍の高速化です。利点を得るためにインスタンスタイプやサイズを変更する必要はありません。

図 2: 13 種類のインスタンスタイプとサイズの組み合わせにおける 13.2 GB イメージのダウンロード時間

結果 3: 最初のエポックのペナルティを除き、トレーニング性能に影響はありません

Docker でプルしたイメージと SOCI でインデックス化したイメージの 2 種類からフルトレーニングジョブを実行し、すべてのエポックを比較しました:

• ResNet-50 による画像分類 Imagenette (9,469 枚の画像)
  • コンテナイメージ: PyTorch(13.2 GB DLC)および TensorFlow(9.9 GB DLC)
  • インスタンス: ml.p5.48xlarge、ml.p4d.24xlarge、ml.g7e.48xlarge、ml.g7.48xlarge、ml.g6e.12xlarge、ml.g6e.24xlarge、および ml.g6e.48xlarge。
• 大規模言語モデル(LLM)ファインチューニング: Qwen2.5-0.5B on alpaca-cleaned （51,760例）をHugging Face Trainerで使用しました。
  • コンテナイメージ：9.4 GB Hugging Face PyTorch DLC。
  • インスタンス：ml.p5.48xlarge、ml.p4d.24xlarge、ml.g6.48xlarge。

以下のチャートは、ジョブの起動時間をイメージのダウンロード時間と最初のトレーニングエポックの時間に分解し、両方のワークロードについてGPUおよびCPUの各構成において、標準的なDocker pullとSOCIインデックス付きイメージを比較したものです。

図3: ジョブ起動をイメージダウンロード時間と最初のトレーニングエポック時間に分割した比較、標準の Docker pull と SOCI インデックス付きイメージの比較

• 2エポック目以降は、どちらのワークロードでも、すべてのインスタンスでエポックあたりの時間は同一です（実行ごとのばらつきの範囲内）。
• 最初のエポックでは、レイジーロードされたイメージからは 3〜16 秒長くかかります。これは、フレームワークが初めて触れるファイルが取得されるためです。この最初のエポックのコストは、どのインスタンスサイズでも同じです。
• 結果: 各ジョブは 1〜6 分早く完了しました。これはダウンロードの節約時間から最初のエポックの一時的なコストを差し引いたもので、ジョブあたりの固定の節約額です。

自身のワークロードで試してみる

これらの数値は、前述の 2 つのワークロードによるものです。最初のエポックのコストは、イメージの内容と、起動時にコードが何にアクセスするかによって異なります。

• 最初のトレーニングステップでイメージの一部だけが必要で、残りがトレーニング中にストリーミングされる場合は、レイジーロードが有利です。
• エントリポイントがトレーニング開始前にイメージの大部分を読み込む場合（例: スキャン、チェックサム計算、内容のプリロード）、最初のエポックはフル pull で取得するのとほぼ同じバイト数を待つことになります。正味の改善は小さくなります。
• それでも、先読み型の起動の方が優れていました:
  • torch.compile とジャストインタイム（JIT）ビルドの CUDA 拡張により最初のエポックに 7〜15 秒追加されましたが、それでも正味の節約は 3 分 25 秒〜3 分 56 秒でした
  • イメージに組み込まれた 6 GB のモデルはロード時にウェイトをストリーミングし、2〜3.5 分追加されましたが、ジョブ全体ではインデックス付きイメージで依然として 2〜2.5 分早く完了しました。

はじめに

SOCI インデックスの採用には 2 つの方法があります: 事前インデックス済みの AWS Deep Learning Container を使用するか、独自のカスタムイメージ用にインデックスを作成します。

オプション 1: 事前インデックス済みの AWS Deep Learning Container を使用する

AWS は、いくつかの人気のある Deep Learning Container の SOCI インデックス付きバリアントを公開しています。イメージタグの -soci サフィックスで識別されます:

763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.9-gpu-py312-cu130-ubuntu22.04-sagemaker-v1-soci

事前インデックス済みバリアントは、PyTorch Training、TensorFlow Training、vLLM、SGLang 向けに用意されており、一部のフレームワークでは arm64 バリアントも含まれます。レイジーロードを採用するには、ジョブ定義のイメージタグを -soci バリアントに更新します。全リストはこちらで参照できます: AWS Deep Learning Containers Gallery.

注: DLC ベースイメージを FROM として独自のイメージをビルドする場合は、最終的なイメージに対して SOCI インデックスを作成する必要があります。ベースイメージのインデックスは派生イメージには引き継がれません。

オプション 2: 独自のイメージ用に SOCI インデックスを作成する

カスタムイメージをインデックス化するには、Amazon Linux 2023 を実行する Amazon Elastic Compute Cloud (Amazon EC2) インスタンスなど、次のものを備えた Linux ホストが必要です。 containerd、 SOCI snapshotter、および Finch または nerdctl がインストールされた Linux ホスト（Amazon Linux 2023 を実行する Amazon Elastic Compute Cloud (Amazon EC2) インスタンスなど）が必要です。

1. 環境変数を設定する.

   ACCOUNT_ID="111122223333"
   REGION="us-east-1"
   REPOSITORY_NAME="my-training-image"
   ORIGINAL_IMAGE_TAG="v1.0"
   SOCI_IMAGE_TAG="v1.0-soci"
   REGISTRY="$ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com"

2. 必要なツールをインストールする.

   sudo dnf install -y containerd fuse soci-snapshotter runfinch-finch
   sudo systemctl start soci-snapshotter
   sudo systemctl restart containerd

3. Amazon ECR に対して認証する.

   aws ecr get-login-password --region "$REGION" \
     | sudo nerdctl login -u AWS --password-stdin "$REGISTRY"

4. 元のイメージをプルする.

   sudo nerdctl pull "$REGISTRY/$REPOSITORY_NAME:$ORIGINAL_IMAGE_TAG"

5. SOCIインデックス付きイメージを作成する.

   sudo nerdctl image convert --soci \
     "$REGISTRY/$REPOSITORY_NAME:$ORIGINAL_IMAGE_TAG" \
     "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG"

6. インデックス付きイメージをプッシュする.

   sudo nerdctl push --platform linux/amd64 \
     "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG"

これにより、リポジトリに2つのアーティファクトが作成されます:SOCIインデックスと、イメージをインデックスにリンクするOCI準拠のイメージインデックスマニフェストです。元のイメージとそのダイジェストは変更されません。

CI/CDでインデックス作成を自動化する

プッシュのたびにインデックスを自動生成するには、 CFN ECR SOCI Index Builderをデプロイします。これは Amazon EventBridge ルールを使用してAmazon ECRのイメージプッシュイベントを検出し、インデックスを構築します。ビルドパイプラインへの変更は不要です。

SageMaker AIジョブでインデックス付きイメージを使用する

SOCIインデックス付きイメージは、他のイメージとまったく同じように参照します。APIの変更は不要で、タグだけが異なります。

トレーニングの例(SageMaker Python SDK)

from sagemaker.train import ModelTrainer
from sagemaker.train.configs import Compute

trainer = ModelTrainer(
    training_image="111122223333.dkr.ecr.us-east-1.amazonaws.com/my-training-image:v1.0-soci",
    role="arn:aws:iam::111122223333:role/SageMakerExecutionRole",
    compute=Compute(
        instance_type="ml.g5.12xlarge",
        instance_count=2,
    ),
)

trainer.train()

イメージにSOCIインデックスが含まれていることを確認する

起動する前に、タグがインデックス済みのイメージに解決されることを確認してください:

aws ecr batch-get-image \
  --repository-name "$REPOSITORY_NAME" \
  --image-ids imageTag="$SOCI_IMAGE_TAG" \
  --accepted-media-types \
  "application/vnd.oci.image.index.v1+json" \
  --query 'images[0].imageManifest' --output text \
  | jq '.manifests[] | {mediaType, artifactType}'

プラットフォームマニフェストと併せて、artifactType が application/vnd.amazon.soci.index. で始まる子マニフェストを探してください。存在しない場合、そのタグはインデックス化されたイメージを指していません。

留意すべき点

• Amazon ECR プライベートリポジトリのみ対象です。 他のレジストリ内のイメージや、プライベート Docker レジストリから取得したイメージ(TrainingRepositoryAccessMode が Vpc に設定されている場合)は、遅延ロードの対象になりません。
• 改善幅は画像サイズに比例して大きくなります。 おおよそ4GB未満のイメージは、比較的小さな改善にとどまります。
• Docker 固有の前提がないか確認してください。 SOCIインデックス化されたイメージは、containerdベースのランタイム上で動作します。Docker CLIを呼び出すコード、/var/run/docker.sockを読み取るコード、cgroupパスから識別子を導出するコードについて、エントリーポイントとセットアップスクリプトを確認してください。SageMaker AIは後方互換性のために/.dockerenvマーカーをマウントしていますが、これに新たな依存関係を追加することは避けてください。
• SageMaker AI ジョブの SOCI インデックス化は、すべての AWS リージョンで利用できます。

結論

SOCIインデックスを使用すると、SageMaker AIのトレーニングジョブや処理ジョブでのコンテナイメージのダウンロード時間を最小化できます。イメージにSOCIインデックスを追加するか、事前にインデックス済みのAWS Deep Learning Containerに切り替えてください。どちらの場合でも、起動時間を数分から数秒に短縮し、ジョブの総コストを削減し、機械学習(ML)のイテレーションサイクルを加速できます。さらに、既存のジョブ定義、API呼び出し、データパイプラインはまったく同じままです。

詳細については Amazon SageMaker AI SOCIインデックスの使用を開始するには、以下のリソースを参照してください:

1. SOCIインデックス対応のDeep Learning Containerを選択してください。 次を閲覧する ディープラーニングコンテナギャラリー 画像タグに -soci を追加してください。フレームワーク固有のガイダンスについては、次を参照してください: AWS Deep Learning Containers の SOCI インデックス.
2. 独自のカスタム画像をインデックスしましょう。 以下のステップバイステップの解説に従ってください。 SageMaker AI 開発者ガイド、またはこちらを参照してください SOCI snapshotter スタートガイド 詳細なオプションについては。
3. CI/CDパイプラインでのインデックス作成を自動化しましょう。 デプロイする CFN ECR SOCIインデックスビルダー 各イメージのプッシュ時にインデックスを自動生成します。パイプラインの変更は不要です。
4. Finchを使用してローカルでテストします。 次を使用します: Finchの遅延ローディングに関するチュートリアル で、ジョブを提出する前にcontainerdベースのランタイム上でイメージを検証できます。
5. この技術についてさらに詳しく調べるには、 AWS Machine Learning Blogの 「Reducing container cold start times using SOCI index on DLAMI and DLC」 を読んでベンチマークを確認するか、AWS Containers Blogの 内部構造の解説: Seekable OCI と AWS Fargate によるコンテナイメージの遅延読み込み を読んで遅延ローディングの仕組みを詳細に理解してください。

コンテナの起動を高速化するのは、イメージタグを1つ変更するだけで実現できます。次回のSageMaker AIジョブでぜひ試して、その違いを体感してください。

著者について