攻克強(qiáng)化學(xué)習(xí)「最慢一環(huán)」！交大字節(jié)聯(lián)手，RL訓(xùn)練速度飆升2.6倍

鷺羽 2025-09-13 16:45:23 來(lái)源：量子位

允中 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

強(qiáng)化學(xué)習(xí)的訓(xùn)練效率，實(shí)在是太低了！

隨著DeepSeek、GPT-4o、Gemini等模型的激烈角逐，大模型“深度思考”能力的背后，強(qiáng)化學(xué)習(xí)（RL）無(wú)疑是那把最關(guān)鍵的密鑰。

然而，這場(chǎng)競(jìng)賽的背后，一個(gè)巨大的瓶頸正悄然限制著所有玩家的速度——相較于預(yù)訓(xùn)練和推理，RL訓(xùn)練更像一個(gè)效率低下的“手工作坊”，投入巨大但產(chǎn)出緩慢。

其中，占據(jù)超過(guò)80%時(shí)間的Rollout（響應(yīng)生成）階段，由于其內(nèi)存帶寬限制和自回歸特性，成為了整個(gè)AI基礎(chǔ)設(shè)施中公認(rèn)的阿喀琉斯之踵。

如何攻克這塊AI基建的最后高地？現(xiàn)在，上海交通大學(xué)與字節(jié)跳動(dòng)研究團(tuán)隊(duì)給出了一個(gè)全新的答案。

該團(tuán)隊(duì)聯(lián)手推出的RhymeRL，從一個(gè)被忽視的現(xiàn)象入手，巧妙地將歷史數(shù)據(jù)變廢為寶，在不犧牲精度的前提下，將RL訓(xùn)練吞吐量提升了2.6倍。

模型生成的答案存在兩大“歷史相似性”

該研究團(tuán)隊(duì)深入分析了大量RL訓(xùn)練過(guò)程，發(fā)現(xiàn)在相鄰的兩個(gè)訓(xùn)練周期中，盡管模型權(quán)重已經(jīng)更新，但對(duì)于同一個(gè)問(wèn)題（Prompt），模型生成的答案（Rollout）存在兩大“歷史相似性”：

第一，序列相似性。

新答案“繼承”了舊答案的思路，高達(dá)95%的歷史Token都可以直接復(fù)用。

第二，長(zhǎng)度分布相似性。

上一輪里，哪些問(wèn)題讓模型“思來(lái)想去”，這一輪大概率依舊如此。響應(yīng)長(zhǎng)度的排序驚人地穩(wěn)定。

該研究團(tuán)隊(duì)認(rèn)為，這種相似性，主要源于PPO/GRPO等主流RL算法為了保證訓(xùn)練穩(wěn)定而采用的梯度裁剪（Clipping）機(jī)制，它讓模型的進(jìn)化變得平滑而有跡可循。

這就好比一個(gè)學(xué)生，雖然每天都在進(jìn)步，但他解決同一類(lèi)型問(wèn)題的思考路徑和草稿篇幅，在短期內(nèi)是高度相似的。

既然如此，歷史的舊草稿能否成為新一輪學(xué)習(xí)的模版？

針對(duì)Rollout階段的低效現(xiàn)狀，提出新框架RhymeRL

基于這一洞察，RhymeRL框架應(yīng)運(yùn)而生。

針對(duì)Rollout階段的低效現(xiàn)狀，它包含兩大核心利器：

HistoSpec和HistoPipe。

RhymeRL的核心設(shè)計(jì)示意圖

HistoSpec

傳統(tǒng)的Rollout，是一個(gè)Token一個(gè)Token往外蹦的自回歸過(guò)程，速度慢且GPU利用率較低。

而HistoSpec獨(dú)創(chuàng)性地將投機(jī)解碼（Speculative Decoding）技術(shù)引入RL，它不再需要一個(gè)額外的小模型來(lái)“猜”草稿，而是直接把上一輪的歷史響應(yīng)作為“最佳劇本”。

