AWS、NVIDIA Cosmos 3 で物理 AI モデルファクトリを構築

ロボットや自律型車両（AV）など、現実世界のデータを物理的な行動に変換する Physical AI システムは、単一のトレーニングジョブで構築できるものではありません。代わりに必要なのは、合成データの生成、事後学習による知覚モデルとポリシーモデルの構築、そしてシステムが周囲を理解し行動できるように評価を行うという一連のプロセスをループさせる継続的なパイプラインです。このパイプラインを継続的に実行するのが「Physical AI モデルファクトリー」の役割であり、新しい現実世界のデータストリームからより優れたモデルを次々と生み出すことを目指します。

本記事では、Amazon SageMaker HyperPod 上で NVIDIA Cosmos 3 を活用して Physical AI モデルファクトリーを構築する方法を紹介します。具体的には以下の内容について解説します。

• Cosmos 3 の独自性：レイヤーごとの共同アテンション（joint attention）を採用し、トレーニングと推論で意図的な非対称性を設けた Mixture-of-Transformers (MoT) デザインの特徴
• この設計がなぜ Amazon Elastic Kubernetes Service (Amazon EKS) を備えた Amazon SageMaker HyperPod に適しているか
• クラスター構築と、共有マルチテラバイトストレージ層のセットアップ方法
• 3 つの代表的なワークロードに対する分散事後学習の実施。特に、公開データセットである DROID データセットを用いたロボットポリシー段階の完全なエンドツーエンド手順について詳しく解説します

同梱のリポジトリには、各ステージのマニフェストと設定ファイルが含まれています。この設計を実際の稼働クラスタに変換するインフラストラクチャテンプレートやジョブマニフェストを含む実行可能なコードは、awsome-distributed-ai GitHub リポジトリ で確認できます。

ループの実行には、リソース容量の確保が不可欠です。GPU を段階的に取得すると、この規模では変動要因が生じます。利用可能な台数や納期は状況によって異なり、確保できた容量も、データ所在地から離れたアベイラビリティゾーンや AWS リージョンに割り当てられる可能性があります。ループ全体に対して容量を事前に確保しておけば、こうした不安定さを回避できます。具体的には、期間限定のキャンペーンでは柔軟なトレーニング計画を立てるか、不特定期間のプロジェクトには容量予約を活用します。

パイプラインがその容量に対して進捗を出しているかどうかに関わらず、確保した容量に対しては課金が発生するため、コストを支配する指標は特定のジョブのピークスループットではありません。重要なのは「GPU goodput（有効スループット）」です。これは、ループ全体を通じて予約された GPU 時間あたりに達成される、実質的なパイプラインの進捗量を指します。

Physical AI パイプラインでは、従来、各工程ごとに個別の GPU 容量を割り当てるのが一般的でした。例えば、合成データの生成用ノード群、微調整（ポストトレーニング）用のノード群、評価用のノード群などがあり、それぞれが独自のライフサイクルを持って起動・停止されていました。

しかし、NVIDIA Cosmos 3 はこうした非効率な運用を不要にします。オープンでオムニモーダルな世界基盤モデルである Cosmos 3 では、動画、画像、行動、音声を単一のトークンストリームとして扱います。同じトランスフォーマーの trunk（幹）が、3 つのモードで動作します。

1. 合成動画生成のための順方向ダイナミクス・ワールドモデル
2. 逆方向ダイナミクスに基づく行動ラベリング
3. デプロイ可能な行動ポリシー

生成、微調整、評価というすべての工程を一つのモデルファミリーがカバーするため、これらの工程は単一のクラスター制御プレーン下にある、永続的で耐障害性の高い GPU ノードプールに割り当てられる複数のワークロードとなります。これは各工程ごとに別々のプールを用意するのではなく、時間共有された容量として機能します。

NVIDIA は Cosmos 3 を Linux Foundation の OpenMDW-1.1 ライセンス のもとで公開しており、そのアーキテクチャの詳細は Cosmos 3 技術レポート で解説されています。

