【編者按】線上程式設計平臺 Replit 分享了一篇如何使用 Databricks、Hugging Face 和 MosaicML 訓練大型語言模型(LLMs)的文章,分享了他們在訓練自家程式設計大模型(LLM)的過程、經驗和教訓,他們把整個過程分為三個階段,資料管道、模型訓練和推理。CSDN藉助 ChatGPT 和 Claude 聯合翻譯,以饗讀者!
原文地址:https://blog.replit.com/llm-training
責編 | 王子彧


介紹
大型語言模型(如 OpenAI 的 GPT-4 或 Google 的 PaLM )已經席捲了人工智慧領域。然而,大多數公司目前都沒有能力訓練這些模型,並且完全依賴於只有少數大型科技公司提供的技術。
在 Replit,我們投入了大量的基礎設施來從頭開始訓練自己的大型語言模型。在本文中,我們將概述如何訓練 LLMs,從原始資料到在使用者面向生產環境中的部署。我們將討論我們在此過程中面臨的工程挑戰,以及我們如何利用我們認為構成現代 LLM 堆疊的供應商:Databricks、Hugging Face 和 MosaicML。
雖然我們的模型主要是用於程式碼生成的用例,但所討論的技術和經驗教訓適用於所有類型的 LLMs,包括通用語言模型。在未來的幾周和幾個月中,我們計劃深入探討我們的流程的細節。

為什麼要訓練自己的 LLMs?
Replit 的 AI 團隊最常見的問題之一是「為什麼要訓練自己的模型?」公司決定訓練自己的 LLMs 的原因有很多,從資料隱私和安全到對更新和改進的更多控制。
在 Replit,我們主要關心的是定製化、減少依賴和成本效益。
個性化定製。訓練自定義模型允許我們根據特定的需求和要求進行定製,包括平臺特定的功能、術語和上下文,這些在通用模型如 GPT-4 或僅限程式碼的模型如 Codex 中都無法很好地覆蓋。例如,我們的模型經過訓練,可以更好地處理在 Replit 上流行的特定基於 Web 的語言,包括 Javascript React (JSX) 和 Typescript React (TSX)。
降低依賴性。雖然我們始終會根據需要使用適當的模型,但我們認為減少對少數人工智慧供應商的依賴是有好處的。這不僅適用於 Replit,也適用於更廣泛的開發者社區。這就是為什麼我們計劃開源一些我們的模型,這是在沒有訓練它們的手段的情況下無法實現的。
成本效益。雖然成本將繼續下降,但 LLMs 對於全球開發者社區來說仍然成本過高。在 Replit,我們的使命是將下一個億的軟體創作者帶上線。我們認為,印度手機上編寫程式碼的學生應該能夠訪問與矽谷專業開發人員相同的人工智慧。為了實現這一目標,我們訓練定製的模型,這些模型更小、更高效,並可以以大大降低的成本進行託管。
資料管道、來源及處理
資料管道
LLM 模型需要大量的資料進行訓練。為了訓練它們,需要構建強大的資料管道,這些管道高度最佳化,同時又足夠靈活,能夠輕鬆包含新的公共和專有資料來源。
資料來源
我們首先以「資料來源」作為我們的主要資料來源,該資料來源可在 Hugging Face 上獲得。Hugging Face 是一個極好的資料集和預訓練模型資源。他們還提供了許多有用的工具,包括用於分詞、模型推斷和程式碼評估的工具,這些工具作為 Transformers 庫的一部分提供。
「資料來源」是由 BigCode 項目提供的。有關資料集構建的詳細資訊可參考 Kocetkov et al. (2022)。在去重之後,版本 1.2 的資料集包含大約 2.7TB 的可許可證源程式碼,涵蓋了超過 350 種程式語言。
Transformers 庫在抽象出許多與模型訓練相關的挑戰方面做得非常好,包括處理大規模資料。但是,我們發現它對我們的過程來說不夠充分,因為我們需要對資料進行額外的控制,並能夠以分散式方式處理資料。

