作者 |  Yi Tay

編譯 | 云昭

出品 | 51CTO技術(shù)棧（微信號：blog51cto）

你敢相信嗎？一位前谷歌大佬，離職成立公司，不到一年，從頭訓(xùn)練出了“GPT3.5”/“Gemini Pro”，注意，后者是多模態(tài)大模型！ 

本文主人公Yi Tay，是一位市面上非常搶手的高性能大模型的大拿。他曾在谷歌Google Brain擔任高級研究科學家，專注于大型語言模型和人工智能的研究。在Google任職期間，曾經(jīng)為業(yè)內(nèi)許多知名的大型語言模型做出了貢獻，例如PaLM、UL2、Flan-{PaLM/UL2/T5}、LaMDA/Bard、MUM等。

另外，Yi還參與了大型多模態(tài)模型如ViT-22B和PaLI-X的研究，負責了新模型PaLM-2和PaLM API的建模工作。

去年3月，Yi離開了谷歌，創(chuàng)辦了一家大模型公司Reka，一直追求打造出令人驚嘆的前沿生成模型。

不到一年的時間，從一張卡都沒有，到推出了可以匹敵GPT3.5/Gemini Pro的大模型Reka。大模型訓(xùn)練、多模態(tài)大模型何其艱難？這期間，究竟發(fā)生了怎樣的離奇的事情呢？

Yi Tay 就此分享了幾點挑戰(zhàn)和教訓(xùn)，比如GPU問題百出，不如TPU、野雞代碼的折磨，“多一些YOLO，少一些原則”等等，非常有意思，值得諸君深思。

1.買算力如同買彩票！

訓(xùn)練模型的首要條件是獲取計算能力。這看起來再簡單不過了，但事實上這就好比買彩票一樣。

你的算力供應(yīng)商是不固定的，集群和加速器及其連接的質(zhì)量也隨著他們各自不同的廠商而帶來巨大差異。

你可能會問，那我們就選擇同一型號的GPU、TPU加速器等，集群也配置對等，不就完事了嗎？

事實總是啪啪打臉，當我們對不同的服務(wù)提供商進行采樣時，我們被結(jié)果震驚到了：

即使對于相同的硬件，即 GPU（我們用的H100），質(zhì)量的差異也很大。請注意，這里的硬件指的是整體集群質(zhì)量，而不一定是芯片或加速器本身。

這就跟買彩票一樣?；旧希?/p>

并非所有硬件都是一樣的。不同硬件提供商的集群質(zhì)量差異如此之大，以至于這實際上是一場彩票，與訓(xùn)練好模型需要經(jīng)歷多少痛苦有關(guān)。簡而言之，算力就是大模型時代的硬件彩票。

事情具體是這樣的，我們從多家計算提供商那里租用了一些集群，每個集群都有數(shù)百到數(shù)千個芯片。

集群問題百出，有的還說得過去，只是一些煩人的問題，只需花費一些 SWE 時間就可以解決，而有的集群則完全不可用，每隔幾個小時就會失敗，而且讓人抓馬的是，原因也各種不一樣。

比如，一些集群的節(jié)點每 N 小時就會出現(xiàn)故障，其問題就包括到底是布線問題（其中 N 過?。┻€是GPU 硬件出錯誤等。

最為讓人百思不得其解地是，同一廠商的每個集群在穩(wěn)健性方面也可能存在巨大差異。

同時，即使其他一些集群可能擁有更穩(wěn)定的節(jié)點，它們也可能會受到 I/O 和文件系統(tǒng)較差的影響，甚至保存檢查點也可能導(dǎo)致超時或大量時間減少集群利用率。

除了這些，有的不同供應(yīng)商來源的算力甚至需要完全不同的軟件層才能運行，并且對于自帶代碼庫的團隊非常不友好 ，因為這就意味著需要額外的遷移成本來運行實驗或大型作業(yè)。

最令人沮喪的部分？幾乎不可能真正提前知道，尤其是在一切逼瘋你的問題全都來臨時，究竟需要獲得什么樣的硬件？體驗的魯棒性/容錯性又該如何預(yù)估和保證？

