MaxText 支援 1,400 萬 token 長序列訓練：TPU 怎麼分攤上下文與記憶體？
MaxText 0.2.4 提供 TPU 長上下文訓練，解析 Allgather、A2A、Ring 與 USP 的記憶體取捨，以及一千四百萬 token 測試的利用率與限制。
訓練模型讀更長的文件，困難常常先出現在記憶體。即使模型參數放得進晶片，為長篇內容保存的中間資料，仍可能在正式計算前就超出容量。
Google 在 2026 年 11 月 9 日公布 MaxText 長上下文訓練更新，以 0.2.4 版提供多種上下文平行方法。官方示範在 TPU 上處理最長一千四百萬 token 的訓練序列。
Token 是模型處理內容的切分單位，無法直接當成相同數量的中文字。這裡公布的是訓練框架的能力，也沒有替一般聊天產品新增一千四百萬 token 的輸入欄位。
MaxText 是開源的大型模型訓練框架，這次重點是把同一段長序列分配給多顆晶片，並安排它們交換計算所需資料。切法不同，會改變記憶體需求與傳輸成本。
我會建議已有 TPU 訓練需求的團隊，先看最小可行的切分規模。長度紀錄有參考價值，但只有在自己的模型與硬體上減少等待、保住吞吐量，才值得增加架構複雜度。
要選出合適的切法，得先弄清楚長序列為什麼會快速占滿記憶體，以及不同晶片到底需要互相傳送什麼。
長序列為何同時拉高計算量與記憶體需求？
語言模型要判斷某一段內容的意思，往往得參考同一段序列中的其他位置。注意力機制就是安排這種關係的計算，讓模型為相關內容分配不同的重要程度。
序列越長，位置之間需要比較的關係也越多。在完整注意力的設定下，這類計算隨長度呈平方成長，因此把輸入加長一倍，工作量可能遠超過一倍。
訓練還需要保存中間結果，供更新參數時回頭計算。這些結果稱為啟用值，連同注意力需要的資訊，會占用加速器的高頻寬記憶體。容量不足時，工作便無法依原設定完成。
TPU 是 Google 用於機器學習的加速器，一組訓練通常由多顆 TPU 合作。增加晶片可分攤負擔，但模型仍必須知道每一顆負責哪一段，並在必要時拿到其他位置的資料。
上下文平行就是把同一序列的不同位置分到不同晶片。它和把不同樣本分給各晶片的資料平行，以及把模型運算拆開的張量平行，分攤的是不同維度，實務上可以搭配。
這樣的分工能延長可處理序列，也多了溝通成本。如果切得過細，每顆晶片拿到的計算太少，反而可能花更多時間等待傳輸，因此增加切分數量本身不等於加速。
Google 公布的更新以不同傳輸安排處理這個取捨。先區分它們各自保存多少資料，就能理解為何同一個工作可能需要不同方法。
Allgather 先湊齊資料，適合容量仍足夠的工作
Allgather 的作法，是讓各晶片先交換資料，取得注意力所需的完整鍵與值。鍵和值常以 K、V 表示，負責讓查詢位置找到相關內容，並取出對應資訊。
當每顆晶片都有完整的 K、V，就可以對自己負責的查詢位置進行計算。這種安排較直接，也讓每個位置看得到整段序列，方便理解與建立基準。
代價是完整資料會在多顆晶片各自保存。序列很長時，即使查詢已經切開，K、V 仍可能占滿本地記憶體，使這種方法先碰到容量上限。
因此，若你的序列長度尚可放入記憶體，Allgather 可以作為容易比對的起點。若它因容量不足而無法執行，問題就已超出調整一般批次大小能解決的範圍，需要改變資料如何流動。
官方測試圖也顯示，較長序列可能讓不同方法的表現拉開。讀圖時需要保留模型、晶片與序列設定，不能把其中一條較高的曲線當成對所有訓練工作都成立的排行。
為了減少每顆晶片都保存完整資料的負擔，其他方法會改用注意力的不同維度重新分配，或讓資料分批經過設備。
A2A 與 Ring 改變資料分配，兩者各有成本
A2A 會重新排列資料，把原本依序列位置切分的內容，改分配到不同注意力頭。注意力頭可以理解為同一模型平行觀察內容的不同通道，各自計算部分關係。
每顆晶片改為負責較少的頭，就能在該範圍取得完整序列。這減少部分資料重複，但可切分的程度受到鍵值頭數與其他張量切分安排限制，沒有無限增加晶片的空間。
資料重排本身也要傳輸。若傳輸無法和注意力計算重疊，就會成為步驟中需要等待的部分。因此實際效率要一起考慮模型的頭數、分組與晶片間連線。
Ring 則讓每顆晶片保留自己的 K、V 區塊，再把區塊沿環狀順序傳給其他晶片。一顆設備處理眼前的區塊時，可以接收後續資料，逐段完成注意力計算。
這種方式避免每顆晶片同時放入全部 K、V，較容易支援更長的序列。若計算時間足以遮住傳輸時間，也有機會減少整段步驟的等待。
但區塊太小時，注意力核心程式可能難以維持高效率。核心程式是直接在加速器上執行的運算實作，工作過度零碎，就可能花更多時間啟動與交換，而非持續處理資料。
不同序列位置的工作量也不一定相同。環狀分配需要安排負載，避免有些設備提早做完，卻仍要等其他設備完成較重的部分，才能一起前進。
這些條件解釋了為何 Ring 並非在所有短序列都更快。它解決的是容量與溝通重疊的組合問題，要和實際長度一起比較，才能看出改善來自哪裡。
USP 結合兩種切分，仍要選對平行規模
USP 把不同切分方式放到兩個維度：一部分設備使用類似 A2A 的安排，另一部分使用 Ring。這讓團隊能在注意力區塊大小與環狀傳輸長度之間取得不同組合。
若單一長環讓傳輸步驟太多，可以把環縮短，同時利用另一個維度分攤資料。若 A2A 受注意力頭數限制，也能透過環狀維度增加可支援的序列範圍。
這項彈性仍要付出配置與測試成本。不同模型的頭數、晶片拓撲與記憶體狀況不同，能成功啟動的組合，也不一定就是完成每一步最快的組合。
官方建議的一個實用原則，是先找出剛好能放進高頻寬記憶體的最小上下文平行度。平行度表示有多少設備共同處理這個維度，過大的數量會增加交換，未必繼續受益。
MaxText 的這些方法使用 Tokamax Splash Attention 等注意力實作。對使用者而言，更新框架與採用相應設定只是起點，還需要確認依賴版本、硬體類型與模型安排是否吻合。
特別是官方也展示特定 TPU 配置中，單純點對點環狀方式已有較佳結果。混合方法提供更多調整空間，並沒有讓所有場合自動出現同一個贏家。
因此，首次比較應保留能執行的簡單基準，再逐次改一個切分條件。若模型、序列與批次量同時改變，就很難判斷速度差異是來自平行方法，還是工作本身已經不同。
一千四百萬 token 成績，要連利用率與條件一起看
官方的長度示範使用 Llama 3.1 8B 與固定六十四顆晶片。在一千四百萬 token 的設定，TPU 7 報告的模型浮點運算利用率為 48.7%，TPU 5p 為 37.1%。
模型浮點運算利用率常簡稱 MFU，用來描述模型有用計算相對於硬體能力的利用程度。它能協助觀察資源是否有效投入，但和模型回答準確率沒有直接對應。
較高的利用率也不能直接改寫成同樣比例的端到端加速。晶片世代、單顆能力與完整工作流程都會影響耗時，若要比較成本，仍須量測整段訓練需要的時間及資源。
這組資料的意義，是顯示極長序列在指定配置中仍能維持一定的運算效率。它沒有證明所有參數規模都能使用相同長度，也沒有直接展示模型學會更多內容的品質改善。
官方亦討論把部分資料移到主機記憶體。主機指的是連接加速器的伺服器，長序列帶來的計算時間，有時能和搬移工作重疊，但這仍受傳輸與運算比例限制。
若搬移時間超過可遮蔽的範圍，就會拉長步驟。團隊應同時觀察加速器記憶體高峰、主機使用量及資料傳輸等待，才能確認節省容量是否換來另一段瓶頸。
這也讓官方原文裡的版本與未來方向需要分開。0.2.4 提供的是本次長上下文訓練能力，原文談到未來要延伸的模型或注意力型態，不能直接當成這版已完整支援。
第一次比較，先讓模型與資料保持一致
評估長上下文時，可以先用同一模型、同一序列與相同批次建立比較。這讓每一步需處理的內容接近，觀察到的耗時差異才比較容易對應到資料交換方式。
接著逐步增加長度，記錄哪一個設定先碰到記憶體上限。若 Allgather 已無法執行，而 Ring 還能完成，得到的是可行範圍擴大。若兩者都能執行，就應比較完整步驟耗時。
兩種成果有不同用途。可行範圍擴大能讓原本無法訓練的序列開始運作，速度改善則影響相同工作要占用硬體多久，不能只用其中一個指標代替全部收益。
更新方式也需要可重現。官方列出 MaxText 0.2.4 的套件、容器及原始程式資源，團隊應記下實際採用的版本與設定，避免比較途中相依程式更新，使前後結果失去基準。
對既有模型而言，還要保留一組較短序列的品質與效率結果。新的安排若只在極長輸入受益，日常短輸入仍可沿用較簡單的設定，不必為了統一配置承擔額外傳輸。
當效能範圍確認，再使用保留的長文件任務檢查模型品質。這讓採用判斷能同時回答「訓練得動」與「學到需要的能力」，而非只留下最長序列紀錄。
MaxText 長上下文訓練常見問題
聊天使用者能直接輸入一千四百萬 token 嗎？
這次更新屬於訓練框架與指定硬體的示範。一般產品的可輸入容量，要以該產品或模型服務的規格為準，訓練序列長度不會自動變成所有聊天入口的容量。
哪一種上下文平行方法一定最快？
需要由模型、長度、設備連線與記憶體共同決定。先使用能放得下的最小切分規模，再比較完整步驟時間，較能避免過度切分造成傳輸成本高於收益。
這些測試已由本站重跑嗎？
本文查核的是 Google 公開的原文、配置與圖表，沒有自行執行六十四顆 TPU 的訓練。利用率與長度成績均屬官方報告，正式採用前仍需在自己的工作上重現。
長上下文是否代表模型一定更會理解長文件？
能完成計算只是必要條件。模型是否正確利用遠處資訊，還要用長文件任務與保留資料評估，不能僅由記憶體容得下或利用率較高推導出品質保證。
已有長序列訓練瓶頸的團隊，可以從超出容量的那一組設定開始，記錄切分前後能否執行、耗時與品質。若工作尚未碰到容量限制，先量現有瓶頸，再決定是否需要新增這層平行配置。
延伸閱讀：Google TPU 8i Max 與 8tx 發表：雙晶片長脈絡、機密推論為何是兩種系統設計？
延伸閱讀：Google TPU 將進 Oracle Cloud：2027 年跨雲部署，訓練與 AI 推論有何差別？
