艾倫人工智能研究所等 5 機(jī)構(gòu)最近公布了史上最全的開(kāi)源模型「OLMo」，公開(kāi)了模型的模型權(quán)重、完整訓(xùn)練代碼、數(shù)據(jù)集和訓(xùn)練過(guò)程，為以后開(kāi)源社區(qū)的工作設(shè)立了新的標(biāo)桿。

多年來(lái)，語(yǔ)言模型一直是自然語(yǔ)言處理（NLP）技術(shù)的核心，考慮到模型背后的巨大商業(yè)價(jià)值，最大最先進(jìn)的模型的技術(shù)細(xì)節(jié)都是不公開(kāi)的。

現(xiàn)在，真?完全開(kāi)源的大模型來(lái)了！

來(lái)自艾倫人工智能研究所、華盛頓大學(xué)、耶魯大學(xué)、紐約大學(xué)和卡內(nèi)基梅隆大學(xué)的研究人員，聯(lián)合發(fā)表了一項(xiàng)足以載入 AI 開(kāi)源社區(qū)史冊(cè)的工作 ——

他們幾乎將從零開(kāi)始訓(xùn)練一個(gè)大模型過(guò)程中的一切數(shù)據(jù)和資料都開(kāi)源了！

具體來(lái)說(shuō)，艾倫人工智能研究所推出的這個(gè)開(kāi)放大語(yǔ)言模型（Open Language Model，OLMo）實(shí)驗(yàn)和訓(xùn)練平臺(tái)，則提供了一個(gè)完全開(kāi)源的大模型，以及所有和訓(xùn)練開(kāi)發(fā)這個(gè)模型有關(guān)的數(shù)據(jù)和技術(shù)細(xì)節(jié) ——

訓(xùn)練和建模：它包括完整的模型權(quán)重、訓(xùn)練代碼、訓(xùn)練日志、消融研究、訓(xùn)練指標(biāo)和推理代碼。

預(yù)訓(xùn)練語(yǔ)料：一個(gè)包含了高達(dá) 3T token 的預(yù)訓(xùn)練開(kāi)源語(yǔ)料庫(kù)，以及產(chǎn)生這些訓(xùn)練數(shù)據(jù)的代碼。

模型參數(shù)：OLMo 框架提供了四個(gè)不同架構(gòu)、優(yōu)化器和訓(xùn)練硬件體系下的 7B 大小的模型，以及一個(gè) 1B 大小的模型，所有模型都在至少 2T token 上進(jìn)行了訓(xùn)練。

同時(shí)，也提供了用于模型推理的代碼、訓(xùn)練過(guò)程的各項(xiàng)指標(biāo)以及訓(xùn)練日志。

評(píng)估工具：公開(kāi)了開(kāi)發(fā)過(guò)程中的評(píng)估工具套件，包括每個(gè)模型訓(xùn)練過(guò)程中每 1000 step 中包含的超過(guò) 500 個(gè)的檢查點(diǎn)以及評(píng)估代碼。

所有數(shù)據(jù)都在 apache 2.0 下授權(quán)使用（免費(fèi)商用）。

如此徹底的開(kāi)源，似乎是給開(kāi)源社區(qū)打了個(gè)樣 —— 以后不像我這樣開(kāi)源的，就別說(shuō)自己是開(kāi)源模型了。

性能評(píng)估

從核心的評(píng)估結(jié)果來(lái)看，OLMo-7B 與同類(lèi)開(kāi)源模型相比略勝一籌。在前 9 項(xiàng)評(píng)測(cè)中，OLMo-7B 有 8 項(xiàng)排名前三，其中有 2 項(xiàng)超越了其他所有模型。

在很多生成任務(wù)或閱讀理解任務(wù)（例如 truthfulQA）上，OLMo-7B 都超過(guò)了 Llama 2，但在一些熱門(mén)的問(wèn)答任務(wù)（如 MMLU 或 Big-bench Hard）上表現(xiàn)則要差一些。

前 9 個(gè)任務(wù)是研究人員對(duì)預(yù)訓(xùn)練模型的內(nèi)部評(píng)估標(biāo)準(zhǔn)，而下面三個(gè)任務(wù)則是為了完善 HuggingFace Open LLM 排行榜而加入的

