AI Engineer Europe 2026 · 演講 · 2026-06-08 · 16 分
把訓練長度推到 500 萬 token
Road to 5 Million Tokens: Breaking Barriers in Long Context Training · Max Ryabinin，Together AI
Together AI 的研發副總 Max Ryabinin 說明長 context 訓練的記憶體都花在哪，並疊加多種既有技巧，再加上自家的 UPipe，把序列長度從 300 萬推到 500 萬 token。
概要
- 長 context 訓練為什麼受關注Max 說原因有兩個，一是 agent 愈來愈普遍，應用會想塞很多 token 進 context；二是影片生成要追蹤很多畫面，還要有時間上的一致性。
- 兩個瓶頸標準 transformer 的計算量是二次方，因為序列裡每個元素之間都要兩兩互動；記憶體則隨 context 線性成長，沒那麼糟，但不用特殊技巧還是很難處理。
- 一層一層疊技巧Max 以 300 萬 token 的訓練為例，依序用 fully sharded data parallelism、DeepSpeed Ulysses 的 context parallelism、activation checkpointing、把部分輸入搬到 CPU，以及沿序列長度切塊處理 element-wise 的運算，最後才做到 300 萬。
- UPipe 再往前推他們發現一次只算一組 attention heads 就足以吃滿 GPU 的算力，所以把 heads 分批依序算，後一批重複使用前一批配置的緩衝區，省下 activation 記憶體。
- 結果與取捨在 8 billion 與 32 billion 規模上，記憶體用量貼近最省記憶體的實作，而且能擴到 500 萬 token；chunk 愈大，記憶體用量愈高，但整個模型跑得愈快。
- 瓶頸常在意想不到的地方Max 建議用 PyTorch profiler 這類工具找，論文裡對此著墨很多。
原話
▶ 03:27“it's still quite important to understand where the memory goes”
即使沒到百萬 token 的規模，搞清楚記憶體花到哪裡去仍然很重要
▶ 04:03“pairwise interactions across all the elements in your sequence”
序列裡所有元素之間的兩兩互動，這是計算量二次方的來源
▶ 04:18“your memory keeps growing linearly”
context 愈長，記憶體就線性成長
▶ 13:05“the bottlenecks might appear where you least expect”
瓶頸可能出現在你最意想不到的地方
延伸
- Max 提到除了論文，之後還會有一串貼文把方法講得更深入。
- Max 說把輸入搬到 CPU 對效能影響不大，因為只有反向傳播到對應那一層時才需要用到。
章節
點時碼直接跳到原片的那一段。
- 開場與題目Max 自我介紹是 Together AI 的研發副總，要講的研究專案是把 context parallelism 的記憶體障礙打破，推到 500 萬序列長度。
- Together AI 簡介Max 說公司是 AI 原生雲，涵蓋 GPU 叢集、微調與強化學習等模型客製化，以及開源模型的推論服務，組合裡有 200 多個模型。
- 為什麼要訓練長 context社群對長 context 訓練的興趣來自 agent 與影片生成；即使沒到百萬 token，搞清楚記憶體去哪了也值得，省下來的可以轉投資在別處。
- 二次方計算與線性記憶體延伸 context 會碰到兩個瓶頸：兩兩互動帶來的二次方計算，以及隨長度線性成長的記憶體。
- 實驗設定目標是看現有技巧加上額外優化能走多遠；Max 拿 300 萬訓練 token 放進單一 H100 節點當例子，光放模型參數就已經塞不下。
- 先用 FSDP 切參數參數分散到 8 張 GPU 後，模型的記憶體用量明顯下降，但 attention 的 activation 仍然讓記憶體爆掉。
- DeepSpeed Ulysses每張 GPU 只負責一個 attention head，但仍對整條序列算 attention，還能沿用 flash attention 等最佳化的實作；用量降了不少，離單節點目標還很遠。
- Activation checkpointing到 backward 時才重算 activation，可再把 activation 用量降到約八分之一，但還不夠。
- 把輸入搬到 CPU每個 transformer block 的部分輸入暫存到 CPU，用不到時就不佔 GPU；據 Max 所知這項優化最早是另一個專案實作的，能大幅擴張 context window。
- 沿序列長度切塊loss 與 MLP 這類 element-wise 的計算可以分塊，避免產生一個某個維度有 300 萬那麼大的緩衝區，這樣才終於做到 300 萬。
- 主要貢獻 UPipe他們的核心優化，是對 context parallelism 的進一步改法；發現一次算一組 heads 就能吃滿 GPU 算力。
- 分批算 heads排在同一張 GPU 的多組 heads 依序分批處理，存下部分結果，下一批重用上一批配置的緩衝區，不必一次配置整塊大緩衝區。
- 量測結果8 billion 與 32 billion 規模上，與最省記憶體的實作相當接近，並能擴到 500 萬 token。
- chunk 大小的取捨chunk 大則記憶體用量高、整體跑得較快；省下的記憶體可以轉投資到別處，或是用來訓練更長的 context。
- 結語與提問瓶頸可能出現在最意想不到的地方，PyTorch profiler 等工具能幫上忙；論文已公開，之後還會有一串貼文講得更深入。
- 問答回答聽眾提問時，Max 說指的只是 transformer 層的 query、key、value 矩陣；序列有 300 萬長時，一般做法會配置一個某軸達 300 萬的大張量，所以得用 UPipe 加上其他方法才不爆記憶體。
名詞
- Together AI
- Max 任職的 AI 原生雲，提供 GPU 叢集、模型客製化與開源模型推論
- UPipe
- 這場介紹的主要優化，把 attention heads 分批依序算並重用緩衝區
- DeepSpeed Ulysses
- Microsoft 最先提出的 context parallelism 技術，每張 GPU 負責一個 attention head，但仍對整條序列算 attention
- FSDP
- fully sharded data parallelism，把模型參數切分到多張 GPU
- Activation checkpointing
- 不存 activation，到 backward 時才重算，用算力換記憶體
- PyTorch profiler
- Max 建議用來找出訓練瓶頸的工具，論文裡有詳細討論
