近日，階躍星辰研究團(tuán)隊(duì)通過大規(guī)模實(shí)證探索，耗費(fèi)了近 100 萬 NVIDIA H800 GPU 小時(shí)（約百萬美元），從頭訓(xùn)練了 3,700 個(gè)不同規(guī)模，共計(jì)訓(xùn)了 100 萬億個(gè) token，揭示了 LLM 超參數(shù)優(yōu)化的全新的普適性縮放規(guī)律，為更好地提升 LLM 性能，提供了開箱即用的工具。該研究也是第一個(gè)全面研究模型最優(yōu)超參隨著 Model Shape、Dense/MoE、預(yù)訓(xùn)練數(shù)據(jù)分布的變化，是否穩(wěn)定不變的工作。研究中凸顯出 Step Law 的魯棒性，大大增加了該工具的實(shí)用性和普適性。同時(shí)該團(tuán)隊(duì)正在逐步開源相關(guān)資料，包括模型、訓(xùn)練日志等，期待更多相關(guān)領(lǐng)域的人基于海量的實(shí)驗(yàn)結(jié)果作出更加深入的研究與解釋。

• 論文標(biāo)題：Predictable Scale: Part Ⅰ — Optimal Hyperparameter Scaling Law in Large Language Model Pretraining
• 論文鏈接：https://arxiv.org/abs/2503.04715
• 工具鏈接：https://step-law.github.io/
• 開源地址：https://github.com/step-law/steplaw
• 訓(xùn)練過程：https://wandb.ai/billzid/predictable-scale

海量實(shí)驗(yàn)，實(shí)證為王

研究團(tuán)隊(duì)從頭訓(xùn)練了 3,700 個(gè)不同規(guī)模、不同超參數(shù)組合、不同形狀、不同數(shù)據(jù)配比、不同稀疏度 (含 MoE、Dense) 的大語言模型（LLM），共訓(xùn)練了超 100 萬億個(gè) token，對(duì)超參數(shù)進(jìn)行了全面的網(wǎng)格搜索，研究團(tuán)隊(duì)發(fā)現(xiàn)了一條普適的縮放法則 (簡(jiǎn)稱 Step Law)：最優(yōu)學(xué)習(xí)率隨模型參數(shù)規(guī)模與數(shù)據(jù)規(guī)模呈冪律變化，而最優(yōu)批量大小主要與數(shù)據(jù)規(guī)模相關(guān)。

圖一：在 400M 的 Dense LLM 上訓(xùn)練 40B Token（左）和在 1B 的 Dense LLM 上訓(xùn)練 100B Token（右）的超參 - 損失等高線圖，并且對(duì)業(yè)內(nèi)不同方法進(jìn)行比較，所有方法都轉(zhuǎn)換成了預(yù)測(cè) Optimal Token Wise BatchSize。這里所有的等高線都是從頭訓(xùn)練的小模型所得的真實(shí)收斂后的 Train Smooth Loss。左右兩張圖的所有等高線，分別來自于兩組共 240 個(gè)采用不同超參（Grid Search）的端到端訓(xùn)練的小模型。Global Mimimum 是來 120 個(gè)小模型中最終 Train Smooth Loss 最小的那個(gè)。等高線表示距離 Global Mimimum 的從最終 loss 角度的相對(duì)距離。而超越 + 2% 的點(diǎn)位，并沒有體現(xiàn)在圖中。

值得關(guān)注的是，實(shí)驗(yàn)表明，在固定模型大小與數(shù)據(jù)規(guī)模時(shí)，超參數(shù)優(yōu)化的 Landscape 呈現(xiàn)出明顯的凸性特征，這意味著存在一個(gè)穩(wěn)定且易尋的最優(yōu)超參數(shù)區(qū)域。

圖二：Learning Rate 與 Batch Size 在 1B 模型訓(xùn)練 100BToken 上的損失分布：散點(diǎn)圖（左）與 3D 曲面（右）圖中的每一個(gè)實(shí)心點(diǎn)都是真實(shí)值，是 120 個(gè)從頭訓(xùn)練的一個(gè)小模型，在訓(xùn)練結(jié)束之后的收斂 Loss。為了展示這樣的凸性，研究員們構(gòu)造了如右圖一樣的 3 維空間，空間的橫軸為 Learning Rate，縱軸為 Batch-size，高度軸為 Loss。對(duì)于這個(gè)三維空間我們進(jìn)行橫面和豎面的切割，如左上圖得到固定不同的 Learning Rate 情況下，最終收斂的 Train Smoothed Loss 隨著 Batchsize 的變化。而左下圖是固定不同的 Batchsize 情況下，最終收斂的 Train Smoothed Loss 隨著 Learning Rate 的變化。可以顯著的觀測(cè)到一種凸性，且在凸性的底端，是一個(gè)相對(duì)平坦的區(qū)域。這意味著 Optimal Learning Rate 和 Batchsize 很可能是一個(gè)比較大區(qū)域。

