Context Parallelism

コンテキスト並列処理は、長い入力シーケンスを複数のアクセラレータに分割する分散コンピューティング技術です。各GPUはシーケンスの一部のみを処理し、アテンションの計算時に他のGPUと連携します。これにより、デバイスごとのアクティベーションメモリが削減され、非常に長いドキュメント、拡張ビデオ、多数の画像パッチのコレクションなど、1台のGPUのメモリを超える可能性のある入力に対してtransformerのトレーニングが可能になります。

モデルのcontext windowを単純に拡大するのとは異なり、コンテキスト並列処理はアーキテクチャが理論的に受け入れられる情報の量を変更しません。その代わりに、シーケンス次元を分散させることで、そのコンテキストの処理を計算上現実的なものにします。

コンテキスト並列処理の仕組み#

シーケンスに32,000トークンが含まれており、コンテキスト並列処理で4つのGPUを使用すると仮定します。各デバイスは最初に約8,000トークンを受け取り、対応する中間アクティベーションを保存します。

正規化やフィードフォワード層などのほとんどの操作は、これらのローカルなシーケンスチャンクを独立して処理できます。課題はattention mechanismです。ローカルのクエリは、他のすべてのデバイスが保持するキーおよびバリューにアテンションを向ける必要がある場合があります。

そのため、実装ではGPU間でキーバリュー（KV）ブロックを交換します。ring attentionでは、各デバイスはローカルデータを使用して部分的なアテンションを計算し、KVブロックを次のデバイスに渡し、完全なシーケンスを処理し終えるまでこれを繰り返します。PyTorch context parallel tutorialは、分散スケーリングドットプロダクトアテンションを通じてこの動作を示し、一方NVIDIA context parallel packageはall-gather、reduce-scatter、およびリングベースの通信オプションについて説明しています。

最終的な結果は、通常の数値的な差異を条件として、完全なシーケンスのアテンションと数学的に同等ですが、単一のGPUがすべてのシーケンスアクティベーションを保持する必要はありません。

コンテキスト並列処理が重要な理由#

長いシーケンスは2つの大きなスケーリングの問題を引き起こします。第1に、シーケンス長が長くなるにつれて、保存されたアクティベーションがより多くのメモリを消費します。第2に、標準的なセルフアテンションはシーケンス全体でトークンを比較し、大量の計算と一時データを生成します。

コンテキスト並列処理は、デバイス間でアクティベーションを分割することでメモリの問題に対処します。また、通信によって新たなコストが発生しますが、アテンションの作業を分割することもできます。効率的なシステムは、NCCL collective operations guideに記載されているような操作を使用して、KV転送と計算をオーバーラップさせます。

この技術は、モデルの重みやバッチサイズではなく、シーケンス長が原因でアウトオブメモリ（OOM）エラーが発生する場合に最も価値があります。これはdistributed trainingのより広い分野に属し、複数の次元を同時にスケーリングするために他の戦略と一般的に組み合わされます。

コンテキスト並列処理と関連技術の比較#

• **Tensor parallelism**は、個々の層内の操作または重み行列を分割します。一方、コンテキスト並列処理は、シーケンス次元に沿ってトークンを分割します。
• **Pipeline parallelism**は、モデル層の異なるグループを異なるデバイスに割り当てます。これはシーケンス長ではなくモデルの深さを分割します。
• Data parallelismはモデルを複製し、各レプリカに異なるトレーニング例を与えます。PyTorch distributed overviewは、完全なモデルと各サンプルが1つのGPUに収まる場合にこれを推奨しています。
• Sequence parallelismは、多くの場合、テンソル並列処理に関連する選択された操作のアクティベーションをシャードします。コンテキスト並列処理は、ネットワークの入力とアクティベーション全体にわたってシーケンス分割をより広範に適用します。

これらのアプローチは補完関係にあります。NVIDIA parallelism strategies guideは、コンテキスト、テンソル、パイプライン、データ並列処理がどのように多次元のデバイスレイアウトを形成できるかを示しています。

実社会での応用#

• 長文ドキュメントAI： 法律や医療の言語モデルは、訴訟記録、患者の履歴、または技術マニュアル全体を処理する必要がある場合があります。コンテキスト並列処理は、何千ものドキュメントトークンをアクセラレータに分散させ、離れたセクション間のアテンションを維持しながら、アクティベーションメモリの負荷を軽減します。

• 長尺ビデオおよびマルチモーダル理解： ビデオトランスフォーマーとlarge vision modelsは、フレーム、画像パッチ、音声セグメント、およびテキストを1つの長いトークンシーケンスとして表現する場合があります。そのシーケンスを分散させることで、モデルはフレーム数や空間的詳細を積極的に削減することなく、長時間の録画を分析できるようになります。AWS Neuron context parallelism overviewは、このような長いコンテキストのワークロードに対して、アクセラレータグループがどのようにKVシャードを交換できるかを示しています。

Ultralytics YOLO26などのコンパクトなコンピュータビジョンアーキテクチャの場合、コンテキスト並列処理は通常不要です。Ultralytics model training workflowを通じた標準的なマルチGPUデータ並列処理は、一般的にトレーニングを加速するためのより適切な方法です。

実践的な使用法とトレードオフ#

次の最小限の例では、PyTorchの実験的なコンテキスト並列APIとscaled dot-product attention controlsを使用しています。これを cp_example.py として保存し、torchrun --standalone --nproc-per-node=2 cp_example.py を使用して2つのGPU上で起動します。

import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(rank)
torch.cuda.manual_seed(0)
dist.init_process_group("nccl")

mesh = init_device_mesh("cuda", (world_size,))
qkv = [torch.randn(1, 4, 4096, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

with sdpa_kernel(SDPBackend.FLASH_ATTENTION), context_parallel(mesh, buffers=tuple(qkv), buffer_seq_dims=(2, 2, 2)):
    output = F.scaled_dot_product_attention(*qkv, is_causal=True)

output.float().square().mean().backward()
dist.destroy_process_group()

ここで、次元2はシーケンス次元であるため、各プロセスはシーケンスシャードを受け取り、アテンションはデバイスメッシュ全体で調整されます。

実際には、エンジニアはメモリ削減が通信オーバーヘッドを上回ることを確認し、高速なインターコネクトを使用し、代表的なシーケンス長でベンチマークを行う必要があります。標準的なビジョンプロジェクトの場合、Ultralytics Platformは、手動でのコンテキスト並列構成を必要とせずに、データセットのアノテーション、トレーニング、デプロイ、およびモニタリングのためのよりシンプルなクラウドおよびローカルのワークフローを提供します。