鑒于我們?cè)赗eka成功地培訓(xùn)了相當(dāng)強(qiáng)大的多模態(tài)語(yǔ)言模型，許多人對(duì)從零開始建立基礎(chǔ)設(shè)施并訓(xùn)練大型語(yǔ)言和多模態(tài)模型的經(jīng)驗(yàn)特別感興趣。

我在社交媒體上經(jīng)常抱怨外部（Google之外）的基礎(chǔ)設(shè)施和代碼，這讓人們對(duì)我在荒野中錯(cuò)過(guò)了什么，以及我對(duì)什么討厭/喜歡非常好奇。所以終于有了這篇文章。這篇博客文章揭示了挑戰(zhàn)和經(jīng)驗(yàn)教訓(xùn)。

我希望這篇文章對(duì)許多人來(lái)說(shuō)既有趣又有教育意義。

在荒野中訓(xùn)練LLMs（圖片由 Dall-E 生成）

在LLMs時(shí)代的硬件抽獎(jiǎng)

?訓(xùn)練模型的第一個(gè)必要條件是獲取計(jì)算資源。這看起來(lái)似乎很簡(jiǎn)單，也足夠容易。然而，最大的驚喜是計(jì)算資源提供商的不穩(wěn)定性，以及集群、加速器及其連接的質(zhì)量取決于來(lái)源的巨大差異。

人們總是認(rèn)為這只是一個(gè)加速器選擇（例如TPUs vs GPUs等）的問(wèn)題/爭(zhēng)論，而所有GPU集群都是平等的。但對(duì)我們來(lái)說(shuō)，這很快就被證明是錯(cuò)誤的。當(dāng)我們嘗試不同的服務(wù)提供商時(shí)，我們發(fā)現(xiàn)即使是相同的硬件，例如GPU（H100），硬件質(zhì)量的差異也是巨大的。請(qǐng)注意，這里的硬件指的是整個(gè)集群的質(zhì)量，而不一定是芯片或加速器本身。就像彩票一樣?；旧希?/p>

并非所有硬件都是一樣的。不同硬件提供商的集群質(zhì)量差異如此之大，以至于如何訓(xùn)練出良好模型會(huì)成為一場(chǎng)真正的抽獎(jiǎng)。簡(jiǎn)而言之，這是LLM時(shí)代的硬件抽獎(jiǎng)。

更具體地說(shuō)，我們從多家計(jì)算資源提供商租用了幾個(gè)集群，每個(gè)集群都有數(shù)百至數(shù)千個(gè)芯片。我們看到的集群范圍從尚可（只是一些煩人的問(wèn)題，只需花費(fèi)一些 工程師的時(shí)間就可以解決）到完全無(wú)法使用的集群，由于各種原因每隔幾個(gè)小時(shí)就會(huì)失敗。具體來(lái)說(shuō)，一些集群的節(jié)點(diǎn)每隔N個(gè)小時(shí)就會(huì)出現(xiàn)故障，問(wèn)題涉及到布線問(wèn)題（其中N是不合理小的數(shù)字）、GPU硬件錯(cuò)誤等。更令人驚訝的是，同一提供商的每個(gè)集群在穩(wěn)健性方面也可能截然不同。

與此同時(shí)，即使其他一些集群可能擁有明顯更穩(wěn)定的節(jié)點(diǎn)，它們可能會(huì)遭受到I/O和文件系統(tǒng)糟糕的影響，甚至保存檢查點(diǎn)都可能導(dǎo)致超時(shí)或花費(fèi)大量時(shí)間降低集群利用率。其他一些計(jì)算資源可能需要完全不同的軟件層才能運(yùn)行，對(duì)于帶有自己代碼庫(kù)的團(tuán)隊(duì)來(lái)說(shuō)并不友好，需要額外的遷移成本來(lái)運(yùn)行實(shí)驗(yàn)或大型作業(yè)。

沒(méi)有完美的東西！但有些肯定比其他的要糟糕得多。

最令人沮喪的部分是什么呢？幾乎不可能事先真正了解，尤其是在一切都在瘋狂進(jìn)行的情況下，人們將得到什么樣的硬件以及體驗(yàn)的魯棒性/容錯(cuò)性。

除此之外，你還無(wú)法知道供應(yīng)商是否會(huì)延遲交貨幾個(gè)月，讓你處于無(wú)法從其他來(lái)源獲取資源的困境中數(shù)周甚至數(shù)月。有些提供商還會(huì)意外刪除你的檢查點(diǎn)ˉ\_(ツ)_/ˉ。

