PrimeIntellect-ai/prime-diloco

解決的問題

Prime 是一個專為在互聯網上高效且全球分散式訓練 AI 模型而設計的框架。它解決了在地理上分散的節點之間訓練大型模型時所面臨的挑戰，例如網路不穩定、高延遲，以及檢查點與恢復的高成本。

工作原理

Prime 採用 DiLoCo（分散式低通訊）方法來最小化通訊開銷。它實現了多項關鍵優化：

• ElasticDeviceMesh：一種容錯抽象，可管理動態程序群組，允許節點在訓練執行期間加入或離開，而無需冷啟動。
• 非同步檢查點：為避免阻塞主訓練流程，檢查點首先儲存至基於 RAM 的檔案系統（/dev/shm），然後非同步上傳至磁碟或遠端儲存。
• 即時檢查點恢復：加入的節點可透過側車 HTTP 伺服器從對等節點快速取得最新模型狀態，從而於訓練過程中中途加入。
• 自訂 Int8 All-Reduce 內核：使用 C++ 實作的環狀歸約內核，將偽梯度量化為 Int8，使網路負載減少 4 倍，且不影響損失曲線。
• 頻寬優化：透過分片偽梯度（多個同時連接）與 VPN 技術優化點對點路由，最大化網路使用率。
• 記憶體管理：使用 PyTorch FSDP2/DTensor ZeRO-3 在節點內 GPU 之間分片模型權重、梯度與最佳化器狀態，並將 DiLoCo 最佳化器張量卸載至 CPU 記憶體。

適用對象

需要在多個資料中心或公共互聯網上而非單一高速叢集內分散訓練大型 AI 模型的研究人員與工程師。

主要亮點

• 透過 ElasticDeviceMesh 與心跳機制實現容錯訓練。
• 使用自訂 Int8 All-Reduce 內核將通訊負載減少 4 倍。
• 透過基於 RAM 的非同步上傳實現零阻塞檢查點。
• 實現高頻寬互聯網訓練，美國境內資料中心間最高可達 4Gb/s。
• 集成 FSDP2 實現高效的記憶體分片（ZeRO-3）。