本文將介紹Replit如何利用Databricks、Hugging Face和MosaicML訓(xùn)練大型語言模型(LLMs)。Replit是美國一家編碼平臺,提供了一個可從瀏覽器訪問的IDE,無需設(shè)置即可開始編碼,此外還提供了一系列簡化編碼任務(wù)的其他功能。

引言

大型語言模型，如OpenAI的GPT-4或谷歌的PaLM，已經(jīng)在人工智能領(lǐng)域掀起了狂潮。然而，大多數(shù)公司目前還沒有能力訓(xùn)練這些模型，完全依賴于少數(shù)幾個大型科技公司作為技術(shù)提供者。

Replit大力投入從零開始訓(xùn)練自己的大型語言模型所需的基礎(chǔ)設(shè)施。本文中，將介紹Replit如何訓(xùn)練LLMs，從原始數(shù)據(jù)到部署到面向用戶的生產(chǎn)環(huán)境。還會介紹過程中遇到的工程挑戰(zhàn)，以及如何利用LLM服務(wù)例如：Databricks、Hugging Face和MosaicML。

雖然模型主要用于代碼生成，但討論的技術(shù)和經(jīng)驗(yàn)教訓(xùn)適用于所有類型的LLMs，包括通用語言模型。后面還會陸續(xù)緊跟他們的進(jìn)展，敬請關(guān)注。

為什么要訓(xùn)練自己的LLM？

Replit人工智能團(tuán)隊(duì)最常見的問題之一是：“為什么要訓(xùn)練自己的模型？”公司可能決定訓(xùn)練自己的LLMs的原因有很多，從數(shù)據(jù)隱私和安全到對更新和改進(jìn)的更多控制。

在Replit，主要關(guān)注定制、降低依賴性和成本效益。

定制化需求：訓(xùn)練自定義模型允許根據(jù)特定需求和要求進(jìn)行調(diào)整，包括平臺特定功能、術(shù)語和上下文，這些在通用模型如GPT-4甚至代碼特定模型如Codex中都不會被很好支持的。例如，模型經(jīng)過訓(xùn)練，可以更好地處理在Replit上受歡迎Web語言，包括JAVAscript React（JSX）和Typescript React（TSX）。
降低依賴性：雖然需要根據(jù)手頭的任務(wù)選擇合適的模型，但減少對少數(shù)AI提供商的依賴是有益的。這不僅適用于Replit，也適用于更廣泛的開發(fā)者社區(qū)。
成本效益：盡管成本會繼續(xù)降低，但LLM在全球開發(fā)者社區(qū)中仍然昂貴得令人望而卻步。Replit的使命是讓下一個十億軟件創(chuàng)作者上線。為了實(shí)現(xiàn)這一目標(biāo)，訓(xùn)練定制的模型，使其更小、更高效，并且可以大幅降低托管成本。

數(shù)據(jù)管道

LLM需要大量的數(shù)據(jù)進(jìn)行訓(xùn)練。訓(xùn)練它們需要構(gòu)建健壯的數(shù)據(jù)管道，這些管道需要高度優(yōu)化，同時又足夠靈活，以便輕松地包含來自公共和專有數(shù)據(jù)的新數(shù)據(jù)源。

堆棧

從作為主要數(shù)據(jù)源的堆棧開始，它可以在Hugging Face上找到。Hugging Face是數(shù)據(jù)集和預(yù)訓(xùn)練模型的極好資源。他們還提供了Transformers庫中的許多有用工具，包括用于標(biāo)記化、模型推理和代碼評估的工具。

堆棧由BigCode項(xiàng)目提供。在去重后，數(shù)據(jù)集的1.2版包含了大約2.7 TB的許可范圍內(nèi)的源代碼，涵蓋了350多種編程語言。

Transformers庫在處理許多與模型訓(xùn)練相關(guān)方面做得很好，包括處理大規(guī)模數(shù)據(jù)。然而，它對我們的過程還不夠，因?yàn)樾枰獙?shù)據(jù)進(jìn)行更多的控制，并能夠以分布式的方式處理數(shù)據(jù)。

