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

引言

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

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

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

為什么要訓練自己的LLM？

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

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

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

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

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

堆棧

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

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

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

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

當需要進行更高級的數(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ù)據(jù)。

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

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

在預處理過程中，采取以下步驟：

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

分詞和詞匯訓練

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

這個步驟是整個過程中最重要的一個，因為它被用在三個階段中（數(shù)據(jù)管道、模型訓練、推理）。這凸顯了為模型訓練過程擁有一個強大且完全集成的基礎設施的重要性。

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

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

模型訓練

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

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

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

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

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

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

評估

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

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

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

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

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

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

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

反饋與迭代

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

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

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

下一步

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