1. Cosmos 3 の仕組み

世界モデルにおける一般的なパターンとして、拡散トランスフォーマーによる動画生成器と、テキスト条件付けを提供する別のビジョン・ランゲージモデルを組み合わせる手法があります。Cosmos 3 はこれとは異なるアプローチを採用しています。両方の機能を一つの trunk で扱い、すべての層で統合されています。

この統合こそが、エンドツーエンドの Physical AI モデルファクトリーのエンジンとして Cosmos 3 を有用にする理由です。その特徴は、以下の 3 つのアーキテクチャ上の選択によって定義されます。

一つのトークンストリーム。すべてのモダリティが単一の共有シーケンスに流入し、1 つのモデルが複数のモダリティ間で読み取りと生成を両立します。理解のためにモデルが読み込む画像、生成されるピクセル、姿勢の差分と把持状態を表すコンパクトなエンボディメントごとのベクトルは、それぞれ独自のエンコーダーを持ちます。ビジョントランスフォーマー（ViT）が画像理解を担当し、凍結された Wan2.2 動画変分オートエンコーダー（VAE）がピクセル生成を担います。この 1 つのアクションベクトルこそが、同じモデルで自動運転車もロボットアームも制御できる理由です。シーケンスは、読み込むテキストとビジョンを含む自己回帰（AR）領域を、生成する動画・音声・アクションを含む拡散領域の前に配置します。

2 つのエクスパートが各層で結合（MoT）。各層では、次のトークンを予測する推論器と、動画・音声・アクションのノイズを除去する生成器が動作します。双ストリームアテンションにより、生成は最終段階でのみではなく、すべての層において推論器の出力に根ざした状態を保ちます。一般的な代替案では、拡散トランスフォーマー（DiT）をビジョン・ランゲージモデル（VLM）に結合し、その最終出力に対して一度だけクロスアテンションを行います。Cosmos 3 は、生成を推論器に底まで根ざさせることで実現しています。

推論時の非対称性

トレーニングとデプロイメントでは、実行される処理量が異なります。トレーニングでは、予測された動画が損失の一部となるため、完全なノイズ除去スケジュールを実行し、動画をピクセルへ復元します。一方、ロボット上では同じモデルが数回のノイズ除去ステップのみを実行し、動画の復元は完全にスキップされます。動画の潜在変数は内部で生成され、行動の根拠として機能しますが、デコードされるのはロボットの関節位置に直接対応する行動トークンのみです。

以下の図では、共有トークンストリームとアテンションによって結合された 2 つのエキスパートを一つのビューで示しています。自己回帰（AR）サブシーケンス（テキストおよびモデルが理解のために読み取るビジョントークン）と拡散モデル（DM）サブシーケンス（動画、音声、行動トークンを生成する部分）は、共有の Reasoner および Generator タワーを通過します。右側のアテンションマスクは 2 つのエキスパートの違いを示しています：DM のクエリは AR キーと DM キーの両方にアテンションを向け（フルアテンション）、AR のクエリは因果関係を保ち、拡散トークンを決して見ません。

図 1: Cosmos 3 の共有トークンストリームと、層ごとのアテンションで結合された 2 つのエキスパート

(出典：Cosmos 3: Omnimodal World Models for Physical AI)

3 つの動作モード、1 つのアーキテクチャ

中間学習済みベースチェックポイントは、どのトークンからノイズとして開始するかを変更することで 3 つのジョブを実行します。その後、ポストトレーニングによりチェックポイントを単一のモードと制御周波数に特化させます。

フォワードダイナミクス（世界モデル）では、アクションは明確で動画にはノイズが含まれます。「このフレームとこのアクションが与えられたとき、次は何が起こるか？」という問いに答える仕組みです。これは合成データ生成エンジンとして機能し、現実の収集ではコスト的に実現できないようなロングテイルの運転シーンや稀な操作インタラクションを生成するために活用されます。