為了便于學(xué)界和業(yè)界應(yīng)用，團(tuán)隊(duì)推出了一款通用的最優(yōu)超參數(shù)估算工具（https://step-law.github.io），其預(yù)測(cè)結(jié)果與窮舉搜索的全局最優(yōu)超參數(shù)相比，性能僅有 0.09% 的差距。同時(shí)，研究團(tuán)隊(duì)還在該網(wǎng)站上公開了所有超參數(shù)組合的 loss 熱力圖，以進(jìn)一步推動(dòng)相關(guān)研究。

圖三：1B 模型、100BToken 訓(xùn)練上的 LR 與 BS 熱力圖。每一個(gè)點(diǎn)上的數(shù)字都是從頭訓(xùn)練的一個(gè)小模型（共訓(xùn)練了 120 個(gè)小模型），在訓(xùn)練結(jié)束之后的收斂真實(shí) Train Smoothed Loss。紅點(diǎn)是上述公式的預(yù)估值所對(duì)應(yīng)的 BS、LR 位置。其中空白的部分，是因?yàn)榉N種原因訓(xùn)練失敗的點(diǎn)位。所有熱力圖見https://step-law.github.io/

研究亮點(diǎn)：Step Law 的普適性，從三個(gè)不同角度

相較于現(xiàn)有的大模型最優(yōu)超參數(shù)估算公式，團(tuán)隊(duì)的研究進(jìn)行了極其充分的、覆蓋模型參數(shù)規(guī)模（N）、訓(xùn)練數(shù)據(jù)規(guī)模（D）、批量大?。˙S）和學(xué)習(xí)率（LR）的網(wǎng)格搜索，最終得到的 Step Law 則展現(xiàn)出顯著的優(yōu)越性，在適用性和準(zhǔn)確度方面均有大幅提升。

表一：不同方法的最佳超參數(shù)縮放定律比較，其中 Data Recipe 是指是否有在不同的預(yù)訓(xùn)練語料的配比下的最優(yōu)超參進(jìn)行研究。Model Sparsity 是指是否同時(shí)支持 MoE Model 和 Dense Model，以及不同的稀疏度下的 MoE 模型。LR 指的是 learning Schedule 中的峰值 Learning Rate，其中 BS 值得是 Token Wise 的 Batch Size。

角度一：跨模型形狀 (Model Shape) 的穩(wěn)定性

研究還深入探討了不同模型形狀（如寬度與深度的不同組合）對(duì)縮放規(guī)律的影響，發(fā)現(xiàn)無論模型是以寬度為主還是深度為主，抑或是寬深平衡的設(shè)計(jì)，Step Law 均表現(xiàn)出了高度的穩(wěn)定性。這表明，縮放規(guī)律不僅適用于特定類型的模型結(jié)構(gòu)，在更廣泛的架構(gòu)設(shè)計(jì)空間中依然適用，為復(fù)雜模型架構(gòu)的設(shè)計(jì)和優(yōu)化提供了指導(dǎo)意義。

圖四：最優(yōu)超參在不同 Model Shape 下的拓?fù)洳蛔冃?。這里的是固定了模型的非詞表參數(shù)量的大小，且固定了模型的訓(xùn)練 Token 數(shù)，但研究團(tuán)隊(duì)使用了不同的 Model Shape。例如變換了層數(shù)，從左到右分別是 14/10/8 層；變換了 Model Hidden Dimension，分別包括 1280/1536/2048 這三種；同時(shí)變換了 6 種不同的 FFN 倍數(shù) (FFN_media_dim/model_dim)，從 1.1 倍到～6.25 倍。其中紅色五角星的點(diǎn)是 Step Law 預(yù)測(cè)的點(diǎn)位，可以觀察到 Step_law 在 6 個(gè)不同的 Shape 上都預(yù)測(cè)到了 Global Minimum 附近。然而也可以同時(shí)觀察到，不同的 Model Shape，Bottom 的一片區(qū)域的位置是會(huì)發(fā)生 shift 的。

