NVIDIA の Cosmos フレームワークチュートリアル：オムニモーダル混合トランスフォーマーを用いた Colab 対応ミニ版 Cosmos 3 ワールドモデルの設計
本記事は、NVIDIA の Cosmos 3 ワールドモデルの完全な推論が Colab 環境では困難であることを示しつつ、そのアーキテクチャを模倣した軽量版の実装とトレーニング手法を詳述し、開発者がマルチモーダル世界モデルの核心技術を体験する実践的なガイドを提供している。
キーポイント
Colab 環境における Cosmos 3 の実用性限界
記事は、標準的な Colab ハードウェア（GPU メモリや計算能力）では本格的な Cosmos 3 チェックポイントの推論が現実的ではないことを明確に示し、環境プロbing の重要性を説いている。
オムニモーダル混合トランスフォーマーの実装
Cosmos の核となる「共有クロスモーダルアテンション」と「モダリティ固有の専門家ルーティング」を持つ、テキスト・ビジョン・アクションを扱うコンパクトな世界モデルを実際に構築・トレーニングする手順を解説している。
合成データを用いた学習とロールアウト
物理世界の関係性を学習させるため合成データを使用し、トレーニング損失の追跡や自己回帰的なロールアウトを通じて、モデルが未来の潜在状態を予測するプロセスを実証している。
Cosmos 3 の実行要件と Colab の限界
本コードは、Nano-16B モデルを実行するために Ampere 以降の GPU 構造、80GB 以上のメモリ、および約 150GB のディスク容量が必要であることを検証し、多くの環境ではこれらの条件を満たしていないことを示しています。
教育的な代替パスの提示
ハードウェア要件が満たされない場合でも、リアルな Cosmos 3 推論は不可能であるため、代わりに教育目的のためのミニチュア版や代替アプローチへの道筋を示唆しています。
環境要件の検証と Colab の限界
Python, PyTorch, CUDA、GPU メモリや計算能力などのシステム情報を確認し、標準的な Colab ハードウェアでは 160 億パラメータを超える本番 Cosmos チェックポイントを実行できない理由を明確に提示します。
公式フレームワークのクローンとマッピング
真のソースとなる `cosmos_framework` パッケージをクローンし、実行環境に適切にマッピングする手順がセクション 1 の主要なタスクとして定義されています。
影響分析・編集コメントを表示
影響分析
この記事は、大規模な AI モデルへのアクセス制限がある開発者や研究者にとって、最先端のマルチモーダル技術（Cosmos）の内部構造と学習原理を低コストで体験できる貴重なリソースとなります。特に、理論的な解説だけでなく、実際のコードベースでの実装とトレーニングプロセスに焦点を当てているため、世界モデルの実践的開発における重要な指針となるでしょう。
編集コメント
Cosmos の完全版は高価なインフラが必要ですが、このチュートリアルはその「骨格」を学ぶための優れた練習問題です。実機での大規模学習が不可能な環境でも、アーキテクチャの理解を深めるには最適なアプローチと言えます。
本チュートリアルでは、実際の Cosmos 3 チェックポイントを実行する際のハードウェア制限について正直に言及しつつ、実用的な Colab フレンドリーな観点から NVIDIA の cosmos-framework を探ります。まず、現在のランタイム、GPU の機能、CUDA の利用可能性、メモリ、ディスク容量を確認し、なぜ標準的な Colab ハードウェア上で完全な Cosmos 3 推論が現実的ではないのかを理解します。そこで立ち止まるのではなく、フレームワークの実際の構造、CLI サフェース、入力スキーマ、モデルモードを基礎として、手作業によるミニチュア実装を行います。その後、合成された物理世界データ、トレーニング損失の追跡、自己回帰ロールアウトを用いて、テキスト・ビジョン・アクションストリームに対してモダリティ固有のエキスパートルーティングを備えた共有クロスモーダルアテンションという Cosmos の核となるアイデアを反映した、コンパクトなオムニモーダル Mixture-of-Transformers（混合トランスフォーマー）世界モデルを構築・訓練します。これにより、モデルがどのようにモダリティ間の関係を学習し、簡略化されつつも技術的に意味のある方法で将来の潜在状態を予測するかを示します。
Colab ハードウェア制限の探求
コードをコピーしました（コピー済み）
別のブラウザを使用してください
import os, sys, json, time, math, textwrap, subprocess, shutil, platform
from pathlib import Path
def rule(title=""):
line = "=" * 86
print("\n" + line + ("\n " + title if title else "") + "\n" + line)
def spark(vals, width=60):
"""Tiny ASCII sparkline for a 1-D sequence (works with no plotting libs)."""
if not vals: return ""
blocks = "▁▂▃▄▅▆▇█"
lo, hi = min(vals), max(vals)
rng = (hi - lo) or 1.0
step = max(1, len(vals) // width)
s = "".join(blocks[min(len(blocks) - 1, int((v - lo) / rng * (len(blocks) - 1)))]
for v in vals[::step])
return s
rule("SECTION 0 — Environment probe: what you have vs. what Cosmos 3 actually needs")
IN_COLAB = "google.colab" in sys.modules
print(f"Running inside Google Colab : {IN_COLAB}")
print(f"Python : {platform.python_version()} ({platform.system()})")
try:
import torch
except ModuleNotFoundError:
print("torch not found — installing CPU build (a few seconds)...")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "torch"], check=False)
import torch
print(f"PyTorch : {torch.__version__}")
CUDA_OK = torch.cuda.is_available()
DEVICE = torch.device("cuda" if CUDA_OK else "cpu")
gpu_name, gpu_mem_gb, cc = "None (CPU)", 0.0, (0, 0)
if CUDA_OK:
p = torch.cuda.get_device_properties(0)
gpu_name = p.name
gpu_mem_gb = p.total_memory / 1024**3
cc = torch.cuda.get_device_capability(0)
print(f"CUDA build : {torch.version.cuda}")
print(f"GPU : {gpu_name}")
print(f"GPU memory : {gpu_mem_gb:.1f} GiB")
print(f"Compute capability : sm_{cc[0]}{cc[1]}")
try:
free_gb = shutil.disk_usage('/').free / 1024**3
print(f"Free disk : {free_gb:.0f} GiB")
except Exception:
free_gb = 0.0
AMPERE = cc[0] >= 8
reqs = [
("GPU architecture", "Ampere+ (sm_80+, A100/RTX30xx)", "OK" if AMPERE else "TOO OLD (T4=sm_75)"),
("GPU memory", ">=80 GiB for Nano-16B (single H100)", "OK" if gpu_mem_gb >= 79 else f"{gpu_mem_gb:.0f} GiB — insufficient"),
("CUDA toolkit", ">=12.8", "check" ),
("Free disk", "~150 GiB first run (~1 TB HF cache)", "OK" if free_gb >= 150 else f"{free_gb:.0f} GiB — insufficient"),
("Attention kernels","FlashAttn-3 (Hopper) / FA2 (Ampere)", "needs Ampere+"),
]
print("\n Can this machine run the REAL Cosmos 3 checkpoints?")
print(" " + "-" * 82)
print(f" {'Requirement':<18}{'Cosmos 3 needs':<38}{'You have'}")
print(" " + "-" * 82)
for k, need, have in reqs:
print(f" {k:<18}{need:<38}{have}")
print(" " + "-" * 82)
VERDICT = AMPERE and gpu_mem_gb >= 79 and free_gb >= 150
print(f" VERDICT: {'This machine could attempt Nano-16B.' if VERDICT else 'NO — real Cosmos 3 inference is not possible here. Educational path below.'}")
まず、ランタイムユーティリティの準備を行い、現在のマシンが Cosmos 3 の推論を現実的にサポートできるかどうかを確認します。Python、PyTorch、CUDA、GPU メモリ、計算能力（compute capability）、利用可能なディスク容量を検査し、私たちの環境が実際のハードウェア要件とどう比較されるかを確認します。その後、標準的な Colab ハードウェアでは通常 16B+ の Cosmos チェックポイントを実行できない理由を説明する明確な結論を表示します。
コピーコード コピー済み 別のブラウザを使用してください
ルール（セクション 1 — 真の cosmos_framework パッケージをクローンしてマッピングする（信頼できるソース））
Cosmos-Framework パッケージのマッピング
コピーコード コピー済み 別のブラウザを使用してください
REPO = "https://github.com/NVIDIA/cosmos-framework.git"
DST = Path("/content/cosmos-framework") if Path("/content").exists() else Path("cosmos-framework")
cloned = False
try:
if not DST.exists():
print(f"Shallow-cloning {REPO} ...")
subprocess.run(["git", "clone", "--depth", "1", REPO, str(DST)],
check=True, capture_output=True, text=True, timeout=180)
cloned = DST.exists()
except Exception as e:
print(f"(Clone skipped/failed — offline is fine, tutorial continues.) {e}")
if cloned:
print(f"Repo at: {DST}\n")
pkg = DST / "cosmos_framework"
if pkg.exists():
print("cosmos_framework/ subpackages (the real code layout):")
for child in sorted(pkg.iterdir()):
if child.is_dir() and not child.name.startswith(("_", ".")):
n_py = len(list(child.rglob("*.py")))
print(f" • {child.name:<20} ({n_py:>3} .py files)")
example = DST / "inputs" / "omni" / "t2v.json"
if example.exists():
print(f"\nReal example input spec ({example.relative_to(DST)}):")
print(textwrap.indent(example.read_text().strip(), " "))
else:
print("Proceeding without a local clone (we already extracted the real schema/CLI).")
print("""
Real CLI surface (docs/inference.md):
Single GPU : python -m cosmos_framework.scripts.inference \\
--parallelism-preset=latency -i "inputs/omni/t2v.json" \\
-o outputs/omni_nano --checkpoint-path Cosmos3-Nano --seed 0
Multi GPU : torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \\
--parallelism-preset=throughput -i "inputs/omni/*.json" \\
-o outputs/omni_super --checkpoint-path Cosmos3-Super --seed 0
Models : Cosmos3-Nano (16B, all modes) | Cosmos3-Super (65B, t2i/t2v/i2v)
Modes : text2image · text2video · image2video · video2video ·
forward_dynamics · inverse_dynamics · policy
Parallelism: FSDP dp-shard / dp-replicate · context (cp) · CFG (cfgp)
presets {latency, throughput}
Guardrails : Cosmos-Guardrail1 + Qwen3Guard-Gen-0.6B + RetinaFace (on by default)
""")
rule("SECTION 2 — Omnimodal Mixture-of-Transformers (MoT) world model — the idea")
print(r"""
Cosmos 3 unifies language, image, video, audio and ACTION in ONE model. The key trick
is a Mixture-of-Transformers: every modality is turned into tokens placed on a SINGLE
interleaved sequence; SELF-ATTENTION is SHARED across all modalities (so vision can be
conditioned on text, actions on vision, etc.), but each token is processed by a
MODALITY-SPECIFIC expert feed-forward block ("Mixture-of-Transformers" routing).
text tokens vision tokens action tokens
[t0 t1 t2 ...] [v0 v1 v2 ...] [a0 a1 ...]
\ | /
\ | /
+----------- one sequence -----------+
|
┌─────────── shared causal self-attention (RoPE) ───────────┐
│ every token attends to all earlier tokens, ANY modality │
└───────────────────────────────────────────────────────────┘
|
route each token to its modality's EXPERT FFN (SwiGLU):
text→Expert0 vision→Expert1 action→Expert2
|
per-modality heads: next-token / next-latent / next-action
Physical-AI modes fall right out of this one model:
text2video = generate the vision-token stream from a text prompt
image2video = condition vision stream on a first frame + text
forward_dynamics= given frames + ACTIONS, roll future frames forward (a world model)
inverse_dynamics= given frames, infer the ACTIONS that caused them
policy = given an observation + goal, emit ACTIONS (+ imagined rollout)
Below we build a faithful ~4M-param miniature of exactly this and train it live.
(The real model uses flow-matching/diffusion for the continuous vision stream; our toy
uses a simple MSE next-latent objective so it trains in seconds — the ROUTING and
SHARED-ATTENTION structure are the same.)
""")
私たちは、ソースから直接パッケージ構造、入力スキーマ、および CLI ワークフローを理解するために、実際の cosmos-framework リポジトリをクローンして調査します。また、テキストから動画への変換、画像から動画への変換、順方向ダイナミクス、逆方向ダイナミクス、ポリシーなどのモードを含む、シングル GPU およびマルチ GPU 起動用の公式推論コマンドパターンも表示します。その後、テキスト、ビジョン、アクショントークンがアテンションを共有しつつも、モダリティ固有のエキスパートフィードフォワードブロックを使用するオムニモーダル Mixture-of-Transformers の概念を紹介します。
Copy CodeCopiedUse a different Browser
rule("SECTION 3 — Implement & train the omnimodal MoT from scratch")
Building The Omnimodal MoT
Copy CodeCopiedUse a different Browser
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
torch.manual_seed(0)
@dataclass
class Cfg:
d_model: int = 192
n_head: int = 6
n_layer: int = 4
ffn_mult: int = 2
n_mod: int = 3
text_vocab:int = 16
vis_dim: int = 8
act_dim: int = 4
Lt: int = 8
Lv: int = 8
La: int = 6
cfg = Cfg()
class RMSNorm(nn.Module):
def __init__(self, d, eps=1e-6):
super().__init__(); self.w = nn.Parameter(torch.ones(d)); self.eps = eps
def forward(self, x):
return self.w * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def build_rope(T, hd, device, base=10000.0):
pos = torch.arange(T, device=device, dtype=torch.float32)[:, None]
idx = torch.arange(0, hd, 2, device=device, dtype=torch.float32)[None, :]
freq = 1.0 / (base ** (idx / hd))
ang = pos * freq
cos = torch.cos(ang).repeat(1, 2)[None, None]
sin = torch.sin(ang).repeat(1, 2)[None, None]
return cos, sin
def rotate_half(x):
hd = x.shape[-1]; x1, x2 = x[..., :hd // 2], x[..., hd // 2:]
return torch.cat([-x2, x1], -1)
def apply_rope(q, k, cos, sin):
return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
class Attention(nn.Module):
"""回転埋め込みを備えた共有型クロスモーダル因果自己アテンション (causal self-attention)。"""
def __init__(self, c: Cfg):
super().__init__()
self.H, self.hd = c.n_head, c.d_model // c.n_head
self.qkv = nn.Linear(c.d_model, 3 * c.d_model, bias=False)
self.proj = nn.Linear(c.d_model, c.d_model, bias=False)
def forward(self, x, cos, sin, mask):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, -1)
q = q.view(B, T, self.H, self.hd).transpose(1, 2)
k = k.view(B, T, self.H, self.hd).transpose(1, 2)
v = v.view(B, T, self.H, self.hd).transpose(1, 2)
q, k = apply_rope(q, k, cos, sin)
att = (q @ k.transpose(-2, -1)) / math.sqrt(self.hd)
att = att.masked_fill(mask, float("-inf")).softmax(-1)
o = (att @ v).transpose(1, 2).reshape(B, T, D)
return self.proj(o)
class Expert(nn.Module):
""」各モーダルごとの SwiGLU フィードフォワード型「トランスフォーマーエキスパート」。"""
def __init__(self, d, mult):
super().__init__(); h = d * mult
self.w1 = nn.Linear(d, h, bias=False)
self.w3 = nn.Linear(d, h, bias=False)
self.w2 = nn.Linear(h, d, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class MoTBlock(nn.Module):
""」共有型アテンションとモーダル別エキスパートによるルーティングを行う Mixture-of-Transformers (MoT)。"""
def __init__(self, c: Cfg):
super().__init__()
self.attn_norm = RMSNorm(c.d_model)
self.attn = Attention(c)
self.ffn_norm = nn.ModuleList([RMSNorm(c.d_model) for _ in range(c.n_mod)])
self.experts = nn.ModuleList([Expert(c.d_model, c.ffn_mult) for _ in range(c.n_mod)])
def forward(self, x, cos, sin, mask, mod_id):
x = x + self.attn(self.attn_norm(x), cos, sin, mask)
out = torch.zeros_like(x)
for i, exp in enumerate(self.experts):
sel = (mod_id == i).view(1, -1, 1).to(x.dtype)
out = out + sel * exp(self.ffn_normi)
return x + out
class OmniMoT(nn.Module):
def __init__(self, c: Cfg):
super().__init__(); self.c = c
self.text_emb = nn.Embedding(c.text_vocab, c.d_model)
self.vis_in = nn.Linear(c.vis_dim, c.d_model)
self.act_in = nn.Linear(c.act_dim, c.d_model)
self.mod_emb = nn.Embedding(c.n_mod, c.d_model)
self.blocks = nn.ModuleList([MoTBlock(c) for _ in range(c.n_layer)])
self.norm = RMSNorm(c.d_model)
self.text_head = nn.Linear(c.d_model, c.text_vocab, bias=False)
self.vis_head = nn.Linear(c.d_model, c.vis_dim, bias=False)
self.act_head = nn.Linear(c.d_model, c.act_dim, bias=False)
ids = torch.cat([torch.zeros(c.Lt), torch.ones(c.Lv), torch.full((c.La,), 2)]).long()
self.register_buffer("mod_id", ids, persistent=False)
def forward(self, text, vis, act):
c = self.c
x = torch.cat([self.text_emb(text), self.vis_in(vis), self.act_in(act)], 1)
x = x + self.mod_emb(self.mod_id)[None]
B, T, D = x.shape
cos, sin = build_rope(T, D // c.n_head, x.device)
mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)[None, None]
for blk in self.blocks:
x = blk(x, cos, sin, mask, self.mod_id)
x = self.norm(x)
ht = self.text_head(x[:, :c.Lt])
hv = self.vis_head(x[:, c.Lt:c.Lt + c.Lv])
ha = self.act_head(x[:, c.Lt + c.Lv:])
return ht, hv, ha
model = OmniMoT(cfg).to(DEVICE)
n_params = sum(p.numel() for p in model.parameters())
print(f"Model built: OmniMoT | {n_params/1e6:.2f}M params | {cfg.n_layer} MoT blocks "
f"x {cfg.n_mod} experts | device={DEVICE}")
PyTorch を用いて、ミニチュア版のオムニモーダル Mixture-of-Transformers モデルをゼロから実装します。RMSNorm、回転埋め込み（rotary embeddings）、共有因果自己注意機構（shared causal self-attention）、モダリティ固有の SwiGLU エキスパート、そして完全な OmniMoT アーキテクチャを定義します。その後、利用可能なデバイス上でモデルを初期化し、パラメータ数、層数、エキスパート構造、および実行時のデバイスを報告します。
Copy CodeCopiedUse a different Browser
K = 4
g = torch.Generator().manual_seed(1)
合成データ上でのトレーニング
Copy CodeCopiedUse a different Browser
K = 4
g = torch.Generator().manual_seed(1)
合成データでのトレーニング
TEXT_TRANS = torch.stack([torch.softmax(torch.randn(cfg.text_vocab, cfg.text_vocab, generator=g), -1)
for _ in range(K)])
VIS_DYN = torch.stack([0.9 * torch.linalg.qr(torch.randn(cfg.vis_dim, cfg.vis_dim, generator=g))[0]
for _ in range(K)])
ACT_MAP = torch.randn(cfg.act_dim, cfg.vis_dim, generator=g) * 0.5
def make_batch(B):
codes = torch.randint(0, K, (B,), generator=g)
text = torch.zeros(B, cfg.Lt, dtype=torch.long)
text[:, 0] = torch.randint(0, cfg.text_vocab, (B,), generator=g)
for t in range(1, cfg.Lt):
probs = TEXT_TRANS[codes, text[:, t-1]]
text[:, t] = torch.multinomial(probs, 1, generator=g).squeeze(1)
vis = torch.zeros(B, cfg.Lv, cfg.vis_dim)
vis[:, 0] = torch.randn(B, cfg.vis_dim, generator=g)
for t in range(1, cfg.Lv):
vis[:, t] = torch.einsum("bij,bj->bi", VIS_DYN[codes], vis[:, t-1]) \
+ 0.02 * torch.randn(B, cfg.vis_dim, generator=g)
vis_state = vis.mean(1)
act = torch.zeros(B, cfg.La, cfg.act_dim)
for t in range(cfg.La):
act[:, t] = (ACT_MAP @ (vis_state * (0.8 ** t)).T).T \
+ 0.02 * torch.randn(B, cfg.act_dim, generator=g)
return text.to(DEVICE), vis.to(DEVICE), act.to(DEVICE), codes
def loss_fn(model, text, vis, act):
ht, hv, ha = model(text, vis, act)
l_text = F.cross_entropy(ht[:, :-1].reshape(-1, cfg.text_vocab), text[:, 1:].reshape(-1))
l_vis = F.mse_loss(hv[:, :-1], vis[:, 1:])
l_act = F.mse_loss(ha[:, :-1], act[:, 1:])
return l_text + l_vis + l_act, (l_text.item(), l_vis.item(), l_act.item())
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
STEPS, BATCH = 400, 64
hist, t0 = [], time.time()
print(f"\nTraining for {STEPS} steps (batch={BATCH})...")
model.train()
for step in range(1, STEPS + 1):
text, vis, act, _ = make_batch(BATCH)
loss, parts = loss_fn(model, text, vis, act)
opt.zero_grad(); loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
hist.append(loss.item())
if step % 50 == 0 or step == 1:
print(f" step {step:4d} total {loss.item():6.3f} "
f"| text {parts[0]:5.3f} vision {parts[1]:6.4f} action {parts[2]:6.4f}")
print(f"Trained in {time.time()-t0:.1f}s | loss {hist[0]:.3f} -> {hist[-1]:.3f}")
print(" loss curve: " + spark(hist))
try:
import matplotlib.pyplot as plt
plt.figure(figsize=(7, 3))
plt.plot(hist); plt.title("OmniMoT training loss"); plt.xlabel("step"); plt.ylabel("loss")
plt.grid(alpha=0.3); plt.tight_layout(); plt.show()
except Exception:
pass
テキスト、ビジョン、アクションのストリームが隠されたシーンコードに依存する合成物理世界データセットを作成します。結合されたクロスエントロピーと平均二乗誤差 (MSE) 目的関数を用いて、ミニチュア型ワールドモデルを訓練し、次のテキストトークン、未来のビジョン潜在変数、および未来のアクションベクトルを予測させます。複数のステップにわたるトレーニング損失を追跡し、必要に応じて曲線をプロットして、モデルがどのようにクロスモーダルダイナミクスを学習するかを示します。
Copy CodeCopiedUse a different Browser
rule("SECTION 4 — Autoregressive world-model rollout (forward_dynamics analog)")
print(textwrap.dedent("""
A world model predicts the FUTURE. Here we give the trained model a partial vision
trajectory and let it roll the vision latents forward one step at a time — exactly
the loop Cosmos 3 runs for forward_dynamics (predict future frames) and policy
(predict future frames + actions). We compare the model's imagined trajectory to the
ground-truth physics (the hidden VIS_DYN map) it never saw explicitly.
"""))
@torch.no_grad()
Autoregressive World-Model Rollout
Copy CodeCopiedUse a different Browser
def rollout(model, text, vis_prefix, act, n_future):
"""ビジョンプレフィックスから視覚潜在変数を自己回帰的に予測し、n_future 分先を生成する。"""
model.eval()
vis = vis_prefix.clone()
preds = []
for _ in range(n_future):
pad = cfg.Lv - vis.shape[1]
vis_in = vis if pad <= 0 else torch.cat(
[vis, vis[:, -1:].repeat(1, pad, 1)], 1)
_, hv, _ = model(text, vis_in[:, :cfg.Lv], act)
nxt = hv[:, min(vis.shape[1], cfg.Lv) - 1:min(vis.shape[1], cfg.Lv)]
preds.append(nxt)
vis = torch.cat([vis, nxt], 1)
return torch.cat(preds, 1)
text, vis, act, codes = make_batch(4)
prefix_len, n_future = 3, 5
pred = rollout(model, text, vis[:, :prefix_len], act, n_future)
gt = vis[:, prefix_len-1:prefix_len].clone()
true_steps = [gt]
cur = gt
for _ in range(n_future):
cur = torch.einsum("bij,bj->bi", VIS_DYN[codes].to(DEVICE), cur[:, -1]).unsqueeze(1)
true_steps.append(cur)
gt_traj = torch.cat(true_steps[1:], 1)
err = F.mse_loss(pred, gt_traj).item()
print(f"Rolled out {n_future} future vision latents for {text.shape[0]} scenes.")
print(f"Imagined-vs-true-physics MSE : {err:.4f} (a small number = it learned the dynamics)")
print(f"Example (scene 0) latent[0] over time:")
print(f" predicted : {pred[0,:,0].detach().cpu().numpy().round(3).tolist()}")
print(f" true : {gt
原文を表示
In this tutorial, we explore NVIDIA’s cosmos-framework from a practical Colab-friendly angle while staying honest about the hardware limits of running real Cosmos 3 checkpoints. We begin by checking the current runtime, GPU capabilities, CUDA availability, memory, and disk space to understand why full Cosmos 3 inference is not realistic on standard Colab hardware. Instead of stopping there, we use the framework’s real structure, CLI surface, input schema, and model modes as the foundation for a hands-on miniature implementation. We then build and train a compact omnimodal Mixture-of-Transformers world model that mirrors the core Cosmos idea: shared cross-modal attention with modality-specific expert routing for text, vision, and action streams. Using synthetic physical-world data, training-loss tracking, and an autoregressive rollout, we show how the model learns relationships across modalities and predicts future latent states in a simplified yet technically meaningful way.
Probing Colab Hardware Limits
Copy CodeCopiedUse a different Browser
import os, sys, json, time, math, textwrap, subprocess, shutil, platform
from pathlib import Path
def rule(title=""):
line = "=" * 86
print("\n" + line + ("\n " + title if title else "") + "\n" + line)
def spark(vals, width=60):
"""Tiny ASCII sparkline for a 1-D sequence (works with no plotting libs)."""
if not vals: return ""
blocks = "▁▂▃▄▅▆▇█"
lo, hi = min(vals), max(vals)
rng = (hi - lo) or 1.0
step = max(1, len(vals) // width)
s = "".join(blocks[min(len(blocks) - 1, int((v - lo) / rng * (len(blocks) - 1)))]
for v in vals[::step])
return s
rule("SECTION 0 — Environment probe: what you have vs. what Cosmos 3 actually needs")
IN_COLAB = "google.colab" in sys.modules
print(f"Running inside Google Colab : {IN_COLAB}")
print(f"Python : {platform.python_version()} ({platform.system()})")
try:
import torch
except ModuleNotFoundError:
print("torch not found — installing CPU build (a few seconds)...")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "torch"], check=False)
import torch
print(f"PyTorch : {torch.__version__}")
CUDA_OK = torch.cuda.is_available()
DEVICE = torch.device("cuda" if CUDA_OK else "cpu")
gpu_name, gpu_mem_gb, cc = "None (CPU)", 0.0, (0, 0)
if CUDA_OK:
p = torch.cuda.get_device_properties(0)
gpu_name = p.name
gpu_mem_gb = p.total_memory / 1024**3
cc = torch.cuda.get_device_capability(0)
print(f"CUDA build : {torch.version.cuda}")
print(f"GPU : {gpu_name}")
print(f"GPU memory : {gpu_mem_gb:.1f} GiB")
print(f"Compute capability : sm_{cc[0]}{cc[1]}")
try:
free_gb = shutil.disk_usage('/').free / 1024**3
print(f"Free disk : {free_gb:.0f} GiB")
except Exception:
free_gb = 0.0
AMPERE = cc[0] >= 8
reqs = [
("GPU architecture", "Ampere+ (sm_80+, A100/RTX30xx)", "OK" if AMPERE else "TOO OLD (T4=sm_75)"),
("GPU memory", ">=80 GiB for Nano-16B (single H100)", "OK" if gpu_mem_gb >= 79 else f"{gpu_mem_gb:.0f} GiB — insufficient"),
("CUDA toolkit", ">=12.8", "check" ),
("Free disk", "~150 GiB first run (~1 TB HF cache)", "OK" if free_gb >= 150 else f"{free_gb:.0f} GiB — insufficient"),
("Attention kernels","FlashAttn-3 (Hopper) / FA2 (Ampere)", "needs Ampere+"),
]
print("\n Can this machine run the REAL Cosmos 3 checkpoints?")
print(" " + "-" * 82)
print(f" {'Requirement':<18}{'Cosmos 3 needs':<38}{'You have'}")
print(" " + "-" * 82)
for k, need, have in reqs:
print(f" {k:<18}{need:<38}{have}")
print(" " + "-" * 82)
VERDICT = AMPERE and gpu_mem_gb >= 79 and free_gb >= 150
print(f" VERDICT: {'This machine could attempt Nano-16B.' if VERDICT else 'NO — real Cosmos 3 inference is not possible here. Educational path below.'}")
We begin by preparing the runtime utilities and checking whether the current machine can realistically support Cosmos 3 inference. We inspect Python, PyTorch, CUDA, GPU memory, compute capability, and available disk space to compare our environment against the actual hardware requirements. We then print a clear verdict explaining why the real 16B+ Cosmos checkpoints cannot usually run on standard Colab hardware.
Copy CodeCopiedUse a different Browser
rule("SECTION 1 — Clone & map the real cosmos_framework package (source of truth)")
Mapping The Cosmos-Framework Package
Copy CodeCopiedUse a different Browser
REPO = "https://github.com/NVIDIA/cosmos-framework.git"
DST = Path("/content/cosmos-framework") if Path("/content").exists() else Path("cosmos-framework")
cloned = False
try:
if not DST.exists():
print(f"Shallow-cloning {REPO} ...")
subprocess.run(["git", "clone", "--depth", "1", REPO, str(DST)],
check=True, capture_output=True, text=True, timeout=180)
cloned = DST.exists()
except Exception as e:
print(f"(Clone skipped/failed — offline is fine, tutorial continues.) {e}")
if cloned:
print(f"Repo at: {DST}\n")
pkg = DST / "cosmos_framework"
if pkg.exists():
print("cosmos_framework/ subpackages (the real code layout):")
for child in sorted(pkg.iterdir()):
if child.is_dir() and not child.name.startswith(("_", ".")):
n_py = len(list(child.rglob("*.py")))
print(f" • {child.name:<20} ({n_py:>3} .py files)")
example = DST / "inputs" / "omni" / "t2v.json"
if example.exists():
print(f"\nReal example input spec ({example.relative_to(DST)}):")
print(textwrap.indent(example.read_text().strip(), " "))
else:
print("Proceeding without a local clone (we already extracted the real schema/CLI).")
print("""
Real CLI surface (docs/inference.md):
Single GPU : python -m cosmos_framework.scripts.inference \\
--parallelism-preset=latency -i "inputs/omni/t2v.json" \\
-o outputs/omni_nano --checkpoint-path Cosmos3-Nano --seed 0
Multi GPU : torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \\
--parallelism-preset=throughput -i "inputs/omni/*.json" \\
-o outputs/omni_super --checkpoint-path Cosmos3-Super --seed 0
Models : Cosmos3-Nano (16B, all modes) | Cosmos3-Super (65B, t2i/t2v/i2v)
Modes : text2image · text2video · image2video · video2video ·
forward_dynamics · inverse_dynamics · policy
Parallelism: FSDP dp-shard / dp-replicate · context (cp) · CFG (cfgp)
presets {latency, throughput}
Guardrails : Cosmos-Guardrail1 + Qwen3Guard-Gen-0.6B + RetinaFace (on by default)
""")
rule("SECTION 2 — Omnimodal Mixture-of-Transformers (MoT) world model — the idea")
print(r"""
Cosmos 3 unifies language, image, video, audio and ACTION in ONE model. The key trick
is a Mixture-of-Transformers: every modality is turned into tokens placed on a SINGLE
interleaved sequence; SELF-ATTENTION is SHARED across all modalities (so vision can be
conditioned on text, actions on vision, etc.), but each token is processed by a
MODALITY-SPECIFIC expert feed-forward block ("Mixture-of-Transformers" routing).
text tokens vision tokens action tokens
[t0 t1 t2 ...] [v0 v1 v2 ...] [a0 a1 ...]
\ | /
\ | /
+----------- one sequence -----------+
|
┌─────────── shared causal self-attention (RoPE) ───────────┐
│ every token attends to all earlier tokens, ANY modality │
└───────────────────────────────────────────────────────────┘
|
route each token to its modality's EXPERT FFN (SwiGLU):
text→Expert0 vision→Expert1 action→Expert2
|
per-modality heads: next-token / next-latent / next-action
Physical-AI modes fall right out of this one model:
text2video = generate the vision-token stream from a text prompt
image2video = condition vision stream on a first frame + text
forward_dynamics= given frames + ACTIONS, roll future frames forward (a world model)
inverse_dynamics= given frames, infer the ACTIONS that caused them
policy = given an observation + goal, emit ACTIONS (+ imagined rollout)
Below we build a faithful ~4M-param miniature of exactly this and train it live.
(The real model uses flow-matching/diffusion for the continuous vision stream; our toy
uses a simple MSE next-latent objective so it trains in seconds — the ROUTING and
SHARED-ATTENTION structure are the same.)
""")
We clone and inspect the real cosmos-framework repository to understand its package structure, input schemas, and CLI workflow directly from the source. We also print the official inference command patterns for single-GPU and multi-GPU launches, including modes such as text-to-video, image-to-video, forward dynamics, inverse dynamics, and policy. We then introduce the omnimodal Mixture-of-Transformers idea, where text, vision, and action tokens share attention while still using modality-specific expert feed-forward blocks.
Copy CodeCopiedUse a different Browser
rule("SECTION 3 — Implement & train the omnimodal MoT from scratch")
Building The Omnimodal MoT
Copy CodeCopiedUse a different Browser
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
torch.manual_seed(0)
@dataclass
class Cfg:
d_model: int = 192
n_head: int = 6
n_layer: int = 4
ffn_mult: int = 2
n_mod: int = 3
text_vocab:int = 16
vis_dim: int = 8
act_dim: int = 4
Lt: int = 8
Lv: int = 8
La: int = 6
cfg = Cfg()
class RMSNorm(nn.Module):
def __init__(self, d, eps=1e-6):
super().__init__(); self.w = nn.Parameter(torch.ones(d)); self.eps = eps
def forward(self, x):
return self.w * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def build_rope(T, hd, device, base=10000.0):
pos = torch.arange(T, device=device, dtype=torch.float32)[:, None]
idx = torch.arange(0, hd, 2, device=device, dtype=torch.float32)[None, :]
freq = 1.0 / (base ** (idx / hd))
ang = pos * freq
cos = torch.cos(ang).repeat(1, 2)[None, None]
sin = torch.sin(ang).repeat(1, 2)[None, None]
return cos, sin
def rotate_half(x):
hd = x.shape[-1]; x1, x2 = x[..., :hd // 2], x[..., hd // 2:]
return torch.cat([-x2, x1], -1)
def apply_rope(q, k, cos, sin):
return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
class Attention(nn.Module):
"""Shared cross-modal causal self-attention with rotary embeddings."""
def __init__(self, c: Cfg):
super().__init__()
self.H, self.hd = c.n_head, c.d_model // c.n_head
self.qkv = nn.Linear(c.d_model, 3 * c.d_model, bias=False)
self.proj = nn.Linear(c.d_model, c.d_model, bias=False)
def forward(self, x, cos, sin, mask):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, -1)
q = q.view(B, T, self.H, self.hd).transpose(1, 2)
k = k.view(B, T, self.H, self.hd).transpose(1, 2)
v = v.view(B, T, self.H, self.hd).transpose(1, 2)
q, k = apply_rope(q, k, cos, sin)
att = (q @ k.transpose(-2, -1)) / math.sqrt(self.hd)
att = att.masked_fill(mask, float("-inf")).softmax(-1)
o = (att @ v).transpose(1, 2).reshape(B, T, D)
return self.proj(o)
class Expert(nn.Module):
"""A per-modality SwiGLU feed-forward 'transformer expert'."""
def __init__(self, d, mult):
super().__init__(); h = d * mult
self.w1 = nn.Linear(d, h, bias=False)
self.w3 = nn.Linear(d, h, bias=False)
self.w2 = nn.Linear(h, d, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class MoTBlock(nn.Module):
"""Shared attention + Mixture-of-Transformers (per-modality expert) routing."""
def __init__(self, c: Cfg):
super().__init__()
self.attn_norm = RMSNorm(c.d_model)
self.attn = Attention(c)
self.ffn_norm = nn.ModuleList([RMSNorm(c.d_model) for _ in range(c.n_mod)])
self.experts = nn.ModuleList([Expert(c.d_model, c.ffn_mult) for _ in range(c.n_mod)])
def forward(self, x, cos, sin, mask, mod_id):
x = x + self.attn(self.attn_norm(x), cos, sin, mask)
out = torch.zeros_like(x)
for i, exp in enumerate(self.experts):
sel = (mod_id == i).view(1, -1, 1).to(x.dtype)
out = out + sel * exp(self.ffn_normi)
return x + out
class OmniMoT(nn.Module):
def __init__(self, c: Cfg):
super().__init__(); self.c = c
self.text_emb = nn.Embedding(c.text_vocab, c.d_model)
self.vis_in = nn.Linear(c.vis_dim, c.d_model)
self.act_in = nn.Linear(c.act_dim, c.d_model)
self.mod_emb = nn.Embedding(c.n_mod, c.d_model)
self.blocks = nn.ModuleList([MoTBlock(c) for _ in range(c.n_layer)])
self.norm = RMSNorm(c.d_model)
self.text_head = nn.Linear(c.d_model, c.text_vocab, bias=False)
self.vis_head = nn.Linear(c.d_model, c.vis_dim, bias=False)
self.act_head = nn.Linear(c.d_model, c.act_dim, bias=False)
ids = torch.cat([torch.zeros(c.Lt), torch.ones(c.Lv), torch.full((c.La,), 2)]).long()
self.register_buffer("mod_id", ids, persistent=False)
def forward(self, text, vis, act):
c = self.c
x = torch.cat([self.text_emb(text), self.vis_in(vis), self.act_in(act)], 1)
x = x + self.mod_emb(self.mod_id)[None]
B, T, D = x.shape
cos, sin = build_rope(T, D // c.n_head, x.device)
mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)[None, None]
for blk in self.blocks:
x = blk(x, cos, sin, mask, self.mod_id)
x = self.norm(x)
ht = self.text_head(x[:, :c.Lt])
hv = self.vis_head(x[:, c.Lt:c.Lt + c.Lv])
ha = self.act_head(x[:, c.Lt + c.Lv:])
return ht, hv, ha
model = OmniMoT(cfg).to(DEVICE)
n_params = sum(p.numel() for p in model.parameters())
print(f"Model built: OmniMoT | {n_params/1e6:.2f}M params | {cfg.n_layer} MoT blocks "
f"x {cfg.n_mod} experts | device={DEVICE}")
We implement the miniature omnimodal Mixture-of-Transformers model from scratch using PyTorch. We define RMSNorm, rotary embeddings, shared causal self-attention, modality-specific SwiGLU experts, and the full OmniMoT architecture. We then initialize the model on the available device and report its parameter count, layer count, expert structure, and runtime device.
Copy CodeCopiedUse a different Browser
K = 4
g = torch.Generator().manual_seed(1)
Training On Synthetic Data
Copy CodeCopiedUse a different Browser
K = 4
g = torch.Generator().manual_seed(1)
Training On Synthetic Data
TEXT_TRANS = torch.stack([torch.softmax(torch.randn(cfg.text_vocab, cfg.text_vocab, generator=g), -1)
for _ in range(K)])
VIS_DYN = torch.stack([0.9 * torch.linalg.qr(torch.randn(cfg.vis_dim, cfg.vis_dim, generator=g))[0]
for _ in range(K)])
ACT_MAP = torch.randn(cfg.act_dim, cfg.vis_dim, generator=g) * 0.5
def make_batch(B):
codes = torch.randint(0, K, (B,), generator=g)
text = torch.zeros(B, cfg.Lt, dtype=torch.long)
text[:, 0] = torch.randint(0, cfg.text_vocab, (B,), generator=g)
for t in range(1, cfg.Lt):
probs = TEXT_TRANS[codes, text[:, t-1]]
text[:, t] = torch.multinomial(probs, 1, generator=g).squeeze(1)
vis = torch.zeros(B, cfg.Lv, cfg.vis_dim)
vis[:, 0] = torch.randn(B, cfg.vis_dim, generator=g)
for t in range(1, cfg.Lv):
vis[:, t] = torch.einsum("bij,bj->bi", VIS_DYN[codes], vis[:, t-1]) \
+ 0.02 * torch.randn(B, cfg.vis_dim, generator=g)
vis_state = vis.mean(1)
act = torch.zeros(B, cfg.La, cfg.act_dim)
for t in range(cfg.La):
act[:, t] = (ACT_MAP @ (vis_state * (0.8 ** t)).T).T \
+ 0.02 * torch.randn(B, cfg.act_dim, generator=g)
return text.to(DEVICE), vis.to(DEVICE), act.to(DEVICE), codes
def loss_fn(model, text, vis, act):
ht, hv, ha = model(text, vis, act)
l_text = F.cross_entropy(ht[:, :-1].reshape(-1, cfg.text_vocab), text[:, 1:].reshape(-1))
l_vis = F.mse_loss(hv[:, :-1], vis[:, 1:])
l_act = F.mse_loss(ha[:, :-1], act[:, 1:])
return l_text + l_vis + l_act, (l_text.item(), l_vis.item(), l_act.item())
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
STEPS, BATCH = 400, 64
hist, t0 = [], time.time()
print(f"\nTraining for {STEPS} steps (batch={BATCH})...")
model.train()
for step in range(1, STEPS + 1):
text, vis, act, _ = make_batch(BATCH)
loss, parts = loss_fn(model, text, vis, act)
opt.zero_grad(); loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
hist.append(loss.item())
if step % 50 == 0 or step == 1:
print(f" step {step:4d} total {loss.item():6.3f} "
f"| text {parts[0]:5.3f} vision {parts[1]:6.4f} action {parts[2]:6.4f}")
print(f"Trained in {time.time()-t0:.1f}s | loss {hist[0]:.3f} -> {hist[-1]:.3f}")
print(" loss curve: " + spark(hist))
try:
import matplotlib.pyplot as plt
plt.figure(figsize=(7, 3))
plt.plot(hist); plt.title("OmniMoT training loss"); plt.xlabel("step"); plt.ylabel("loss")
plt.grid(alpha=0.3); plt.tight_layout(); plt.show()
except Exception:
pass
We create a synthetic physical-world dataset where text, vision, and action streams depend on hidden scene codes. We train the miniature world model to predict next text tokens, future vision latents, and future action vectors using a combined cross-entropy and MSE objective. We track the training loss over multiple steps and optionally plot the curve to show how the model learns the cross-modal dynamics.
Copy CodeCopiedUse a different Browser
rule("SECTION 4 — Autoregressive world-model rollout (forward_dynamics analog)")
print(textwrap.dedent("""
A world model predicts the FUTURE. Here we give the trained model a partial vision
trajectory and let it roll the vision latents forward one step at a time — exactly
the loop Cosmos 3 runs for forward_dynamics (predict future frames) and policy
(predict future frames + actions). We compare the model's imagined trajectory to the
ground-truth physics (the hidden VIS_DYN map) it never saw explicitly.
"""))
@torch.no_grad()
Autoregressive World-Model Rollout
Copy CodeCopiedUse a different Browser
def rollout(model, text, vis_prefix, act, n_future):
"""Predict n_future vision latents autoregressively from a vision prefix."""
model.eval()
vis = vis_prefix.clone()
preds = []
for _ in range(n_future):
pad = cfg.Lv - vis.shape[1]
vis_in = vis if pad <= 0 else torch.cat(
[vis, vis[:, -1:].repeat(1, pad, 1)], 1)
_, hv, _ = model(text, vis_in[:, :cfg.Lv], act)
nxt = hv[:, min(vis.shape[1], cfg.Lv) - 1:min(vis.shape[1], cfg.Lv)]
preds.append(nxt)
vis = torch.cat([vis, nxt], 1)
return torch.cat(preds, 1)
text, vis, act, codes = make_batch(4)
prefix_len, n_future = 3, 5
pred = rollout(model, text, vis[:, :prefix_len], act, n_future)
gt = vis[:, prefix_len-1:prefix_len].clone()
true_steps = [gt]
cur = gt
for _ in range(n_future):
cur = torch.einsum("bij,bj->bi", VIS_DYN[codes].to(DEVICE), cur[:, -1]).unsqueeze(1)
true_steps.append(cur)
gt_traj = torch.cat(true_steps[1:], 1)
err = F.mse_loss(pred, gt_traj).item()
print(f"Rolled out {n_future} future vision latents for {text.shape[0]} scenes.")
print(f"Imagined-vs-true-physics MSE : {err:.4f} (a small number = it learned the dynamics)")
print(f"Example (scene 0) latent[0] over time:")
print(f" predicted : {pred[0,:,0].detach().cpu().numpy().round(3).tolist()}")
print(f" true : {gt
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み