作者 |  Yi Tay

編譯 | 云昭

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

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

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

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

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

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

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

1.買算力如同買彩票！

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

4.野雞代碼的痛苦

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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