資料處理
當需要進行更高級的資料處理時,我們使用 Databricks 來構建資料管道。這種方法還使我們能夠輕鬆地將其他資料來源(例如 Replit 或 Stack Overflow)引入我們的處理過程中,這是我們計劃在未來的迭代中做的。
首先,我們需要從 Hugging Face 下載原始資料。我們使用 Apache Spark 將資料集構建過程並行化到每種程式語言中。然後,我們重新分區資料,並使用最佳化設置將其以 parquet 格式重新編寫,以便進行下游處理。
接下來,我們轉向資料清理和預處理。通常,重要的是對資料進行去重和修復各種編碼問題,但是 The Stack 已經使用 Kocetkov 等人(2022)提出的一種近似去重技術為我們完成了這項工作。然而,一旦我們開始將 Replit 資料引入我們的管道中,就必須重新運行去重處理過程。這就是使用 Databricks 這樣的工具的好處,我們可以將 The Stack、Stackoverflow 和 Replit 資料視為較大資料湖中的三個源,並根據需要在下游處理中利用它們。
使用 Databricks 的另一個好處是,我們可以對底層資料運行可擴展和可追蹤的分析。我們對資料來源運行各種摘要統計資訊,檢查長尾分佈,並診斷過程中的任何問題或不一致性。所有這些都在 Databricks 筆記本中完成,還可以與 MLFlow 集成以跟蹤和重現我們一路走來的所有分析。這一步,相當於定期對我們的資料進行一次 X 光檢查,還有助於指導我們進行預處理的各種步驟。
對於預處理,我們採取以下步驟:
我們通過刪除包括電子郵件、IP 地址和金鑰在內的任何個人可識別資訊(PII)來匿名化資料。
我們使用一些啟發式方法來檢測和刪除自動生成的程式碼。
對於一些語言的子集,我們會刪除不能編譯或無法使用標準語法解析器進行解析的程式碼。
我們基於平均行長度、最大行長度和包含字母數字字符的百分比來過濾檔案。


Tokenization and vocabulary training
在進行分詞 (tokenization) 之前,我們使用相同的資料集的隨機子樣本來訓練自己的自定義詞彙表 (vocabulary)。自定義詞彙表可以使我們的模型更好地理解和生成程式碼內容,提高模型性能並加速模型訓練和推斷。
這一步是整個過程中最重要的步驟之一,因為它在我們的整個過程中的所有三個階段(資料流水線、模型訓練和推斷)中都被使用到。這凸顯了擁有一個強大而完整的基礎架構對模型訓練過程的重要性。
我們計劃在以後的部落格文章中深入探討分詞 (tokenization)。在高層次上,我們需要考慮的一些重要事項是詞彙表大小、特殊標記 (special tokens) 以及保留用於 sentinel 標記 (sentinel tokens) 的空間。
一旦我們訓練好了自定義詞彙表,我們就對資料進行分詞 (tokenize)。最後,我們構建我們的訓練資料集,並將其寫成分片格式,以便最佳化輸入到模型訓練過程中。

模型訓練
我們使用 MosaicML 訓練我們的模型。在以前部署自己的訓練集群之後,我們發現 MosaicML 平臺為我們提供了幾個關鍵的好處。
多個雲提供商。Mosaic 允許我們利用不同雲提供商的 GPU,而無需設置賬戶和所有必要的集成的開銷。
LLM 訓練配置。Composer 庫有許多針對訓練各種模型和不同類型的訓練目標最佳化的配置。
託管基礎架構。他們託管的基礎架構為我們提供編排、效率最佳化和容錯性(即從節點故障中恢復)。
在確定我們模型的參數時,我們考慮模型大小、上下文窗口、推理時間、記憶體佔用量等方面的各種權衡。較大的模型通常提供更好的性能,並且更容易進行遷移學習。然而,這些模型在訓練和推理方面都有更高的計算要求。後者對我們來說特別重要。Replit 是一個性能像本地應用程序的雲原生 IDE,因此我們的程式碼補全模型需要非常快速。因此,我們通常傾向於選擇記憶體佔用量和低延遲推理的較小模型。
除了模型參數外,我們還可以選擇各種各樣的訓練目標,每個目標都有其自己獨特的優點和缺點。最常見的訓練目標是下一個令牌預測。這通常對程式碼補全有效,但未考慮文件中下游的上下文。這可以通過使用「填補中間」目標來緩解,在該目標中,文件中的一系列令牌被遮擋,模型必須使用周圍的上下文預測它們。另一種方法是 UL2 (無監督隱語言學習),它將不同的語言模型訓練目標框架為去噪任務,模型必須恢復給定輸入的丟失子序列。

一旦我們決定模型配置和訓練目標,我們就會在多節點 GPU 集群上啟動訓練運行。我們可以根據要訓練的模型大小和訓練過程的完成速度調整為每個運行分配的節點數。運行大型 GPU 集群非常昂貴,所以確保以最有效的方式利用它們非常重要。我們密切監控 GPU 利用率和記憶體,以確保我們的計算資源獲得最大可能的利用。
我們使用 Weights & Biases 監控訓練過程,包括資源利用率和訓練進度。我們監控我們的損失曲線,以確保模型在訓練過程的每一步中都在有效學習。我們還注意損失峰值。這些是損失值突然增加,通常表示訓練資料或模型架構的潛在問題。因為這些情況通常需要進一步調查和潛在調整,我們在流程中執行資料確定性,以便更 easily 重現、診斷和解決任何這種損失峰值的潛在來源。

