Road to 5 Million Tokens: Breaking Barriers in Long Context Training — Max Ryabinin, Together AI
三句話摘要
透過組合多種優化技術，在標準硬體上實現 500 萬令牌序列長度的大模型訓練。 訓練長文脈模型不是單一技術能解決，而是需要針對意想不到的內存瓶頸，系統性地堆疊多個已知與新創優化方法。 長文脈訓練的必要性在於代理應用與視頻生成等場景需要模型理解更多歷史脈絡，但標準 Transformer 的注意力機制會產生令牌數平方級的計算量與線性的內存消耗。
重點整理
重點- 1
長文脈訓練的必要性在於代理應用與視頻生成等場景需要模型理解更多歷史脈絡，但標準 Transformer 的注意力機制會產生令牌數平方級的計算量與線性的內存消耗。
- 2
DeepSpeed Ulysses 是關鍵技術突破，透過將多頭注意力分散到不同 GPU，每個 GPU 只計算一個注意力頭卻覆蓋整個序列，可以使用 Flash Attention 等最優實現，將內存使用降低約 8 倍。
- 3
U-Pipe 是 Together AI 的核心貢獻，進一步發現單個 GPU 計算一組注意力頭已足以飽和，因此可分塊迭代計算並重用緩衝區，在保持吞吐量的前提下進一步降低激活內存。
- 4
內存優化是堆疊式的：先用 FSDP 分散模型參數，再用上下文並行削減注意力激活，透過檢查點重計算、CPU 卸載輸入、序列瓦片化 MLP 等進一步壓縮，最終 U-Pipe 實現質的突破。
實用技巧與重點
乾貨- 基礎硬體配置：8×H100 GPU 節點
- 模型架構：Llama 3B
- 訓練序列長度目標：300 萬令牌，最終達到 500 萬令牌
- 關鍵技術堆疊順序及效果：
- 完全分片數據並行 (FSDP) — 模型參數分散到 8 個 GPU
- DeepSpeed Ulysses 上下文並行 — 內存使用降低約 8 倍
- 激活檢查點 (Activation Checkpointing) — 進一步降低 8 倍
- CPU 卸載 (Offloading) — 達到 37GB 數據可卸載到 CPU
- 序列瓦片化 (Tiling) — 分塊執行元素級操作 (MLP、損失函數)
- U-Pipe (新方法) — 注意力頭分塊迭代，緩衝區重用
- 效果指標：在 8B 與 32B 模型規模上記憶體使用與最優實現相當，甚至在較短序列下效能更佳
- 性能調參：塊大小越大，內存利用越高但吞吐量越低；可透過調整找到平衡點
結論
結論“訓練長文脈模型不是單一技術能解決，而是需要針對意想不到的內存瓶頸，系統性地堆疊多個已知與新創優化方法。”
完整解析
詳細大模型訓練正面臨新的需求浪潮。隨著 AI 代理應用與視頻生成等領域的發展，開發者需要訓練能理解更長文脈的模型，以便在推理時充分利用提供的上下文資訊。然而，標準 Transformer 架構面臨兩道難以逾越的障礙：首先是注意力機制的計算複雜度——查詢與鍵之間的兩兩交互導致計算量隨序列長度平方增長；其次是線性但仍可觀的內存消耗，隨著序列擴展，GPU 記憶體快速耗盡。
Together AI 的研究團隊以 Llama 3B 模型為基準，嘗試在單個 8×H100 GPU 節點上訓練 300 萬令牌序列。他們採用漸進式優化策略。首先應用完全分片數據並行將模型權重均勻分散在 8 個 GPU 上，但這只解決了模型參數的內存壓力，注意力激活依然是瓶頸。接著導入 Microsoft 開發的 DeepSpeed Ulysses 上下文並行技術，其核心思路是將多頭注意力計算解耦：每個 GPU 專責一個注意力頭，卻計算該頭對整個序列的注意力，透過頭間通信同步激活。這一步驟將內存用量銳減 8 倍，並允許使用 Flash Attention 等最優實現。
儘管如此，激活記憶體仍是瓶頸。團隊繼續堆疊優化：激活檢查點重計算機制再度削減 8 倍激活記憶體；CPU 卸載將暫時不需的層輸入移至主記憶體，反向傳播時按需取回；序列瓦片化將 MLP 與損失計算分塊執行，避免分配指數級緩衝區。最後，他們引入 U-Pipe——一個新的優化方法，將注意力頭分成更小的塊，逐塊迭代計算並重用緩衝區，既保持吞吐量又進一步降低峰值內存。透過這套組合拳，他們成功在同一硬體配置上訓練 500 萬令牌序列，並在 8B 與 32B 模型規模驗證了方法的可擴展性。
關鍵時刻
Pipeline v2帶時間戳的重點，會在逐字稿層級分析上線後產生。目前請先透過原始影片觀看。
事實查核
Pipeline v2說法查證是下一次管線升級的一部分。KeyFrame 只會顯示它真正能驗證的內容。
