ウルリセス解放：通信と計算の重なりに関する実験
fal.aiは、ByteDanceのVeOmniで導入された「Async Ulysses」手法を用い、QKV計算とAll-to-All通信のオーバーラップを実現することで、大規模コンテキストにおけるビデオ拡散モデルのトレーニング効率を大幅に向上させる技術的進歩を発表した。
キーポイント
Ulyssesのボトルネックと解決策
従来のUlysses（コンテキスト並列処理の標準手法）は、QKV投影後のAll-to-All通信がシリアライズされ、パフォーマンスの主要なボトルネックとなっていた。
Async Ulyssesによる通信・計算の重なり
Q、K、Vブランチが独立している性質を利用し、あるブランチの通信中に次のブランチの計算を実行する非同期スケジューリングを採用し、通信待機時間を削減した。
実装手法とベンチマーク結果
NCCLの非同期操作またはPyTorchの機能的集合通信を用いて実装し、8台のB200 GPUでのベンチマークにより、事前SDPAチャンクのレイテンシとエンドツーエンドのステップレイテンシを改善した。
非同期通信の実装アプローチの比較と選択
従来のNCCL（async_op=True）とPyTorchの機能的なcollectives（torch.distributed._functional_collectives）の2つの実装方法があり、後者はtorch.compileとの親和性が高くスケジュールが明確なため採用された。
Symmetric MemoryによるSMリソース競合の解消
非同期通信でもGPUストリーミングマクロ（SM）リソースを共有するため競合が発生し得るため、Copy Engine経由のSymmetric Memoryを使用してデータ転送を分離し、GEMMやアテンションカーネルへのSM容量の割り当てを改善した。
GPU数によるパフォーマンスのトレードオフ
Symmetric Memoryを使用するAsync Ulyssesは2および4 GPUで性能が向上するが、8 GPUではペイロードが小さくなり固定オーバーヘッド（バッファリングやシグナリング）を相殺しきれないため、性能が低下する。
Fused QKVのスケール特性
2〜4 GPUでは最も高い最適化効果をもたらすが、8 GPU以上では通信メッセージが小さくなるため固定オーバーヘッドが支配し、恩恵はほぼ消失する。
影響分析・編集コメントを表示
影響分析
この記事は、大規模言語モデルやビデオ生成モデルのトレーニングコストを削減するための具体的な最適化手法を示しており、インフラエンジニアリングの現場において即座に適用可能な知見を提供している。特に通信ボトルネックを解消する手法は、マルチGPU環境でのスケーラビリティ向上に直結するため、AIインフラ分野の標準的なベストプラクティスを更新する重要な事例となる。
編集コメント
通信と計算のオーバーラップは高性能計算の古典的な課題だが、それをDiffusion Modelsの文脈で具体的に実装・ベンチマークした事例は貴重である。ByteDanceの研究がfal.aiのインフラにどう反映されているかの具体例として注目すべきだ。
動画拡散モデルの規模が大きくなるにつれ、シーケンス長が不都合になるほど巨大化します。実用的な解決策として文脈並列化（context parallelism）があり、その中でも Ulysses が標準的なアプローチです。核となる考え方はシンプルで、「完全なシーケンスをシャードされたヘッドに分割する」ことです。これは高スループットのアールツーオール通信（all-to-all communication）のために構築された現代の GPU クラスターに適しており、大規模な文脈長でも密着型アテンション（dense attention）を実行し続けることができます。
Ulysses は以下の straightforward な実行フローに従います：
QKV 射影：各ランクが局所的に Q、K、V を計算します。
アテンション前アールツーオール：ランク間でシャードを交換し、各 GPU がアテンションに必要なシーケンススライスを受け取ります。
アテンション計算：集約された Q/K/V 上で SDPA（Scaled Dot-Product Attention）が実行されます。
アテンション後パス：出力は再び交換され、出力射影が適用されます。
このフローを具体的に理解するために、まず Baseline Ulysses のアテンション前チャンクから始めます。各ランクは 3 つの局所的な GEMM（Q、K、V）を実行した後、SDPA を実行する前に 3 回のアールツーオール交換を行います。このチャンク段階が主な通信ボトルネックであり、本稿の後半で最適化の対象となる部分です。
imageAs video diffusion models scale, sequence lengths get uncomfortably large. A practical answer is context parallelism, and within that family, Ulysses is the canonical approach. The core idea is simple: full sequence, sharded heads. It maps well to modern GPU clusters built for high-throughput all-to-all communication, while still letting you run dense attention at large context lengths.
Ulysses follows a straightforward execution flow:
QKV projection: each rank computes local Q, K, and V.
Pre-attention all-to-all: ranks exchange shards so each GPU has the sequence slice needed for attention.
Attention compute: SDPA runs on gathered Q/K/V.
Post-attention path: outputs are exchanged back and the output projection is applied.
To make the flow concrete, we start with the pre-attention chunk in baseline Ulysses. Each rank runs three local GEMMs (Q, K, V), then performs three all-to-all exchanges before SDPA. This chunk stage is the main communication hotspot and the optimization target for the rest of the post.
ベンチマーク設定。8 枚の B200 GPU（bf16、B=1、H=40、D=128）上で、固定された GPU クロックと cuDNN アテンションバックエンドを用いた合成 Ulysses プリアテンションベンチマークを実行します。各設定は 15 回のウォームアップ反復と 40 回の計測反復を 3 回繰り返すもので、最大ランクの中央値レイテンシを報告します。ここでは、プリアテンションチャンクレイテンシ（QKV プロジェクション + プリ SDPA コミュニケーション）とエンドツーエンドステップレイテンシの両方を報告します。ストロングスケーリングでは GPU 数が増加してもグローバルシーケンス長は固定され、ウィークスケーリングではローカルシーケンス長を固定してワールドサイズに応じてグローバルシーケンス長が拡大するように設定されます。
非同期 Ulysses
従来の Baseline Ulysses は、プリアテンションパスが主に逐次実行されるため、性能を十分に引き出せていません。各ランクで Q、K、V を計算した後、3 回の all-to-all 交換を実行し、その後 SDPA（Scaled Dot-Product Attention）に入ります。
ByteDance が VeOmni の一部として導入した Async Ulysses は、アテンション前の Q/K/V ブランチが互いに独立しているという単純な観察に基づいています。これにより、数学的計算を変更することなく、あるブランチからの通信を次のブランチの計算と重畳（オーバーラップ）させることが可能になります。
その結果得られるスケジューリングは以下の通りです：
Q を計算し、Q の all-to-all を起動
Q の all-to-all が実行されている間に K を計算し、K の all-to-all を起動
K の all-to-all が実行されている間に V を計算し、V の all-to-all を起動
集約されたテンソルを待機/復元した後、SDPA を実行
この結果生じるオーバーラップは以下の図に示されています
このオーバーラップを実装する実用的な方法は 2 つあります。1 つ目は古典的な NCCL のアプローチで、非同期操作 (async_op=True) を使用して集合演算を開始し、Work ハンドルを保持した上で後で明示的な待機処理を行う方法です。2 つ目は PyTorch 関数型集合演算 (torch.distributed._functional_collectives) で、これは AsyncTensor 値を返すため、開始と待機の動作がテンソルフロー内で直接表現されます。ここでは新しい機能であり、torch.compile との親和性が高く、オーバーラップスケジュールの記述と推論がより明確になるため、関数型パスを採用しています。
非同期 Ulysses は期待通りの動作を示します：2/4/8 GPU 環境でチャンクあたりのレイテンシが約 23–25% 低下し、エンドツーエンドの性能は約 3% 向上します。オーバーラップにより事前アテンションの直列化が除去されましたが、SDPA と出力処理が依然としてステップ全体の時間を支配しています。
対称メモリを備えた非同期 Ulysses
現在の Async Ulysses パスには、微妙な制限があります。通信を別々のストリーム上で非同期に起動しているにもかかわらず、それらは依然として同じ基盤となる SM（Streaming Multiprocessor）リソース上で実行される可能性があります。NCCL コレクティブは、ロード/ストアおよび還元演算を発行する GPU カーネルとして実装されているため、SM サイクルを GEMM と競合します。実際には、これはタイムライン上の重なりが必ずしもスループット上の重なりに対応しないことを意味します：通信と計算は時間的に重なることがあっても、互いにスロットリングし合う可能性があります。
この競合を軽減するために、Async Ulysses は変更せず、トランスポート層のみを変更します。プレアテンション転送をコピーエンジン経由でルーティングすることで、データ移動を SM 計算から分離し、GEMM やアテンションカーネルにより多くの中断のない SM 容量を残します。PyTorch Symmetric Memory を使用してこの経路にアクセスしつつ、同じアテンション数式と実行スケジュールを維持しています。
async_symm は 2 台および 4 台の GPU では優れていますが、8 台の GPU では性能が低下します。ペイロードが大きくなると、コピーエンジンによるトランスポートは SM の競合を軽減し、実質的な重なりを改善しますが、8 台の GPU ではペイロードが小さくなるため、Symmetric Memory の固定オーバーヘッド（バッファリング/シグナリング/ブックキーピング）を分散することが難しくなります。
Fused QKV Projections
非同期 Ulysses は、オーバーラップが機能することを証明しました。オーバーラップの自然な次のステップは融合です：プレアテンションパスで起動されるカーネルと集合演算の数を減らすことです。個別の Q/K/V 射影と個別の通信交換を実行する代わりに、各ランクは 1 つのパックされたローカル重みシャード（そのローカルヘッドの Q|K/V 行）を構築し、torch.ops.symm_mem.fused_all_gather_matmul を呼び出します。この単一の演算は、シーケンスのアールゲザーとローカルヘッドの行列積を同時に実行し、パックされた (B, S_global, 3*H_local*D) の出力を返します。その後、パックされた出力を Q/K/V に分割し、SDPA 用に集約されたローカルヘッドテンソルへ再整形します。
これは 2 GPU と 4 GPU で最も強力なチャンク最適化であり（それぞれ -37.3%、-33.4%）、そこでのエンドツーエンドの向上も最大となります（それぞれ -5.0%、-4.8%）。8 GPU ではその恩恵はほとんど消滅します（チャンクで -4.6%、全体で -0.3%）：強いスケーリング下ではメッセージサイズがすでに小さいため、固定された融合オーバーヘッドが支配的になります。
弱スケーリング
強いスケーリングでは、GPU 数が増えるにつれてローカルシーケンスが縮小するため、通信コストを隠蔽できてしまいます。そのため、弱スケーリングも実行します：ローカルシーケンスを固定し、グローバルシーケンスを増大させます。
ローカルシーケンス長 16K における弱スケーリングでは、Fused QKV は 2〜4 GPU で最も性能を発揮しますが、Async Ulysses はより滑らかにスケールし、8 GPU で最適となります。Symmetric Memory を備えた Async Ulysses は小規模な環境でも競争力がありますが、規模が大きくなるにつれて性能が急速に低下します。
教訓：オーバーラップ（重なり）処理は高スケーリングにおける最も堅牢なデフォルト戦略であり、融合（Fusion）は中低スケーリング領域において最も強力です。
結論
これらの結果は 8xB200 ノードから得られたものであり、ここでは all-to-all 通信がすでに非常に高速です。より低速な相互接続環境では、隠蔽可能な通信の余裕（slack）がさらに存在する可能性があります。
より広範な視点では、スタックはより緊密な通信・計算融合へと移行しています：デバイス起動型の対称メモリ/NVSHMEM スタイルの通信、Triton カーネル、そして永続的な実行です。Kraken はこの方向性に対する優れたレシピスタイルの参照資料となります。
単一の戦略がすべての領域で勝利するわけではないため、軽量なランタイムポリシーがシーケンス長、ワールドサイズ、およびハードウェア/相互接続の特徴に基づいて、パッキング（詰め込み）、オーバーラップ、融合のいずれを選択するかを決定できるべきです。
通信集約型ワークロード（例えば MoE ルーティング）は、この密着アテンション設定よりもさらに大きな恩恵を受けるはずです。
参考文献
DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
VeOmni: Scaling Any Modality Model Training with Model-Centric Distributed Recipe Zoo
原文を表示
imageAs video diffusion models scale, sequence lengths get uncomfortably large. A practical answer is context parallelism, and within that family, Ulysses is the canonical approach. The core idea is simple: full sequence, sharded heads. It maps well to modern GPU clusters built for high-throughput all-to-all communication, while still letting you run dense attention at large context lengths.
Ulysses follows a straightforward execution flow:
QKV projection: each rank computes local Q, K, and V.
Pre-attention all-to-all: ranks exchange shards so each GPU has the sequence slice needed for attention.
Attention compute: SDPA runs on gathered Q/K/V.
Post-attention path: outputs are exchanged back and the output projection is applied.
To make the flow concrete, we start with the pre-attention chunk in baseline Ulysses. Each rank runs three local GEMMs (Q, K, V), then performs three all-to-all exchanges before SDPA. This chunk stage is the main communication hotspot and the optimization target for the rest of the post.
Benchmark setup. We run a synthetic Ulysses pre-attention benchmark on 8 B200 GPUs (bf16, B=1, H=40, D=128) with fixed GPU clocks and the cuDNN attention backend for attention compute. Each configuration uses 15 warmup iterations and 40 timed iterations, repeated 3 times; we report max-rank median latency. We report both pre-attention chunk latency (QKV projection + pre-SDPA communication) and end-to-end step latency. Strong scaling keeps global sequence fixed as GPU count grows, while weak scaling keeps local sequence fixed so global sequence grows with world size.
Async Ulysses
Baseline Ulysses leaves performance on the table because the pre-attention path is mostly serialized: each rank computes Q, K, and V, then runs three all-to-all exchanges, then enters SDPA.
Async Ulysses (as introduced by ByteDance as part of VeOmni) uses a simple observation: the Q/K/V branches are independent before attention. That lets us overlap communication from one branch with compute from the next, without changing the math.
The resulting schedule is:
Compute Q, launch Q all-to-all
Compute K while Q all-to-all is in flight, then launch K all-to-all
Compute V while K all-to-all is in flight, then launch V all-to-all
Wait/restore gathered tensors, then run SDPA
The resulting overlap is shown in the diagram below
There are two practical ways to implement this overlap. The first is the classic NCCL route: launch collectives with async_op=True, carry Work handles, and place explicit waits later. The second is PyTorch functional collectives (torch.distributed._functional_collectives), which return AsyncTensor values so launch/wait behavior is expressed directly in the tensor flow. We use the functional path here because it is newer, composes better with torch.compile, and makes the overlap schedule cleaner to write and reason about.
Async Ulysses does what we want: chunk latency drops by about 23–25% at 2/4/8 GPUs, while end-to-end improves by ~3%. The overlap removes pre-attention serialization, but SDPA and output still dominate total step time.
Async Ulysses with Symmetric Memory
There is a subtle limitation in our current Async Ulysses path. Even though we launch communication asynchronously on separate streams, they can still run on the same underlying SM resources. NCCL collectives are implemented as GPU kernels that issue loads/stores and reduction work, so they compete with GEMMs for SM cycles. In practice, this means timeline overlap does not always translate into throughput overlap: communication and computation can overlap in time while still throttling each other.
To reduce this contention, we keep Async Ulysses unchanged and change only the transport layer. Routing pre-attention transfers through the Copy Engine separates data movement from SM compute, which leaves more uninterrupted SM capacity for GEMMs and attention kernels. We use PyTorch Symmetric Memory to access that path while keeping the same attention math and execution schedule.
async_symm is better at 2 and 4 GPUs, but regresses at 8 GPUs. When payloads are larger, copy-engine transport reduces SM contention and improves effective overlap; at 8 GPUs payloads are smaller, so Symmetric Memory fixed overheads (buffering/signaling/bookkeeping) are harder to amortize.
Fused QKV Projections
Async Ulysses proved overlap works. A natural follow-up to overlap is fusion: reduce how many kernels and collectives the pre-attention path launches. Instead of running separate Q/K/V projections and separate communication exchanges, each rank builds one packed local weight shard (Q|K|V rows for its local heads) and calls torch.ops.symm_mem.fused_all_gather_matmul. That single op performs sequence all-gather and local-head matmul together, returning a packed (B, S_global, 3*H_local*D) output. We then split the packed output into Q/K/V and reshape back to gathered local-head tensors for SDPA.
This is the strongest chunk optimization at 2 and 4 GPUs (-37.3%, -33.4%) and gives the best end-to-end gains there (-5.0%, -4.8%). At 8 GPUs the benefit mostly vanishes (-4.6% chunk, -0.3% total): messages are already small under strong scaling, so fixed fusion overheads take over.
Weak Scaling
Strong scaling can hide communication costs because local sequence shrinks with more GPUs. So we also run weak scaling: keep local sequence fixed and let global sequence grow.
For weak scaling at local-seq 16K, Fused QKV is best at 2-4 GPUs, while Async Ulysses scales more smoothly and is best at 8 GPUs. Async Ulysses with Symmetric Memory is competitive at lower scale, then degrades faster.
Takeaway: overlap is the most robust high-scale default; fusion is strongest in lower/mid-scale regimes.
Conclusion
These results come from an 8xB200 node, where all-to-all is already very strong; on slower interconnects, there is likely more communication slack to hide.
More broadly, the stack is moving toward tighter comm-comp fusion: device-initiated symmetric-memory/NVSHMEM-style communication, Triton kernels, and persistent execution. Kraken is a good cookbook-style reference for that direction.
Since no single strategy wins every regime, a lightweight runtime policy could choose between packing, overlap, and fusion based on sequence length, world size, and hardware/interconnect characteristics.
Communication-heavy workloads (for example, MoE routing) should benefit even more than this dense-attention setup.
References
DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
VeOmni: Scaling Any Modality Model Training with Model-Centric Distributed Recipe Zoo
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み