インバースダイナミクス（アクションラベラー）では、動画は明確でアクションにはノイズが含まれます。「この 2 つのフレームが与えられたとき、どの変化を引き起こしたアクションか？」という問いに答える仕組みです。これにより、ラベル付けされていない動画（生テレオペレーション記録やサードパーティのロボット映像、YouTube の運転映像など）を、アクション付きのトレーニングデータに変換します。

ポリシー（実際に稼働するロボット）では、両方にノイズが含まれつつも、3 視点の画像と固有感覚（プロプリオセプション）に基づいて条件付けられます。32 フレーム先の関節位置を出力し、予測された動画フレームはアクション予測を裏付ける副産物として生成されます。

このモデルファミリーには 2 つのティアがあります。Cosmos3-Nano は 16B パラメータで、密集型（dense）の 8B パラメータ Qwen3-VL ベースカーン上に構築されています。一方、Cosmos3-Super は 64B パラメータで、密集型の 32B パラメータ Qwen3-VL ベースカーン上に構築されています。Cosmos3-Nano-Policy-DROID のようなタスク別バリアントはこれらのティアを基盤にしています。また NVIDIA は Cosmos3-Edge もリリースしており、これはオンデバイス展開用のコンパクトな 4B ティアです（Jetson Thor および Orin でベンチマーク済み）。Edge は Nano や Super と同じ物理世界事前学習データを使用しますが、Qwen3-VL から初期化されるのではなく、ゼロから訓練された密集型約 2B パラメータのベースカーン上に構築されているため、重みの系統は別物です。つまり、Nano チェックポイントを縮小して Edge に適用するのではなく、ターゲットハードウェア向けに Edge を直接ポストトレーニングします。

以下の図は、3 つのモードを並べて示しています。実線のボックスはクリーン（既知）なトークンを、点線のボックスはモデルがノイズ除去するノイジーなトークンを表します。

フォワードダイナミクスでは、アクションと現在のフレームをクリーンに保ちつつ、未来の動画をノイズ除去します。インバースダイナミクスでは、動画をクリーンに保ちながら、アクションをノイズ除去します。ポリシーは最初のフレームのみをクリーンとして受け取り、ロボットが実行するアクションをノイズ除去します。

これら 3 つのモードには同じアーキテクチャが使われており、異なる点はクリーンなトークンとノイジーなトークンのパターンだけです。ベースチェックポイントでは 3 つのモードすべてが利用可能ですが、Cosmos3-Nano-Policy-DROID のようなポストトレーニング済みバリアントは、15 Hz で 32 ステップのホライズンを備えたポリシーモードに特化しています。

Figure 2: 一つのチェックポイントにおける 3 つのアクションモード。どのトークンがノイズとして始まるかで設定されます

（出典：Cosmos 3: Omnimodal World Models for Physical AI）

2. 単一モデルから、永続的なモデルファクトリーへ

ロボットや自動運転車（AV）を開発するチームは、一度きりの微調整ワークロードを実行するわけではありません。彼らは継続的なループを回します。

「実データを収集し、キュレーションを行い、合成データで拡張し、ポストトレーニングを実施し、クローズドループシミュレーションで評価し、ポリシーを展開し、さらにデータを収集して、また繰り返す」というサイクルです。

Figure 3: 4 つの段階からなる物理 AI モデルファクトリーのフライホイール

このループには 4 つの段階があります。(1) DROID、BridgeData2、自動運転車のセンサーログといった実世界の Physical AI データを、Amazon Simple Storage Service (Amazon S3) や Amazon FSx for Lustre にある共有コーパスに収集・整備します。(2) Cosmos3-Super という教師モデルが、そのコーパスを補完する合成データを生成します。(3) 合成データと実データを組み合わせたコーパスで、展開可能な Cosmos3-Nano ポリシーのポストトレーニングを行い、Nano と Super の両階層に対してビジョン微調整を適用します。(4) ポリシーはクローズドループシミュレーション内で評価され、その失敗事例が次の世代のターゲットとして抽出され、次ラウンドのために再びコーパスに組み込まれます。