數(shù)據(jù)處理

當(dāng)需要進(jìn)行更高級的數(shù)據(jù)處理時，使用Databricks來構(gòu)建管道。這種方法也能夠輕松地將其他數(shù)據(jù)源（如Replit或Stack Overflow）引入。

第一步是從Hugging Face下載原始數(shù)據(jù)。使用Apache Spark將數(shù)據(jù)集構(gòu)建器過程在每種編程語言之間并行化。然后，重新分區(qū)數(shù)據(jù)，并以parquet格式使用針對下游處理優(yōu)化的設(shè)置重寫數(shù)據(jù)。

接下來，清理和預(yù)處理我們的數(shù)據(jù)。通常，重要的是要對數(shù)據(jù)進(jìn)行去重并修復(fù)各種編碼問題，The Stack完成了這項(xiàng)工作。然而，將Replit數(shù)據(jù)引入管道，就需要重新運(yùn)行去重過程。這時，擁有Databricks這樣的工具就顯得非常有價值，可以將The Stack、Stackoverflow和Replit數(shù)據(jù)視為更大數(shù)據(jù)湖中的三個源，并根據(jù)需要在下游過程中使用它們。

使用Databricks的另一個好處是可以在底層數(shù)據(jù)上運(yùn)行可擴(kuò)展且可追蹤的分析。對數(shù)據(jù)源運(yùn)行各種摘要統(tǒng)計(jì)信息，檢查長尾分布，并診斷過程中的任何問題或不一致。所有這些都是在Databricks notebook中完成的，這些notebook還可以與MLFlow集成，以跟蹤和重現(xiàn)在過程中進(jìn)行的所有分析。這個步驟相當(dāng)于對我們的數(shù)據(jù)進(jìn)行定期X光檢查，它還有助于了解在預(yù)處理中采取的各種步驟。

在預(yù)處理過程中，采取以下步驟：

通過刪除任何個人身份信息（PII），包括電子郵件、IP地址和密鑰，對數(shù)據(jù)進(jìn)行匿名處理。
使用一些啟發(fā)式方法來檢測和刪除自動生成的代碼。
對于某些語言，刪除無法編譯或無法使用標(biāo)準(zhǔn)語法解析器解析的代碼。
根據(jù)平均行長度、最大行長度和字母數(shù)字字符百分比來過濾文件。

分詞和詞匯訓(xùn)練

在分詞之前，使用與模型訓(xùn)練相同的數(shù)據(jù)的隨機(jī)樣本來訓(xùn)練自己的自定義詞匯表。自定義詞匯表使模型更好地理解和生成代碼內(nèi)容。這提高了模型性能，加速了模型訓(xùn)練和推理。

這個步驟是整個過程中最重要的一個，因?yàn)樗挥迷谌齻€階段中（數(shù)據(jù)管道、模型訓(xùn)練、推理）。這凸顯了為模型訓(xùn)練過程擁有一個強(qiáng)大且完全集成的基礎(chǔ)設(shè)施的重要性。

在較高層次上，還需要考慮的一些重要事項(xiàng)包括詞匯量大小、特殊標(biāo)記和哨兵標(biāo)記的保留空間。

完成自定義詞匯表的訓(xùn)練后，對數(shù)據(jù)進(jìn)行分詞。最后，構(gòu)建訓(xùn)練數(shù)據(jù)集，并將其寫入一個分片格式，以便優(yōu)化地輸入到模型訓(xùn)練過程中。

模型訓(xùn)練

使用MosaicML來訓(xùn)練模型。之前已經(jīng)部署了訓(xùn)練集群，發(fā)現(xiàn)MosaicML平臺帶來了一些關(guān)鍵優(yōu)勢。

