500 万トークンへの道 — Max Ryabinin (Together AI) の Untied Ulysses

AI Engineer 2026 / 講演約 16 分

マックス・リャビニン / Max Ryabinin · 13:09 「大きなコンテキスト長でモデルを訓練するのは面白く挑戦的な目標だが、ボトルネックは思いがけない場所に現れる」

AI Engineer 2026 での講演「Road to 5 Million Tokens: Breaking Barriers in Long Context Training」 (スライド上の表題は「Road to 5 Million Sequence Length: Breaking Memory Barriers in Context Parallelism」、講演約 16 分、末尾に Q&A、動画公開 2026-06-08、 AI Engineer 公式チャンネル)。講師は Max Ryabinin (Together AI、 VP of R&D)。基となる論文は「Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking」 (arXiv 2602.21196、 2026-02、 Together AI / Ravi Ghadia・Maksim Abraham・Sergei Vorobyov・Max Ryabinin)。

Together AI は AI ネイティブクラウドで、 GPU クラスタ、ファインチューニング、 200 以上のモデルを抱える推論基盤を提供する。この talk が扱うのは学習側 — 標準的な transformer を 500 万トークンのコンテキスト長まで、 1 ノード (8 基の H100) で訓練するために、メモリ最適化をどう積み上げるか。

なぜ長コンテキストか

長コンテキスト学習への関心が高まる理由を Max Ryabinin は二つ挙げる。一つはエージェントの普及 — コンテキストに入れたいトークンが増え、モデルにそれを実際に活用させたい。もう一つは動画生成のような応用 — 複数フレーム (時に毎秒複数フレーム) を追う必要があり、トークンをすぐ大量に占有する。加えて「数秒前、理想的には数分前に何が起きていたか」を見られる時間的一貫性 (temporal consistency) も要る。これらを成立させるには、学習時にモデルがその長いコンテキストを正しく処理できる必要がある。そして数百万トークン規模でなくても、メモリの行き先を理解しておけば、浮いた分を別のことに再投資して学習を速められる。

二つのボトルネック

標準的な transformer のコンテキストを延ばそうとすると、二つのボトルネックにぶつかる。一つ目は二次計算量 — transformer は系列中の全要素どうしの総当たりの相互作用を計算するので、計算が系列長の 2 乗で膨らむ (N 人が全員と握手すると約 N² 回になるのと同じ)。二つ目はより厄介で、コンテキストを延ばすほどメモリが線形に増え続ける。線形は 2 乗ほど悪くはないが、特定の техник群を当てないと扱いづらい。 Hugging Face の学習ブログの例でも、系列長の増加がメモリ上限をかなり圧迫することが示されている。

メモリを削る積み重ね

Llama3-8B を 8 基の H100 ノードで 300 万トークンに載せる、という想定で、既知の手法を順に積む様子を Ryabinin は示す。

最初はモデルのパラメータを置くだけでメモリが溢れる。次に FSDP (Fully Sharded Data Parallelism) でパラメータを 8 基の GPU に分割する — モデルのメモリは大きく下がるが、 attention の活性 (activation) でまだ溢れる。そこで context parallelism / DeepSpeed Ulysses を導入する。 Microsoft の DeepSpeed Ulysses は、全 GPU が全系列の attention を別々に計算するのではなく、 head ごとに別の GPU で計算し、必要な活性を通信し合う。 1 つの GPU が 1 つの head を担当しつつ、全系列の attention を計算する。 FlashAttention の最良実装も活かせる。これで利用メモリは約 8 分の 1 に落ちるが、まだ 1 ノードに収める目標には遠い。

続いて活性チェックポインティング (activation checkpointing) で、逆伝播時に活性を再計算する — 活性メモリをさらに約 8 分の 1 に。さらに、各 transformer ブロックの入力の一部を GPU ではなく CPU に退避し、必要になる直前にプリフェッチする CPU オフロードを当てる (約 37GB を退避)。この最適化は Unsloth が最初に実装したものだと紹介される。最後に、要素ごとの計算 (損失や MLP) を系列長方向にタイル分割し、「系列長 300 万」という巨大なバッファを作らずに済ませる。ここまで積んで、ようやく 300 万トークンが可能になる。

Untied Ulysses (UPipe) — head 単位のチャンク化

では、さらに先へ行くには。ここで本研究の主要な最適化 Untied Ulysses (UPipe) が登場する。 context parallelism のより深い分析と拡張、と Ryabinin は位置付ける。

着想はこうだ — 1 セットの head を計算するだけで、 1 反復のうちに GPU の計算能力はすでに飽和する。つまり 1 つの GPU に複数の head が割り当てられているなら、それをチャンクに分けて時間方向に反復すればよい。あるグループの head を再計算 → その attention を計算 → 部分結果を保存 → 次の段が前段で確保したバッファを再利用する。以前のように巨大なバッファを 1 つ確保する代わりに、小さなバッファを 2 回以上の反復で使い回す。小規模ではスループットを大きく損なわずに活性メモリをさらに節約できる。小さな作業机を 1 つ用意して head のグループごとに使い回す、という発想に近い。

結果と教訓

