跳到主要內容
AI News HubLIVE
站內改寫3 分鐘閱讀

在Modal上構建強化學習定理證明工作流

文章摘要

AE Studio利用Modal平台,通過進化策略(ES)和GRPO兩種強化學習方法訓練語言模型進行數學定理證明。他們使用Lean驗證器,並藉助Modal的並行GPU、沙盒隔離和卷存儲功能高效運行實驗。結果顯示ES在某些場景下媲美甚至超越GRPO,且成本顯著降低。

在Modal上構建強化學習定理證明工作流
報告錯誤

更正渠道尚未開通,可先複製下方文章資訊留存。

查看更正說明
直接讀正文

AE Studio最近在Modal平台開展了一項引人注目的研究,旨在通過強化學習訓練語言模型進行數學定理證明。他們系統比較了兩種方法:當前主流的組相對策略優化(GRPO)和一種受自然選擇啓發的進化策略(ES)。研究使用了Lean形式化驗證語言作為獎勵信號——模型生成的Lean代碼由編譯器驗證正確性,從而提供明確的獎勵信號。

為了高效運行實驗,AE Studio選擇了Modal平台,該平台提供了並行GPU計算、沙盒隔離和持久化卷存儲等關鍵功能。具體工作流程分為三個部分:證明生成使用vLLM在GPU上運行,證明驗證使用Lean編譯器在CPU上通過Modal沙盒隔離執行,協調進程管理整個訓練循環。ES方法通過創建一羣略有不同的模型副本,評估它們並更新原始模型,而GRPO則基於組內相對性能進行梯度更新。

在實現細節上,AE Studio充分利用了Modal的特色功能。他們為每個計算角色使用獨立的鏡像,而不是將所有依賴打包到一個臃腫的環境中。GPU工作節點專注於推理,Lean沙盒專注於驗證,協調器保持輕量。並行GPU扇出通過Modal的.map()實現,每個擾動評估只需要當前檢查點、定理批次、一個擾動種子和生成參數。由於ES的權重擾動完全由種子決定,更新步驟只需種子和獎勵即可聚合所有個體,避免了GPU間昂貴的權重傳輸。驗證環節使用Modal沙盒進行隔離,每個驗證批次啓動一個獨立的Lean服務器,處理完即關閉,防止錯誤的證明掛起或崩潰整個系統。這種設計使得每輪迭代可以創建3840個證明嘗試,並按批次並行驗證。

ES的一個優雅特性是,整個模型狀態可以由基模型加上一組(種子,獎勵)對來描述。協調器維護一個種子/獎勵歷史列表,並傳遞給每個工作節點。GPU工作節點從卷存儲加載原始基模型,然後重放完整歷史以重構當前權重,再應用自己的擾動。這個重放過程遍歷過去每次迭代的種子和獎勵,在GPU上使用確定性種子重新生成同樣的噪聲,並就地應用加權更新。整個模型狀態作為純Python列表傳遞,小到可以作為函數參數傳給每次遠程調用。

性能方面,Modal的實現僅需約250行平台設置代碼,不到其他平台典型600行的一半。複雜度降低使得整個項目從啓動到成功訓練運行僅用了不到兩天時間,比替代平台快60%。運行時效率也顯著提升:生成步驟(需GPU)平均每輪147秒,但完整循環因驗證步驟(僅CPU)而耗時538秒。Modal的彈性計費意味着只需為GPU實際使用時間付費,將浪費的GPU時間減少了約3.7倍。基於這些時間數據,估算整個運行成本在Modal上僅為122美元,而同等硬件上的替代方案成本在180至480美元之間。

早期結果令人鼓舞。在多次定理證明運行中,ES在每輪驗證的證明數量上匹配甚至超越了GRPO基線,並且當訓練數據有限時,ES通常展現出更高的樣本效率。但在其他設置中增益較小,且尚未分離超參數選擇、數據集差異或種羣規模等因素帶來的變化。ES在這一領域的縮放行為仍未被充分理解,需要更多實驗來確定其一致優勢以及在大規模訓練中能否與GRPO競爭或超越。

研究問題仍然開放,但Modal上的基礎設施設置將繼續使用,以便運行、檢查和迭代這個基於驗證器的訓練循環,而無需投入時間構建定製基礎設施。接下來的實驗包括方差分析(對sigma、種羣規模和定理選擇進行系統掃描)和更大模型測試(使用開源的7B蒸餾Kimina模型)。

值得注意的是,這一基礎設施設置並非定理證明所獨有。它在許多機器學習系統中普遍適用:在GPU上生成輸出,對外部驗證器或編譯器運行結果,將結果轉化為訓練信號,並在大量獨立候選方案上重複。關鍵的技術要點是並行化GPU推斷、隔離驗證以及保持訓練狀態的可移植性——全部在一個工作流中完成。.map()分佈了證明生成,沙盒控制了驗證器故障,卷存儲使檢查點在迭代之間可用,而無需移回本地機器。

通過Modal,團隊能夠在三種截然不同的運行時中運行相同的稀疏獎勵強化學習工作流,而無需每次重新搭建。他們從本地機器啓動運行,將基模型保持在遠程,記錄到MLflow,並根據實驗進展更改種羣大小、定理批次大小和驗證並行度。由於Modal平台的特性和易用性,項目啓動數天內就能看到實驗結果。其他基礎設施平台可能需要數週測試和迭代才能獲得任何信號。一旦搭建完畢,團隊可以完全專注於實驗設計,無需擔心繫統可靠性。

完整的實驗代碼已在GitHub上公開(github.com/agencyenterprise/modal-rl-theorem-case-study),供深入研究或適配到自己的工作流。

展開要點與分析

文章情報

工程師進階

要點

  • AE Studio在Modal上實現了進化策略(ES)用於定理證明,並與GRPO進行了對比。
  • 利用Modal的.map()並行GPU推斷、沙盒隔離驗證和卷存儲,大幅簡化了基礎設施搭建。
  • ES在多項測試中匹配或超越GRPO,尤其在數據有限時展現出更高的樣本效率。
  • 整體運行成本約122美元,比替代平台降低60%以上,且代碼量僅為250行。

要點與分析由自動化流程生成,可能有誤,請結合原始來源核實。

在Modal上構建強化學習定理證明工作流 | AI News Hub