艾倫人工智能研究所等 5 機構(gòu)最近公布了史上最全的開源模型「OLMo」,公開了模型的模型權(quán)重,、完整訓(xùn)練代碼,、數(shù)據(jù)集和訓(xùn)練過程,,為以后開源社區(qū)的工作設(shè)立了新的標(biāo)桿。
多年來,,語言模型一直是自然語言處理(NLP)技術(shù)的核心,,考慮到模型背后的巨大商業(yè)價值,最大最先進的模型的技術(shù)細節(jié)都是不公開的,。
現(xiàn)在,,真?完全開源的大模型來了!
來自艾倫人工智能研究所,、華盛頓大學(xué),、耶魯大學(xué)、紐約大學(xué)和卡內(nèi)基梅隆大學(xué)的研究人員,,聯(lián)合發(fā)表了一項足以載入 AI 開源社區(qū)史冊的工作 ——
他們幾乎將從零開始訓(xùn)練一個大模型過程中的一切數(shù)據(jù)和資料都開源了,!
具體來說,,艾倫人工智能研究所推出的這個開放大語言模型(Open Language Model,OLMo)實驗和訓(xùn)練平臺,,則提供了一個完全開源的大模型,,以及所有和訓(xùn)練開發(fā)這個模型有關(guān)的數(shù)據(jù)和技術(shù)細節(jié) ——
訓(xùn)練和建模:它包括完整的模型權(quán)重、訓(xùn)練代碼,、訓(xùn)練日志,、消融研究、訓(xùn)練指標(biāo)和推理代碼,。
預(yù)訓(xùn)練語料:一個包含了高達 3T token 的預(yù)訓(xùn)練開源語料庫,,以及產(chǎn)生這些訓(xùn)練數(shù)據(jù)的代碼。
模型參數(shù):OLMo 框架提供了四個不同架構(gòu),、優(yōu)化器和訓(xùn)練硬件體系下的 7B 大小的模型,,以及一個 1B 大小的模型,所有模型都在至少 2T token 上進行了訓(xùn)練,。
同時,,也提供了用于模型推理的代碼、訓(xùn)練過程的各項指標(biāo)以及訓(xùn)練日志,。
評估工具:公開了開發(fā)過程中的評估工具套件,,包括每個模型訓(xùn)練過程中每 1000 step 中包含的超過 500 個的檢查點以及評估代碼。
所有數(shù)據(jù)都在 apache 2.0 下授權(quán)使用(免費商用),。
如此徹底的開源,,似乎是給開源社區(qū)打了個樣 —— 以后不像我這樣開源的,就別說自己是開源模型了,。
性能評估
從核心的評估結(jié)果來看,,OLMo-7B 與同類開源模型相比略勝一籌。在前 9 項評測中,,OLMo-7B 有 8 項排名前三,,其中有 2 項超越了其他所有模型。
在很多生成任務(wù)或閱讀理解任務(wù)(例如 truthfulQA)上,,OLMo-7B 都超過了 Llama 2,,但在一些熱門的問答任務(wù)(如 MMLU 或 Big-bench Hard)上表現(xiàn)則要差一些,。
前 9 個任務(wù)是研究人員對預(yù)訓(xùn)練模型的內(nèi)部評估標(biāo)準,,而下面三個任務(wù)則是為了完善 HuggingFace Open LLM 排行榜而加入的
下圖展示了 9 個核心任務(wù)準確率的變化趨勢。除了 OBQA 外,,隨著 OLMo-7B 接受更多數(shù)據(jù)的訓(xùn)練,,幾乎所有任務(wù)的準確率都呈現(xiàn)上升趨勢。
與此同時,,OLMo 1B 與其同類模型的核心評估結(jié)果表明,,OLMo 與它們處于同一水平,。
通過使用艾倫 AI 研究所的 Paloma(一個基準測試)和可獲取的檢查點,研究人員分析了模型預(yù)測語言能力與模型規(guī)模因素(例如訓(xùn)練的 token 數(shù)量)之間的關(guān)系,。
可以看到,,OLMo-7B 在性能上與主流模型持平。其中,,每字節(jié)比特數(shù)(Bits per Byte)越低越好,。
通過這些分析,研究人員發(fā)現(xiàn)模型在處理不同數(shù)據(jù)源時的效率差異較大,,這主要取決于模型訓(xùn)練數(shù)據(jù)與評估數(shù)據(jù)的相似度,。
特別地,OLMo-7B 在主要基于 Common Crawl 的數(shù)據(jù)源上表現(xiàn)出色(比如 C4),。
不過,,在與網(wǎng)絡(luò)抓取文本關(guān)系不大的數(shù)據(jù)源上,如 WikiText-103,、M2D2 S2ORC 和 M2D2 Wikipedia,,OLMo-7B 與其他模型相比效率較低。
RedPajama 的評估也體現(xiàn)了相似的趨勢,,可能是因為它的 7 個領(lǐng)域中只有 2 個來源于 Common Crawl,,且 Paloma 對每個數(shù)據(jù)源中的各個領(lǐng)域給予了相同的權(quán)重。
鑒于像 Wikipedia 和 arXiv 論文這樣的精選數(shù)據(jù)源提供的異質(zhì)數(shù)據(jù)遠不如網(wǎng)絡(luò)抓取文本豐富,,隨著預(yù)訓(xùn)練數(shù)據(jù)集的不斷擴大,,維持對這些語言分布的高效率會很更加困難。
OLMo 架構(gòu)
在模型的架構(gòu)方面,,團隊基于的是 decoder-only 的 Transformer 架構(gòu),,并采用了 PaLM 和 Llama 使用的 SwiGLU 激活函數(shù),引入了旋轉(zhuǎn)位置嵌入技術(shù)(RoPE),,并改進了 GPT-NeoX-20B 的基于字節(jié)對編碼(BPE)的分詞器,,以減少模型輸出中的個人可識別信息。
此外,,為了保證模型的穩(wěn)定性,,研究人員沒有使用偏置項(這一點與 PaLM 的處理方式相同)。
預(yù)訓(xùn)練數(shù)據(jù)集:Dolma
雖然研究人員在獲取模型參數(shù)方面取得了一定的進展,,但開源社區(qū)目前預(yù)訓(xùn)練數(shù)據(jù)集的開放程度還遠遠不夠,。
之前的預(yù)訓(xùn)練數(shù)據(jù)往往不會隨著模型的開源而公開(閉源模型就更不用說了)。
而且有關(guān)這些數(shù)據(jù)的說明文檔也常常缺乏足夠的細節(jié),,但是這些細節(jié)對于想要復(fù)現(xiàn)研究或完全理解相關(guān)工作至關(guān)重要,。
這一情況加大了語言模型研究的難度 —— 比如,了解訓(xùn)練數(shù)據(jù)如何影響模型能力和其局限性,。
為了推動語言模型預(yù)訓(xùn)練領(lǐng)域的開放研究,,研究人員構(gòu)建并公開了預(yù)訓(xùn)練數(shù)據(jù)集 Dolma,。
這是一個包含了從 7 種不同數(shù)據(jù)來源獲取的 3 萬億個 token 的多樣化、多源語料庫,。
這些數(shù)據(jù)源一方面在大規(guī)模語言模型預(yù)訓(xùn)練中常見,,另一方面也能被普通大眾所接觸。
Dolma 的構(gòu)建過程包括六個步驟:語言過濾,、質(zhì)量過濾,、內(nèi)容過濾、去重,、多源混合和 token 化,。
在整理和最終發(fā)布 Dolma 過程中,研究人員確保各數(shù)據(jù)源的文檔保持獨立,。
他們還開源了一套高效的數(shù)據(jù)整理工具,,這套工具能夠幫助進一步研究 Dolma、復(fù)制成果,,并簡化預(yù)訓(xùn)練語料庫的整理工作,。
此外,研究人員也開源了 WIMBD 工具,,以助于數(shù)據(jù)集分析,。
訓(xùn)練 OLMo
分布式訓(xùn)練框架
研究人員利用 PyTorch 的 FSDP 框架和 ZeRO 優(yōu)化器策略來訓(xùn)練模型。這種方法通過將模型的權(quán)重和它們對應(yīng)的優(yōu)化器狀態(tài)在多個 GPU 中進行分割,,從而有效減少了內(nèi)存的使用量,。
在處理高達 7B 規(guī)模的模型時,這項技術(shù)使研究人員能夠在每個 GPU 上處理 4096 個 token 的微批大小,,以實現(xiàn)更高效的訓(xùn)練,。
對于 OLMo-1B 和 7B 模型,研究人員固定使用大約 4M token(2048 個數(shù)據(jù)實例,,每個實例包含 2048 個 token 的序列)的全局批大小,。
而對于目前正在訓(xùn)練中的 OLMo-65B 模型,研究人員采用了一個批大小預(yù)熱策略,,起始于大約 2M token(1024 個數(shù)據(jù)實例),,之后每增加 100B token,批大小翻倍,,直至最終達到大約 16M token(8192 個數(shù)據(jù)實例)的規(guī)模,。
為了加快模型訓(xùn)練的速度,研究人員采用了混合精度訓(xùn)練的技術(shù),,這一技術(shù)是通過 FSDP 的內(nèi)部配置和 PyTorch 的 amp 模塊來實現(xiàn)的,。
這種方法特別設(shè)計,,以確保一些關(guān)鍵的計算步驟(例如 softmax 函數(shù))始終以最高精度執(zhí)行,,以保證訓(xùn)練過程的穩(wěn)定性,。
與此同時,其他大部分計算則使用一種稱為 bfloat16 的半精度格式,,以減少內(nèi)存使用并提高計算效率,。
在特定配置中,每個 GPU 上的模型權(quán)重和優(yōu)化器狀態(tài)都以最高精度保存,。
只有在執(zhí)行模型的前向傳播和反向傳播,,即計算模型的輸出和更新權(quán)重時,每個 Transformer 模塊內(nèi)的權(quán)重才會臨時轉(zhuǎn)換為 bfloat16 格式,。
此外,,各個 GPU 間同步梯度更新時,也會以最高精度進行,,以確保訓(xùn)練質(zhì)量,。
優(yōu)化器
研究人員采用了 AdamW 優(yōu)化器來調(diào)整模型參數(shù)。
無論模型規(guī)模大小如何,,研究人員都會在訓(xùn)練初期的 5000 步(大約處理 21B 個 token)內(nèi)逐漸增加學(xué)習(xí)率,,這一過程稱為學(xué)習(xí)率預(yù)熱。
預(yù)熱結(jié)束后,,學(xué)習(xí)率將按線性規(guī)律逐漸減少,,直到降至最高學(xué)習(xí)率的十分之一。
此外,,研究人員還會對模型參數(shù)的梯度進行裁剪,,確保其總的 L1 范數(shù)不會超過 1.0。
數(shù)據(jù)集
研究人員利用開放數(shù)據(jù)集 Dolma 中的一個 2T token 的樣本,,構(gòu)建了他們的訓(xùn)練數(shù)據(jù)集,。
研究人員將每篇文檔的 token 連接起來,每篇文檔的末尾都會加上一個特殊的 EOS token,,接著將這些 token 分成每組 2048 個,,形成訓(xùn)練樣本。
這些訓(xùn)練樣本在每次訓(xùn)練時都會以同樣的方式進行隨機打亂,。研究人員還提供了一些工具,,使得任何人都可以復(fù)原每個訓(xùn)練批次的具體數(shù)據(jù)順序和組成。
研究人員已經(jīng)發(fā)布的所有模型至少都經(jīng)過了一輪(2T token)的訓(xùn)練,。其中一些模型還進行了額外的訓(xùn)練,,即在數(shù)據(jù)上進行第二輪訓(xùn)練,但采用了不同的隨機打亂順序,。
根據(jù)之前的研究,,這樣重復(fù)使用少量數(shù)據(jù)的影響是微乎其微的。
英偉達和 AMD 都要 YES,!
為了確保代碼庫能夠同時在英偉達和 AMD 的 GPU 上都能高效運行,,研究人員選擇了兩個不同的集群進行了模型訓(xùn)練測試:
利用 LUMI 超級計算機,,研究人員部署了最多 256 個節(jié)點,每個節(jié)點搭載了 4 張 AMD MI250X GPU,,每張 GPU 擁有 128GB 內(nèi)存和 800Gbps 的數(shù)據(jù)傳輸速率,。
通過 MosaicML (Databricks) 的支持,研究人員使用了 27 個節(jié)點,,每個節(jié)點配備了 8 張英偉達 A100 GPU,,每張 GPU 擁有 40GB 內(nèi)存和 800Gbps 的數(shù)據(jù)傳輸速率。
雖然研究人員為了提高訓(xùn)練效率對批大小進行了微調(diào),,但在完成 2T token 的評估后,,兩個集群的性能幾乎沒有差異。
總結(jié)
與以往大多數(shù)僅僅提供模型權(quán)重和推理代碼的模型不同,,研究人員開源了 OLMo 的全部內(nèi)容,,包括訓(xùn)練數(shù)據(jù)、訓(xùn)練和評估代碼,,以及訓(xùn)練日志,、實驗結(jié)果、重要發(fā)現(xiàn)以及 Weights & Biases 的記錄等等,。
此外,,團隊正在研究如何通過指令優(yōu)化和不同類型的強化學(xué)習(xí)(RLHF)來改進 OLMo。而這些微調(diào)代碼,、數(shù)據(jù)和經(jīng)過微調(diào)后的模型也都會被開源,。
研究人員致力于持續(xù)支持和發(fā)展 OLMo 及其框架,推動開放語言模型(LM)的發(fā)展,,助力開放研究社區(qū)的發(fā)展,。為此,研究人員計劃引入更多不同規(guī)模的模型,、多種模態(tài),、數(shù)據(jù)集、安全措施和評估方法,,豐富 OLMo 家族,。
他們希望通過今后持續(xù)進行的徹底開源工作,增強開源研究社區(qū)的力量,,并引發(fā)新一輪的創(chuàng)新浪潮,。