NVIDIA Cosmos 3: Как запустить мир-модель на Colab

import os
import sys
import json
import time
import math
import textwrap
import subprocess
import shutil
import 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(len(vals) // width, 1)
    return "".join(
        blocks[min(len(blocks) - 1, int(lo + rng * (v / max(vals, 1))) if max(vals, 1) > 0 else 0)]
        for v in vals[::step]
    )

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, ""
if CUDA_OK:
    gpu_name = torch.cuda.get_device_properties(0).name
    gpu_mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
    cc = torch.cuda.get_device_capability()

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 >= (8, 0)
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.items():
    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.'}")

rule("SECTION 1 — Clone & map the real cosmos_framework package (source of truth)")

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 pkg.is_dir() else []:
            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("")

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 3 — Implement & train the omnimodal MoT from scratch")

import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

torch.manual_seed(42)

@dataclass
class Cfg:
    d_model: int = 192
    n_head: int = 3
    n_layer: int = 4
    ffn_mult: int = 4
    n_mod: int = 3  # text, vis, act
    text_vocab: int = 16
    vis_dim: int = 32
    act_dim: int = 16
    Lt: int = 5
    Lv: int = 5
    La: int = 3

cfg = Cfg()

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return self.weight * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

def build_rope(hd, device, base=10000.0):
    pos = torch.arange(device=device, dtype=torch.float32)
    idx = torch.arange(hd, device=device, dtype=torch.float32)
    freq = 1.0 / (base ** (idx / hd))
    ang = pos.unsqueeze(1) * freq.unsqueeze(0)
    cos = torch.cos(ang).repeat(1, 2)
    sin = torch.sin(ang).repeat(1, 2)
    return cos, sin

def rotate_half(x):
    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(cos, sin):
    return (cos * rotate_half(sin)), (sin * rotate_half(cos))

class Attention(nn.Module):
    """Shared cross-modal causal self-attention with rotary embeddings."""
    def __init__(self, cfg: Cfg):
        super().__init__()
        self.hd = cfg.d_model // cfg.n_head
        self.qkv = nn.Linear(cfg.d_model, cfg.d_model * 3, bias=False)
        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)

    def forward(self, x, cos, sin, mask):
        b, t, c = x.size()
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(b, t, self.n_head, self.hd).transpose(1, 2)
        k = k.view(b, t, self.n_head, self.hd).transpose(1, 2)
        v = v.view(b, t, self.n_head, self.hd).transpose(1, 2)
        
        # Apply RoPE
        q = apply_rope(cos, sin) * q
        k = apply_rope(cos, sin) * k
        
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.hd))
        att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        out = (att @ v).transpose(1, 2).reshape(b, t, c)
        return self.proj(out)

class Expert(nn.Module):
    """A per-modality SwiGLU feed-forward 'transformer expert'."""
    def __init__(self, dim, mult):
        super().__init__()
        self.w1 = nn.Linear(dim, dim * mult, bias=False)
        self.w3 = nn.Linear(dim, dim * mult, bias=False)
        self.w2 = nn.Linear(dim * mult, dim, 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, cfg: Cfg):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model)
        self.attn = Attention(cfg)
        self.ffn_norm = nn.ModuleList([RMSNorm(cfg.d_model) for _ in range(cfg.n_mod)])
        self.experts = nn.ModuleList([Expert(cfg.d_model, cfg.ffn_mult) for _ in range(cfg.n_mod)])

    def forward(self, x, cos, sin, mask, mod_id):
        h = 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).unsqueeze(-1).to(x.dtype)
            out = out + sel * exp(self.ffn_norm[i](x))
        return out

class OmniMoT(nn.Module):
    def __init__(self, cfg: Cfg):
        super().__init__()
        self.text_emb = nn.Embedding(cfg.text_vocab, cfg.d_model)
        self.vis_in = nn.Linear(cfg.vis_dim, cfg.d_model)
        self.act_in = nn.Linear(cfg.act_dim, cfg.d_model)
        self.mod_emb = nn.Embedding(cfg.n_mod, cfg.d_model)
        self.blocks = nn.ModuleList([MoTBlock(cfg) for _ in range(cfg.n_layer)])
        self.norm = RMSNorm(cfg.d_model)
        self.text_head = nn.Linear(cfg.d_model, cfg.text_vocab, bias=False)
        self.vis_head = nn.Linear(cfg.d_model, cfg.vis_dim, bias=False)
        self.act_head = nn.Linear(cfg.d_model, cfg.act_dim, bias=False)
        
        # Create modality IDs for the sequence
        ids = torch.cat([torch.zeros(cfg.Lt), torch.ones(cfg.Lv), torch.full((cfg.La,), 2, dtype=torch.long)])
        self.register_buffer("mod_id", ids, persistent=False)

    def forward(self, text, vis, act):
        x = torch.cat([
            self.text_emb(text),
            self.vis_in(vis),
            self.act_in(act),
            self.mod_emb(self.mod_id[:len(text)+len(vis)+len(act)].unsqueeze(0))
        ], dim=1)
        
        b, t, c = x.shape
        cos, sin = build_rope(self.blocks[0].attn.hd, x.device)
        mask = torch.triu(torch.ones(t, t, dtype=torch.bool, device=x.device), diagonal=1)
        
        for blk in self.blocks:
            x = blk(x, cos, sin, mask, self.mod_id[:t].unsqueeze(0))
        
        x = self.norm(x)
        ht = self.text_head(x[:, :cfg.Lt])
        hv = self.vis_head(x[:, cfg.Lt:cfg.Lt+cfg.Lv])
        ha = self.act_head(x[:, cfg.Lt+cfg.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  x {cfg.n_mod} experts  |  device={DEVICE}")

torch.manual_seed(42)

# Synthetic Data Generation
TEXT_TRANS = torch.softmax(torch.randn(cfg.text_vocab, cfg.text_vocab), dim=-1)
VIS_DYN = torch.linalg.qr(torch.randn(cfg.vis_dim, cfg.vis_dim))[0]
ACT_MAP = torch.randn(cfg.act_dim, cfg.vis_dim)

def make_batch(batch_size):
    codes = torch.randint(0, 10, (batch_size,))
    text = torch.zeros(batch_size, cfg.Lt, dtype=torch.long)
    for t in range(cfg.Lt):
        probs = TEXT_TRANS[codes] if t == 0 else TEXT_TRANS[text[:, t-1]]
        text[:, t] = torch.multinomial(probs, 1).squeeze()
    
    vis = torch.zeros(batch_size, cfg.Lv, cfg.vis_dim)
    vis_state = torch.randn(batch_size, cfg.vis_dim)
    for t in range(cfg.Lv):
        vis[:, t] = torch.einsum('bij,bj->bi', VIS_DYN, vis_state) + 0.02 * torch.randn(batch_size, cfg.vis_dim)
        vis_state = vis[:, t]
    
    act = torch.zeros(batch_size, cfg.La, cfg.act_dim)
    for t in range(cfg.La):
        act[:, t] = torch.einsum('ij,bj->bi', ACT_MAP, vis_state) + 0.8 * torch.randn(batch_size, cfg.act_dim)
        
    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.reshape(-1, cfg.text_vocab), text.reshape(-1))
    l_vis = F.mse_loss(hv, vis)
    l_act = F.mse_loss(ha, act)
    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(STEPS):
    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 == STEPS-1:
        print(f"  step {step:4d}  total {loss.item():6.3f}  | 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(f"  loss curve: {spark(hist)}")