理想的にはこのループは継続し、新しいデータが入ってくるたびに各段階が繰り返されます。このサイクルによってコストの主要因は、単一ジョブのピークスループットではなく、GPU のグッドプット（予約された GPU 時間あたりの実質的なパイプライン進捗）となります。グッドプットを最大化するには、各段階を一つのプールで共有し、別々のクラスター間で再プロビジョニングやデータ移動による時間を失わないようにする必要があります。Cosmos 3 はこれを可能にします。世界シミュレーション生成器、ポリシー、知覚モデルという 3 つのモデルクラスを、異なるモードで動作する単一のモデルとして統合するからです。このフライングホイールを支えるには、その下にあるクラスターも同じ形状である必要があります。つまり、ジョブごとにバラバラな計算環境を用意するのではなく、一つの管理平面に一つの永続的なプールを持つ構成です。

Amazon SageMaker HyperPod on Amazon EKS は、まさにその形を実現します。Cosmos 3 の背後にある各アーキテクチャ上の選択が、具体的なクラスター需要を生み出しています。単一のトークンストリームと 64B MoT 構造により、トレーニングは低遅延な相互接続を必要とする長系列・マルチノードジョブとなります。トレーニングと推論の非対称性によって、生成、ポストトレーニング、評価がすべて同一モデルおよび同一ストレージ層上で完結するため、各工程で容量を細分化するのではなく、一つのコミットされたキャパシティプールを時間共有して運用できます。このサイクルを継続的に回し続けるには、予約されかつ常時監視されたキャパシティが必要です。Amazon SageMaker HyperPod の 4 つの特性が、これらの要件に順次応えます。

すべての工程を 1 つのクラスターで完結。Amazon SageMaker HyperPod は EKS を介してクラスターをオーケストレーションするため、ループ内の 3 つのエンジン（生成、ポストトレーニング、評価）は、単一の共有 GPU プール上で通常の Kubernetes ワークロードとして実行されます。

生成処理は vLLM-Omni サーバーで、ポストトレーニングは torchrun 下での cosmos-framework（Fully Sharded Data Parallel (FSDP2) と Ulysses コンテキスト並列化を併用）で、評価は単一 GPU のポリシーサーバーで行います。また、これらは同じストレージ層も共有します。

Amazon FSx for Lustre ファイルシステムが Elastic Fabric Adapter (EFA) を介してアクセスされ、その背後には Amazon S3 バケットがデータリポジトリ連携（DRA）を介してバックアップされています。このファイルシステムは 1 回マウントするだけで、同じパスからすべての工程にデータを供給します。

生成処理で合成クリップを書き込み、ポストトレーニングでそれを読み取り、ポリシーサーバーも同じボリュームからチェックポイントを読み出します。工程間で再プロビジョニングやテラバイト規模のデータ移行は不要です。また、地域ロックされた AV データも 1 つのリージョン内のクラスター内に留まります。

「健全性を確認し、自動復旧するキャパシティ」：生成処理はバースト状で GPU 使用時間を支配し、ポストトレーニングは多数のノードにわたって数日間に及ぶため、事前にプロビジョニングされ継続的に監視されたキャパシティが求められます。Amazon SageMaker HyperPod は故障したノードを常時検知して自動的に再起動または交換し、柔軟なトレーニングプランでキャパシティを事前に確保できます。また、管理されたジョブの自動再開機能により、ワーカーの障害が発生しても回復範囲を限定できます。Kubeflow PyTorchJob がポッド群を再構築し、NCCL が通信経路を再形成する一方、cosmos-framework は最新の PyTorch Distributed Checkpoint (DCP) から処理を再開します。その結果、ノードの障害による損失は、1 チェックポイント間隔分の再計算と、ノード交換・スケジューリング遅延に限定され、ジョブ全体が失われることはありません。

