如何使用Databricks、Hugging Face和MosaicML訓練大型語言模型
大型語言模型，如OpenAI的GPT-4或谷歌的PaLM，已經席卷了人工智能世界。然而，大多數公司目前沒有能力培訓這些模型，完全依賴少數大型科技公司作為技術提供商。在本篇文章當中，我們將用Replit團隊的實際經驗，來概述如何訓練LLM，從原始數據到面向用戶的生產環(huán)境中的部署，在這一過程中面臨的工程挑戰(zhàn)，以及我們如何利用構成現代LLM堆棧的供應商：Databricks、Hugging face和MosaicML。雖然我們的模型主要用于代碼生成的用例，但所討論的技術和經驗教訓適用于所有類型的LLM，包括通用語言模型。為什么要培訓自己的LLM？Replit人工智能團隊最常見的問題之一是“為什么要訓練自己的模型？”一家公司可能決定訓練自己的LLM有很多原因，從數據隱私和安全，到加強對更新和改進的控制。
- 定制
訓練自定義模型使我們能夠根據自己的特定需求和要求進行定制，包括GPT-4等通用模型甚至Codex等代碼特定模型無法很好地涵蓋的平臺特定功能、術語和上下文。例如，我們的模型經過培訓，可以更好地使用Replit上流行的特定基于web的語言，包括Javascript React（JSX）和Typescript React（TSX）。
- 減少依賴性
雖然我們總是會根據手頭任務情況，而選用對應的模型，但我們相信，減少對少數人工智能提供商的依賴是有好處的。這不僅適用于Replit，也適用于更廣泛的開發(fā)人員社區(qū)。這就是為什么我們計劃開源一些模型，如果沒有訓練它們的手段，我們就無法做到這一點。
- 成本效益
盡管成本將繼續(xù)下降，但LLM在全球開發(fā)者群體中的使用成本仍然高得令人望而卻步。在Replit，我們的使命是讓下十億軟件創(chuàng)作者在線。我們認為，在印度用手機編碼的學生應該能夠像硅谷的專業(yè)開發(fā)人員一樣使用相同的人工智能。為了實現這一點，我們訓練了更小、更高效的自定義模型，并且可以大幅降低成本。數據管道LLM需要大量的數據進行訓練。培訓他們需要建立強大的數據管道，這些管道經過高度優(yōu)化，但足夠靈活，可以輕松地包括公共和專有數據的新來源。我們從The Stack開始，它是主要數據源，可在Hugging Face上獲得。Hugging Face是數據集和預訓練模型的一個很好的資源。作為Transformers庫的一部分，它們還提供了各種有用的工具，包括用于標記化、模型推理和代碼評估的工具。該堆棧由BigCode項目提供。數據集構建的細節(jié)可在Kocetkov等人（2022）中獲得。重復數據消除后，數據集的1.2版包含約2.7 TB的許可源代碼，這些代碼是用350多種編程語言編寫的。Transformers庫在抽象與模型訓練相關的許多挑戰(zhàn)方面做得很好，包括大規(guī)模處理數據。然而，我們發(fā)現這對流程來說是不夠的，因為需要對數據進行額外的控制，以及以分布式方式處理數據的能力。
當需要更高級的數據處理時，我們使用Databricks來構建我們的管道。這種方法還可以很容易地將額外的數據源（如Replit或Stack Overflow）引入到我們的流程中，計劃在未來的迭代中這樣做。第一步是從Hugging Face下載原始數據。使用ApacheSpark在每種編程語言中并行化數據集生成器過程。然后，重新劃分數據，并將其改寫為鑲木地板格式，并為下游處理進行優(yōu)化設置。接下來，我們轉向數據的清理和預處理。通常，消除數據重復并解決各種編碼問題很重要，但the Stack已經使用Kocetkov等人（2022）中概述的近乎重復的技術做到了這一點。但是，一旦開始將Replit數據引入我們的管道，就必須重新運行重復數據消除過程。這就是擁有Databricks這樣的工具的回報所在，在這里，可以將堆棧、堆棧溢出和Replit數據視為一個更大的數據湖中的三個源，并根據需要在下游流程中使用它們。使用Databricks的另一個好處是，可以對底層數據運行可擴展和可處理的分析。在數據源上運行所有類型的匯總統(tǒng)計數據，檢查長尾分布，并診斷過程中的任何問題或不一致。所有這些都是在Databricks筆記本電腦中完成的，它還可以與MLFlow集成，以跟蹤和復制一路上的所有分析。這一步驟相當于對數據進行周期性的x射線檢查，也有助于為我們采取的各種預處理步驟提供信息。對于預處理，采取以下步驟：
- 通過刪除任何個人身份信息（PII）來匿名化數據，包括電子郵件、IP地址和密鑰。
- 使用許多啟發(fā)式方法來檢測和刪除自動生成的代碼。
- 對于一個子集的語言，刪除無法編譯或無法使用標準語法解析器解析的代碼。
- 根據平均行長度、最大行長度和字母數字字符的百分比來過濾文件。
標記化與詞匯訓練
在標記化之前，使用用于模型訓練的相同數據的隨機子樣本，來訓練我們自己的自定義詞匯表。自定義詞匯表使模型能夠更好地理解和生成代碼內容。這導致了模型性能的提高，并加快了模型訓練和推理的速度。這一步驟是流程中最重要的步驟之一，因為它用于流程的所有三個階段（數據管道、模型訓練、推理）。它強調了為您的模型培訓過程提供強大且完全集成的基礎設施的重要性。一旦訓練了我們的自定義詞匯，就能夠對數據進行標記。最后，構建訓練數據集，并將其寫入經過優(yōu)化的分片格式，以供模型訓練過程使用。模型培訓使用MosaicML訓練我們的模型。在之前部署了訓練集群之后，發(fā)現MosaicML平臺帶來了一些關鍵的好處。
- 多個云提供商。Mosaic使我們能夠利用來自不同云提供商的GPU，而無需設置帳戶和所有所需集成的開銷。
- LLM培訓配置。Composer庫有許多經過良好調整的配置，用于訓練各種模型和不同類型的訓練目標。
- 托管基礎設施。他們的托管基礎設施為我們提供了協(xié)調、效率優(yōu)化和容錯（即從節(jié)點故障中恢復）。
在確定模型的參數時，我們考慮了模型大小、上下文窗口、推理時間、內存占用等之間的各種權衡。較大的模型通常提供更好的性能，并且更有能力進行遷移學習。然而，這些模型對訓練和推理都有更高的計算要求。后者尤其重要。Replit是一個云原生IDE，其性能就像桌面原生應用程序，因此我們的代碼完成模型需要閃電般快。因此，我們通常會選擇內存占用較小、推理延遲較低的較小模型。除了模型參數，我們還從各種訓練目標中進行選擇，每個目標都有自己獨特的優(yōu)點和缺點。最常見的訓練目標是下一個令牌預測。這通常適用于代碼完成，但沒有考慮文檔下游的上下文。這可以通過使用“中間填充”目標來緩解，在該目標中，文檔中的一系列標記被屏蔽，并且模型必須使用周圍的上下文來預測它們。另一種方法是UL2（無監(jiān)督潛在語言學習），它將用于訓練語言模型的不同目標函數框定為去噪任務，其中模型必須恢復給定輸入的缺失子序列。
一旦決定了模型配置和訓練目標，我們就在GPU的多節(jié)點集群上啟動訓練運行。我們能夠根據正在訓練的模型的大小，以及希望以多快的速度完成訓練過程，從而來調整為每次跑步分配的節(jié)點數量。運行一個大型GPU集群是昂貴的，所以要以最有效的方式利用它們是很重要的。并且密切監(jiān)控GPU的利用率和內存，以確保計算資源得到最大可能的利用。此外，使用權重和偏差來監(jiān)控培訓過程，包括資源利用率和培訓進度。在監(jiān)控損失曲線，以確保模型在訓練過程的每一步都能有效學習。同樣，也注意損失的峰值。這些是損失值的突然增加，通常表明底層訓練數據或模型架構存在問題。由于這些事件通常需要進一步的調查和潛在的調整，我們在流程中強制執(zhí)行數據決定論，因此可以更容易地復制、診斷和解決任何此類損失峰值的潛在來源。評價為了測試我們的模型，使用了Chen等人（2021）中描述的HumanEval框架的變體。我們使用該模型生成一個給定函數簽名和文檔字符串的Python代碼塊。然后，對生成的函數運行一個測試用例，以確定生成的代碼塊是否按預期工作。這種方法最適合Python，有現成的評估器和測試用例。但由于Replit支持許多編程語言，需要評估各種其他語言的模型性能。然而這很難做到，而且沒有廣泛采用的工具或框架可以提供全面的解決方案。兩個具體的挑戰(zhàn)包括在任何編程語言中創(chuàng)造一個可復制的運行時環(huán)境，以及在沒有廣泛使用的測試用例標準（如HTML、CSS等）的編程語言中的模糊性。部署到生產模型一旦訓練并評估完畢，就到了將其部署到生產中的時候了。正如我們前面提到的，我們的代碼完成模型應該感覺很快，請求之間的延遲非常低。使用NVIDIA的FasterTransformer和Triton Server加速我們的推理過程。FasterTransformer是一個實現基于transformer的神經網絡推理加速引擎的庫，Triton是一個穩(wěn)定快速的推理服務器，配置簡單。這種組合在轉換器模型和底層GPU硬件之間為我們提供了一個高度優(yōu)化的層，并允許對大型模型進行超快速的分布式推理。在將模型部署到生產中后，能夠使用Kubernetes基礎設施自動擴展它以滿足需求。托管推理服務器帶來了一系列獨特的挑戰(zhàn)。這些包括大的偽影（即，模型權重）和特殊的硬件要求（即，不同的GPU大小/計數）。我們設計了部署和集群配置，以便能夠快速可靠地發(fā)貨。例如，集群旨在解決單個區(qū)域的GPU短缺問題，并尋找最便宜的可用節(jié)點。在把一個模型放在實際用戶面前之前，通常會先自己測試它，并了解模型的“氛圍”。在之前計算的HumanEval測試結果很有用，但沒有什么比使用模型來了解它更好的了，包括它的延遲、建議的一致性和一般的幫助性。
而后繼續(xù)監(jiān)控模型性能和使用指標。對于模型性能，監(jiān)控請求延遲和GPU利用率等指標。對于使用，跟蹤代碼建議的接受率，并在包括編程語言在內的多個維度上進行分解。這也允許我們對不同的模型進行A/B test，并獲得一個模型與另一個模型比較的定量度量。反饋和迭代我們的模型訓練平臺使我們能夠在不到一天的時間內從原始數據到部署在生產中的模型。但更重要的是，它能夠去訓練和部署模型，收集反饋，然后根據反饋快速迭代。同樣重要的是，流程要對底層數據源、模型訓練目標或服務器架構的任何變化保持穩(wěn)健。這使我們能夠在一個快速發(fā)展的領域利用新的進步和能力。雖然目前已經取得了很大的進步，但仍處于LLM培訓的早期階段，還有很多改進要做，還有很多難題要解決。隨著語言模型的不斷發(fā)展，這一趨勢只會加速。數據、算法和模型評估將面臨一系列新的挑戰(zhàn)。
