Together AI の Max Ryabinin氏、500 万トークン到達へ長文コンテキスト訓練の壁を突破

500 万トークンの壁を破った Together AI の技術：長文コンテキスト訓練の現実解

Together AI は、標準的なトランスフォーマーアーキテクチャが直面する「計算量の二乗爆発」という物理的限界を突破し、8 GPU ノードで 500 万トークンという驚異的な長文コンテキストの訓練に成功しました。この成果は、AI エージェントや動画生成など、長大な文脈を必要とする次世代アプリケーションの実現に向けた重要な転換点となります。

標準モデルが直面する「計算とメモリ」の二重の壁

現在、コミュニティで長文コンテキストを持つモデルへの関心が高まっている背景には、AI エージェントの普及や動画生成といった新しいユースケースがあります。これらは数秒前、あるいは数分前の状況を理解するために、コンテキスト内に数百万トークンを含むことを要求します。

しかし、標準的なトランスフォーマーベースの言語モデルを拡張しようとすると、2 つの重大なボトルネックに直面します。

1. 計算量の二乗増加: シーケンス内のすべての要素間でペアごとの相互作用が発生するため、コンテキスト長が増えると計算量が指数関数的（二次関数）に増大します。
2. メモリ使用量の直線的増大: コンテキストを拡張するにつれ、活性化値の保存に必要なメモリ量が線形に増加し、ハードウェアの物理的制限にすぐにぶつかります。

Together AI の Max Ryabinin 氏はこう指摘します。

「標準的な方法では、300 万トークンのシーケンス長を持つ巨大なテンソル全体を割り当てる必要があり、これは想像できるほど非常に重要です。単一の技術に頼るだけでは不十分です。」

既存最適化手法の組み合わせで「300 万」を実現

Together AI はまず、既存の高度な最適化技術を組み合わせて、8 つの H100 GPU ノード上で Llama 3B アーキテクチャを用いた 300 万トークンの訓練を可能にしました。このアプローチは以下の 3 つの柱で構成されています。

1. DeepSpeed Ulysses によるコンテキスト並列化

マイクロソフトが開発した「DeepSpeed Ulysses」を活用し、シーケンス全体を各 GPU で個別に計算するのではなく、マルチヘッドアテンションの計算を複数の GPU に分散させました。これにより、単一の GPU が全シーケンスに対するアテンションを計算する必要がなくなり、メモリ効率を劇的に改善しています。

2. アクティベーションチェックポイント（Activation Checkpointing）

逆伝播時に必要な活性化値を再計算するこの手法は、現在ほぼすべての深層学習フレームワークで利用可能です。これにより、活性化のメモリ使用量をさらに約 8 分の 1 に削減できます。

3. CPU オフロードとタイル化

不要になった活性化データを GPU から CPU メモリにオフロードし、必要時にプリフェッチする手法です。これは Unsloth によって実装された技術で、コンテキストウィンドウを劇的に拡張します。さらに、巨大なバッファを避けるために計算要素を「タイル（チャンク）」に分割して処理することで、300 万トークンの訓練が可能になりました。

「Untitled Ulysses」がもたらす 500 万トークンへの突破

しかし、Together AI はここで止まりませんでした。彼らは独自の研究開発により、「Untitled Ulysses」と名付けた新たな手法を開発し、訓練可能トークン数をさらに 500 万まで引き上げました。

この手法の核心は、アテンションヘッドの計算チャンクを時間軸で分割・再利用する点にあります。従来の並列処理では一度にすべてのヘッドを計算するために巨大なバッファが必要でしたが、Untitled Ulysses は以下のプロセスを採用しました。

1. GPU の計算容量が飽和するレベルまで、一度に 1 つのセットのヘッドのみを計算します。
2. 複数の異なるヘッドをチャンクに分割し、時間経過とともに順次処理します。
3. 前のステージで割り当てたバッファを次のステージで再利用することで、メモリ使用量を最小化します。

このアプローチにより、Together AI はより小さなバッファを複数回再利用する形を実現しました。その結果、80 億パラメータおよび 320 億パラメータのスケールにおいて、既存の最もメモリ効率の高い実装に匹敵する性能を維持しながら、500 万トークンまでのスケーリングが可能となりました。

「チャンクが大きければメモリ利用率は高くなりますが、同時にモデル全体を少し高速で実行できます。これらの手法を積み重ねることで、トレーニング時に追加のメモリを解放し、別の場所で再投資することが可能になります。」

結論：ハードウェア制約を超えた未来へ

Together AI のこの研究は、LLM が扱う情報の量（コンテキストウィンドウ）の限界を物理的なハードウェア制約を超えて拡張する道筋を示しています。特に、既存の高性能 GPU クラスターを活用しながらメモリ効率を劇的に改善した手法は、AI エンタープライズにおけるコスト削減とスケーラビリティ向上に直結します。

長文ドキュメント分析や複雑なマルチモーダル処理の実用化に向けたこの技術的転換点は、今後の AI アプリケーション開発において不可欠な基盤となるでしょう。詳細な手法については、Together AI が公開している論文および技術ブログでさらに深く解説されています。