NVIDIA’s Cosmos-Framework Tutorial: Designing a Colab-Friendly Miniature of Cosmos 3 World Models with Omnimodal Mixture-of-Transformers

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 / 10243 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 / 10243 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':= 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: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_norm[i](x)) 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 = sel

[truncated for AI cost control]