Max Ryabinin, Together AI · 15:50
這支影片在講：5M token 長上下文不是只把 window 拉大
Max Ryabinin 拆解 Together AI 怎麼把 LLaMA 類模型訓練推到 3M、5M tokens。重點不是單一神招，而是一整層記憶體工程：sharding、context parallelism、checkpointing、CPU offload、chunked sequence training，最後再用 Untied Ulysses 繼續壓 activation buffer。
Published: 2026-06-08Video ID: TUnPNY4E2fw
01 / 問題
3M tokens 先卡在記憶體
標準 LLaMA 3B 在單台 8xH100 上訓練 3M context，甚至還沒開始跑，模型參數就已經吃掉 GPU memory。
02 / 解法
一層一層把記憶體拆掉
FSDP 分攤參數；Ulysses context parallelism 把 activation 降 8x；activation checkpointing 再降 8x；CPU offload 和 chunking 避免巨大 buffer。
03 / 反轉
5M 還是不夠
到 5M tokens 時，上述堆疊仍不足。Untied Ulysses 把 attention head buffer 再切細並跨迭代重用，才繼續往前推。
Watch for this
影片最值得抓的四個訊號
Context target
3M → 5M
3M 已經會讓標準訓練流程在啟動前爆掉；5M 需要新的 buffer 策略。
Node shape
8xH100
討論基準是單台 8 張 H100，不是無限 cluster 下的抽象問題。
Two reductions
8x + 8x
Ulysses 先降 activation；checkpointing 再用重算換記憶體。
Sequence gain
+25%
Untied Ulysses 相對既有 Ulysses 實作，把可訓練 sequence length 再推遠。
Core stack
把 5M token 訓練撐起來的六層工程
Layer 1
Fully Sharded Data Parallelism：先別讓每張卡都背完整模型
3B 模型在長上下文下已經很快碰到記憶體牆。FSDP 的角色是把參數、梯度與 optimizer state 分散，避免每張 GPU 都保留完整副本。
Layer 2
DeepSpeed Ulysses：把 context 維度切開
Ulysses context parallelism 針對 attention / sequence 維度切分工作，RSS 描述中點出它帶來約 8x activation reduction，是長上下文訓練的第一個大幅降記憶體步驟。
Layer 3
Activation checkpointing：用重算換第二個 8x
不把所有中間 activation 都存著，而是在 backward 需要時重算。這不是免費午餐，會增加 compute，但在 long context 下是把訓練跑起來的必要交易。
Layer 4
CPU offloading：把 transformer block inputs 移出 GPU
當 GPU memory 是硬限制，某些 block inputs 可以被搬到 CPU，讓 GPU 留給真正高壓的 attention 與 activation。代價是資料搬移與 throughput trade-off。
Layer 5
Chunked sequence training：不要配置 3M token 寬的 buffer
如果一次配置完整 sequence 寬度的 buffer，記憶體會直接失控。chunking 的核心是把「看起來是一條超長序列」的工作切成可管理片段。
Layer 6
Untied Ulysses：把 head group buffer 再切細並重用
到 5M tokens 時，前面五層仍不夠。Untied Ulysses 不再為每個 attention head group 配一大塊 buffer，而是把 heads 進一步 chunk，跨迭代重用 buffer，降低 activation memory，且 throughput 影響很小。
Why it matters
這支影片對 AI infra 團隊的啟發
長 context 是系統問題，不只是模型規格把 context window 標成百萬 tokens 很容易；要能訓練，必須同時處理參數、activation、attention buffer、通訊與 offload。
記憶體最佳化通常是堆疊，不是單點FSDP、Ulysses、checkpointing、CPU offload、chunking 各自處理不同記憶體來源；少一層就可能跑不起來。
throughput trade-off 要被明確量化checkpointing 與 offload 會用 compute / IO 換 memory；Untied Ulysses 的價值在於降低 memory 同時維持近似 throughput。
5M context 會逼出新的 parallelism 介面既有 Ulysses 到極端長度會撞到 buffer 形狀限制；Untied Ulysses 的重點是改掉 buffer allocation 與 reuse 方式。
Use this checklist
如果你要評估長上下文訓練方案，先問這五題
你的 memory budget 有拆成參數、activation、optimizer state、buffer 嗎？只看模型大小會低估 long context 的壓力；activation 和 attention buffer 才常是爆點。
context parallelism 是否真的降低 activation，而不是只把 batch 拆開？長序列瓶頸在 sequence 維度，parallelism 策略要對準 sequence / attention 的 memory shape。
checkpointing、offload、chunking 的代價有沒有被量到？跑得起來只是第一步；production 訓練還要看重算成本、CPU-GPU 搬移與 end-to-end throughput。
buffer allocation 是否會隨 sequence length 線性爆炸？如果某個 buffer 仍必須一次配置 3M 或 5M token 寬度，後面再多 GPU 也可能只是延後撞牆。
你的方法在 8B、32B 等不同規模是否仍成立？影片提到 Untied Ulysses 在 8B 與 32B scale 都接近最省記憶體 transformer training baseline，這比單一 toy setting 更有參考價值。