下圖展示了 9 個(gè)核心任務(wù)準(zhǔn)確率的變化趨勢(shì)。除了 OBQA 外，隨著 OLMo-7B 接受更多數(shù)據(jù)的訓(xùn)練，幾乎所有任務(wù)的準(zhǔn)確率都呈現(xiàn)上升趨勢(shì)。

與此同時(shí)，OLMo 1B 與其同類(lèi)模型的核心評(píng)估結(jié)果表明，OLMo 與它們處于同一水平。

通過(guò)使用艾倫 AI 研究所的 Paloma（一個(gè)基準(zhǔn)測(cè)試）和可獲取的檢查點(diǎn)，研究人員分析了模型預(yù)測(cè)語(yǔ)言能力與模型規(guī)模因素（例如訓(xùn)練的 token 數(shù)量）之間的關(guān)系。

可以看到，OLMo-7B 在性能上與主流模型持平。其中，每字節(jié)比特?cái)?shù)（Bits per Byte）越低越好。

通過(guò)這些分析，研究人員發(fā)現(xiàn)模型在處理不同數(shù)據(jù)源時(shí)的效率差異較大，這主要取決于模型訓(xùn)練數(shù)據(jù)與評(píng)估數(shù)據(jù)的相似度。

特別地，OLMo-7B 在主要基于 Common Crawl 的數(shù)據(jù)源上表現(xiàn)出色（比如 C4）。

不過(guò)，在與網(wǎng)絡(luò)抓取文本關(guān)系不大的數(shù)據(jù)源上，如 WikiText-103、M2D2 S2ORC 和 M2D2 Wikipedia，OLMo-7B 與其他模型相比效率較低。

RedPajama 的評(píng)估也體現(xiàn)了相似的趨勢(shì)，可能是因?yàn)樗?7 個(gè)領(lǐng)域中只有 2 個(gè)來(lái)源于 Common Crawl，且 Paloma 對(duì)每個(gè)數(shù)據(jù)源中的各個(gè)領(lǐng)域給予了相同的權(quán)重。

鑒于像 Wikipedia 和 arXiv 論文這樣的精選數(shù)據(jù)源提供的異質(zhì)數(shù)據(jù)遠(yuǎn)不如網(wǎng)絡(luò)抓取文本豐富，隨著預(yù)訓(xùn)練數(shù)據(jù)集的不斷擴(kuò)大，維持對(duì)這些語(yǔ)言分布的高效率會(huì)很更加困難。

OLMo 架構(gòu)

在模型的架構(gòu)方面，團(tuán)隊(duì)基于的是 decoder-only 的 Transformer 架構(gòu)，并采用了 PaLM 和 Llama 使用的 SwiGLU 激活函數(shù)，引入了旋轉(zhuǎn)位置嵌入技術(shù)（RoPE），并改進(jìn)了 GPT-NeoX-20B 的基于字節(jié)對(duì)編碼（BPE）的分詞器，以減少模型輸出中的個(gè)人可識(shí)別信息。

此外，為了保證模型的穩(wěn)定性，研究人員沒(méi)有使用偏置項(xiàng)（這一點(diǎn)與 PaLM 的處理方式相同）。

預(yù)訓(xùn)練數(shù)據(jù)集：Dolma

雖然研究人員在獲取模型參數(shù)方面取得了一定的進(jìn)展，但開(kāi)源社區(qū)目前預(yù)訓(xùn)練數(shù)據(jù)集的開(kāi)放程度還遠(yuǎn)遠(yuǎn)不夠。

之前的預(yù)訓(xùn)練數(shù)據(jù)往往不會(huì)隨著模型的開(kāi)源而公開(kāi)（閉源模型就更不用說(shuō)了）。

而且有關(guān)這些數(shù)據(jù)的說(shuō)明文檔也常常缺乏足夠的細(xì)節(jié)，但是這些細(xì)節(jié)對(duì)于想要復(fù)現(xiàn)研究或完全理解相關(guān)工作至關(guān)重要。

這一情況加大了語(yǔ)言模型研究的難度 —— 比如，了解訓(xùn)練數(shù)據(jù)如何影響模型能力和其局限性。

為了推動(dòng)語(yǔ)言模型預(yù)訓(xùn)練領(lǐng)域的開(kāi)放研究，研究人員構(gòu)建并公開(kāi)了預(yù)訓(xùn)練數(shù)據(jù)集 Dolma。