我提到了不同集群的模型浮點(diǎn)操作（Model Flop Utilisation，MFU）也會(huì)不同嗎？如果不幸地找到布線不良或其他問(wèn)題的提供商，這將是一筆不可忽略的計(jì)算資源浪費(fèi)。具有非常次優(yōu)文件系統(tǒng)的系統(tǒng)在團(tuán)隊(duì)成員開始在集群之間傳輸大量數(shù)據(jù)時(shí)，訓(xùn)練運(yùn)行的MFU會(huì)瞬間下降。

每個(gè)服務(wù)提供商也提供不同級(jí)別的支持。這些支持從禮貌到漠不關(guān)心，從“ChatGPT風(fēng)格”的套話回復(fù)再到責(zé)怪用戶每一件事情都出錯(cuò)。

總的來(lái)說(shuō)，我們嘗試過(guò)的每個(gè)集群都感覺(jué)有著自己的氛圍、掙扎和失敗模式。幾乎每個(gè)集群似乎都需要為自己的一系列問(wèn)題提供自己的修復(fù)措施 - 有些問(wèn)題比其他問(wèn)題更容忍。也就是說(shuō)，我們已經(jīng)學(xué)到了備用方案的重要性，快速為任何集群找到修復(fù)措施可能是關(guān)鍵。

在過(guò)去的幾個(gè)月里，我們已經(jīng)做了很多工作，只是為了確保事物能夠使用，例如，圍繞監(jiān)控工具、高效的檢查點(diǎn)，以及各種其他優(yōu)化，甚至安裝我們的自定義文件系統(tǒng)以實(shí)現(xiàn)可擴(kuò)展的數(shù)據(jù)存儲(chǔ) - 并且這只是實(shí)際需求的冰山一角。

這些工具組合帶來(lái)了MFU 的顯著改進(jìn)，同時(shí)還最大限度地減少了糟糕硬件帶來(lái)的停機(jī)時(shí)間。

GPU與TPU的對(duì)比

在Reka，我們大部分時(shí)間都是在GPU上訓(xùn)練我們的模型。就我個(gè)人而言，在Reka之前的Google生活中，我一直使用TPU進(jìn)行大型語(yǔ)言模型訓(xùn)練。CUDA和nccl對(duì)我來(lái)說(shuō)是最陌生的東西。（我只是從一個(gè)曾在Nvidia工作過(guò)的同事那里學(xué)到它的發(fā)音是“Nickel” 哈哈）

與我在Google使用TPU的經(jīng)驗(yàn)相比，我完全驚訝于GPU的故障率。事實(shí)上，我實(shí)際上不太記得在Google使用TPU時(shí)出現(xiàn)過(guò)太多故障，即使是在大規(guī)模運(yùn)行時(shí)，盡管我不確定我是否被這個(gè)簡(jiǎn)直太好的基礎(chǔ)設(shè)施和專門的硬件團(tuán)隊(duì)的強(qiáng)大保護(hù)著。事實(shí)上，UL2 20B模型（在Google）是通過(guò)不小心讓作業(yè)運(yùn)行了一個(gè)月來(lái)訓(xùn)練的。它從未失敗過(guò)。如果這是在GPU領(lǐng)域，它肯定會(huì)在最初的幾天內(nèi)就失敗了。

話雖如此，我認(rèn)為這更多地與管理加速器的硬件團(tuán)隊(duì)的能力有關(guān)，而不是與底層芯片有關(guān)。擁有良好的硬件支持（來(lái)自計(jì)算資源提供商）很重要。很多事情取決于他們是否真的很有能力，這加強(qiáng)了“硬件抽獎(jiǎng)”的概念。

GPU領(lǐng)域感覺(jué)很奇怪。與在TPU架構(gòu)中作為分布式訓(xùn)練的首要選擇不同，在GPU領(lǐng)域，多節(jié)點(diǎn)訓(xùn)練似乎更像是事后想到的。在GPU領(lǐng)域，似乎不同的提供商以不同的方式將它們連接起來(lái)以實(shí)現(xiàn)多節(jié)點(diǎn)訓(xùn)練，這導(dǎo)致了在不同地方進(jìn)行事情的方式存在很大差異。雖然我不是硬件方面的專家，但這是我得到的印象。

多集群設(shè)置的痛苦

我職業(yè)生涯的大部分時(shí)間都花在了谷歌的基礎(chǔ)設(shè)施上，它主要運(yùn)行在Borg、Xmanager和Colossus上，從任何地方都可以訪問(wèn)所有內(nèi)容。因此，實(shí)際上需要在不同的集群中設(shè)置新環(huán)境的概念對(duì)我來(lái)說(shuō)是陌生的。

