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框架的核心機制，並在資源受限的環境中動手實踐。