單張GPU搞定GPT-3超參數(shù)!先訓(xùn)練小模型,再“一鍵遷移”
本文經(jīng)AI新媒體量子位(公眾號(hào)ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請(qǐng)聯(lián)系出處。
不不不,現(xiàn)在情況有變——
在單個(gè)GPU上完全可以調(diào)整大規(guī)模模型的超參數(shù)。
怎么說?
原來有人發(fā)現(xiàn)了一種新的調(diào)參方法,無論模型規(guī)模怎么變化,所得的最優(yōu)超參數(shù)都能保持性能穩(wěn)定。
由此一來,我們可以先訓(xùn)練一個(gè)小版本的模型,在上面間接調(diào)整好超參,然后以零樣本的方式直接將它們復(fù)制到全尺寸模型上,就能獲得相當(dāng)不錯(cuò)的性能。
這對(duì)手里GPU資源不夠的人來說簡直不要太好了吧。
目前,相關(guān)帖子也在Reddit上引起熱議,得到了300+的點(diǎn)贊支持。
在一個(gè)GPU上調(diào)參GPT-3大模型
方法名叫muP (Maximal Update Parametrization),作者分別來自微軟和OpenAI。
想法很簡單,利用的是他們?cè)谥暗墓ぷ髦邪l(fā)現(xiàn)的一個(gè)叫做μP的特殊參數(shù)化思路:
窄神經(jīng)網(wǎng)絡(luò)和寬神經(jīng)網(wǎng)絡(luò)共享同一組最優(yōu)超參數(shù),即使寬度無限大(width->∞)時(shí)也一樣。
具體原理可以參見論文《Feature Learning in Infinite-Width Neural Networks》。
可共享的超參數(shù)包括學(xué)習(xí)率learning rate、learning rate schedule、initialization、parameter multipliers……甚至可以單獨(dú)針對(duì)每個(gè)參數(shù)張量。
作者在寬度高達(dá)4096的Transformer以及ResNet上驗(yàn)證了這一結(jié)論。
因此,資源匱乏的煉丹er就可以在單張GPU上對(duì)一個(gè)小版本的GPT-3模型進(jìn)行超參數(shù)調(diào)整:
如果在這個(gè)小模型上得到的參數(shù)接近最優(yōu),那么在大模型上也可以得到一樣的結(jié)果。
ps. 這種調(diào)參方式又被命名為“μ遷移(μTransfer)”。
具體效果如何?
作者訓(xùn)練了一個(gè)只有4000萬參數(shù)規(guī)模的小型GPT-3,它小到可以直接跑在一張GPU上。
然后把它的超參數(shù)“μ遷移”到一個(gè)有67億參數(shù)的大規(guī)模GPT-3上,結(jié)果發(fā)現(xiàn)它的性能和原始的GPT-3完全相當(dāng)——盡管原GPT-3的參數(shù)規(guī)模還是它的兩倍!
而這一調(diào)整成本只占整個(gè)預(yù)訓(xùn)練成本的7%。
由于模型規(guī)模增大,直接調(diào)整小型模型的成本仍大致相同,如果用該方式來調(diào)參175億規(guī)模的GPT-3,其成本可能最多只有總預(yù)訓(xùn)練成本的0.3%。
好了,這時(shí)你可能會(huì)問:能不能只縮小模型的寬度呢?
作者表示,對(duì)于“non-width stuff”,沒有理論保證。
不過好消息是,他們?cè)趐reLN Transformer的合理范圍內(nèi)對(duì)depth、batch size、sequence length和timestep的遷移效果進(jìn)行了測試。
其中,他們將BERT-base和BERT-large在寬度和深度上縮小到同樣的規(guī)模, 然后同時(shí)進(jìn)行超參數(shù)調(diào)整后發(fā)現(xiàn):
相比已經(jīng)調(diào)優(yōu)的megatron BERT基線,兩者的性能都得到了改善,尤其是BERT-large提升更大。
由此也總結(jié)出一個(gè)道理:
遷移后的模型規(guī)模越大,收益越高。
所以作者還調(diào)侃道,雖然我們沒有測試175億規(guī)模的GPT-3,但保證結(jié)果能讓你“流口水”。
說了這么多,到底如何實(shí)現(xiàn)?
下表概括了如何通過fan-in或fan-out調(diào)整你的模型的initialization和learning rate。
其中粉色文本為μP,括號(hào)中的灰色文本為pytorch默認(rèn)值。
當(dāng)然,如果你不想自己手動(dòng)操作,作者也開源了Pytorch實(shí)現(xiàn),通過pip install mup就可以應(yīng)用到你的模型中。
關(guān)于作者
一作名叫Greg Yang,微軟高級(jí)研究員。
通訊作者為微軟研究院深度學(xué)習(xí)技術(shù)中心合伙人研究經(jīng)理、IEEE Fellow高劍峰。
還有兩位華人作者分別為來自微軟的Liu Xiaodong(北京郵電大學(xué)校友)和Chen Weizhu (已在微軟工作16年)。
他們的這篇成果已被NeurIPS 2021接收。
GitHub鏈接:
??https://github.com/microsoft/mup??
論文地址:
??https://arxiv.org/abs/2203.03466??