在當(dāng)前世界中，除非專門為單個(gè)位置的大量集群構(gòu)建，否則擁有多個(gè)加速器池集群似乎是不可避免的。更具體地說(shuō)，GPU供應(yīng)（或缺乏）也自然導(dǎo)致了這種以集群采購(gòu)的模式，其性質(zhì)是分散的。訓(xùn)練大型模型還需要大量的數(shù)據(jù)，即使只是將它們移動(dòng)一下也會(huì)帶來(lái)很多不便。與此同時(shí)，通常復(fù)制數(shù)據(jù)也不是直接的，在極大規(guī)模上是難以承受的。

顯然，理想情況下是有一種專門建立的編排層，可以將作業(yè)發(fā)送到不同的服務(wù)器。我相信許多大型AI優(yōu)先公司通常都有某種基礎(chǔ)設(shè)施來(lái)改善AI研究人員的生活質(zhì)量。然而，對(duì)于一個(gè)剛起步的精益的新創(chuàng)企業(yè)來(lái)說(shuō)，構(gòu)建這種復(fù)雜和花哨的ML訓(xùn)練基礎(chǔ)設(shè)施實(shí)際上是不太可能的。

目前，我們開發(fā)了許多內(nèi)部工作流來(lái)減輕許多這些問(wèn)題，并且正在繼續(xù)朝著世界一流的實(shí)驗(yàn)基礎(chǔ)設(shè)施的黃金標(biāo)準(zhǔn)邁進(jìn)。

（有人告訴我，這種雜亂的設(shè)置或多或少是非頂級(jí)/大公司的常態(tài)）。

代碼在野外

毫無(wú)疑問(wèn)，我有史以來(lái)最喜歡的代碼庫(kù)是T5X和Mesh Tensorflow（命名張量太贊了），但這些選項(xiàng)很快就變得不可行，原因是：1）它們?cè)诠雀柚獾貌坏教嘀С郑?）它們已經(jīng)有點(diǎn)被棄用了，3）對(duì)于我們團(tuán)隊(duì)中的非谷歌人員來(lái)說(shuō)并不友好。

我們最終選擇了一些常規(guī)的、看起來(lái)比較穩(wěn)定且更受歡迎的選項(xiàng)（即，pytorch），這對(duì)大多數(shù)團(tuán)隊(duì)成員來(lái)說(shuō)更易于接觸（除了我哈哈）。在我最初的幾個(gè)月里，我在pip、git、docker和所有這些野生的東西感到困惑。不過(guò)話說(shuō)回來(lái)，我不太確定使用谷歌的代碼庫(kù)外部會(huì)有多穩(wěn)定或用戶友好（我猜這會(huì)非常令人討厭）。

坦率地說(shuō)，我不得不說(shuō)外部代碼庫(kù)的質(zhì)量明顯落后于我在谷歌習(xí)慣的那些。主要是因?yàn)楣雀鑳?nèi)部的代碼庫(kù)往往是由機(jī)器學(xué)習(xí)界的明星們親自編寫的（例如，Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung等等），并且與我在外部嘗試過(guò)的代碼庫(kù)相比，感覺(jué)更好（例如，優(yōu)越的氛圍）。特別是，當(dāng)我嘗試使用其他公司構(gòu)建的東西時(shí)，我發(fā)現(xiàn)自己對(duì)代碼質(zhì)量感到非常惱火（有些比其他更糟糕??）。

另外，我從來(lái)不知道更改模型并行性的能力不是自動(dòng)（免費(fèi)）的，直到某些代碼庫(kù)要求我編寫一個(gè)轉(zhuǎn)換器來(lái)更改模型的并行性。對(duì)我來(lái)說(shuō)這絕對(duì)是一個(gè)WTF時(shí)刻。

另一個(gè)引人注目的事情是這些代碼庫(kù)對(duì)大規(guī)模編碼器-解碼器訓(xùn)練甚至 prefixLM 訓(xùn)練的支持非常少。為此，即使 Flash Attention 也一直拒絕為 prefixLM 訓(xùn)練（即自定義掩碼）提供支持，盡管出于某種原因?qū)ζ?github 問(wèn)題有合理的需求。