「EFA はマルチノード NCCL 用に既に配線済み」：Cosmos 3 のトレーニングでは、動画の潜在表現（latents）、テキスト、行動データを組み合わせた非常に長いシーケンス（各々数万トークン）を処理するため、64B モデル級では FSDP2 に加えてコンテキスト並列化が必須となります。これにより、すべてのレイヤーでノード間通信（クロスノード・コレクティブ）が発生します。これを手動で構築するのは通常、EFA スタックの整合、NCCL プラグインの設定、および cosmos-framework が指定する Torch と NCCL の正確なバージョン合わせなど、マルチノード環境における時間的なボトルネックとなります。Amazon SageMaker HyperPod はこれらを事前設定済みの状態で提供しており、AWS Deep Learning Containers (DLC) イメージと組み合わせることで、Torch や aws-ofi-nccl のバージョンがフレームワークの要件と一致しているため、NCCL over EFA がすぐに動作するようになります。

複数のロボット種別や自動運転車（AV）のバリエーションを同時に扱うファクトリの場合、Amazon SageMaker HyperPod のタスクガバナンス機能 task governance が役立ちます。これは Kueue を基盤としており、リソースプールを名前空間ごとのキューに分割し、クォータや優先度、プリエンプション（割り込み）を設定できます。これにより、多様なジョブが数十個あっても、競合してリソースを奪い合うのではなく、一つの容量予約を共有できます。その結果、プロジェクト間で otherwise-idle だった GPU を活用でき、実効スループット（goodput）が向上します。

単一のロボット種別のみを対象とするプログラムでは、この機能は必須ではありません。ただし、多数のジョブが同じプールを争うようになれば、タスクガバナンスを導入するメリットは明確になります。

Amazon SageMaker HyperPod と軽量な選択肢を選ぶ基準は、「作業の単位」にあります。一度きりのファインチューニングであれば、HyperPod クラスター特有の耐障害性や永続性は必ずしも必要ありません。短時間のランでノード障害に遭遇することは稀なので、一時的な管理型トレーニングジョブ（例：Amazon SageMaker AI のトレーニングジョブ）で十分です。

一方、Cosmos 3 のフィードバックループのように、生成処理が継続的に行われ、ポストトレーニングがマルチノードかつ長時間にわたり、評価も同じストレージ層上で共起し、統計的に頻繁な障害が発生する環境では、Amazon SageMaker HyperPod が最適です。

3. 私たちが構築しているもの

本ソリューションは、フライングホイールの各段階を担う3 つの代表的なワークロードに対してポストトレーニングを実施し、生成と評価の両工程をエンドツーエンドで検証します。これら 3 つのワークロードはすべて、p5en.48xlarge ノード（NVIDIA H200 GPU を 8 基搭載）上で実運用可能なチェックポイントを用いて実行されます。具体的には、ロボット操作ポリシーと、2 つのビジョン知覚タスク向けのファインチューニングワークロードです。

AV（自動運転車）のポストトレーニングについては本稿では直接取り上げませんが、Cosmos 3 のベースモデルは、パブリックな合成ドライビングコーパス（SDG-DriveSim、Hugging Face 上の nvidia/PhysicalAI-WorldModel-Synthetic-Autonomous-Driving-Scenarios データセット）を用いて訓練されています。各エンボディメントごとのアクション投影は、AV の自己位置（ego-pose）アクション空間へ拡張できるように設計されているため、本稿で解説するレシピやクラスター設定は、AV のポストトレーニングにもそのまま適用可能です。

このトレーニングスタックには NVIDIA 製の cosmos-framework を使用しており、フレームワークパッケージへのフォークやソースコードの改変は一切行われていません。そのため、アップストリームの更新を問題なく取り込むことができます。訓練では FSDP2 を採用し、シーケンス長やノード数が増大するにつれて、ハイブリッド・シャーデッド・データ並列処理（HSDP）およびコンテキスト並列処理へとスケーリングします。これらの並列化戦略を各ティアでどのように設定するかについては、セクション 6.2 で詳しく解説しています。

本ガイドでは、参照インスタンスタイプとして p5en.48xlarge を通して使用します。