多個云提供商：Mosaic能夠利用不同云提供商的GPU，而無需設(shè)置帳戶和所有必需的集成。
LLM訓(xùn)練配置：Composer庫有許多為訓(xùn)練各種模型以及不同類型的訓(xùn)練目標(biāo)進(jìn)行了優(yōu)化的配置。
托管基礎(chǔ)設(shè)施：他們的托管基礎(chǔ)設(shè)施提供了編排、效率優(yōu)化和容錯（即從節(jié)點(diǎn)故障中恢復(fù)）。

在確定模型參數(shù)時，考慮模型大小、上下文窗口、推理時間、內(nèi)存占用等之間的各種權(quán)衡。更大的模型通常具有更好的性能，且更適于遷移學(xué)習(xí)。然而，這些模型對訓(xùn)練和推理都有更高的計(jì)算要求。后者尤為重要。Replit是一個云原生IDE，其性能感覺像一個桌面原生應(yīng)用程序，因此代碼補(bǔ)全模型需要非常快速。因此，通常傾向于選擇具有較小內(nèi)存占用和低延遲推理的較小模型。

除了模型參數(shù)外，還從各種訓(xùn)練目標(biāo)中進(jìn)行選擇，每個目標(biāo)都具有其獨(dú)特的優(yōu)點(diǎn)和缺點(diǎn)。最常見的訓(xùn)練目標(biāo)是下一個token預(yù)測。這對于代碼補(bǔ)全通常效果很好，但無法考慮到文檔后續(xù)上下文中的信息。這可以通過使用“完形填空”目標(biāo)來彌補(bǔ)，其中文檔中的一系列tokens被屏蔽，模型必須使用周圍的上下文來預(yù)測它們。另一種方法是 UL2（無監(jiān)督潛在語言學(xué)習(xí)），將訓(xùn)練語言模型的不同目標(biāo)函數(shù)構(gòu)建為去噪任務(wù)，其中模型必須恢復(fù)給定輸入的缺失子序列。

一旦確定了模型配置和訓(xùn)練目標(biāo)，就會在多節(jié)點(diǎn)GPU集群上啟動訓(xùn)練。根據(jù)正在訓(xùn)練的模型的大小以及希望完成訓(xùn)練過程的速度，可以調(diào)整為每次運(yùn)行分配的節(jié)點(diǎn)數(shù)量。運(yùn)行大型GPU集群的成本很高，因此需要以盡可能高效的方式利用它們。同時密切監(jiān)控GPU利用率和內(nèi)存，以確保充分利用了計(jì)算資源。

使用Weights & Biases來監(jiān)控訓(xùn)練過程，包括資源利用情況以及訓(xùn)練進(jìn)度。監(jiān)控?fù)p失曲線，以確保模型在訓(xùn)練過程的每個階段都能有效地學(xué)習(xí)。另外還要關(guān)注損失值的突然增加，這通常表示訓(xùn)練數(shù)據(jù)或模型架構(gòu)存在問題。因?yàn)檫@些情況通常需要進(jìn)一步調(diào)查和可能的調(diào)整，在過程中實(shí)施數(shù)據(jù)確定性，以便更容易地重現(xiàn)、診斷和解決任何這種損失激增的潛在來源。

評估

為了測試模型，使用了HumanEval框架的一個變體。使用模型根據(jù)函數(shù)簽名和文檔字符串生成Python/ target=_blank class=infotextkey>Python代碼塊。然后，對生成的函數(shù)運(yùn)行一個測試用例，以確定生成的代碼塊是否按預(yù)期工作。運(yùn)行多個樣本并分析相應(yīng)的Pass@K值。

這種方法對于Python效果最好，有現(xiàn)成的評估器和測試用例。但是，因?yàn)镽eplit支持許多編程語言，需要評估更多其他語言的模型性能。這很難做到，而且沒有廣泛采用的工具或框架提供全面的解決方案。兩個具體的挑戰(zhàn)包括在任何編程語言中制作可復(fù)現(xiàn)的運(yùn)行時環(huán)境，以及在沒有廣泛使用的測試用例標(biāo)準(zhǔn)的編程語言中存在的歧義（例如html、css等）。萬幸的是，一個“可在任何編程語言中復(fù)現(xiàn)的運(yùn)行時環(huán)境”正是Replit的特長！目前正在構(gòu)建一個評估框架，允許任何研究人員接入并測試他們的多語言基準(zhǔn)。