測定では、 8B (Llama3-8B) でも 32B (Qwen3-32B) でも、最もメモリ効率の良い transformer 学習の実装に近い性能を保ちつつ、 500 万トークンまでスケールでき、短い系列ではむしろ高速な場合もある。論文によれば、 Llama3-8B を 1 ノード (8×H100) で 500 万トークンまで (従来 SOTA の FPDT の 400 万から 25% 拡張)、 2 ノード (16×H100) で 800 万トークンまで訓練でき、 32B では attention の中間テンソルのメモリを最大 87.5% 削減する。チャンクサイズ (同時に計算する head 数) とスループットの関係は素直で、チャンクが大きいほどメモリ利用は高いが少し速く回せる。全部を積んで UPipe を乗せれば、浮いたメモリを別の用途 (例えばパイプラインの段間) に再投資することも、 500 万トークン級の学習を成立させることもできる。

教訓として Ryabinin が挙げるのは「ボトルネックは思いがけない場所に現れる」こと。 PyTorch profiler のようなツール (論文で詳述) が大きく助けになる、と締めくくる。 Q&A では QKV (transformer 層の query / key / value 行列) を確認し、系列長 300 万では総当たりの相互作用のために巨大なテンソルを確保することになるため、 UPipe 単体ではなく複数の手法を組み合わせて初めてメモリ不足を避けられる、と補足した。

編集所見

この talk の価値は「魔法の一手」ではなく「積み重ね」を見せたこと。 FSDP → DeepSpeed Ulysses → 活性チェックポインティング → CPU オフロード (Unsloth) → タイル化、という既知の段を順に積み、そこで止まらずに「head 単位のチャンク化 (UPipe)」という最後の一段を足す。各段がメモリのどこを削るかを 1 枚ずつ示す構成は、長コンテキスト学習が単一の発明ではなく工学的な layering で前進する分野であることを伝える。推論側で「速さ = 能力」を説いた Tanishq Kumar と対をなすように、こちらは学習側で「メモリの会計を理解すれば、同じハードでより遠くへ行ける」を実演している。一次情報 (公開論文 + コード) が揃っている点も、 archive 価値を高くする。

着眼点

「GPU はすでに飽和している」という観察が鍵

UPipe の核心は新しい数学ではなく、「1 セットの head を計算するだけで GPU の計算能力は 1 反復のうちに飽和する」という観察にある。飽和しているなら、残りの head を同時に展開する意味は薄く、時間方向に小さく回してバッファを使い回したほうがメモリを節約できる。計算資源の実測 (profiler) から最適化の余地を見つける、という Ryabinin の締めの教訓を、手法自体が体現している。

動画の構成

(00:00) 自己紹介 (Together AI、 VP of R&D)、 Together AI の概要 (クラウド / fine-tuning / 推論)
(01:53) 本題 — 学習・カスタマイズ・ファインチューニング、長コンテキストへの関心
(02:15) なぜ長コンテキストか — エージェントと動画生成、時間的一貫性
(03:42) 二つのボトルネック — 二次計算量と線形に増えるメモリ (Hugging Face ブログ例)
(04:46) 想定 — Llama3-8B / 300 万トークン / 8×H100、まず OOM
(05:43) FSDP でパラメータを分割、それでも活性で溢れる
(06:16) DeepSpeed Ulysses (context parallelism)、約 8 倍の削減
(07:48) 活性チェックポインティングでさらに約 8 倍
(08:38) CPU オフロード (Unsloth)、約 37GB を退避
(09:30) 要素ごと計算のタイル化、 300 万トークンが可能に
(10:16) Untied Ulysses (UPipe) — head 単位のチャンク化
(11:43) 結果 — 8B / 32B、 500 万トークン、チャンクサイズと速度
(13:09) 教訓 — ボトルネックは思いがけない場所に、 PyTorch profiler
(13:58) Q&A — QKV と巨大テンソルの確保

関連リンク

マックス・リャビニン

Max Ryabinin

Together AI VP of R&D / 長コンテキスト学習・分散学習システム

用語集

コンテキスト並列 (context parallelism): 長い系列を複数 GPU で分担して処理する並列化。 Microsoft の DeepSpeed Ulysses は、全 GPU が全系列の attention を計算する代わりに head ごとに別 GPU で計算し、必要な活性を通信し合う。 1 GPU が 1 head を担当しつつ全系列の attention を計算するため、 FlashAttention 等の最良実装も活かせる。
Untied Ulysses (UPipe): Together AI による context parallelism の発展 (arXiv 2602.21196)。 1 GPU に複数の head が割り当てられているとき、それらを小さなチャンクに分けて時間方向に反復処理し、小さなバッファを使い回す。巨大なバッファを 1 つ確保する従来比で活性メモリを節約し、 32B モデルで attention の中間テンソルを最大 87.5% 削減する。
FSDP (Fully Sharded Data Parallelism): PyTorch の完全シャード化データ並列。モデルのパラメータ・勾配・オプティマイザ状態を複数 GPU に分割して保持する。 8 基の GPU で 1 冊の巨大な本を 1/8 ずつ分担するイメージ。モデルのメモリは大きく下がるが、 attention の活性で依然溢れることがある。
活性チェックポインティング / CPU オフロード: 活性チェックポインティングは、順伝播の中間活性を保持せず逆伝播時に再計算してメモリを削る手法。 CPU オフロードは、活性の一部を GPU から CPU へ退避し必要直前に取り戻す手法 (この talk では Unsloth が最初に実装と紹介)。どちらもメモリを稼ぐが、再計算や転送のコストと引き換えになる。
二次計算量と線形メモリ: transformer の attention は系列中の全要素どうしの総当たりを計算するため、計算量が系列長の 2 乗で増える。一方メモリは系列長に線形に増える。長コンテキスト学習では、この二つを各種手法 (並列化・再計算・オフロード・タイル化) で抑え込む必要がある。

comment is stripped from the HTML output. */}