NVIDIA Cosmos-Framework: Jak stworzyć miniaturowy model świata Cosmos 3 w Google Colab

NVIDIA udostępniła szczegółowy samouczek dotyczący swojego Cosmos-Framework, koncentrujący się na praktycznym podejściu do tworzenia miniaturowej wersji modelu świata Cosmos 3, która jest kompatybilna ze środowiskiem Google Colab. Celem jest zrozumienie architektury i działania pełnowymiarowego modelu, jednocześnie uwzględniając ograniczenia sprzętowe typowe dla Colab.

Na początku samouczka przeprowadzana jest analiza środowiska wykonawczego, weryfikująca możliwości GPU, dostępność CUDA, pamięć oraz przestrzeń dyskową. Ma to na celu wyjaśnienie, dlaczego pełna inferencja modelu Cosmos 3 nie jest możliwa na standardowym sprzęcie dostępnym w Colab. Zamiast rezygnować, NVIDIA proponuje budowę kompaktowego, omnimodalnego modelu świata typu Mixture-of-Transformers, który odzwierciedla kluczowe idee Cosmos 3.

Wymagania sprzętowe a możliwości Colab

Samouczek rozpoczyna się od szczegółowej „sondy środowiskowej”, która porównuje dostępne zasoby sprzętowe z rzeczywistymi potrzebami Cosmos 3. Weryfikowane są takie parametry jak wersja Pythona, PyTorch, dostępność i wersja CUDA, nazwa i pamięć GPU, a także zdolność obliczeniowa (compute capability) oraz wolne miejsce na dysku. Okazuje się, że pełne uruchomienie Cosmos 3 Nano-16B wymaga:

• Architektury GPU Ampere+ (np. A100/RTX30xx, sm_80+).
• Co najmniej 80 GB pamięci GPU (dla Nano-16B na pojedynczym H100).
• Zestawu narzędzi CUDA w wersji 12.8 lub nowszej.
• Około 150 GB wolnego miejsca na dysku dla pierwszego uruchomienia (i około 1 TB dla pamięci podręcznej Hugging Face).

Standardowe środowiska Google Colab, często wyposażone w karty T4 (sm_75) i znacznie mniejszą pamięć GPU, nie spełniają tych wymagań. W związku z tym, samouczek jasno wskazuje, że rzeczywista inferencja Cosmos 3 nie jest możliwa w takim środowisku, co prowadzi do „ścieżki edukacyjnej” polegającej na budowie miniaturowej wersji.

Architektura Cosmos 3: Omnimodalny Mixture-of-Transformers

Kluczową ideą Cosmos 3 jest unifikacja różnych modalności – tekstu, obrazu, wideo, dźwięku i akcji – w ramach jednego modelu. Samouczek wyjaśnia, że osiąga się to poprzez architekturę Mixture-of-Transformers (MoT). W tym podejściu każda modalność jest konwertowana na tokeny, które są następnie umieszczane w jednej, przeplatanej sekwencji. Współdzielona uwaga własna (self-attention) działa na wszystkich modalnościach, co pozwala na warunkowanie wizji tekstem, akcji wizją itd.

Jednakże, każdy token jest przetwarzany przez specyficzny dla modalności blok feed-forward (ekspert), co stanowi element „Mixture-of-Transformers”. Oznacza to, że tokeny tekstowe są kierowane do eksperta tekstowego, wizyjne do eksperta wizyjnego, a akcji do eksperta akcji. Takie podejście umożliwia modelowi efektywne uczenie się relacji między różnymi typami danych.

Samouczek demonstruje, jak zbudować i wytrenować miniaturową wersję tego modelu (około 4 milionów parametrów) od podstaw. Mimo że rzeczywisty model wykorzystuje bardziej zaawansowane techniki (np. flow-matching/dyfuzję dla strumienia wizyjnego), miniaturowa wersja zachowuje kluczową strukturę routingu i współdzielonej uwagi. Wykorzystując syntetyczne dane ze świata fizycznego, śledzenie strat treningowych i autoregresywną generację, model uczy się relacji między modalnościami i przewiduje przyszłe stany ukryte w uproszczony, ale technicznie znaczący sposób.

Implementacja i trening miniaturowego modelu MoT

Samouczek szczegółowo opisuje implementację miniaturowego modelu OmniMoT (Omnimodal Mixture-of-Transformers). Przedstawia kod PyTorch, który definiuje kluczowe komponenty architektoniczne, takie jak:

• RMSNorm: normalizacja warstwowa.
• RoPE (Rotary Positional Embeddings): osadzanie pozycji obrotowych.
• Attention: współdzielona, krzyżowo-modalna uwaga własna z osadzeniami obrotowymi.
• Expert: blok feed-forward specyficzny dla modalności, wykorzystujący architekturę SwiGLU.
• MoTBlock: główny blok modelu, łączący współdzieloną uwagę z routingiem do ekspertów.
• OmniMoT: kompletny model, zawierający osadzenia dla tekstu, wizji i akcji, a także głowy wyjściowe dla każdej modalności.

Model jest konfigurowany z parametrami takimi jak d_model=192, n_head=6, n_layer=4 i n_mod=3 (dla tekstu, wizji i akcji). Trening na syntetycznych danych fizycznych pozwala zademonstrować, jak model uczy się przewidywać przyszłe stany ukryte, co jest kluczowe dla modeli świata. Choć jest to uproszczona wersja, pozwala na praktyczne zrozumienie mechanizmów leżących u podstaw bardziej złożonych systemów, takich jak pełny Cosmos 3.

Samouczek NVIDIA Cosmos-Framework stanowi cenne narzędzie edukacyjne, umożliwiające programistom i badaczom eksplorację zaawansowanych koncepcji modeli świata, nawet przy ograniczonych zasobach sprzętowych. Pokazuje, jak fundamentalne idee, takie jak Mixture-of-Transformers i współdzielona uwaga krzyżowo-modalna, mogą być skalowane i adaptowane do różnych środowisk, otwierając drogę do dalszych eksperymentów i rozwoju w dziedzinie multimodalnej sztucznej inteligencji.

Źródło: marktechpost.com