斯坦福大學基礎模型研究中心(CRFM)主導的開放基礎模型專案Marin,於上週正式啟動總引數達5350億的MoE模型Marin 535B-A23B的訓練。與多數模型公司保密訓練細節不同,該專案將訓練曲線、資料配方、模型配置和技術討論即時公開,甚至允許外部研究者參與審查。專案負責人、斯坦福電腦科學副教授Percy Liang在社交平台X上宣佈了這一訊息,帖子瀏覽量已突破80萬,並獲吳恩達轉發,吳恩達稱其為捍衛AI開放性的“珍貴示範”。

Marin 535B-A23B計劃處理18.75萬億Token,其中約80%用於預訓練,20%用於中期訓練,訓練執行在11套NVIDIA GB200 NVL72系統上,預計持續約3個月,總計算量約2.7×10^24 FLOPs。該模型採用混合專家(MoE)架構,總引數5350億,但每個Token僅啟用約230億引數,通過路由模組將輸入分配給部分專家網路,以控制計算成本。Marin團隊在技術說明中透露,模型每層保留2個共享專家和8個路由專家,均採用半寬結構,路由專家還使用2倍壓縮,約三分之一的專家計算來自始終工作的共享專家,這種設計旨在降低MoE訓練中“Token丟棄”的風險。

Marin專案並非首個直播訓練的大模型專案,此前Hugging Face的BigScience在訓練BLOOM時公開了TensorBoard日誌,Pythia、LLM360和Ai2的OLMo也開放了訓練資料、程式碼和日誌。但Marin的獨特之處在於將開放從模型釋出行為擴充套件為實驗室的預設工作方式,從提出假設、提交程式碼到訓練失敗都儘量即時公開。專案發起公告承認了EleutherAI、AI2、Hugging Face等先行專案的貢獻,但認為開放模型仍缺少類似軟體開源的協作機制,因此提出“開放實驗室”機制,通過GitHub Issue宣告實驗目標和假設,外部研究者可參與Review。

訓練MoE模型面臨通訊複雜、專家負載不均等挑戰。Marin團隊披露,在之前的實驗中,當上下文長度從4K擴充套件到65K時,Token Dropping比例從約7%升至約40%,因此535B模型選擇以4K上下文啟動預訓練,並測試了新型pooled/wave專家並行方案,將Token Dropping降至約3%。團隊還專門為JAX環境手寫了專家並行實現,因為未找到效能足夠的現成方案。訓練棧基於JAX、XLA和Levanter,此前Marin 8B和32B模型主要執行在Google TPU上,此次轉向NVIDIA GB200 NVL72,需重新處理GPU叢集上的通訊效率。

在正式訓練前,Marin團隊先訓練了一組規模遞增的小型MoE模型,形成四級“縮放梯”,覆蓋從1.6B總引數到27.7B總引數的模型,用於預測535B模型的損失水平、暴露訓練穩定性問題,並區分正常波動與失控前兆。這些縮放實驗僅佔最終計算量的約1%,但能幫助團隊在訓練早期發現數據、路由或最佳化器問題。團隊曾通過縮放實驗發現梯度範數增長到4以上,可能引發訓練發散,最終加入logit z-loss以限制輸出分佈數值幅度。

儘管Marin 535B已進入大規模訓練區間,但專案仍處於早期,預訓練損失不能直接證明其前沿效能。MoE的535B不能與稠密模型等價,計算成本和表達能力取決於專家分工有效性,且程式碼、數學等能力還受資料質量、中期訓練及後訓練影響。專案在GitHub上的標題寫作“535B-A23B on 18T tokens”,而Percy Liang公告使用18.75T Token口徑,說明訓練計劃仍在持續對齊。這場實驗最重要的產出可能是沿途留下的訓練記錄,包括專家路由、Token Dropping變化、梯度異常處理、長上下文擴充套件以及JAX在GB200上的專家並行實現,這些經驗可被其他規模的模型複用。