這是一個(gè)包含了從 7 種不同數(shù)據(jù)來(lái)源獲取的 3 萬(wàn)億個(gè) token 的多樣化、多源語(yǔ)料庫(kù)。

這些數(shù)據(jù)源一方面在大規(guī)模語(yǔ)言模型預(yù)訓(xùn)練中常見(jiàn)，另一方面也能被普通大眾所接觸。

Dolma 的構(gòu)建過(guò)程包括六個(gè)步驟：語(yǔ)言過(guò)濾、質(zhì)量過(guò)濾、內(nèi)容過(guò)濾、去重、多源混合和 token 化。

在整理和最終發(fā)布 Dolma 過(guò)程中，研究人員確保各數(shù)據(jù)源的文檔保持獨(dú)立。

他們還開(kāi)源了一套高效的數(shù)據(jù)整理工具，這套工具能夠幫助進(jìn)一步研究 Dolma、復(fù)制成果，并簡(jiǎn)化預(yù)訓(xùn)練語(yǔ)料庫(kù)的整理工作。

此外，研究人員也開(kāi)源了 WIMBD 工具，以助于數(shù)據(jù)集分析。

訓(xùn)練 OLMo

分布式訓(xùn)練框架

研究人員利用 PyTorch 的 FSDP 框架和 ZeRO 優(yōu)化器策略來(lái)訓(xùn)練模型。這種方法通過(guò)將模型的權(quán)重和它們對(duì)應(yīng)的優(yōu)化器狀態(tài)在多個(gè) GPU 中進(jìn)行分割，從而有效減少了內(nèi)存的使用量。

在處理高達(dá) 7B 規(guī)模的模型時(shí)，這項(xiàng)技術(shù)使研究人員能夠在每個(gè) GPU 上處理 4096 個(gè) token 的微批大小，以實(shí)現(xiàn)更高效的訓(xùn)練。

對(duì)于 OLMo-1B 和 7B 模型，研究人員固定使用大約 4M token（2048 個(gè)數(shù)據(jù)實(shí)例，每個(gè)實(shí)例包含 2048 個(gè) token 的序列）的全局批大小。

而對(duì)于目前正在訓(xùn)練中的 OLMo-65B 模型，研究人員采用了一個(gè)批大小預(yù)熱策略，起始于大約 2M token（1024 個(gè)數(shù)據(jù)實(shí)例），之后每增加 100B token，批大小翻倍，直至最終達(dá)到大約 16M token（8192 個(gè)數(shù)據(jù)實(shí)例）的規(guī)模。

為了加快模型訓(xùn)練的速度，研究人員采用了混合精度訓(xùn)練的技術(shù)，這一技術(shù)是通過(guò) FSDP 的內(nèi)部配置和 PyTorch 的 amp 模塊來(lái)實(shí)現(xiàn)的。

這種方法特別設(shè)計(jì)，以確保一些關(guān)鍵的計(jì)算步驟（例如 softmax 函數(shù)）始終以最高精度執(zhí)行，以保證訓(xùn)練過(guò)程的穩(wěn)定性。

與此同時(shí)，其他大部分計(jì)算則使用一種稱為 bfloat16 的半精度格式，以減少內(nèi)存使用并提高計(jì)算效率。

在特定配置中，每個(gè) GPU 上的模型權(quán)重和優(yōu)化器狀態(tài)都以最高精度保存。

只有在執(zhí)行模型的前向傳播和反向傳播，即計(jì)算模型的輸出和更新權(quán)重時(shí)，每個(gè) Transformer 模塊內(nèi)的權(quán)重才會(huì)臨時(shí)轉(zhuǎn)換為 bfloat16 格式。

此外，各個(gè) GPU 間同步梯度更新時(shí)，也會(huì)以最高精度進(jìn)行，以確保訓(xùn)練質(zhì)量。

優(yōu)化器

研究人員采用了 AdamW 優(yōu)化器來(lái)調(diào)整模型參數(shù)。

無(wú)論模型規(guī)模大小如何，研究人員都會(huì)在訓(xùn)練初期的 5000 步（大約處理 21B 個(gè) token）內(nèi)逐漸增加學(xué)習(xí)率，這一過(guò)程稱為學(xué)習(xí)率預(yù)熱。