角度二：跨模型架構(gòu)的泛化性

研究結(jié)果發(fā)現(xiàn)，這一縮放規(guī)律不僅適用于 Dense 模型，還能很好地推廣到不同稀疏度的 MoE 模型（Mixture-of-Experts），展示了極強(qiáng)的泛化能力。

圖五：不同稀疏比下 MoE 模型的超參 - 損失等高線圖。左：低稀疏度（N_a/N=0.27），中間：中等稀疏度（N_a/N=0.58，D/N=10），右：中等稀疏度、較少訓(xùn)練 Token 數(shù)（N_a/N=0.58，D/N=4）。研究員們?cè)诓煌∈瓒?，不?D/N 的 MoE 模型配置，每一種配置都從頭訓(xùn)練了 45 個(gè)小模型，來做最優(yōu)超參搜取。共計(jì)從頭訓(xùn)練了 495 個(gè)不同稀疏度、不同超參、不同 D/N 的 MoE 模型。從而得到了不同配置下的基于真實(shí)值的 Global Minimum Train Smoothed Loss。其中除了一組 D/N=1 的實(shí)驗(yàn)，其余實(shí)驗(yàn) Step Law 預(yù)測(cè)位置都在 Global Minimum+0.5% 的范圍之內(nèi)。并且大多數(shù)配置下都在 Global Minimum+0.25% 的范圍之內(nèi)。充分的驗(yàn)證了 Step Law 的魯棒性。詳細(xì)結(jié)果可以參考論文的附錄部分。

角度三：跨數(shù)據(jù)分布的穩(wěn)定性

那么 Step Law 是否能兼容不同的預(yù)訓(xùn)練數(shù)據(jù)分布呢？研究團(tuán)隊(duì)進(jìn)一步驗(yàn)證了不同數(shù)據(jù)分布下的規(guī)律一致性：無論是英語主導(dǎo)、中英雙語、Code 和英語混合，還是代碼主導(dǎo)的數(shù)據(jù)分布，Step Law 都表現(xiàn)出了穩(wěn)定的性能。這為多語言、多任務(wù)場(chǎng)景下的實(shí)際應(yīng)用提供了可靠支持。

表二：實(shí)驗(yàn)的不同數(shù)據(jù)分布。其中 Baseline 是得出 Step Law 的訓(xùn)練 Recipe。而 Code-Math，是壓縮英文 web-data 的配比近一半，擴(kuò)大 code-math 的比例至近 40%。而 More Code-Math 比例更加極端，將英文 web-data 的配比壓縮為之前的 1/4，將 Code-math 擴(kuò)大為近 2/3。EN-CN 是下調(diào)英文 web-data 的配比近一半，將余量的部分都轉(zhuǎn)化為中文網(wǎng)頁數(shù)據(jù)。

圖六：不同數(shù)據(jù)分布下的超參 - 損失等高線。左：雙語數(shù)據(jù)（表格中 En-CN），中間：加入 Code 數(shù)據(jù)（表格中的 Code+Math），右：主要為 Code 數(shù)據(jù)（表格中的 More Code+Math）。每一個(gè)圖都是從頭訓(xùn)練了 45 個(gè)模型，每一個(gè)模型除了 Bs/lr 不同以外，其他設(shè)置完全相同?？偣灿?xùn)練了 135 個(gè)在三種數(shù)據(jù)分布下的模型。其中 Global Minimum 是通過這種 grid search 的方法得到的最低 Final Train Smoothed Loss 的真實(shí)值。Step Law 預(yù)測(cè)出來的最優(yōu) Batch Size/Learning Rate 都在最低 Loss +0.125%/0.25% 的范圍內(nèi)。

 研究細(xì)節(jié)解讀

1. 學(xué)習(xí)率調(diào)度策略優(yōu)化

