AWS環境でNVIDIA Cosmos 3 Superを動かす
はじめに
NVIDIAのCosmos 3には、16BのNanoと64BのSuperがあります。前回はNanoのセットアップをしたので、今回はSuperを触っていきます。
この記事ではAWS環境上でCosmos3 Superを動かすまでの記録を紹介します。
基本的に公式の手順通りに進めればOKですが、いくつかつまづいた点やポイントがあったので、それについて紹介します。
今回の個人的なつまづきとポイントは以下の3つでした。
- 8GPUインスタンスの在庫が取れない（オンデマンドもオンデマンドキャパシティ予約も）。
- 重み133GBを起動のたびにダウンロードすると、料金が高いインスタンスの時間を溶かす。
- p5起動直後にNCCL/Libfabricでsegfault。
後で解説します。
なぜH100×8基の構成にしたのか
Super（64B）はReasoner（32B）とGenerator（32B）の2つ（Mixture-of-Transformers）でできていて、BF16の重みは合わせて約130GB。単一のH100（80GB）には載りません。そこで生成側を含め8枚のH100にFSDPで分割して動かします。実行は公式のcosmos-frameworkで、torchrun --nproc-per-node=8（8プロセス＝8GPU）で起動します。
| 項目 | 内容 |
|---|---|
| インスタンス | p5.48xlarge（NVIDIA H100 80GB × 8） |
| リージョン | 東京（ap-northeast-1） |
| AMI | Deep Learning Base OSS Nvidia Driver GPU AMI（Ubuntu 24.04） |
| セットアップ | uv sync --all-extras --group=cu130-train |
セットアップの全体像（まずは公式手順）
冒頭にも書いたとおり、基本は公式ドキュメント通りで大丈夫です。まずはこの2つに目を通しておけば、セットアップから実行までの流れは追えます。
- セットアップ手順（公式）: cosmos-framework / docs/setup.md
- 推論の実行（公式）: cosmos-framework / docs/inference.md
全体の流れとしてはざっくりこんな感じです。
- p5.48xlarge（8×H100）を起動して、重みのボリュームをアタッチする
- システムの依存パッケージを入れる。
sudo apt-get install -y --no-install-recommends curl ffmpeg git-lfs libx11-dev tree wget - cosmos-frameworkを
uvで入れる（CUDAに合わせてグループを選ぶ。今回はcu130）。uv sync --all-extras --group=cu130-train source .venv/bin/activate && export LD_LIBRARY_PATH= - 推論を実行する（Superは8GPUに分割するので
torchrun。← 初回のsegfaultがポイント3。コマンドは後述）。
この流れ自体は公式通りなので、この記事では私が詰まったところ・公式に書いていないところの3つのポイントを中心に紹介していきます。
ポイント1: GPU8基の在庫が取れない → Capacity Blockで解決
まずオンデマンドでp5.48xlargeを起動しようとしましたが、東京の両AZ・us-east-1・us-west-2のいずれもInsufficientInstanceCapacityで取れませんでした。オンデマンドキャパシティ予約も、p5・p4de（8×A100）とも同様に枯渇していました。
まあ当然ですよね。大量のGPUインスタンスを使いたい時は突発的には使えないものだと思って、事前に段取っておくよいです。
そこでEC2 Capacity Blocks for MLを使いました。
- 最小1日分・前払い。今回は1ノード（8×H100）を予約。
- 予約は開始・終了時刻が決まっていて、その枠内で起動します。
注意: 終了30分前にAWSが自動終了する
これは実際に遭遇したのですが、Capacity Blockは終了時刻の30分前に、AWSが動いているインスタンスを自動的に終了します（AWS公式ドキュメント）。つまり実質のデッドラインは終了時刻の30分前。作業中でも問答無用で切れるので、逆算して片付ける必要があります。
私は最後の検証をしている最中にこれで終了してしまいました泣
ポイント2: 重み133GBを毎回DLしない → EBSスナップショットで事前ステージ
Superの重み（HuggingFaceのnvidia/Cosmos3-Super、実測で124GiB）を、高価なp5の上で毎回ダウンロードすると、ダウンロード待ちの時間だけで結構な金額（私の検証時は$68/h）が溶けます。そこで安いインスタンスで先に重みを落としておき、EBSスナップショット化しておきました。
- 安い
m7i.2xlargeにgp3ボリューム（250GB）をアタッチし、hf download nvidia/Cosmos3-Superで重みを取得。 - アンマウントしてEBSスナップショット化（残るのはスナップショットだけなので安価）。
- p5を起動するとき、このスナップショットからボリュームを復元してアタッチ。復元時にProvisioned Init Rate（
--volume-initialization-rate、今回は300MiB/s）で事前ウォームしておくと、初回アクセスのレイテンシを避けられます。復元先のAZはCapacity BlockのAZに合わせます。
これでp5側では重みの再ダウンロードが不要になり、高額なインスタンスの時間を計算・生成に使えます。
ポイント3: NCCL/Libfabricのsegfault → NCCL_NET_PLUGIN=none
p5はEFA（高速ネットワーク）が有効で、初回の分散初期化時にNCCL/Libfabric周りでsegfaultしました。これは公式のinference.mdのTroubleshootingにあるとおり、
export NCCL_NET_PLUGIN=none
で解決します。単一ノードの推論ではEFAのネットワークプラグインは不要なので、無効化してしまえば問題ありません。
起動
セットアップを終えれば、あとは8つのGPUに分割して起動するだけです。
export NCCL_NET_PLUGIN=none
torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \
--parallelism-preset=throughput --dp-shard-size=8 --dp-replicate-size=1 --cp-size=1 --cfgp-size=1 \
-i "inputs/xxx.json" -o outputs/yyy \
--checkpoint-path /mnt/weights/Cosmos3-Super --no-guardrails --seed=0
--parallelism-preset=throughput＋--dp-shard-size=8で、Superを8枚のH100にFSDPで分割します。
まとめ
- Super（64B）は単一GPUに載らないので8×H100が必要。オンデマンドは品薄なのでCapacity Blockで確保する（前払い・最小24h・終了30分前に自動terminate）。
- 重み133GBは安いインスタンスで先に取ってEBSスナップショット化し、p5には復元してアタッチする。これで高い箱の時間を再ダウンロードで溶かさない。
- p5はEFA有効で初回segfault →
NCCL_NET_PLUGIN=noneで解決。デバイス名やスナップショット実行場所の細かい落とし穴に注意。
次回以降、このSuper環境で実際に試した内容を書いていきます。