總之一句話，情況沒有最差的，只有更差的！

就比如，供應(yīng)商也有延遲和放鴿子的情況，你無法判斷供應(yīng)商是否能按時交貨，而且僅僅只是發(fā)貨延遲了幾個月，供應(yīng)商自己也尷尬，他們也無法從其他來源采購，這種延遲情況從數(shù)周到數(shù)月不等。此外，某些提供商還會意外地刪除你的檢查點文件。

這還沒完，對于不同的集群，您還會獲得不同的模型失敗率利用率 (MFU)！如果不幸找到一個節(jié)點布線不良或存在其他問題的提供商，這就會導(dǎo)致一場價值不菲的算力浪費。

再有，當團隊成員開始跨集群傳輸大量數(shù)據(jù)時，具有非常次優(yōu)文件系統(tǒng)的系統(tǒng)的訓(xùn)練運行 MFU 將會下降。

此外，所有的服務(wù)供應(yīng)商的服務(wù)態(tài)度也都分三六九等，都有著不同級別的支持，從禮貌到漠不關(guān)心，從“chatgpt 式”的預(yù)設(shè)回復(fù)到將每一件出錯的事情歸咎于用戶。

一整套搞下來，最大的感觸就是，我們嘗試過的每個集群，都感覺它們有自己的氛圍、掙扎和失敗模式。幾乎每個集群都需要針對自己的一系列問題進行熱修復(fù)——有些問題比其他問題更容易忍受。

總之，就是為故障做保險非常重要，但關(guān)鍵之處在于如何為所有集群找到快速修復(fù)的方法。

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

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

2.跟TPU相比，GPU簡直菜雞

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

與我在谷歌使用 TPU 的經(jīng)歷相比，GPU 的故障率讓我完全大吃一驚。事實上，我印象中TPU 即便在大規(guī)模運行時也沒有失敗過，可能的確是谷歌基礎(chǔ)設(shè)施非常出色，擁有著絕對穩(wěn)健性，而且谷歌擁有著一支專門的硬件團隊。

事實上，UL2 20B模型（在 Google）是通過讓任務(wù)無特別維護地情況下運行一個月來進行訓(xùn)練的，它從未失敗過。如果這要是換成 GPU ，那么它肯定會在最初幾天內(nèi)就會宕掉。

也就是說，我認為這可能更多地取決于管理加速器的硬件團隊的能力，而不是底層芯片。擁有良好的硬件支持（來自計算提供商）非常重要。這在很大程度上取決于他們的實際能力，這強化了“硬件彩票”的概念。

GPU 領(lǐng)域感覺很奇怪。感覺多節(jié)點訓(xùn)練更像是事后的想法，而不是作為 TPU pod 上的一等公民的分布式訓(xùn)練。在 GPU 領(lǐng)域，感覺好像不同的提供商以不同的方式連接它們以實現(xiàn)多節(jié)點訓(xùn)練，這導(dǎo)致不同地方的工作方式存在很大差異。

雖然我不是硬件專家，但這就是我的印象。

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

我職業(yè)生涯的大部分時間都花在了 Google 基礎(chǔ)設(shè)施上，它主要運行在Borg、Xmanager和Colossus上?；旧先魏渭憾际沁@樣的配置。然而，創(chuàng)業(yè)后，才意識到不同集群必須要實際設(shè)置新環(huán)境，這種概念對我來說是陌生的。

在當今世界，擁有多個加速器池集群似乎是不可避免的，除非專門在一個位置為大量集群構(gòu)建。更具體地說，GPU 供應(yīng)（或缺乏）也自然導(dǎo)致了這種集群采購模式，其中事物本質(zhì)上是碎片化的。訓(xùn)練大型模型還需要大量 TB 的數(shù)據(jù)，即使只是移動它們也會帶來很多不便。同時，在超大規(guī)模的情況下復(fù)制數(shù)據(jù)通常也不讓人望而卻步的。