預(yù)熱結(jié)束后，學(xué)習(xí)率將按線性規(guī)律逐漸減少，直到降至最高學(xué)習(xí)率的十分之一。

此外，研究人員還會(huì)對(duì)模型參數(shù)的梯度進(jìn)行裁剪，確保其總的 L1 范數(shù)不會(huì)超過(guò) 1.0。

數(shù)據(jù)集

研究人員利用開(kāi)放數(shù)據(jù)集 Dolma 中的一個(gè) 2T token 的樣本，構(gòu)建了他們的訓(xùn)練數(shù)據(jù)集。

研究人員將每篇文檔的 token 連接起來(lái)，每篇文檔的末尾都會(huì)加上一個(gè)特殊的 EOS token，接著將這些 token 分成每組 2048 個(gè)，形成訓(xùn)練樣本。

這些訓(xùn)練樣本在每次訓(xùn)練時(shí)都會(huì)以同樣的方式進(jìn)行隨機(jī)打亂。研究人員還提供了一些工具，使得任何人都可以復(fù)原每個(gè)訓(xùn)練批次的具體數(shù)據(jù)順序和組成。

研究人員已經(jīng)發(fā)布的所有模型至少都經(jīng)過(guò)了一輪（2T token）的訓(xùn)練。其中一些模型還進(jìn)行了額外的訓(xùn)練，即在數(shù)據(jù)上進(jìn)行第二輪訓(xùn)練，但采用了不同的隨機(jī)打亂順序。

根據(jù)之前的研究，這樣重復(fù)使用少量數(shù)據(jù)的影響是微乎其微的。

英偉達(dá)和 AMD 都要 YES！

為了確保代碼庫(kù)能夠同時(shí)在英偉達(dá)和 AMD 的 GPU 上都能高效運(yùn)行，研究人員選擇了兩個(gè)不同的集群進(jìn)行了模型訓(xùn)練測(cè)試：

利用 LUMI 超級(jí)計(jì)算機(jī)，研究人員部署了最多 256 個(gè)節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)搭載了 4 張 AMD MI250X GPU，每張 GPU 擁有 128GB 內(nèi)存和 800Gbps 的數(shù)據(jù)傳輸速率。

通過(guò) MosaicML (Databricks) 的支持，研究人員使用了 27 個(gè)節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)配備了 8 張英偉達(dá) A100 GPU，每張 GPU 擁有 40GB 內(nèi)存和 800Gbps 的數(shù)據(jù)傳輸速率。

雖然研究人員為了提高訓(xùn)練效率對(duì)批大小進(jìn)行了微調(diào)，但在完成 2T token 的評(píng)估后，兩個(gè)集群的性能幾乎沒(méi)有差異。

總結(jié)

與以往大多數(shù)僅僅提供模型權(quán)重和推理代碼的模型不同，研究人員開(kāi)源了 OLMo 的全部?jī)?nèi)容，包括訓(xùn)練數(shù)據(jù)、訓(xùn)練和評(píng)估代碼，以及訓(xùn)練日志、實(shí)驗(yàn)結(jié)果、重要發(fā)現(xiàn)以及 Weights & Biases 的記錄等等。

此外，團(tuán)隊(duì)正在研究如何通過(guò)指令優(yōu)化和不同類(lèi)型的強(qiáng)化學(xué)習(xí)（RLHF）來(lái)改進(jìn) OLMo。而這些微調(diào)代碼、數(shù)據(jù)和經(jīng)過(guò)微調(diào)后的模型也都會(huì)被開(kāi)源。

研究人員致力于持續(xù)支持和發(fā)展 OLMo 及其框架，推動(dòng)開(kāi)放語(yǔ)言模型（LM）的發(fā)展，助力開(kāi)放研究社區(qū)的發(fā)展。為此，研究人員計(jì)劃引入更多不同規(guī)模的模型、多種模態(tài)、數(shù)據(jù)集、安全措施和評(píng)估方法，豐富 OLMo 家族。

他們希望通過(guò)今后持續(xù)進(jìn)行的徹底開(kāi)源工作，增強(qiáng)開(kāi)源研究社區(qū)的力量，并引發(fā)新一輪的創(chuàng)新浪潮。