評估
為了測試我們的模型,我們使用類似於 Chen 等人(2021)中描述的 HumanEval 框架的變體。我們使用模型根據給定的函數簽名和文件字串生成一個 Python 程式碼塊。然後我們對產生的函數運行一個測試用例來確定生成的程式碼塊是否按預期工作。我們運行多個樣本並分析相應的 Pass@K 數字。
這種方法最適用於 Python,因為有現成的評估器和測試用例。但由於 Replit 支持許多程式語言,我們需要評估模型在更廣泛的程式語言中的性能。我們發現這很難做到,並且沒有廣泛採用的工具或框架提供全面的解決方案。兩個特定的挑戰包括在任何程式語言中產生可重複的運行時環境,以及對於沒有廣泛使用的測試用例標準的程式語言存在歧義(例如,HTML、CSS 等)。幸運的是,「在任何程式語言中產生可重複的運行時環境」是我們在 Replit 的專長領域!我們正在構建一個評估框架,使任何研究人員都可以插入和測試他們的多語言基準測試。我們將在未來的部落格文章中討論這個問題。


部署到生產環境
一旦我們訓練並評估了我們的模型,就該將其部署到生產環境中了。如前所述,我們的程式碼補全模型應該感覺很快,請求之間的延遲非常低。我們使用 NVIDIA 的 FasterTransformer和Triton Server加速我們的推理過程。FasterTransformer 是一個實現加速器引擎的庫,用於變壓器基神經網路的推理,Triton 是一個穩定快速的推理伺服器,配置簡單。這種組合為我們提供了一個高度最佳化的層,位於變壓器模型和底層 GPU 硬體之間,可以進行大型模型的超快分散式推理。
部署我們的模型到生產環境後,我們能夠使用 Kubernetes 基礎架構根據需求自動擴展它。雖然我們在前面的部落格文章中討論了自動擴展,但值得一提的是,託管推理伺服器會帶來一套獨特的挑戰。這些包括大型工件(即模型權重)和特殊的硬體需求(即不同的 GPU 大小/數量)。我們的部署和集群配置已設計為能夠快速可靠地發貨。例如,我們的集群被設計為可以繞過個別區域中的 GPU 短缺,尋找最便宜的可用節點。
在將模型放在實際使用者面前之前,我們喜歡先自己測試它,並對模型的「氛圍」有所了解。我們之前計算的 HumanEval 測試結果雖然很有用,但是親自使用模型沒有什麼能替代的,包括其延遲、建議的一致性和總體的有用性。將模型放在 Replit 工作人員面前就像翻轉開關一樣簡單。一旦我們對它感到舒服,我們再翻轉另一個開關,將其推出給我們的其他使用者。

我們繼續監控模型性能和使用指標。對於模型性能,我們監視諸如請求延遲和 GPU 利用率等指標。對於使用情況,我們跟蹤程式碼建議的接受率,並根據多個維度進行細分,包括程式語言。這也使我們能夠 A/B 測試不同的模型,並獲得定量衡量一個模型與另一個模型的比較。

反饋和迭代
我們的模型訓練平臺讓我們能夠在不到一天的時間內從原始資料訓練出一個在生產環境中部署的模型。但更重要的是,它讓我們能夠訓練和部署模型、收集反饋,並基於這些反饋快速迭代。
我們的流程也需要對底層資料來源、模型訓練目標或伺服器架構的任何變化保持穩健。這使我們能夠利用這個快速發展的領域中每天都會帶來新的令人興奮的進展和能力。
接下來,我們將擴展我們的平臺,使我們能夠利用 Replit 本身來改進我們的模型。這包括使用基於人類反饋的強化學習(RLHF)技術,以及使用從 Replit 賞金任務中收集的資料進行指令調整。

下一步
雖然我們已經取得了巨大進步,但我們仍處於訓練 LLM 的初期階段。我們還有大量改進要做,許多困難問題還有待解決。隨著語言模型的不斷發展,這一趨勢只會加速。與資料、演算法和模型評估相關的一系列新的挑戰將持續存在。
如果您對訓練 LLM 的許多工程挑戰感興趣,我們很樂意與您交談。我們歡迎反饋,並樂意聽取您對我們遺漏的內容及您會如何不同的意見。