SageMaker、GPUを最大7分割 MIGで推論を相乗り

切り分けたチョコレートケーキ AI

AWS は11月6日、Amazon SageMaker AI のリアルタイム推論エンドポイントで、NVIDIA の MIG(Multi-Instance GPU:1枚の GPU を、計算部分とメモリごと最大7つに区切る機能)を使えるようにしました。80GB の H100 1枚に、8B クラスの小さなモデルを別々に載せられます。区切った部分どうしはハードウェアの側で分けられていて、ほかのモデルが混んでも遅くなりにくい作りです。東京リージョンでも使えますが、料金はこれまでどおりインスタンスまるごとの時間単価で、使い方によっては区切らないほうが安くなります。

背景と文脈

小さなモデルを本番で動かすと、GPU が余りがちです。8B のモデルを16ビットで持つと、重みは約16GB。そこに KV キャッシュ(文章を生成するとき、前に読んだ部分の計算結果をためておくメモリ)と1〜2GB の作業用メモリが足されます。それでも 40GB や 80GB の GPU 1枚を丸ごと使うことになり、空いた分にもお金を払っていました。

SageMaker には、これまでも1つのエンドポイントに複数のモデルを置くやり方がありました。マルチモデルエンドポイントは、モデルを同じ GPU に置きますが、モデルどうしを分ける仕組みがありません。しばらく使われなかったモデルは外され、次に呼ばれたときに読みこみ直す時間がかかります。推論コンポーネント(エンドポイントの上に、モデルごとに CPU・メモリ・GPU の割り当てを決めて載せる仕組み)も、GPU は1枚単位でしか割り当てられませんでした。

MIG 自体は、NVIDIA が A100 の世代から入れている機能です。Kubernetes の上で自分で組めば前から使えましたが、ドライバーの版、区切り方の設定、デバイスプラグイン、割り当ての仕組みまで自分で面倒を見る必要がありました。SageMaker の学習・推論用クラスターの HyperPod では先に MIG に対応していて、今回はふつうの推論エンドポイントでも使えるようになりました。

技術/ビジネス面

区切られたロッカーの扉
Photo by Jan Antonin Kolar on Unsplash

設定で増えたのは、省略できる項目2つだけです。新しい API や IAM の権限は増えていません。エンドポイントの設定(CreateEndpointConfig)に、GPU をどう区切るかを書く AcceleratorPartitionLayouts が入りました。推論コンポーネントの側では、ComputeResourceRequirements の中の AcceleratorPartitionType で、どの区切りに載せるかを選びます。どちらも書かなければ、これまでと同じ動きです。

区切り方は GPU ごとに決まっていて、名前の「3g.40gb」は計算部分7つのうち3つと、メモリ 40GB を使うという意味です。AWS のブログと NVIDIA の MIG の案内から、対応しているインスタンスをまとめます。

インスタンス GPU(メモリ) 1枚あたりの最大の区切り数 区切り方の例
P4d A100(40GB) 7 1g.5gb、2g.10gb、3g.20gb、4g.20gb
P5 H100(80GB) 7 1g.10gb、2g.20gb、3g.40gb、4g.40gb
G7 RTX PRO 4500 Blackwell(32GB) 2 1g.16gb、2g.32gb
G7e RTX PRO 6000 Blackwell Server Edition(96GB) 4 1g.24gb、2g.48gb、4g.96gb

区切り方は、そのエンドポイントの全部の GPU に同じように当てはまります。8枚の GPU があるインスタンスで3種類の区切りを指定すると、それぞれが8つずつ、合わせて24個できます。1つのモデルが区切りをまたいで2枚分を使うことはできません。GPU が複数枚いる大きなモデルは、これまでどおり GPU を丸ごと割り当てます。

区切りはインスタンスを用意するときに作られ、AWS の試験では準備の時間が約17秒延びました。指定した区切りが GPU に合わないときや、合計が1枚に収まらないときは、設定を作る時点で断られます。黙って GPU 1枚分に切りかえることはありません。

開発者として試すなら

手順は、推論コンポーネントを使うエンドポイントを作るときとほぼ同じです。AWS のブログの例では、エンドポイントの設定に次のように書きます。

