Cosmos3 Model Parallelism (#14054)

* feat(cosmos3): multi-GPU inference — context + tensor parallelism

Cosmos 3 cannot use diffusers' declarative `_cp_plan` CP path: it is grouped-query
attention (the shared Ulysses kernel assumes K/V share the query head count), its
understanding (causal) and generation (full) streams are separate packed sequences
(gen attends to cat(und, gen)), and per-pathway lengths are ragged. The model carries
no parallelism logic -- it exposes only small, CP-agnostic seams; all sharding lives
outside it, in a reusable example module.

Model (transformer_cosmos3.py): adds two default-None `forward` seams -- `_cp_shard_fn`
(shards und/gen + rotary before the decoder layers) and `_cp_gather_fn` (gathers/unpads
after the final norm) -- and extracts `Cosmos3AttnProcessor._run_attention` as an
override point. The non-parallel path is unchanged.

Helpers (examples/cosmos3/cosmos_parallel.py): one importable module, two orthogonal
and composable axes:
  * Context parallelism (Ulysses) -- `enable_cosmos3_context_parallel`. Shards the
    sequence; brackets the two attention pathways with all-to-all (DTensor redistribute),
    repeats GQA KV heads, pads ragged lengths and masks padded generation keys.
  * Tensor parallelism (Megatron) -- `enable_cosmos3_tensor_parallel`. Column/row-shards
    the attention + MLP weights so a checkpoint that does not fit one GPU (Super, ~120 GB)
    loads across several; weights load to CPU then shard layer by layer.
Both expand KV heads to the query-head count and call SDPA with enable_gqa=False so it
dispatches to the flash kernel; enable_gqa=True forces the math path, which materializes
the full [S, S] score matrix and OOMs on long videos. A dense `Cosmos3FlashAttnProcessor`
(`enable_cosmos3_flash_attention`) provides the same for TP without CP.

CLI (examples/cosmos3/inference_cosmos3.py): imports these helpers, so any modality
(text-to-image/video, image-to-video, sound, action) runs single- or multi-GPU via
`--tp-degree` / `--cp-degree` (their product must equal --nproc_per_node). Single-GPU
behavior is unchanged.

Docs + example README updated. Verified: CP attention core is bit-exact vs non-CP in
fp32 (max|d|=0), and a full 36-layer forward matches CP-on vs CP-off to ~1e-6 in fp32
(bf16 differs only by floating-point rounding).

* refactor(cosmos3): make CP/TP attention processors standalone

Address review feedback on PR #14054: the parallel attention processors no
longer subclass Cosmos3AttnProcessor, so the model file needs no override seam.

- transformer_cosmos3.py: revert Cosmos3AttnProcessor to inline the attention
  in __call__ (remove the _run_attention seam); restores it to its base version.
- cosmos_parallel.py: Cosmos3CPAttnProcessor and Cosmos3FlashAttnProcessor each
  get their own full __call__, sharing a _project_qkv_with_rope prologue helper.

Verified behavior-preserving on 4x RTX PRO 6000: cp_unit_test (fp32) passes at
1e-4; cp_numeric_check is byte-identical to the pre-refactor code; the end-to-end
CLI passes in CP-only, TP-only, and TP+CP modes.

* fix(cosmos3): support modular multi-GPU inference

Author: atharvajoshi10

docs/source/en/api/pipelines/cosmos3.md

@@ -660,6 +660,125 @@ if result.action is not None:
 </hfoption>
 </hfoptions>
 
+## Context parallelism
+
+For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens.
+
+Cosmos 3 supports **Ulysses** context parallelism (all-to-all sequence/head exchange). Ring attention is not supported.
+
+Unlike most diffusers models, Cosmos 3 does **not** wire CP into the transformer or the declarative [`~ModelMixin.enable_parallelism`] path: its grouped-query attention, separate understanding/generation streams (the generation stream attends to both), and ragged per-stream lengths can't be expressed as a `_cp_plan`. Instead, the model exposes small no-op shard/gather seams, and the implementation lives in [`examples/cosmos3/cosmos_parallel.py`](https://github.com/huggingface/diffusers/blob/main/examples/cosmos3/cosmos_parallel.py) — a self-contained module you can read end to end and adapt. It offers two orthogonal, composable sharding axes:
+
+| Helper | Shards | Use for |
+|---|---|---|
+| `enable_cosmos3_context_parallel(transformer, cp_mesh)` | sequence (CP / Ulysses) | latency on a model that fits one GPU (`Nano`) |
+| `enable_cosmos3_tensor_parallel(transformer, tp_mesh)` | weights (TP) | fitting a model that doesn't fit one GPU (`Super`) |
+
+Use either alone or both together on a 2-D `(tp, cp)` mesh (see [Fitting large models with tensor parallelism](#fitting-large-models-with-tensor-parallelism)).
+
+Two requirements are specific to Cosmos 3:
+
+- Use the `native` attention backend. Cosmos 3 uses grouped-query attention (GQA), and the native SDPA backend is the only one that accepts `enable_gqa` (cuDNN and flash reject it). The helpers expand the KV heads to the query-head count and call SDPA with `enable_gqa=False` so it still dispatches to the flash kernel (the math fallback would materialize the full `[S, S]` scores and OOM on long sequences).
+- The CP (Ulysses) degree must divide the query-head count (32 for `Nano`, 64 for `Super`); for TP, the degree must divide the KV heads (8). The understanding (text) and generation (video/sound) streams are sharded independently along the sequence, and ragged lengths are zero-padded internally to a multiple of the world size.
+
+### Run it
+
+The full CLI [`examples/cosmos3/inference_cosmos3.py`](https://github.com/huggingface/diffusers/blob/main/examples/cosmos3/inference_cosmos3.py) uses [`Cosmos3OmniModularPipeline`] and reuses these helpers, so **any modality** (text-to-image/video, image-to-video, sound, action modes) runs multi-GPU via `--tp-degree` / `--cp-degree`. Launch with [torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html); `--tp-degree * --cp-degree` must equal `--nproc_per_node`. Every rank produces the same output; rank 0 writes it.
+
+```bash
+# CP only — Nano (fits one GPU); CP degree must divide 32 query heads.
+torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py --model nano --cp-degree 4 --prompt "..."
+
+# TP only — Super; TP degree must divide 64 query heads and 8 KV heads.
+torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py --model super --tp-degree 4 --prompt "..."
+
+# TP + CP — Super, with sound (TP=2 x CP=2 across 4 GPUs).
+torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py \
+    --model super --tp-degree 2 --cp-degree 2 --enable-sound --prompt "..."
+```
+
+`Super`'s ~120 GB of weights do not fit on one 96 GB GPU, so it needs TP; `Nano` fits on a single GPU, so CP for it is a pure latency optimization. (Omit both flags to run single-GPU.)
+
+### Fitting large models with tensor parallelism
+
+CP shards *activations* but replicates every weight on every rank, so it does not reduce a model's weight footprint — a model that doesn't fit on one GPU still won't fit under CP alone. To shard the **weights**, `enable_cosmos3_tensor_parallel(transformer, tp_mesh)` applies Megatron-style tensor parallelism on a second, orthogonal mesh axis:
+
+- The attention and MLP projections are column/row sharded across the TP group (`to_q/to_k/to_v` + `add_q/k/v` and the MLPs' `gate/up` are column-parallel; `to_out/to_add_out` and the MLPs' `down` are row-parallel with an all-reduce). Each rank ends up owning `query_heads / tp` query heads and `kv_heads / tp` KV heads.
+- TP composes with CP on a 2-D `(tp, cp)` device mesh: TP splits heads/weights persistently, CP shards the sequence on top. The constraints are `tp` divides the KV heads (8), and `tp * cp` divides the query heads (32 for `Nano`, 64 for `Super`).
+- Weights are loaded to CPU and sharded onto the GPUs layer by layer, so the full model is never materialized on a single device.
+
+> [!TIP]
+> TP issues an all-reduce on every attention and MLP block, so it is bandwidth-heavy. On hosts without NVLink it is the dominant cost; prefer the smallest TP degree that makes the weights fit and put the remaining GPUs into CP.
+
+### Use it in your own modular pipeline
+
+The CLI flags are convenient, but you can call the helpers directly with [`Cosmos3OmniModularPipeline`]. Load the pipeline configuration and components on CPU, apply TP *before* moving the pipeline to the rank-local GPU, switch to the `native` backend, and then enable CP. Do not use `device_map` for this flow:
+
+```python
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+from diffusers import Cosmos3OmniModularPipeline
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from torch.distributed.device_mesh import init_device_mesh
+
+# Make the helper module importable.
+sys.path.insert(0, "examples/cosmos3")
+from cosmos_parallel import (
+    enable_cosmos3_context_parallel,
+    enable_cosmos3_flash_attention,
+    enable_cosmos3_tensor_parallel,
+)
+
+# torchrun sets RANK / WORLD_SIZE / LOCAL_RANK. Pick tp_degree * cp_degree == world size.
+local_rank = int(os.environ["LOCAL_RANK"])
+torch.cuda.set_device(local_rank)
+dist.init_process_group("nccl")
+mesh = init_device_mesh("cuda", (tp_degree, cp_degree), mesh_dim_names=("tp", "cp"))
+
+# Load components on CPU first; a TP-sharded model may not fit one GPU.
+pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.enable_safety_checker()
+
+if tp_degree > 1:
+    enable_cosmos3_tensor_parallel(pipe.transformer, mesh["tp"])  # shard weights -> GPUs
+pipe.to(f"cuda:{local_rank}")                                     # move the replicated remainder
+pipe.transformer.set_attention_backend("native")
+if cp_degree > 1:
+    enable_cosmos3_context_parallel(pipe.transformer, mesh["cp"])  # shard the sequence
+elif tp_degree > 1:
+    enable_cosmos3_flash_attention(pipe.transformer)               # GQA-safe dense attention
+
+# Modular pipelines replace components through update_components().
+scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
+)
+pipe.update_components(scheduler=scheduler)
+
+# A single output name returns that value; a list returns a dictionary.
+outputs = pipe(
+    prompt='{"scene":"A robot arm in a kitchen"}',
+    num_frames=189,
+    height=720,
+    width=1280,
+    output=["videos", "sound", "sampling_rate", "action"],
+)
+videos = outputs["videos"]
+sound = outputs["sound"]  # None unless sound generation was requested.
+action = outputs["action"]  # None unless an action workflow produced actions.
+```
+
+For CP only (no weight sharding), use a 1-D mesh: `init_device_mesh("cuda", (world_size,), mesh_dim_names=("cp",))` and just `enable_cosmos3_context_parallel`.
+
+`enable_safety_checker()` loads and enables the default checker; `disable_safety_checker()` explicitly disables it. Use those pipeline methods instead of the task-pipeline `enable_safety_checker=` construction argument or `enable_safety_check=` call argument. Modular pipelines also do not return `Cosmos3OmniPipelineOutput`: use `output="videos"` for frames alone, or an output list and its returned dictionary as shown above instead of `result.video`, `result.sound`, or `result.action`.
+
+> [!TIP]
+> On some multi-GPU topologies the first NCCL all-to-all can hang. If a CP run stalls at the start of the first denoising step, set `NCCL_P2P_DISABLE=1` in the environment before launching `torchrun`.
+
+CP and TP compose with all the workflows above (text-to-video, image-to-video, text-to-video with sound, and action-conditioned generation) and with both the `Nano` and `Super` checkpoints — only the pipeline construction and the parallelism setup lines change.
+
 ## Metadata templates
 
 `tokenize_prompt` appends short metadata sentences inside the user message so the LLM sees the conditioning the model was trained with. The positive prompt gets sentences like *"The video is 7.9 seconds long and is of 24 FPS."* and *"This video is of 720x1280 resolution."*; the negative prompt gets the inverse (*"… is not …"*).

examples/cosmos3/README.md

@@ -1,13 +1,18 @@
 # Cosmos3 — smoke-test runner
 
-The canonical reference for `Cosmos3OmniPipeline` lives in the diffusers docs:
+The canonical reference for `Cosmos3OmniModularPipeline` lives in the diffusers docs:
 [`docs/source/en/api/pipelines/cosmos3.md`](../../docs/source/en/api/pipelines/cosmos3.md). Use the
 examples there as the source of truth for application code — they cover text-to-image,
 text-to-video, image-to-video, and text+sound modes.
 
-This directory provides a small CLI wrapper (`inference_cosmos3.py`) that exercises the full
-load → encode → denoise → decode path against either the Hub release or a local checkpoint
-during development.
+This directory provides two files:
+
+- `inference_cosmos3.py` — the runnable `Cosmos3OmniModularPipeline` CLI (text-to-image/video,
+  image-to-video, sound, action modes). Single-GPU by default; pass `--tp-degree` / `--cp-degree`
+  and launch with `torchrun` to run any modality multi-GPU (see [Multi-GPU inference](#multi-gpu-inference-context-parallelism)
+  below).
+- `cosmos_parallel.py` — the importable multi-GPU helpers (context + tensor parallelism). No
+  `main`; the CLI imports from it. Read it to understand or adapt the sharding.
 
 ## Setup
 
@@ -178,3 +183,82 @@ Pick the tier that matches the native resolution of your conditioning input (`48
 | `--no-duration-template` | off | Skip the duration metadata sentence appended to the prompt and negative prompt. Ignored for `--num-frames 1` and for action modes (which build a structured caption instead). |
 | `--no-resolution-template` | off | Skip the resolution metadata sentence appended to the prompt and negative prompt. Ignored for action modes. |
 | `--output` | `.` | Directory to write `sample.jpg` or `sample.mp4`. |
+
+## Multi-GPU inference (context parallelism)
+
+Cosmos 3 can be sharded across GPUs on two orthogonal axes (implemented in `cosmos_parallel.py`):
+
+- **Context parallelism (CP)** — `enable_cosmos3_context_parallel`. The *sequence* is sharded
+  across GPUs and attention runs with two Ulysses all-to-all collectives per layer, cutting
+  per-step latency for long videos / high resolutions. Weights are replicated, so this is for
+  models that already fit one GPU (`Nano`).
+- **Tensor parallelism (TP)** — `enable_cosmos3_tensor_parallel`. The attention and MLP *weight*
+  matrices are sharded across GPUs (Megatron-style), so a checkpoint that doesn't fit one GPU
+  (`Super`, ~120 GB) loads. The sequence is not sharded.
+- **TP + CP** — both at once on a 2-D `(tp, cp)` mesh: a large model *and* latency.
+
+The model itself carries no parallelism logic — it exposes small no-op shard/gather seams, and
+`cosmos_parallel.py` implements the entire path (collectives, GQA KV-head handling, ragged-length
+padding, the dual-pathway attention, weight sharding) behind those two helpers. It is meant to be
+read end to end and adapted.
+
+The CLI imports these helpers, so you run **any modality** (text-to-image/video, image-to-video,
+sound, action modes) multi-GPU by adding `--tp-degree` / `--cp-degree` and launching with
+[torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html) — `--tp-degree * --cp-degree`
+must equal `--nproc_per_node`:
+
+```bash
+# CP only (Nano): CP degree must divide the 32 query heads.
+torchrun --nproc_per_node 4 examples/cosmos3/inference_cosmos3.py --model nano --cp-degree 4 --prompt "..."
+
+# TP only (Super): TP degree must divide the 64 query heads and 8 KV heads.
+torchrun --nproc_per_node 4 examples/cosmos3/inference_cosmos3.py --model super --tp-degree 4 --prompt "..."
+
+# TP + CP (Super), 4 GPUs as 2 x 2, with sound:
+torchrun --nproc_per_node 4 examples/cosmos3/inference_cosmos3.py \
+    --model super --tp-degree 2 --cp-degree 2 --enable-sound --prompt "A waterfall in a forest."
+```
+
+### Modular pipeline setup
+
+The CLI uses `Cosmos3OmniModularPipeline`, not the legacy task pipeline. Its distributed setup order is important:
+
+1. Construct it with `Cosmos3OmniModularPipeline.from_pretrained(...)` and call
+   `pipe.load_components(torch_dtype=torch.bfloat16)` while it is still on CPU. Do not use
+   `device_map`.
+2. Initialize the NCCL process group and call `torch.cuda.set_device(local_rank)` before building
+   the device mesh or applying TP. If needed, replace the scheduler with
+   `pipe.update_components(scheduler=...)` and apply TP to `pipe.transformer` while it is still on
+   CPU.
+3. Move the pipeline to `cuda:${LOCAL_RANK}`, select the `native` attention backend, and then
+   enable CP (or the GQA-safe dense attention helper for TP-only runs).
+4. Before generation, call `pipe.enable_safety_checker()` to load and enable the default checker,
+   or `pipe.disable_safety_checker()` to explicitly opt out. The task-pipeline
+   `enable_safety_checker=` construction argument and `enable_safety_check=` call argument do not
+   configure the modular pipeline.
+
+Modular calls request outputs explicitly: `output="videos"` returns frames directly, while
+`output=["videos", "sound", "sampling_rate", "action"]` returns a dictionary. Read values such as
+`outputs["videos"]` from that dictionary rather than `result.video`, `result.sound`, or
+`result.action`; sound and action are `None` when their respective workflows are not used. The
+[pipeline documentation](../../docs/source/en/api/pipelines/cosmos3.md#context-parallelism) has a
+complete direct-use example.
+
+Notes:
+
+- The helpers use the `native` attention backend (the only one that supports GQA's `enable_gqa`),
+  and expand the KV heads to the query-head count so SDPA picks the flash kernel — passing
+  `enable_gqa=True` forces the math kernel, which materializes the full `[S, S]` scores and OOMs
+  on long sequences.
+- Only Ulysses is supported (not ring attention).
+- The CP/Ulysses degree must divide the query heads (32 for `Nano`, 64 for `Super`). For TP,
+  `tp` must divide the KV heads (8), and `tp * cp` must divide the query heads.
+- TP all-reduces on every block, so it's bandwidth-heavy — use the smallest TP degree that makes
+  the weights fit and put the remaining GPUs into CP.
+- Generation size is set with the usual CLI flags (`--num-frames` / `--height` / `--width`), and
+  multi-GPU runs require a seed for reproducibility across ranks (the CLI sets one if you omit `--seed`).
+- On some multi-GPU topologies the first NCCL all-to-all can hang; if a run stalls at the first
+  denoising step, set `NCCL_P2P_DISABLE=1` before launching.
+
+See the [pipeline docs](../../docs/source/en/api/pipelines/cosmos3.md#context-parallelism) for how
+to enable CP and TP from your own pipeline code.