部署到生產(chǎn)環(huán)境

在訓(xùn)練和評估模型之后，是時候?qū)⑵洳渴鸬缴a(chǎn)環(huán)境了。如前所述，代碼補(bǔ)全模型應(yīng)具有快速的感覺，在請求之間具有非常低的延遲。使用NVIDIA的FasterTransformer和Triton服務(wù)器加速推理過程。FasterTransformer是一個實(shí)現(xiàn)加速基于Transformer的神經(jīng)網(wǎng)絡(luò)推理引擎的庫，而Triton是一個穩(wěn)定且快速的推理服務(wù)器，配置簡單。這種組合提供了一個高度優(yōu)化的層，位于Transformer模型和底層GPU硬件之間，并允許大型模型的超快分布式推理。

在將模型部署到生產(chǎn)環(huán)境后，能夠使用Kube.NETes基礎(chǔ)設(shè)施根據(jù)需求自動擴(kuò)展它。托管推理服務(wù)器還面臨著一系列獨(dú)特的挑戰(zhàn)。這些挑戰(zhàn)包括大型工件（即模型權(quán)重）和特殊硬件要求（即不同的GPU大小/數(shù)量）。通過設(shè)計(jì)部署和集群配置，以便能夠快速且可靠地發(fā)布。例如，集群設(shè)計(jì)為解決單個區(qū)域中的GPU短缺問題，并尋找最便宜的可用節(jié)點(diǎn)。

在將模型置于實(shí)際用戶之前，還需要測試它，并了解模型的“氛圍”。之前計(jì)算的HumanEval測試結(jié)果很有用，不過只有親自使用模型才能真正了解它，包括其延遲、建議的一致性和總體幫助程度。將模型放在開發(fā)人員面前就像撥動一個開關(guān)一樣簡單。一旦對其感到滿意，就撥動另一個開關(guān)，將其推廣給其他用戶。

持續(xù)監(jiān)控模型性能和使用指標(biāo)。對于模型性能，監(jiān)控諸如請求延遲和GPU利用率之類的指標(biāo)。對于使用情況，跟蹤代碼建議的接受率，并根據(jù)編程語言等多個維度進(jìn)行細(xì)分。這還允許我們對不同模型進(jìn)行A/B測試，并獲得一個模型與另一個模型進(jìn)行比較的定量指標(biāo)。

反饋與迭代

模型訓(xùn)練平臺使開發(fā)人員能夠在不到一天的時間里，從原始數(shù)據(jù)到部署到生產(chǎn)環(huán)境的模型。但更重要的是，它使訓(xùn)練和部署模型，收集反饋，然后根據(jù)這些反饋快速迭代。

這個過程要保持對底層數(shù)據(jù)來源、模型訓(xùn)練目標(biāo)或服務(wù)器架構(gòu)的穩(wěn)定性。這可以快速發(fā)展出領(lǐng)域中的新能力，每天似乎都有新改進(jìn)。

接下來，是擴(kuò)展平臺，使開發(fā)人員能夠使用Replit本身來改進(jìn)模型。這包括基于人類反饋的強(qiáng)化學(xué)習(xí)（RLHF）等技術(shù)，以及使用來自Replit賞金任務(wù)的數(shù)據(jù)進(jìn)行指令調(diào)整。

下一步

盡管Replit已經(jīng)取得了很大的進(jìn)展，但在訓(xùn)練LLMs方面仍處于非常初期階段。還有很多改進(jìn)要做，很多困難問題需要解決。隨著語言模型的不斷發(fā)展，這種趨勢將只會加速。與數(shù)據(jù)、算法和模型評估相關(guān)的新挑戰(zhàn)將持續(xù)出現(xiàn)。