作為第一作者,想從idea來源的視角介紹一下這個工作。在知乎潛水多年,首答就獻給這個問題了。答得不好的地方請大家多多批評指正(鞠躬)。非常感謝劉忠鑫老師和 Binyuan Hui 老師的指導,以及所有合作者的幫助,尤其感謝通義千問的資源支持。
我們都知道,除了拓展數據量以外,現在有兩條主流的 scaling 路線來增強拉大模型的計算量,增強大模型的能力:
那麼問題來了:能不能有一種新的 scaling 路線,不會帶來顯著的內存和時延增加,同時又適用於一切的場景?
我們的核心想法就是:在參數量不變的情況下,同時拉大訓練和推理并行計算量。
動機由來
最初其實是舍友在宿舍學習 diffusion 模型,他對 diffusion model 必用的一個 trick 百思不得其解:Classifier-Free Guidance(CFG)。CFG 在推理階段拿到輸入 x 時,首先做一次正常的 forward 得到 f(x);然后再對 x 進行主動的劣化(比如去除條件)變為 x',再進行一次 forward 得到 f(x')。最終的輸出 g(x) 是 f(x) 和 f(x') 的一個加權組合,它的效果比 f(x) 更好,更能遵循輸入的條件。
這個現象事實上有點反直覺:f(x) 和訓練階段是對齊的,而 g(x) 明顯和訓練階段的目標存在 gap。按照常識,只有訓練目標和推理目標形式相同,推理才能發揮最大效果。另外,f(x) 的參數量和 g(x) 也是相同的,輸入的有效信息量也相同,為什麼 f(x) 反而學不到 g(x) 的能力?這說明背后或許存在更深層次的原因。
我們做出了大膽的猜想:CFG 生效的原因,本質上是用了雙倍的并行計算量,它拉大了模型本身的 capacity。
這就啟發我們,可以進一步 scale:
這個方法非常簡單,可以應用到任何的模型架構、任務和數據上。我們在大語言模型上首先探索了這一想法,如下圖所示:
Parallel Scaling Law
我們首先進行了一系列理論分析,得出一個結論:將一個參數量為 N 的模型并行 P 個流,等價於將參數量變為原來的 倍(詳見論文的分析)。其中 Diversity 與不同流之間的殘差相關系數有關,不好接著分析。但這至少說明,放縮并行計算量和放縮參數之間必然存在著某個聯系
因此,我們接下去做了大量的實驗,最終得到:并行P個流,等價於把參數放大O(logP)倍,但它相較於放大參數有著非常顯著的推理效率優勢:
我們也提供了一個 HuggingFace 空間,來更直觀地感受 Scaling Law 的力量。歡迎大家試用:
兩階段訓練
先前的實驗主要聚焦在 pre-train,由於 batchsize 會拉大 P 倍,訓練開銷較大。因此我們嘗試了一個后訓練策略:首先第一階段訓練 1T token(常數學習率),然后在二階段使用 ParScale(退火學習率)后訓練 20B token。可以發現這個策略也是非常有效的。
之后,我們將 ParScale 應用到 Qwen-2.5 模型上(已經訓了 12 T token),包括全參數持續訓練(CPT)和 PEFT 訓練(凍結主網絡,只微調引入的 prefix 參數)。
值得一提的是,PEFT 訓練展示出了動態并行拓展的前景:我們可以使用相同的模型權重,在不同的場合下使用不同的 P,從而快速地動態調整能力以及推理開銷。這是目前的主流方法比較難做到的。
總結
ParScale 是我們探索 LLM Scaling Law 的新嘗試,目前的研究仍然在進行中。我們相信擴展計算量可以帶來智能的涌現。展望未來,我們計劃進一步在更多的模型架構(比如MoE)以及更大的數據上進行嘗試,從而更好地理解擴大并行計算能帶來的收益。更多的未來方向討論詳見論文。歡迎大家多多批評指正!