鑒于我們在Reka成功地培訓了相當強大的多模態語言模型，許多人對從零開始建立基礎設施并訓練大型語言和多模態模型的經驗特別感興趣。

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

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

在荒野中訓練LLMs（圖片由 Dall-E 生成）

在LLMs時代的硬件抽獎

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

人們總是認為這只是一個加速器選擇（例如TPUs vs GPUs等）的問題/爭論，而所有GPU集群都是平等的。但對我們來說，這很快就被證明是錯誤的。當我們嘗試不同的服務提供商時，我們發現即使是相同的硬件，例如GPU（H100），硬件質量的差異也是巨大的。請注意，這里的硬件指的是整個集群的質量，而不一定是芯片或加速器本身。就像彩票一樣。基本上：

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

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

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

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

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

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

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

每個服務提供商也提供不同級別的支持。這些支持從禮貌到漠不關心，從“ChatGPT風格”的套話回復再到責怪用戶每一件事情都出錯。

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

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

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

GPU與TPU的對比

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

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

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

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

多集群設置的痛苦

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

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

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

目前，我們開發了許多內部工作流來減輕許多這些問題，并且正在繼續朝著世界一流的實驗基礎設施的黃金標準邁進。

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

代碼在野外

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

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

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

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

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

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

少一點原則，多一點Yolo

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

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

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

簡而言之

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

總的來說，這只是我們如何在不到一年的時間里創辦一家公司、籌集一些資金、購買一些芯片并匹配 Gemini pro/GPT 3.5 并超越許多其他公司的故事的一小部分，而一切都從頭開始。

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

關于作者：

Yi Tay是Reka的聯合創始人兼首席科學家。Reka致力于構建先進的生成式 AI 模型。

此前，他曾是Google Brain的高級研究科學家，專注于大型語言模型和人工智能領域。在加入Brain團隊之前，他是Google Research美國研究團隊的技術負責人，致力于Transformer的擴展和架構研究，并為約20個產品發布做出了貢獻。

他在Google任職期間，參與了許多產業級的大語言模型的開發，如PaLM、UL2、Flan-{PaLM/UL2/T5}、LaMDA/Bard、MUM等。他還為大型多模態模型，如ViT-22B和PaLI-X做出了貢獻。他還是新PaLM-2和PaLM API的建模聯合負責人。