AWS 2026-06-14

SageMaker AI のモデルロード最適化 — 遅延ロード・シャーディング・量子化

前提:SageMaker のモデルロードは 4 ステップ

なぜ起動が遅く、メモリを食うのかは、ロードの流れを知ると見通せます。SageMaker AI が推論エンドポイントを立ち上げるとき、モデルは概ね次の順で GPU に載ります。

大規模モデルでは各ステップが重く、特に「ホストメモリへ全部展開してからシャーディング」が時間と RAM を消費します。以下の 3 戦略は、それぞれ別の場所に効きます。

1. 遅延ロード(Lazy Loading) — 起動時間を縮める

遅延ロードは「全部を一度にロードしない」発想です。SageMaker 周辺では主に 2 つの文脈で登場します。

コンテナイメージの遅延ロード(SOCI)

SageMaker Studio などでは SOCI(Seekable OCI) インデックスにより、コンテナイメージ全体のダウンロードを待たずに、起動に必要な部分だけを先に取得してコンテナを立ち上げられます。GB 級の ML イメージほどコールドスタート短縮の効果が大きくなります。関連して、オートスケール時のスケールアウトを速める Container Caching もあり、よく使うコンテナイメージをノード上にキャッシュして再ダウンロードを避けます。

モデル重みの遅延ロード / Fast Model Loading(Weight Streaming)

モデル本体側では、Fast Model Loading が「ホストメモリへの全展開 → シャーディング」という遠回りを省き、S3 から GPU へ重みを直接ストリーミングします。前掲図の 2〜4 を圧縮するイメージです。Inference Optimization Toolkit で最適化した一部モデル(Llama 系など)で利用でき、巨大モデルのコールドスタートを大幅に短縮できます。

また量子化の処理自体でも遅延ロードが使われ、Transformer ブロックを 1 つずつディスクから読み込み、GPU 上で量子化が終わったら即座に CPU へオフロードすることで、量子化中の GPU メモリピークを抑えます。

2. モデルシャーディング(Tensor Parallelism) — メモリに「収める」

1 枚の GPU の VRAM に収まらないモデルは、複数 GPU に分割(シャーディング)して載せます。SageMaker の LMI コンテナは DeepSpeed / TensorRT-LLM / vLLM / Transformers-NeuronX といった分割フレームワークを内蔵しており、環境変数 OPTION_TENSOR_PARALLEL_DEGREE で使用 GPU 数(テンソル並列度)を指定するだけで分割できます。

# LMI コンテナの環境変数例(8 GPU に Tensor Parallel で分割)
OPTION_MODEL_ID=meta-llama/Llama-3-70b
OPTION_TENSOR_PARALLEL_DEGREE=8
OPTION_ROLLING_BATCH=vllm

分割方式には主に 2 種類あります。

• Tensor Parallel(テンソル並列): 各レイヤーの重み行列を GPU 間で横分割する。全 GPU が常に稼働するためハードウェア効率が高く、AWS は基本的にこちらを推奨しています。
• Pipeline Parallel(パイプライン並列): レイヤー単位で GPU を分担する。GPU 間の待ち(バブル)が出やすく、単一ノードでは Tensor Parallel が有利です。

AOT(Ahead-of-Time)パーティショニング

毎回の起動時にシャーディングすると遅いため、事前にモデルを分割(AOT パーティショニング)しておき、分割済みアーティファクトを S3 に保存しておく手法があります。エンドポイント起動時は「分割済みのものをそのまま配るだけ」になるため、起動時間を短縮できます。シャーディング(メモリ対策)が、結果的に起動時間最適化にもつながる例です。

3. 量子化(Quantization) — メモリ使用量そのものを縮める

量子化は、重みや活性化を FP16/BF16 → INT8/INT4 などの低精度データ型で表現し、モデルのメモリフットプリントを物理的に小さくする手法です。これにより、より安価で入手しやすい(VRAM の小さい)GPU にモデルを載せられます。INT4 ならおおよそ重みサイズが 1/4 になります。

SageMaker の Inference Optimization Toolkit / LMI では、no-code または少量設定で主要な量子化手法を利用できます。

• AWQ(Activation-aware Weight Quantization): 活性化の分布を考慮した weight-only 量子化。INT4-AWQ が精度劣化を抑えつつ大幅にサイズ削減できる代表的手法。
• GPTQ: キャリブレーションデータを使う事後量子化(PTQ)。低ビット重みで高精度を狙う。
• SmoothQuant: 重みと活性化の両方を量子化(W8A8 等)し、活性化外れ値を重み側に「移す」ことで精度を保つ。
• FP8: 新しめの GPU(H100 等)が持つ FP8 を活かす量子化。

量子化はメモリと帯域を減らすため スループット向上・コスト削減にも効きますが、低ビットほど精度が落ちうるため、用途に応じて手法とビット数を選び、評価することが重要です。

3 戦略の使い分け(まとめ)

• メモリに収まらない → まず量子化でモデルを縮める。それでも 1 GPU に載らなければシャーディングで複数 GPU に分割。
• 起動が遅い / コールドスタートが痛い → 遅延ロード(SOCI・Container Caching)でイメージを、Fast Model Loading や AOT パーティショニングで重みのロードを短縮。
• これらは排他ではなく併用が前提です。例えば「INT4-AWQ で量子化 → Tensor Parallel で 2 GPU に分割 → Fast Model Loading で起動短縮」のように重ねます。