這就像開(kāi)卷考試，你提前拿到了去年的標(biāo)準(zhǔn)答案作為參考。

1. 起草（Draft）：直接從歷史響應(yīng)中巧妙地總結(jié)出樹(shù)狀草稿。
2. 驗(yàn)證（Verify）：將整段草稿一次性扔給大模型，通過(guò)單次前向傳播，并行驗(yàn)證所有Token的正確性，然后“接收”所有匹配的部分。

由于歷史序列的超高相似性，草稿的接受率極高。

這使得計(jì)算過(guò)程從逐字生成變成了“批量驗(yàn)證”，計(jì)算密度飆升，打破了內(nèi)存帶寬的枷鎖，讓單個(gè)響應(yīng)的生成速度實(shí)現(xiàn)了質(zhì)的飛躍。

HistoSpec采用的基于樹(shù)的歷史響應(yīng)管理，實(shí)現(xiàn)了草稿的高速、準(zhǔn)確生成

HistoPipe

僅僅讓單個(gè)響應(yīng)變快還不夠。

在批處理中，不同任務(wù)的響應(yīng)長(zhǎng)度不一，短任務(wù)總要等待長(zhǎng)任務(wù)，導(dǎo)致大量GPU資源被閑置，產(chǎn)生了巨大的“空泡”（Bubble）。

HistoPipe是一位具有前瞻性的調(diào)度大師，它的目標(biāo)是：

榨干每一滴GPU算力，實(shí)現(xiàn)無(wú)空泡的完美流水線。

基于“長(zhǎng)度分布相似性”的洞察，HistoPipe不再?gòu)?qiáng)求單一步內(nèi)實(shí)現(xiàn)負(fù)載均衡，而是玩起了“跨步互補(bǔ)”——在奇數(shù)步，讓所有GPU由短到長(zhǎng)處理任務(wù)；在偶數(shù)步，則反過(guò)來(lái)，由長(zhǎng)到短處理。

這樣一來(lái)，上一步因?yàn)樘幚黹L(zhǎng)任務(wù)而拖慢的GPU，在下一步會(huì)優(yōu)先處理短任務(wù)，完美填補(bǔ)了時(shí)間差。

通過(guò)這種巧妙的削峰填谷，HistoPipe將整個(gè)集群的資源浪費(fèi)降至最低。

HistoPipe的流水線設(shè)計(jì)，通過(guò)跨步互補(bǔ)調(diào)度實(shí)現(xiàn)了無(wú)空泡調(diào)度

2.6倍加速，精度無(wú)損

當(dāng)模板HistoSpec遇上調(diào)度大師HistoPipe，產(chǎn)生反應(yīng)是驚人的——實(shí)驗(yàn)結(jié)果表明，在數(shù)學(xué)、代碼等任務(wù)上，RhymeRL相比于基礎(chǔ)系統(tǒng)取得了大幅性能提升，端到端訓(xùn)練吞吐量提升高達(dá)2.61倍。

這意味著，研究者和企業(yè)可以用更少的資源、在更短的時(shí)間內(nèi)，完成更強(qiáng)大的模型訓(xùn)練，極大地加速了AI技術(shù)迭代的步伐。

RhymeRL在不同模型大小和不同響應(yīng)長(zhǎng)度下都取得了顯著加速效果

RhymeRL的重要意義在于，它提出了一種新的強(qiáng)化學(xué)習(xí)范式——基于歷史信息來(lái)端到端地加速?gòu)?qiáng)化學(xué)習(xí)效率。

強(qiáng)化學(xué)習(xí)不是簡(jiǎn)單的推理與訓(xùn)練的拼接，通過(guò)深入剖析其任務(wù)特性，RhymeRL能夠充分發(fā)揮系統(tǒng)統(tǒng)籌調(diào)度能力與底層硬件的算力資源，同時(shí)無(wú)損地適配各種已有的訓(xùn)練算法。

arXiv鏈接：
https://www.arxiv.org/abs/2508.18588