圖七：不同學(xué)習(xí)率策略的比較。藍(lán)色等高線 (傳統(tǒng)衰減策略): 學(xué)習(xí)率會(huì)從一個(gè)最大值 (max_lr) 逐漸減小到一個(gè)最小值 (min_lr，常是峰值的十分之一)。紅色等高線 (固定最終學(xué)習(xí)率策略): 保持一個(gè)固定的最小學(xué)習(xí)率 (min_lr = 1e-5)，而不是像傳統(tǒng)方法那樣與最大學(xué)習(xí)率掛鉤。兩張圖都分別為 120 個(gè)從頭訓(xùn)練的模型，在相同的 batch size/learning rate 范圍內(nèi)做的 Grid Search。紅色和藍(lán)線的 Global Minimum 都是各自配置下的真值 - 最小的 Final Train Smoothed Loss?？梢杂^察到改成 max_lr/10 之后，藍(lán)點(diǎn)會(huì)向左上方偏移，即更小的 Learning Rate 和更大的 Batchsize。如果不是對(duì)比相對(duì)值，而是對(duì)比真值，min_lr=1e-5 的最終收斂 loss 普遍小于 max_lr/10。相關(guān)的真值將陸續(xù)開源。相關(guān)的真值開源在https://github.com/step-law/steplaw

研究團(tuán)隊(duì)通過對(duì)比分析發(fā)現(xiàn)，學(xué)習(xí)率調(diào)度策略對(duì)最優(yōu)超參選擇產(chǎn)生顯著影響。研究揭示了傳統(tǒng)學(xué)習(xí)率衰減與固定最小學(xué)習(xí)率方案間的重要差異：

傳統(tǒng)學(xué)習(xí)率衰減方案將最小學(xué)習(xí)率設(shè)為最大值的十分之一 (max_lr/10)，而團(tuán)隊(duì)提出的固定方案則采用恒定的絕對(duì)最小值 (1e-5)。等高線圖分析清晰表明，傳統(tǒng)衰減方法使得最優(yōu)學(xué)習(xí)率區(qū)域出現(xiàn)明顯的左偏分布 —— 即損失最小區(qū)域向較低學(xué)習(xí)率區(qū)間顯著偏移。

團(tuán)隊(duì)通過分析了傳統(tǒng)學(xué)習(xí)率（min_lr=max_lr/10）調(diào)度方案的局限性：采用較高初始峰值學(xué)習(xí)率時(shí)，其退火機(jī)制會(huì)同步抬升最低學(xué)習(xí)率閾值。這種耦合設(shè)計(jì)在訓(xùn)練末期會(huì)使學(xué)習(xí)率超出理想?yún)^(qū)間，過大的參數(shù)更新幅度引發(fā)損失函數(shù)在收斂階段持續(xù)振蕩。相比之下，固定最小學(xué)習(xí)率策略通過解耦初始學(xué)習(xí)率與終值學(xué)習(xí)率的關(guān)聯(lián)，在訓(xùn)練后期始終維持符合梯度下降動(dòng)態(tài)特性的更新步長(zhǎng)。

此外，這種固定最終較小最終學(xué)習(xí)率的策略也與業(yè)界的訓(xùn)練經(jīng)驗(yàn)相匹配，更有實(shí)際應(yīng)用價(jià)值。

2. 訓(xùn)練損失與驗(yàn)證損失的最優(yōu)超參一致性

圖八：平滑訓(xùn)練損失（Final Train Smoothed Loss）的超參 - 損失等高線圖（左）和驗(yàn)證損失（Validation Loss of Final Checkpoint）的超參 - 損失等高線圖（右）。兩張圖都是在同一組實(shí)驗(yàn)下進(jìn)行，對(duì)于相同的模型尺寸，相同的訓(xùn)練 Token 數(shù)，分別采用了 64 組不同的超參進(jìn)行 Grid Search。從而得到 64 個(gè)模型的 Final Train Smoothed Loss、和 Validation Loss。

在 429M 模型上訓(xùn)練 40B 的 Token 進(jìn)行驗(yàn)證，當(dāng)平滑訓(xùn)練損失達(dá)到最優(yōu)時(shí)，學(xué)習(xí)率為 1.95×10^-3，批量大小為 393,216，這一點(diǎn)與驗(yàn)證損失最優(yōu)時(shí)的超參數(shù)完全重合。此外，學(xué)習(xí)率和批量大小的變化趨勢(shì)在平滑訓(xùn)練損失和驗(yàn)證損失中表現(xiàn)出高度一致性。這一發(fā)現(xiàn)表明，平滑的訓(xùn)練損失曲線可以為實(shí)際超參數(shù)選擇提供可靠的指導(dǎo)。盡管采用 Train Smoothed Loss 可以降低實(shí)驗(yàn)成本（節(jié)省了 Final Checkpoint 在 Validation Set 上推理的算力），但仍然具有一定的局限性，例如訓(xùn)練數(shù)據(jù)不能重復(fù)。研究團(tuán)隊(duì)將會(huì)陸續(xù)開源這近 4000 個(gè)模型的 Final Checkpoint，供廣大研究員進(jìn)行進(jìn)一步的分析。

