NVIDIA的Cosmos-Framework教程：用全模態混合Transformer設計Colab友好的Cosmos 3世界模型微型版

本教程從實用的Google Colab角度探索NVIDIA的Cosmos框架，同時誠實地說明了執行真實Cosmos 3檢查點所需的硬體限制。我們首先檢查當前執行時環境，包括GPU能力、CUDA可用性、記憶體和磁碟空間，以理解為何完整的Cosmos 3推理在標準Colab硬體上不現實。但教程並未止步於此，而是利用框架的真實結構、CLI介面、輸入模式和模型模式，作為動手實現微型版本的基礎。

我們構建並訓練了一個緊湊的全模態混合Transformer世界模型，它體現了Cosmos的核心思想：共享的跨模態注意力，以及針對文本、視覺和動作流的模態特定專家路由。使用合成物理世界資料、訓練損失跟蹤和自迴歸滾動，我們展示了模型如何學習跨模態的關係，並以簡化但技術上有意義的方式預測未來的潛在狀態。

硬體限制探測

教程首先評估Colab環境。透過檢查Python版本、PyTorch安裝、CUDA可用性、GPU名稱和記憶體、計算能力以及磁碟空間，將實際硬體與Cosmos 3的要求（如Ampere+架構、至少79 GiB GPU記憶體、150+ GiB磁碟空間）進行對比。結論明確：標準Colab（例如T4 GPU，記憶體約15 GiB）無法執行真實的Cosmos 3檢查點。

克隆並對映框架

教程接著克隆NVIDIA/cosmos-framework倉庫，並檢查其包結構。展示了官方推理命令，包括單GPU和多GPU啟動模式，以及各種模式（文生圖、文生影片、圖生影片、前向動力學、逆向動力學、策略）。

全模態混合Transformer概念

Cosmos 3的關鍵創新是混合Transformer：每個模態被轉換為令牌並放置在一個單一的穿插序列上。自注意力在所有模態之間共享，使視覺可以根據文本條件化，動作可以根據視覺條件化等。但每個令牌由其模態特定的專家前饋網路處理。這種路由機制使得一個模型能夠執行多種物理AI任務。

從零實現並訓練

教程使用PyTorch從零實現了約400萬引數的微型模型。核心元件包括：RMSNorm、旋轉位置編碼（RoPE）、共享因果自注意力，以及三個模態（文本、視覺、動作）各自的SwiGLU專家網路。使用合成資料訓練模型以預測下一個令牌、下一個潛在狀態和下一個動作。雖然真實模型使用流匹配/擴散用於連續視覺流，但本教程使用簡單的MSE目標，使訓練在幾秒內完成，同時保留了路由和共享注意力結構。

透過這個過程，讀者能夠理解Cosmos 3框架的核心機制，並在資源受限的環境中動手實踐。