• <xmp id="om0om">
  • <table id="om0om"><noscript id="om0om"></noscript></table>
  • 對話式人工智能

    通過高效的長上下文大語言模型訓練擴展到數百萬個 Token

    大語言模型 (LLM) 的演變標志著其處理和生成文本的能力有了顯著提升。在這些發展中,上下文長度的概念 (模型可以處理的單個輸入樣本中的 token 數量) 已成為定義這些模型在不同應用中可以實現的目標的關鍵因素。

    例如,這些模型允許處理視頻輸入、總結冗長的文檔、在多輪對話中保持一致性、通過思維鏈進行推理,以及使用大量示例執行詳細的上下文學習。在視頻生成和理解、法律文檔分析、低資源語言翻譯等必須保留和利用全面上下文的場景中,以及在使用 AI 助手時,這種擴展功能尤為重要。

    在本文中,我們將探討長上下文 LLM 的技術基礎,以及如何有效訓練它們的技巧。我們規劃了需求和挑戰,以及如何使用 NVIDIA NeMo 框架通過各種優化技術來解決這些問題,從而提供高吞吐量的高效訓練。

    對擴展上下文長度的需求和相關挑戰

    隨著越來越多的多模態用例涌現,處理長視頻內容需要模型同時處理數千幀,同時保持時間一致性。具有擴展上下文長度的模型 (例如支持多達 100 萬個 token 的模型) 可以在大量視覺輸入幀中保留詳細的時間信息。

    針對復雜推理優化的模型 (例如 DeepSeek-R1 和 Llama Nemotron) 依靠擴展上下文來通過鏈式推理來解決多步問題。如果沒有足夠的上下文窗口大小,這些模型將截斷關鍵邏輯路徑,從而導致錯誤。DeepSeek-R1 的上下文長度超過 128K,而 Llama 4 則將上下文長度的界限擴大到超過 1000 萬個 token。

    使用擴展上下文長度訓練 LLM 會帶來重大的技術障礙,尤其是在內存管理方面。基于 Transformer 的 LLM 會隨著序列長度的增加 (如果使用閃光注意力,則時間復雜度為 O (n^ 2)) 進行計算擴展。這使得使用超長上下文進行訓練的成本非常高昂。

    借助 NVIDIA NeMo 實現長上下文訓練

    作為開發者,您可以在長上下文訓練期間通過以下方式改進內存管理:

    • 激活重新計算
    • 上下文并行
    • 激活卸載

    NeMo 框架通過先進的實施實現了這些功能,并為熱門社區模型提供長上下文方法。

    激活重新計算

    在訓練期間存儲中間激活函數所需的顯存會隨著序列長度和模型深度的增加而增加,甚至會迅速超過最大 GPU 的容量。

    NeMo 框架支持激活重新計算,這是一種可解決此瓶頸的內存節省技術。訓練過程沒有存儲反向傳播所需的所有中間激活函數,而是選擇性地僅檢查子集 (例如每個 Transformer 層的輸入) 。在向后傳遞期間計算梯度時,系統會通過重新執行部分向前傳遞來即時重新計算所需的激活值。

    通過僅存儲一小部分激活函數并重新計算其余部分,激活函數重新計算可顯著減少內存占用。這對于在有限的 GPU 顯存中擬合超長序列和大批量大小至關重要。隨著上下文長度的增加,激活內存甚至可能超過模型權重和優化器狀態所需的內存。重新計算允許將訓練擴展到更長的環境,同時保持成本效益。

    A comparison of two checkpointing strategies. Top flow showcases full checkpointing where all activations are recomputed during backprop. Bottom flow showcases selective checkpointing where only a fraction of activations are recomputed and the rest are saved.
    圖 1。檢查激活函數的子集并重新計算其余部分,以減少設備顯存占用

    上下文并行

    雖然激活重新計算通過在向后傳遞期間丟棄和重新計算激活函數來有效減少內存使用量,但這種方法會引入大量的重新計算用度 (通常每個訓練步驟高達 30%) ,從而減緩訓練過程。

    上下文并行 (CP) 提供了一種更高效的替代方案。CP 在 NeMo 框架中實現,并在針對近無限上下文的塊狀 Transformer 的 Ring Attention 中引入,它將序列維度分割到多個 GPU。每個 GPU 僅處理和存儲序列的一個塊,從而能夠在不超過顯存限制的情況下訓練具有更長輸入序列的模型。

    CP 與序列并行 (SP) 的不同之處在于,SP 僅分割幾個選定層 (例如 LayerNorm 和 Dropout) 的序列,而 CP 分割所有層的序列,通信成本通常與計算重疊。這使 CP 能夠克服單 GPU 顯存容量的限制,同時避免重新計算用度。這種方法為在長序列上訓練大型模型提供了可擴展且高計算效率的解決方案,使其成為大規模深度學習時代的強大工具。

    上下文并行的工作原理

    總的來說,CP 允許標準模組 (如 Linear、LayerNorm 和其他逐點運算) 在不進行修改的情況下運行。這些層不需要令牌間通信,因此自然支持分割序列布局。對于注意力機制,每個令牌的查詢 (Q) 必須關注同一序列中所有令牌的鍵 (K) 和值 (V) 。

    在前向傳遞期間,CP 在每個 GPU 上為其本地序列塊存儲 KV,在后向傳遞期間根據需要再次收集 KV 張量,從而更高效地利用內存。所涉及的通信集合 ( all-gather 和 reduce-scatter) 作為環形拓撲內的優化點對點通信來實現。交換 KV 還可以利用 MQA/ GQA 來減少通信量,因為它們只有一個或幾個 KV 注意力頭。

    例如,在圖 2 中,GPU0 和 GPU1 形成一個張量并行組,GPU0 和 GPU2 形成一個上下文并行組,相互交換 KV 對。GPU1 和 GPU3 之間也會執行相同的操作。CP 通過以下方式進一步提高性能:

    1. 利用最新的開源軟件 (OSS) 和 NVIDIA cuDNN 閃存注意力內核,實現更快、更節省內存的注意力計算。
    2. 消除由低三角形因果遮罩引起的不必要計算,并在 GPU 之間實現最佳負載平衡。
    Illustration of a transformer layer running using a mix of tensor parallelism and context parallelism across four GPUs.
    圖 2。使用 TP2CP2 運行的 Transformer 層

    AG/ RS:全聚在正向,并減少向后散射。RS/ AG:歸約散射在前向,全聚在后向,/ AG:無操作在前向,全聚在后向。

    CP 基準測試

    圖 3 顯示了在 Llama 3 8B 上,上下文并行的序列長度在 16K 到 100 萬個序列之間的效率。從 32K 及以上的序列長度開始,可以看到使用 CP 會產生更高的 teraFLOPS。序列長度為 100 萬時,必須使用 CP 才能運行模型。請注意,盡管序列長度增加,teraFLOPS 仍開始趨于平穩,這表明 CP 實現是以最小的開銷高效完成的。

    Bar chart showing how context parallelism results in more than 2x speedup for Llama 3 8B with sequences ranging from 16K to 1 million tokens.

    激活卸載

    CPU 卸載外,高效管理 GPU 顯存的另一項技術是 CPU 卸載CPU 卸載的工作原理是將中間激活函數和非活動權重卸載到 CPU 內存,從而減少 GPU 顯存占用峰值。NeMo 框架支持 Transformer 層卸載,允許用戶配置使用此策略的層數。在前向傳遞期間,NeMo Framework 會在最佳時間卸載激活函數,而在后向傳遞期間,它會根據需要重新加載這些函數。

    這種動態卸載機制有助于進一步擴展每個 GPU 的內存容量,尤其是在訓練非常深度的模型時,使其成為大型模型訓練中上下文并行的重要補充。

    總結

    雖然您可以實施各種技術來改進模型長上下文長度,但在進行優化時,最好考慮模型架構和硬件選擇。

    NVIDIA NeMo 框架是適用于 LLM、語音模型和多模態模型的 GPU 加速訓練框架,可提供經過測試的長上下文模型訓練方法。這些 recipe 可在 NeMo Framework LLM recipe 目錄中找到。現有方案包括 Llama 3 8B 和 70B、Mixtral 8x7B 以及 Nemotron 4 15B 和 22B,序列長度分別為 16K、64K 和 128K。

    您還可以通過預訓練檢查點擴展上下文窗口。有關更多信息,請參閱長上下文 recipe 文檔

    ?

    0

    標簽

    人人超碰97caoporen国产