我知道我應(yīng)該使用Jax。一個(gè)朋友剛剛因我使用pytorch而羞辱了我，但這是一個(gè)創(chuàng)業(yè)公司，我們決定快速行動(dòng)。對(duì)此我很抱歉，下次我們會(huì)做得更酷一些。我對(duì)這個(gè)事實(shí)并不感到自豪。

少一點(diǎn)原則，多一點(diǎn)Yolo

系統(tǒng)地?cái)U(kuò)展模型通常需要以一種原則性的方式從小規(guī)模到大規(guī)模進(jìn)行，即在多個(gè)階段運(yùn)行實(shí)驗(yàn)（從1B到8B到64B到300B等），然后選擇出獲勝者并持續(xù)擴(kuò)展它們。在創(chuàng)業(yè)公司中，我們擁有的計(jì)算資源遠(yuǎn)遠(yuǎn)不足以進(jìn)行大規(guī)模的參數(shù)掃描。最終，我們不得不依靠許多“Yolo”運(yùn)行（幸運(yùn)的是，結(jié)果還不錯(cuò)）。

最終，我們只需要進(jìn)行了非常少量的小規(guī)模和較短的割裂實(shí)驗(yàn)，就能夠獲得強(qiáng)大的21B Reka Flash模型和7B邊緣模型（以及即將推出的最大核心模型）。在有限數(shù)量的運(yùn)行中找到一個(gè)可靠的配方是具有挑戰(zhàn)性的，因?yàn)樗阉骺臻g非常巨大，必須同時(shí)改變?cè)S多變量。為了做到這一點(diǎn)，人們必須放棄大型科技公司的系統(tǒng)性，大量依賴于“Yolo”、直覺(jué)和本能。

值得慶幸的是，我（以及我們團(tuán)隊(duì)中的許多人）在我們的機(jī)器學(xué)習(xí)職業(yè)生涯中已經(jīng)積累了相當(dāng)多的直覺(jué)，以便在相當(dāng)短的嘗試次數(shù)內(nèi)得到正確的結(jié)果。雖然我們之前在以往的工作中訓(xùn)練過(guò)非常好的模型，但是訓(xùn)練基礎(chǔ)設(shè)施、數(shù)據(jù)、新思想的融入以及其他環(huán)境問(wèn)題的差異仍然可能導(dǎo)致結(jié)果的巨大差異。盡管如此，強(qiáng)大的先驗(yàn)知識(shí)有助于顯著減少搜索空間，并且可能是我們能夠用如此少的嘗試、資源和實(shí)驗(yàn)訓(xùn)練出非常強(qiáng)大的模型的最簡(jiǎn)單的解釋之一。

簡(jiǎn)而言之

探索荒野中的事物是一次有趣的經(jīng)歷。不幸的是，這并不是一帆風(fēng)順的。計(jì)算資源的稀缺以及不可靠的計(jì)算資源提供商會(huì)使事情變得比預(yù)期的要困難得多，但我們很高興用強(qiáng)大的技術(shù)實(shí)力度過(guò)了難關(guān)。

總的來(lái)說(shuō)，這只是我們?nèi)绾卧诓坏揭荒甑臅r(shí)間里創(chuàng)辦一家公司、籌集一些資金、購(gòu)買一些芯片并匹配 Gemini pro/GPT 3.5 并超越許多其他公司的故事的一小部分，而一切都從頭開始。

還有更多要寫的，比如數(shù)據(jù)管道、人工評(píng)估等，但這已經(jīng)是一篇很長(zhǎng)的文章了。下次再寫。

關(guān)于作者：

Yi Tay是Reka的聯(lián)合創(chuàng)始人兼首席科學(xué)家。Reka致力于構(gòu)建先進(jìn)的生成式 AI 模型。

此前，他曾是Google Brain的高級(jí)研究科學(xué)家，專注于大型語(yǔ)言模型和人工智能領(lǐng)域。在加入Brain團(tuán)隊(duì)之前，他是Google Research美國(guó)研究團(tuán)隊(duì)的技術(shù)負(fù)責(zé)人，致力于Transformer的擴(kuò)展和架構(gòu)研究，并為約20個(gè)產(chǎn)品發(fā)布做出了貢獻(xiàn)。

他在Google任職期間，參與了許多產(chǎn)業(yè)級(jí)的大語(yǔ)言模型的開發(fā)，如PaLM、UL2、Flan-{PaLM/UL2/T5}、LaMDA/Bard、MUM等。他還為大型多模態(tài)模型，如ViT-22B和PaLI-X做出了貢獻(xiàn)。他還是新PaLM-2和PaLM API的建模聯(lián)合負(fù)責(zé)人。