Cosmos3 context parallel

[atharvajoshi10]

Cosmos3 multi-GPU inference: context + tensor parallelism

What this PR does

Adds multi-GPU inference for the Cosmos3 (Cosmos3OmniPipeline) family along two orthogonal, composable sharding axes, so the model can be run faster and/or made to fit when a single checkpoint exceeds one GPU's memory. All parallelism logic lives in examples/cosmos3/cosmos_parallel.py; the model itself stays parallelism-free apart from two tiny no-op seams.

The two axes

• Context parallelism (CP / Ulysses) — enable_cosmos3_context_parallel. Shards the sequence across GPUs; attention runs with two all-to-all collectives per layer (gather-seq/scatter-heads → local attention → gather-heads/scatter-seq). Weights are replicated, so it cuts latency but not weight memory.
• Tensor parallelism (TP) — enable_cosmos3_tensor_parallel. Shards the attention and MLP weight matrices (Megatron-style: column-parallel q/k/v + gate/up, row-parallel out + down), so a checkpoint that doesn't fit one GPU (e.g. Cosmos3-Super, ~120 GB) loads across several.

They compose on a 2-D (tp, cp) device mesh (e.g. TP=2 × CP=2 over 4 GPUs).

Design

The model carries no parallelism logic — it exposes two optional, default-None seams on Cosmos3OmniTransformer: _cp_shard_fn / _cp_gather_fn, which shard each pathway's sequence (and rotary embeddings) before the decoder stack and re-gather after the final norm. Attention parallelism lives entirely in standalone attention processors (Cosmos3CPAttnProcessor, Cosmos3FlashAttnProcessor) installed via set_processor — each is self-contained (its own __call__), so the core model file needs no override hooks.

Why a custom CP path (not the declarative _cp_plan): Cosmos3 attention has (1) grouped-query attention — KV heads must be repeated to match query heads; (2) separate understanding (causal) and generation (full) token streams, where generation attends to cat(und, gen); (3) ragged per-stream lengths that are padded independently with the padded generation keys masked. These can't be expressed declaratively.

GQA + flash: SDPA's flash/cuDNN kernels reject enable_gqa, and the native kernel falls back to math (materializing the full [S, S] scores → OOM on long sequences). Both attention paths instead expand KV heads up to the query-head count and call SDPA with enable_gqa=False, so it dispatches to flash (O(S) memory).

Usage

The example runner examples/cosmos3/inference_cosmos3.py works across all modalities (t2i / t2v / i2v / v2v / sound / action) — just launch with torchrun and pass degrees:

# CP=4 (lower latency, Nano):
torchrun --nproc_per_node 4 inference_cosmos3.py --cp-degree 4 --prompt "..."

# TP=2 × CP=2 across 4 GPUs (Super):
torchrun --nproc_per_node 4 inference_cosmos3.py --model super --tp-degree 2 --cp-degree 2 --prompt "..."

--nproc_per_node must equal --tp-degree × --cp-degree.