顯然，這里的理想情況是某種專門用于將任務(wù)發(fā)送到不同服務(wù)器的編排層。我相信許多AI優(yōu)先的大型公司通常都配有某種基礎(chǔ)設(shè)施來改善人工智能研究人員的工作質(zhì)量。然而，對于一個精益的新初創(chuàng)公司來說，一開始就不可能構(gòu)建這種復(fù)雜而精美的訓(xùn)練基礎(chǔ)設(shè)施。

目前，我們最終開發(fā)了許多內(nèi)部工作流程來緩解其中的許多問題，并繼續(xù)朝著世界一流實驗基礎(chǔ)設(shè)施的黃金標準邁進。

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

4.野雞代碼的痛苦

眾所周知，我一直以來最喜歡的代碼庫是T5X和Mesh Tensorflow（名為tensor ftw），但這些選項很快就變得不可行，理由有三點：1）它們在 Google 之外沒有得到那么多的支持，2）它們有點被棄用了， 3）他們對我們團隊中非 xoogler 的人不友好。

我們最終選擇了一些普通的、看似穩(wěn)定且更流行的東西（即 pytorch），對團隊中的大多數(shù)人（除了我）來說更容易使用。在最初的幾個月里，我對 pip、git、docker 和所有這些野生的東西感到困惑。

話又說回來，我并不能 100% 確定在外部使用 Google 代碼庫有多穩(wěn)定或多用戶友好。

坦率地說，我不得不說外部代碼庫的質(zhì)量明顯落后于我在 Google 習慣的代碼庫。主要是因為 Google 內(nèi)部的代碼庫往往是由 “ML 搖滾明星”自己編寫（例如，Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung 等人），并且與我在外部嘗試過的代碼庫相比，感覺更好（例如，優(yōu)越的氛圍） 。特別是，當我涉足其他公司構(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)公司，我們需要行動快速。我并不為此感到自豪。

5.少一點原則，多一點Yolo

系統(tǒng)地擴展模型通常需要有原則地從小到大，即分多個階段（1B->8B->64B->300B等）進行實驗，然后選出獲勝者并不斷擴展它們。在初創(chuàng)公司中，我們執(zhí)行這些大規(guī)模掃描來檢查 hparam 的計算量要少得多。最后，我們不得不進行多次Yolo run（幸運的是結(jié)果很好）。

YOLO，you only live once。即開始就全套運行，而不是一小步一小步地訓(xùn)練。

最終，我們只需要極少數(shù)的較小規(guī)模和較短的燒蝕運行，即可獲得強大的 21B Reka Flash 和 7B 邊緣模型（以及我們即將推出的最大核心模型）。在運行次數(shù)非常有限的情況下找到可靠的配方具有挑戰(zhàn)性，并且考慮到搜索空間極其巨大，需要立即更改許多變量。

為了做到這一點，人們必須放棄大型科技公司的系統(tǒng)性，而在很大程度上依賴“Yolo”、直覺和本能。

我（以及我們團隊中的許多人）在我們的 ML 職業(yè)生涯中已經(jīng)建立了相當多的這種直覺，可以在很短的嘗試時間內(nèi)得到正確的結(jié)果。雖然我們在之前的工作中訓(xùn)練過非常好的模型，但訓(xùn)練基礎(chǔ)設(shè)施、數(shù)據(jù)、新想法的融合和其他環(huán)境問題的差異仍然可能導(dǎo)致結(jié)果的巨大差異。

也就是說，強大的先驗有助于顯著減少搜索空間，這可能是解釋為什么我們能夠通過如此少的試驗、資源和實驗來訓(xùn)練真正強大的模型的最簡單的解釋之一。

6.寫在最后：荒野中起舞，痛并快樂

在荒野中弄清楚事情是一次有趣的經(jīng)歷。不幸的是，這并不是無痛的。計算稀缺和不可靠的計算提供商使事情比預(yù)期困難得多，但我們很高興我們憑借強大的技術(shù)實力渡過了難關(guān)。

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

參考鏈接:https://reka.ai/reka-flash-an-efficient-and-capable-multimodal-language-model/