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

引言

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

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

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

為什么要訓練自己的LLM？

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

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

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

數據管道

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

堆棧

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

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

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

數據處理

當需要進行更高級的數據處理時，使用Databricks來構建管道。這種方法也能夠輕松地將其他數據源（如Replit或Stack Overflow）引入。

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

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

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

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

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

分詞和詞匯訓練

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

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

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

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

模型訓練

使用MosaicML來訓練模型。之前已經部署了訓練集群，發現MosaicML平臺帶來了一些關鍵優勢。

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

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

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

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

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

評估

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

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

部署到生產環境

在訓練和評估模型之后，是時候將其部署到生產環境了。如前所述，代碼補全模型應具有快速的感覺，在請求之間具有非常低的延遲。使用NVIDIA的FasterTransformer和Triton服務器加速推理過程。FasterTransformer是一個實現加速基于Transformer的神經網絡推理引擎的庫，而Triton是一個穩定且快速的推理服務器，配置簡單。這種組合提供了一個高度優化的層，位于Transformer模型和底層GPU硬件之間，并允許大型模型的超快分布式推理。

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

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