"ProductionVariants": [{
  "VariantName": "AllTraffic",
  "InstanceType": "ml.p4d.24xlarge",
  "InitialInstanceCount": 1,
  "AcceleratorPartitionLayouts": [
    { "AcceleratorPartitionTypes": ["mig-4g.20gb", "mig-2g.10gb", "mig-1g.5gb"] }
  ]
}]

推論コンポーネントの側は、"AcceleratorPartitionType": "mig-4g.20gb" と "NumberOfAcceleratorDevicesRequired": 1 を書きます。区切りを指定したときの「1」は GPU 1枚ではなく、区切り1つの意味になります。MinMemoryRequiredInMb は今までどおり、GPU ではなくサーバー側のメモリの量です。呼び出しは InvokeEndpoint で推論コンポーネントを指定するだけで、変わりません。

P4d で mig-1g.10gb を使うときは、ドライバーに気をつけてください。この区切りには R525 以降のドライバーが要りますが、P4d の標準のイメージは 470.256.02 です。ブログによると、このままだと区切りの作成が失敗し、エラーにどの区切りが原因かが出ません。InferenceAmiVersion に al2023-ami-sagemaker-inference-gpu-4-1 を指定すると、580 系のドライバーになります。

費用は、AWS の測定を紹介した Unite.AI の記事に条件まで書かれています。比べたのは、A100(40GB)1枚を丸ごと使う場合と、H100 を mig-3g.40gb で区切った1つです。モデルは Llama-3.1-8B-Instruct で、同時に送る要求の数を2から512まで変えています。

項目 結果
出力100万トークンあたりの費用 区切った H100 のほうが18〜34%安い(混み具合が一定を超えたとき)
安くなり始める同時要求数 100トークンの短い入力で64、8,000トークンの長い入力で4
1秒あたりの出力トークン 53〜89%多い(費用の差が一番大きい負荷のとき)
それより空いているとき A100 を丸ごと使うほうが安い

この比べ方は、同じ GPU を区切った前後の比較ではありません。新しい世代の H100 の約半分が、古い A100 1枚より安くて速かったという結果です。自分の負荷がこの境目のどちら側にあるかは、手元のアクセスの記録で同時要求数を見て確かめてください。

運用で困りそうな点もあります。Unite.AI の記事によると、区切ったエンドポイントは GPU・CPU・メモリの使用率を CloudWatch に出しません。呼び出しの回数と応答時間は見られます。区切り方は動いているエンドポイントでは変えられず、変えるには新しいエンドポイントを作って通信を移します。非同期推論やサーバーレス推論、1つのモデルだけのエンドポイントでは使えません。CPU やネットワークは区切られないので、サーバー側の取り合いは残ります。

これからどうなるか

Unite.AI などによると、使えるリージョンは東京・シンガポール・フランクフルト・アイルランド・米国の3つの、合わせて7か所です。料金はインスタンス単位のままで、区切っても時間単価は下がりません。安くなるのは、1台に載せるモデルの数を増やして、空いていた GPU の分を減らせたときだけです。AWS のブログでは、Salesforce が本番に近い負荷で試す準備を始めたと紹介されていて、使った会社の結果が出てくるのはこれからです。

自分のサービスで効くのは、8B 以下のモデルをお客さまごとや用途ごとに分けて動かしている場合です。いまモデル1つに GPU を1枚ずつ割り当てているなら、H100 1枚を2つや4つに区切ると、エンドポイントの台数を減らせるか計算できます。RAG(検索拡張生成:検索で見つけた文書をモデルに渡して答えさせる仕組み)で、埋め込み用と回答用の小さなモデルを同じ GPU に置く構成なら、片方が混んでももう片方の応答時間が延びにくくなります。一方、使用率が CloudWatch に出ないので、GPU が余っているかどうかは呼び出しの数や応答時間から見るしかありません。

まとめ

SageMaker の推論エンドポイントで、H100 や A100 を最大7つに区切って小さなモデルを並べられるようになりました。設定は2つの項目を足すだけで、東京でも使えます。料金はインスタンス単位のままなので、混み具合が境目を超えないと丸ごと使うほうが安くなります。使用率の指標が出ない点と、区切り方をあとから変えられない点も、試す前に確かめておいてください。

参考リンク

アイキャッチ画像: Photo by Waldemar Brandt on Unsplash

タイトルとURLをコピーしました