最優(yōu)超參的 Scaling Law 擬合

圖九：(a) 散點(diǎn)圖表示模型規(guī)模為 N 時(shí)，經(jīng)驗(yàn)最優(yōu)學(xué)習(xí)率與批量大小的關(guān)系；(b) 散點(diǎn)圖表示數(shù)據(jù)集規(guī)模為 D 時(shí)，經(jīng)驗(yàn)最優(yōu)學(xué)習(xí)率與批量大小的關(guān)系。曲線表示超參數(shù)縮放定律的預(yù)測(cè)結(jié)果，陰影區(qū)域表示基于采樣擬合策略得到參數(shù)不確定性范圍。圖上的每一個(gè)點(diǎn)，背后都代表著 45~120 個(gè)采用了不同的超參的從頭訓(xùn)練的模型。圖上的每一個(gè)點(diǎn)位都在不同的 Model Size、Data Size 下通過 Grid Search 得到的最優(yōu)的超參 (Optimal Learning Rate，Optimal Batch Size)。這張圖總共涉及了 1912 個(gè)從頭訓(xùn)練的 LLM。擬合方法開源在https://github.com/step-law/steplaw

通過從頭訓(xùn)練了 1912 個(gè) LLM，研究團(tuán)隊(duì)分析發(fā)現(xiàn)：

• 最優(yōu)學(xué)習(xí)率：隨模型規(guī)模增大而減小，隨數(shù)據(jù)規(guī)模增大而增大。
• 最優(yōu)批量大?。弘S數(shù)據(jù)規(guī)模增大而增大，與模型規(guī)模弱相關(guān)。

研究通過對(duì)數(shù)變換將冪律關(guān)系轉(zhuǎn)化為線性形式，采用最小二乘法擬合參數(shù)，并通過 Bootstrap 采樣方法提升穩(wěn)健性。最終，團(tuán)隊(duì)提出了一套精確的預(yù)測(cè)公式，為大模型預(yù)訓(xùn)練的超參數(shù)設(shè)置提供了一個(gè)開箱即用的工具。據(jù)了解，在階躍星辰真正的研發(fā)過程中，他們已經(jīng)廣泛使用了 Step Law，主要是在大于 1B 的模型和非極端的 D/N 下適用。

討論與未來工作

研究團(tuán)隊(duì)坦言，盡管付出了很多的算力，和大量的精力來分析相關(guān)的實(shí)驗(yàn)。但是很多子 Topic 分析仍然是值得深究的。研究團(tuán)隊(duì)認(rèn)為他們的這份工作僅僅是一個(gè)開始，他們會(huì)陸續(xù)將實(shí)驗(yàn)的各個(gè)細(xì)節(jié)整理并且開源出來，供整個(gè)社區(qū)的研究人員進(jìn)行進(jìn)一步分析和深入的理論解釋。這些子 Topic 包括以下話題：

• 在給定模型、訓(xùn)練 Token 數(shù)的情況下，(Loss,bs,lr) 這三維空間是否是真正的凸性。
• 是否有更好的 optimal BS LR 的擬合方法，并且可以兼容 BS、LR 的內(nèi)在關(guān)系。
• 盡管 Step Law 在不同 Model Shape、不同稀疏的 MoE 模型是魯棒的，但是次優(yōu)的區(qū)域是在不同配置下是變化的，有無更好的解釋方法。
• 上文中這些基于海量 Grid Search 的數(shù)據(jù)驅(qū)動(dòng)的結(jié)論的理論解釋。
• 不同的超參、不同 Model Size、Model Shape、Model Sparsity 下的 Training Dynamic 研究。

為此研究團(tuán)隊(duì)公開了他們的開源計(jì)劃，以及 Predictable Scale 系列工作的發(fā)布節(jié)奏。據(jù)了解，Predictable Scale 系列后續(xù)可能進(jìn)一步討論超大模型性能預(yù)測(cè)、Code&Math Scaling 性質(zhì)、不同 Attention 類型的 Scaling 性質(zhì)等問題。非常值得期待。