NVIDIA Cosmos требует 80 ГБ видеопамяти: как собрать прототип в бесплатном Colab

Фреймворк NVIDIA Cosmos объединяет генерацию изображений, видео, прямую и обратную динамику в одном пакете. Модели Cosmos 3 выпущены в двух размерах: Nano на 16 миллиардов параметров (все режимы) и Super на 65 миллиардов (текст в изображение, текст в видео, изображение в видео). Для запуска даже младшей Nano нужен один GPU класса H100 с 80 ГБ видеопамяти, архитектура Ampere и выше, около 150 ГБ свободного диска и CUDA не ниже 12.8. В стандартном Google Colab этого нет, но структура фреймворка открыта, и её можно воспроизвести в миниатюре.

Что понадобится?

• Google Colab (бесплатный тариф с GPU T4 подойдёт для миниатюрной модели; для полноценного Cosmos 3 нужен A100 или H100)
• Python 3.10+ и PyTorch (ставится автоматически в Colab)
• Git для клонирования репозитория NVIDIA
• Около 5 ГБ свободного диска для учебного проекта
• 30 минут на весь цикл от проверки железа до обученного прототипа

Пошаговая инструкция

1. Проверьте, на чём вы запускаете. Cosmos 3 требует конкретного железа, и первый шаг определяет, какой путь вам доступен. Запустите в ячейке Colab:

import torch, shutil

cuda_ok = torch.cuda.is_available()
if cuda_ok:
    p = torch.cuda.get_device_properties(0)
    print(f"GPU: {p.name}, память: {p.total_memory / 1024**3:.1f} ГБ")
    print(f"Compute capability: sm_{p.major}{p.minor}")
else:
    print("GPU не обнаружен, только CPU")

free_gb = shutil.disk_usage('/').free / 1024**3
print(f"Свободное место на диске: {free_gb:.0f} ГБ")

Скрипт покажет вашу модель GPU, объём видеопамяти, поколение архитектуры и свободный диск. На стандартном Colab вы увидите T4 с 15 ГБ памяти и архитектурой sm_75, это ниже порога Ampere (sm_80+). Вердикт: полноценный инференс (запуск обученной модели для получения результата) невозможен, переходим к учебному пути.

1. Клонируйте репозиторий NVIDIA Cosmos. Даже без запуска тяжёлых чекпоинтов (файлов с сохранёнными весами модели) репозиторий содержит реальную структуру кода, CLI-команды и входные схемы:

git clone --depth 1 https://github.com/NVIDIA/cosmos-framework.git

Внутри папки cosmos_framework вы найдёте подпакеты для каждого режима генерации, модули параллелизма и защитные фильтры (guardrails).

1. Изучите реальный CLI фреймворка. Запуск полноценной модели на одном 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

Для мульти-GPU (8 карт) используется torchrun с пресетом throughput. Входной файл t2v.json описывает текстовый промпт и параметры генерации видео.

1. Соберите миниатюрную MoT-модель. MoT (Mixture-of-Transformers, смесь трансформеров) повторяет ключевую идею Cosmos: общее внимание (cross-modal attention) между модальностями (текст, изображение, действие) с маршрутизацией к модальностно-специфичным экспертам. Создайте в Colab ячейку с компактной архитектурой:

import torch
import torch.nn as nn

class ModalityExpert(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim)
        )
    def forward(self, x):
        return self.net(x)

class MiniMoT(nn.Module):
    def __init__(self, dim=128, heads=4, n_experts=3):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.experts = nn.ModuleList(
            [ModalityExpert(dim, dim * 2) for _ in range(n_experts)]
        )
        self.norm = nn.LayerNorm(dim)

    def forward(self, tokens, modality_ids):
        # общее кросс-модальное внимание
        attn_out, _ = self.attn(tokens, tokens, tokens)
        x = self.norm(tokens + attn_out)
        # маршрутизация по модальностям
        out = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            mask = (modality_ids == i).unsqueeze(-1)
            out += expert(x) * mask.float()
        return out

Три эксперта обрабатывают три потока: текст (id=0), визуальные токены (id=1), действия (id=2).

1. Сгенерируйте синтетические данные и обучите модель. Для демонстрации не нужны реальные видеоролики. Создайте случайные последовательности с разметкой модальностей и запустите обучение:

dim, seq_len, batch = 128, 24, 32
model = MiniMoT(dim=dim).to("cuda" if cuda_ok else "cpu")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
losses = []

for step in range(200):
    tokens = torch.randn(batch, seq_len, dim).to(model.norm.weight.device)
    mod_ids = torch.randint(0, 3, (batch, seq_len)).to(tokens.device)
    target = torch.randn_like(tokens)

    pred = model(tokens, mod_ids)
    loss = ((pred - target) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(loss.item())
    if step % 50 == 0:
        print(f"Шаг {step}: loss = {loss.item():.4f}")

print("Обучение завершено")

1. Запустите авторегрессивный роллаут. Проверьте, как модель предсказывает следующие латентные состояния (скрытые числовые представления, в которых модель «думает» о мире):

with torch.no_grad():
    seed = torch.randn(1, 1, dim).to(model.norm.weight.device)
    mod_id = torch.zeros(1, 1, dtype=torch.long).to(seed.device)
    trajectory = [seed.squeeze().cpu()]
    for _ in range(10):
        seed = model(seed, mod_id)
        trajectory.append(seed.squeeze().cpu())
    print("Роллаут из 10 шагов завершён, размерность выхода:", trajectory[-1].shape)

Модель берёт начальный токен и последовательно генерирует будущие состояния, как это делает Cosmos при предсказании кадров видео.

Что делать с этим прямо сейчас?

Разработчику. Клонируйте репозиторий NVIDIA Cosmos, изучите структуру входных JSON-схем и CLI-команд. Даже без доступа к мощному GPU вы поймёте архитектуру и подготовите пайплайн, который заработает при появлении ресурсов.

Автору Дзена. Тема генерации видео через ИИ набирает просмотры. Статья «Я собрал мини-Cosmos в Colab за 30 минут» с реальными скриншотами потери и роллаута даст вам уникальный технический контент, которого пока нет на русскоязычном Дзене.

Предпринимателю в РФ. Полноценный Cosmos 3 требует серверного GPU стоимостью от нескольких тысяч долларов в месяц за облачную аренду. Прежде чем тратить, проверьте идею на миниатюрной версии. Из российских облаков с GPU A100 можно использовать Yandex Cloud или Selectel.

Фреймворк NVIDIA Cosmos открыт, документация доступна, а входной порог для экспериментов с архитектурой равен нулю, нужен только Colab и полчаса.