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

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

我希望這篇文章對許多人來說既有趣又有教育意義。

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

在LLMs時代的硬件抽獎

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

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

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

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

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

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

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

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

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

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

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

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

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

GPU與TPU的對比

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

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

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

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

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

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

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

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

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

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

代碼在野外

毫無疑問，我有史以來最喜歡的代碼庫是T5X和Mesh Tensorflow（命名張量太贊了），但這些選項很快就變得不可行，原因是：1）它們在谷歌之外得不到太多支持，2）它們已經(jīng)有點被棄用了，3）對于我們團(tuán)隊中的非谷歌人員來說并不友好。

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

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

另外，我從來不知道更改模型并行性的能力不是自動（免費）的，直到某些代碼庫要求我編寫一個轉(zhuǎn)換器來更改模型的并行性。對我來說這絕對是一個WTF時刻。

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

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

少一點原則，多一點Yolo

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

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

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

簡而言之

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

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

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

關(guān)于作者：

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

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

他在Google任職期間，參與了許多產(chǎn)業(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é)人。