作者 |  Yi Tay

編譯 | 云昭

出品 | 51CTO技術棧（微信號：blog51cto）

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

本文主人公Yi Tay，是一位市面上非常搶手的高性能大模型的大拿。他曾在谷歌Google Brain擔任高級研究科學家，專注于大型語言模型和人工智能的研究。在Google任職期間，曾經(jīng)為業(yè)內許多知名的大型語言模型做出了貢獻，例如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。大模型訓練、多模態(tài)大模型何其艱難？這期間，究竟發(fā)生了怎樣的離奇的事情呢？

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

1.買算力如同買彩票！

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

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

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

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

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

這就跟買彩票一樣。基本上：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

3.多集群設置的痛苦

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

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

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

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

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

4.野雞代碼的痛苦

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

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

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

坦率地說，我不得不說外部代碼庫的質量明顯落后于我在 Google 習慣的代碼庫。主要是因為 Google 內部的代碼庫往往是由 “ML 搖滾明星”自己編寫（例如，Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung 等人），并且與我在外部嘗試過的代碼庫相比，感覺更好（例如，優(yōu)越的氛圍） 。特別是，當我涉足其他公司構建的東西時，我發(fā)現(xiàn)自己對代碼質量非常惱火。

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

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

我知道我應該使用 Jax。一位朋友剛剛因為我使用 pytorch 而羞辱我，但這是一家初創(chuàng)公司，我們需要行動快速。我并不為此感到自豪。

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

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

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

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

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

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

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

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

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

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

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