2025年5月30日,來自清華大學(xué)和螞蟻研究院的研究團(tuán)隊在arXiv上發(fā)布了一篇標(biāo)題為《AREAL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning》的研究論文。這項由Wei Fu、Jiaxuan Gao、Yi Wu等人共同完成的研究,提出了一種全新的異步強(qiáng)化學(xué)習(xí)系統(tǒng),專為提升大語言模型的推理能力而設(shè)計。論文代碼已在GitHub上開源,有興趣的讀者可通過https://github.com/inclusionAI/AReaL/查看完整代碼庫。
一、為什么我們需要異步強(qiáng)化學(xué)習(xí)系統(tǒng)?
想象你正在教一個聰明但經(jīng)驗不足的學(xué)生解決復(fù)雜問題。傳統(tǒng)的教學(xué)方式是:你給學(xué)生一道題,等他完全解答出來后,再根據(jù)他的答案給予反饋,然后修正他的思維方法,再給下一道題。這種教學(xué)方式雖然直觀,但效率低下——尤其是當(dāng)不同學(xué)生解題速度差異巨大時,總是要等最慢的學(xué)生完成,其他人則在空等。
在大語言模型(LLM)的世界里,強(qiáng)化學(xué)習(xí)(RL)就像這樣的教學(xué)過程,而當(dāng)前主流的RL系統(tǒng)大多采用這種"同步"方式工作。這導(dǎo)致了嚴(yán)重的計算資源浪費(fèi)問題,特別是在訓(xùn)練大型推理模型(Large Reasoning Models,LRMs)時,不同輸入可能產(chǎn)生長度從幾十到數(shù)萬不等的思考過程,使得GPU利用率低下。
研究團(tuán)隊發(fā)現(xiàn),現(xiàn)有的強(qiáng)化學(xué)習(xí)系統(tǒng)面臨兩個核心問題:首先,推理設(shè)備利用率低,因為系統(tǒng)必須等待最長輸出完成才能開始訓(xùn)練;其次,同步系統(tǒng)的可擴(kuò)展性差,因為將生成任務(wù)分散到更多設(shè)備上會降低每個GPU的批處理大小,使得解碼過程變成內(nèi)存IO受限的狀態(tài),難以通過增加設(shè)備來提高吞吐量。
二、AREAL:一個全異步的強(qiáng)化學(xué)習(xí)系統(tǒng)
AREAL系統(tǒng)就像是完全重新設(shè)計的教學(xué)流程——不再要求所有學(xué)生同步提交答案,而是讓每個學(xué)生獨(dú)立作答,老師隨時收集已完成的答卷進(jìn)行批改,同時不斷更新教學(xué)方法。這種異步模式讓資源利用率大大提高。
具體來說,AREAL由四個核心組件組成:
首先是"可中斷的推理工作器"(Interruptible Rollout Worker)。這些工作器負(fù)責(zé)接收兩類請求:一是生成響應(yīng),二是更新權(quán)重。當(dāng)新的模型參數(shù)可用時,這些工作器會中斷當(dāng)前正在進(jìn)行的生成任務(wù),加載新參數(shù),然后繼續(xù)未完成的序列生成。這就像學(xué)生正在解題過程中,老師提供了新的解題思路,學(xué)生立即采納并繼續(xù)解題。
其次是"獎勵服務(wù)"(Reward Service),負(fù)責(zé)評估模型生成的回答質(zhì)量。比如在編程任務(wù)中,這個服務(wù)會提取代碼并運(yùn)行單元測試來驗證其準(zhǔn)確性。
第三是"訓(xùn)練工作器"(Trainer Workers),它們從經(jīng)驗回放緩沖區(qū)持續(xù)采樣數(shù)據(jù),直到達(dá)到配置的訓(xùn)練批量大小,然后執(zhí)行PPO(Proximal Policy Optimization)更新并將結(jié)果參數(shù)存儲在分布式存儲中。為確保數(shù)據(jù)新鮮度,回放緩沖區(qū)的數(shù)據(jù)只使用一次。
最后是"推理控制器"(Rollout Controller),作為推理工作器、獎勵服務(wù)和模型工作器之間的關(guān)鍵橋梁。它從數(shù)據(jù)集讀取數(shù)據(jù)并調(diào)用推理工作器的生成請求,然后將收到的響應(yīng)發(fā)送給獎勵服務(wù)以獲取獎勵。軌跡和獎勵存儲在回放緩沖區(qū)中,等待模型工作器進(jìn)行訓(xùn)練。當(dāng)模型工作器更新參數(shù)后,控制器會調(diào)用推理工作器的權(quán)重更新功能。
這種完全異步的設(shè)計確保了生成和訓(xùn)練資源的持續(xù)滿負(fù)荷運(yùn)行,大大提高了系統(tǒng)效率。想象一下,這就像是一個工廠的生產(chǎn)線,不同工位獨(dú)立運(yùn)作,材料流動從不停止。
三、異步學(xué)習(xí)的算法挑戰(zhàn)與解決方案
然而,異步系統(tǒng)設(shè)計雖然提高了效率,但也帶來了幾個技術(shù)挑戰(zhàn)。
首先是"數(shù)據(jù)過時"問題。由于系統(tǒng)的異步特性,每個訓(xùn)練批次包含的數(shù)據(jù)可能來自多個不同版本的策略。想象一下,如果學(xué)生用舊方法解的題被用來改進(jìn)最新的教學(xué)方法,效果可能并不理想。之前在強(qiáng)化學(xué)習(xí)訓(xùn)練系統(tǒng)的研究表明,這種過時性會降低學(xué)習(xí)性能。
其次是"策略版本不一致"問題。由于生成過程可能被中斷并使用新參數(shù)繼續(xù),同一軌跡可能包含由不同策略版本生成的片段,這違反了標(biāo)準(zhǔn)PPO假設(shè)——所有動作都由單一策略生成。
為了解決這些問題,研究團(tuán)隊提出了兩個關(guān)鍵的算法創(chuàng)新:
第一是"過時性感知訓(xùn)練"(Staleness-Aware Training)。他們引入了一個名為η的超參數(shù),代表最大允許過時度。給定最新參數(shù)版本i,總生成軌跡數(shù)Nr和訓(xùn)練批量大小B,系統(tǒng)強(qiáng)制執(zhí)行:?Nr/B? ≤ i + η。當(dāng)η=0時,系統(tǒng)退化為同步RL設(shè)置;當(dāng)η=1時,系統(tǒng)恢復(fù)為之前的一步重疊方法。在訓(xùn)練過程中,系統(tǒng)優(yōu)先處理較舊的軌跡,確保過時度保持在η以下。
第二是"解耦PPO目標(biāo)函數(shù)"(Decoupled PPO Objective)。研究團(tuán)隊?wèi)?yīng)用了一個解耦的PPO目標(biāo)函數(shù),將行為策略和近端策略分離。行為策略πbehav代表用于采樣軌跡的策略,代理策略πprox作為近期目標(biāo),用于規(guī)范πθ的更新。通過對采樣軌跡應(yīng)用重要性采樣,他們推導(dǎo)出適用于異步RL訓(xùn)練的解耦PPO目標(biāo)函數(shù)。
這個解耦的PPO目標(biāo)函數(shù)提供了一個自然的好處:它放寬了所有數(shù)據(jù)在一個訓(xùn)練批次中應(yīng)該由單一策略生成的要求。這對于結(jié)合可中斷生成與策略更新至關(guān)重要。研究團(tuán)隊證明,一個軌跡中不一致的策略版本等同于單一行為策略πbehav。
實際實現(xiàn)中,他們簡單地使用每次模型更新前的參數(shù)作為πprox,并在每個訓(xùn)練步驟中全局批次到達(dá)時重新計算令牌概率。
四、AREAL的系統(tǒng)實現(xiàn)與優(yōu)化
AREAL系統(tǒng)是使用Python和PyTorch實現(xiàn)的,建立在ReaLHF框架基礎(chǔ)上。該系統(tǒng)結(jié)合了SGLang v0.4.6作為生成服務(wù)后端,以及Megatron-Core v0.11.0作為訓(xùn)練后端,由SLURM管理資源調(diào)度。
為了最大化生成和訓(xùn)練階段的吞吐量,研究團(tuán)隊實施了幾項關(guān)鍵的系統(tǒng)級優(yōu)化:
首先,AREAL將GPU計算與CPU操作解耦,包括基于規(guī)則的獎勵計算(如數(shù)學(xué)問題的字符串匹配或代碼的單元測試執(zhí)行)和基于TCP的數(shù)據(jù)傳輸。通過在單獨(dú)的線程中執(zhí)行這些操作并流水線化工作流程,系統(tǒng)可以將獎勵計算和數(shù)據(jù)傳輸與后續(xù)生成請求重疊。他們使用asyncio協(xié)程在推理工作器中并發(fā)運(yùn)行多個請求,避免相互阻塞等待。
其次,為了處理可變長度序列的訓(xùn)練,他們采用了無填充的序列打包策略,結(jié)合動態(tài)分配算法。該算法在固定內(nèi)存約束下平衡微批次間的令牌分布,最大化GPU內(nèi)存利用率,同時最小化所需的前向-后向傳遞次數(shù)。
這些優(yōu)化共同導(dǎo)致了系統(tǒng)性能的顯著提升,使AREAL能夠充分利用可用的計算資源。
五、實驗結(jié)果:AREAL的性能優(yōu)勢
研究團(tuán)隊在具有挑戰(zhàn)性的數(shù)學(xué)和編程任務(wù)上評估了AREAL系統(tǒng),使用了從1.5B到32B參數(shù)不等的模型。他們采用了從DeepSeek-R1蒸餾的Qwen2模型系列作為基礎(chǔ)模型。
在最終的對比實驗中,AREAL與最先進(jìn)的同步RL系統(tǒng)進(jìn)行了比較:用于數(shù)學(xué)推理的DeepScaleR(1.5B模型)和用于代碼生成的DeepCoder(14B模型),兩者都使用verl訓(xùn)練。對于較大的7B和32B模型,由于缺乏可比較的基線,他們使用AREAL的同步變體進(jìn)行了對照實驗。
實驗結(jié)果令人印象深刻:AREAL在不影響(甚至提高)最終性能的情況下,實現(xiàn)了顯著的加速。在數(shù)學(xué)推理任務(wù)中,使用1.5B和7B模型時,AREAL分別將訓(xùn)練時間從41.0小時和57.7小時減少到14.8小時和25.4小時,同時保持了42.2%和63.1%的AIME24基準(zhǔn)準(zhǔn)確率。在代碼生成任務(wù)中,使用14B和32B模型時,AREAL將訓(xùn)練時間從48.8小時和51.1小時減少到21.9小時和31.1小時,同時在LiveCodeBench基準(zhǔn)上分別達(dá)到58.1%和61.0%的準(zhǔn)確率。
總體而言,AREAL實現(xiàn)了高達(dá)2.57倍的訓(xùn)練吞吐量提升,并在多達(dá)512個GPU上展示了線性擴(kuò)展效率。關(guān)鍵的是,這種加速甚至伴隨著這些任務(wù)上的解決方案準(zhǔn)確率提高,說明AREAL在不犧牲(實際上是增強(qiáng))模型性能的同時提供了顯著的效率提升。
六、算法驗證實驗與系統(tǒng)消融研究
為了驗證研究團(tuán)隊在第三部分提出的算法創(chuàng)新,他們進(jìn)行了一系列消融研究,通過在數(shù)學(xué)任務(wù)上訓(xùn)練1.5B LRM來驗證他們的方法。
首先,他們改變了最大允許過時度η,并比較了有無解耦PPO目標(biāo)函數(shù)的配置。結(jié)果顯示,樸素PPO無法匹配同步RL的性能(即η=0時的性能)。即使輕微的過時度也會由于不當(dāng)?shù)牟眉糁行暮涂芍袛嗌蛇^程中的策略變化而顯著降低最終性能。此外,增加數(shù)據(jù)過時度會持續(xù)降低學(xué)習(xí)性能,這與之前在其他領(lǐng)域的研究觀察一致。
然而,通過比較學(xué)習(xí)曲線,他們發(fā)現(xiàn)解耦PPO目標(biāo)函數(shù)在處理過時數(shù)據(jù)時大大提高了訓(xùn)練穩(wěn)定性,這與游戲領(lǐng)域的發(fā)現(xiàn)一致。值得注意的是,即使使用解耦目標(biāo)函數(shù),無界過時度(η→∞)仍然導(dǎo)致比零過時度的效果差。當(dāng)適當(dāng)約束時,適度的過時度(如η≤4)對最終性能影響很小,同時通過異步流水線顯著加速訓(xùn)練。
研究團(tuán)隊還對系統(tǒng)組件進(jìn)行了消融研究。沒有可中斷生成功能時,控制器必須等待最長的響應(yīng),導(dǎo)致1.5B和7B模型在4個節(jié)點上的吞吐量分別降低12%和17%。這驗證了他們的架構(gòu)設(shè)計選擇。
在動態(tài)微批次分配方面,與標(biāo)準(zhǔn)微批次策略相比,動態(tài)批處理在各種模型大小上平均提高了30%的吞吐量。標(biāo)準(zhǔn)微批次策略可能導(dǎo)致多個長序列被分配到同一個微批次,通常需要足夠多的微批次以防止內(nèi)存不足錯誤。
這些消融研究結(jié)果證實了研究團(tuán)隊的設(shè)計選擇,并驗證了他們提出的算法創(chuàng)新在提高異步RL訓(xùn)練穩(wěn)定性和效率方面的有效性。
七、AREAL的局限性與未來研究方向
盡管AREAL取得了令人印象深刻的成果,但這項研究仍存在一些局限性,這也為未來的研究提供了方向。
首先,推理與訓(xùn)練設(shè)備之間的比例可以進(jìn)一步優(yōu)化,特別是針對特定的訓(xùn)練設(shè)置。此外,這一比例可能受益于訓(xùn)練過程中的動態(tài)調(diào)整,尤其是在微調(diào)預(yù)訓(xùn)練基礎(chǔ)模型時,上下文長度通常會增加。
研究團(tuán)隊的評估主要集中在單步數(shù)學(xué)和編碼任務(wù)上,但AREAL架構(gòu)本身并不僅限于這些領(lǐng)域。未來的工作可以探索多輪交互和代理場景,擴(kuò)展AREAL的應(yīng)用范圍。
此外,隨著模型參數(shù)量和上下文長度的持續(xù)增長,研究人員可能需要開發(fā)更高效的異步算法,進(jìn)一步優(yōu)化系統(tǒng)設(shè)計,以應(yīng)對更大規(guī)模的強(qiáng)化學(xué)習(xí)訓(xùn)練挑戰(zhàn)。
八、總結(jié)與展望
AREAL代表了大語言模型強(qiáng)化學(xué)習(xí)訓(xùn)練的一個重要進(jìn)步。通過完全解耦生成和訓(xùn)練過程,結(jié)合創(chuàng)新的算法方法來處理數(shù)據(jù)過時性和策略不一致性,AREAL實現(xiàn)了顯著的效率提升,同時保持或甚至提高了模型性能。
這種方法就像是重新設(shè)計了教學(xué)流程——不再要求所有學(xué)生同步提交答案,而是允許每個學(xué)生以自己的節(jié)奏學(xué)習(xí),教師則不斷收集和評估完成的作業(yè),并持續(xù)更新教學(xué)方法。這種異步流程不僅提高了效率,還可能通過允許更多樣化的學(xué)習(xí)路徑來增強(qiáng)學(xué)習(xí)效果。
AREAL的成功表明,在大規(guī)模強(qiáng)化學(xué)習(xí)訓(xùn)練中,系統(tǒng)設(shè)計與算法創(chuàng)新的結(jié)合至關(guān)重要。通過解決強(qiáng)化學(xué)習(xí)訓(xùn)練中的關(guān)鍵瓶頸,AREAL為未來的研究提供了一個可靠的起點,有望支持更先進(jìn)的AI系統(tǒng)發(fā)展,進(jìn)一步推動機(jī)器智能的邊界。
對于研究人員和工程師來說,AREAL提供了一個寶貴的框架,可以在更短的時間內(nèi)訓(xùn)練更強(qiáng)大的推理模型。對于更廣泛的AI社區(qū)來說,這意味著我們可以更快、更有效地開發(fā)出具有更強(qiáng)推理能力的大語言模型,最終為各種應(yīng)用提供更智能的解決方案。
好文章,需要你的鼓勵
浙江大學(xué)團(tuán)隊提出動態(tài)專家搜索方法,讓AI能根據(jù)不同問題靈活調(diào)整內(nèi)部專家配置。該方法在數(shù)學(xué)、編程等任務(wù)上顯著提升推理準(zhǔn)確率,且不增加計算成本。研究發(fā)現(xiàn)不同類型問題偏愛不同專家配置,為AI推理優(yōu)化開辟新路徑。
清華大學(xué)研究團(tuán)隊提出SIRI方法,通過"壓縮-擴(kuò)張"交替訓(xùn)練策略,成功解決了大型推理模型"話多且準(zhǔn)確率低"的問題。實驗顯示,該方法在數(shù)學(xué)競賽題上將模型準(zhǔn)確率提升43.2%的同時,輸出長度減少46.9%,真正實現(xiàn)了效率與性能的雙重優(yōu)化,為AI模型訓(xùn)練提供了新思路。
南洋理工大學(xué)與騰訊聯(lián)合研究團(tuán)隊開發(fā)出Rolling Forcing技術(shù),實現(xiàn)AI視頻實時流式生成的重大突破。該技術(shù)通過滾動窗口聯(lián)合去噪、注意力錨點機(jī)制和高效訓(xùn)練算法三項創(chuàng)新,解決了長視頻生成中的錯誤累積問題,可在單GPU上以16fps速度生成多分鐘高質(zhì)量視頻,延遲僅0.76秒,質(zhì)量漂移指標(biāo)從傳統(tǒng)方法的1.66降至0.01,為交互式媒體和內(nèi)容創(chuàng)作開辟新可能。
華中科技大學(xué)研究團(tuán)隊發(fā)現(xiàn),通過讓AI模型學(xué)習(xí)解決幾何問題,能夠顯著提升其空間理解能力。他們構(gòu)建了包含約30000個幾何題目的Euclid30K數(shù)據(jù)集,使用強(qiáng)化學(xué)習(xí)方法訓(xùn)練多個AI模型。實驗結(jié)果顯示,幾何訓(xùn)練在四個空間智能測試基準(zhǔn)上都帶來顯著提升,其中最佳模型達(dá)到49.6%準(zhǔn)確率,超越此前最好成績。這項研究揭示了基礎(chǔ)幾何知識對培養(yǎng)AI空間智能的重要價值。