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 / 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 =


