diffusers/examples/cosmos3/cosmos_parallel.py

# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Multi-GPU helpers for Cosmos 3, implemented entirely outside the model.

This module holds the implementation; the runnable entry point is
``inference_cosmos3.py`` (pass ``--tp-degree`` / ``--cp-degree`` and launch with
``torchrun`` to use any of these across all the pipeline's modalities). Two orthogonal
sharding axes are provided, composable on a 2-D ``(tp, cp)`` device mesh:

  * 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).
    Replicates the weights, so it cuts latency but not weight memory.
  * Tensor parallelism (TP) — ``enable_cosmos3_tensor_parallel``. Shards the attention
    and MLP *weight* matrices (Megatron-style), so a checkpoint that doesn't fit one GPU
    (Cosmos3-Super, ~120 GB) loads across several. Attention stays dense; pair it with
    ``enable_cosmos3_flash_attention`` (or with CP) so GQA uses the flash kernel.

The model carries no parallelism logic — it exposes small no-op seams on the transformer
(``_cp_shard_fn`` / ``_cp_gather_fn``, sequence shard/gather around the decoder stack) and
runs whatever attention processor is installed. Attention parallelism lives entirely in the
standalone processors below (each a self-contained ``__call__``, installed via
``set_processor``), so the model file needs no attention seam. The helpers below wire those up.

Why Cosmos 3 needs a custom CP path (not diffusers' declarative ``_cp_plan``):
  1. grouped-query attention — K/V heads must be repeated to match the query heads;
  2. separate understanding (causal) / generation (full) streams, with the generation
     stream attending to ``cat(und, gen)``;
  3. ragged per-stream lengths — each pathway is padded independently and the padded
     generation keys are masked.

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

import torch

from diffusers.models.attention_dispatch import AttentionBackendName, dispatch_attention_fn
from diffusers.models.transformers.transformer_cosmos3 import _rotate_half


try:  # torch >= 2.4
    from torch.distributed.tensor import DTensor, Replicate, Shard
except ImportError:  # pragma: no cover - older torch
    from torch.distributed._tensor import DTensor, Replicate, Shard


def _repeat_kv_heads(x, repeats):
    """Repeat KV heads (for GQA): ``[seq, num_kv_heads, d] -> [seq, num_kv_heads * repeats, d]``.

    Each KV head is repeated ``repeats`` times contiguously, matching GQA grouping
    (query head i pairs with KV group i // repeats).
    """
    if repeats == 1:
        return x
    seq_len, num_kv_heads, head_dim = x.shape
    x = x[:, :, None, :].expand(seq_len, num_kv_heads, repeats, head_dim)
    return x.reshape(seq_len, num_kv_heads * repeats, head_dim)


def _project_qkv_with_rope(attn, und_seq, gen_seq, rotary_emb):
    """Shared processor prologue: project both pathways to q/k/v, RMS-norm q/k, and apply
    rotary embeddings. Returns ``(q_und, k_und, v_und, q_gen, k_gen, v_gen)``, each shaped
    ``[seq, heads, head_dim]``.

    Mirrors the projection + RoPE code in ``Cosmos3AttnProcessor.__call__`` so the parallel
    processors below stay standalone (each writes its own ``__call__``) without the model file
    needing an override seam.
    """
    q_und = attn.to_q(und_seq).view(-1, attn.num_attention_heads, attn.head_dim)
    k_und = attn.to_k(und_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
    v_und = attn.to_v(und_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
    q_gen = attn.add_q_proj(gen_seq).view(-1, attn.num_attention_heads, attn.head_dim)
    k_gen = attn.add_k_proj(gen_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
    v_gen = attn.add_v_proj(gen_seq).view(-1, attn.num_key_value_heads, attn.head_dim)

    q_und = attn.norm_q(q_und)
    k_und = attn.norm_k(k_und)
    q_gen = attn.norm_added_q(q_gen)
    k_gen = attn.norm_added_k(k_gen)

    cos_und, sin_und, cos_gen, sin_gen = rotary_emb
    cos_und = cos_und.unsqueeze(1)
    sin_und = sin_und.unsqueeze(1)
    q_und = q_und * cos_und + _rotate_half(q_und) * sin_und
    k_und = k_und * cos_und + _rotate_half(k_und) * sin_und
    cos_gen = cos_gen.unsqueeze(1)
    sin_gen = sin_gen.unsqueeze(1)
    q_gen = q_gen * cos_gen + _rotate_half(q_gen) * sin_gen
    k_gen = k_gen * cos_gen + _rotate_half(k_gen) * sin_gen
    return q_und, k_und, v_und, q_gen, k_gen, v_gen


# =============================================================================
# Context parallelism (Ulysses)
# =============================================================================
# --- Collective primitives (all-to-all / all-gather via DTensor redistribute) ---
def _cp_all_to_all(local_input, scatter_dim, gather_dim, cp_mesh):
    """All-to-all via DTensor redistribute: gather ``gather_dim``, scatter ``scatter_dim``."""
    dt = DTensor.from_local(local_input, cp_mesh, [Shard(gather_dim)], run_check=False)
    return dt.redistribute(cp_mesh, [Shard(scatter_dim)]).to_local()


def _cp_all_gather(local_input, gather_dim, cp_mesh):
    """All-gather via DTensor redistribute: ``Shard(gather_dim) -> Replicate()``."""
    dt = DTensor.from_local(local_input, cp_mesh, [Shard(gather_dim)], run_check=False)
    return dt.redistribute(cp_mesh, [Replicate()]).to_local()


def _cp_gather_seq_scatter_heads(x, cp_mesh):
    """``[seq/cp, h, d] -> [seq, h/cp, d]``."""
    return _cp_all_to_all(x, scatter_dim=1, gather_dim=0, cp_mesh=cp_mesh)


def _cp_gather_heads_scatter_seq(x, cp_mesh):
    """``[seq, h/cp, d] -> [seq/cp, h, d]``."""
    return _cp_all_to_all(x, scatter_dim=0, gather_dim=1, cp_mesh=cp_mesh)


def _cp_pad_dim0(x, target_len):
    """Zero-pad ``x`` along dim 0 up to ``target_len`` (no-op if already there)."""
    pad = target_len - x.shape[0]
    if pad <= 0:
        return x
    return torch.cat([x, x.new_zeros((pad, *x.shape[1:]))], dim=0)


def _cp_shard_dim0(x, cp_mesh):
    """Keep this rank's contiguous shard along dim 0 (dim 0 must be divisible)."""
    world = cp_mesh.size()
    if world == 1:
        return x
    rank = cp_mesh.get_local_rank()
    shard = x.shape[0] // world
    return x.narrow(0, rank * shard, shard).contiguous()


def _cp_gather_dim0(x, cp_mesh):
    """Reassemble the full sequence (dim 0) from per-rank shards, in rank order."""
    if cp_mesh.size() == 1:
        return x
    return _cp_all_gather(x, gather_dim=0, cp_mesh=cp_mesh)


# --- Sharding / gathering of the dual-pathway packed sequence ---
def shard_cosmos3_sequence(und_seq, gen_seq, rotary_emb, cp_mesh):
    """Pad each pathway to a multiple of the CP world size, then shard the und/gen
    hidden states and their rotary embeddings along the sequence dim, independently
    per pathway. ``rotary_emb`` is ``(cos_und, sin_und, cos_gen, sin_gen)``.

    Returns ``(und_seq, gen_seq, rotary_emb, meta)`` where ``meta`` records the real
    and padded per-pathway lengths so attention can mask padded keys and the caller
    can slice the padding off after gathering.
    """
    world = cp_mesh.size()
    cos_und, sin_und, cos_gen, sin_gen = rotary_emb
    und_real, gen_real = und_seq.shape[0], gen_seq.shape[0]
    und_padded = ((und_real + world - 1) // world) * world
    gen_padded = ((gen_real + world - 1) // world) * world
    meta = {"und_real": und_real, "gen_real": gen_real, "und_padded": und_padded, "gen_padded": gen_padded}

    und_seq = _cp_shard_dim0(_cp_pad_dim0(und_seq, und_padded), cp_mesh)
    gen_seq = _cp_shard_dim0(_cp_pad_dim0(gen_seq, gen_padded), cp_mesh)
    rotary_emb = (
        _cp_shard_dim0(_cp_pad_dim0(cos_und, und_padded), cp_mesh),
        _cp_shard_dim0(_cp_pad_dim0(sin_und, und_padded), cp_mesh),
        _cp_shard_dim0(_cp_pad_dim0(cos_gen, gen_padded), cp_mesh),
        _cp_shard_dim0(_cp_pad_dim0(sin_gen, gen_padded), cp_mesh),
    )
    return und_seq, gen_seq, rotary_emb, meta


def build_gen_key_mask(meta, dtype, device):
    """Additive key mask for the generation pathway's full attention. Padded und/gen
    key positions are set to ``-inf`` so real generation queries ignore them. Returns
    ``None`` when no padding was added. Shape ``[1, 1, 1, S_k]``.
    """
    if meta["und_real"] == meta["und_padded"] and meta["gen_real"] == meta["gen_padded"]:
        return None
    s_k = meta["und_padded"] + meta["gen_padded"]
    mask = torch.zeros(s_k, dtype=dtype, device=device)
    neg_inf = torch.finfo(dtype).min
    mask[meta["und_real"] : meta["und_padded"]] = neg_inf
    mask[meta["und_padded"] + meta["gen_real"] :] = neg_inf
    return mask.view(1, 1, 1, s_k)


def gather_and_unpad(und_out, gen_out, meta, cp_mesh):
    """Gather each pathway to its full padded length, then slice off the padding."""
    und_out = _cp_gather_dim0(und_out, cp_mesh)[: meta["und_real"]]
    gen_out = _cp_gather_dim0(gen_out, cp_mesh)[: meta["gen_real"]]
    return und_out, gen_out


def cosmos3_cp_attention(cp_mesh, q_und, k_und, v_und, q_gen, k_gen, v_gen, gen_key_mask=None):
    """Ulysses context-parallel attention for the dual-pathway packed sequence.

    All inputs are *sequence-sharded*: q ``[S/cp, H, D]``, k/v ``[S/cp, Hkv, D]``
    (no batch dim). Returns ``(causal_out, full_out)`` flattened to ``[S/cp, H*D]``
    and ``[S_gen/cp, H*D]``. The understanding pathway self-attends causally; the
    generation pathway attends to the concatenation of und+gen keys/values.
    """
    world = cp_mesh.size()
    q_heads = q_und.shape[1]
    kv_heads = k_und.shape[1]
    if q_heads % world != 0:
        raise ValueError(f"Query heads ({q_heads}) must be divisible by CP world size ({world}).")

    # GQA, step 1: repeat KV heads up to a multiple of the world size so the head
    # scatter in the all-to-all is valid (each rank must receive an equal share).
    kv_head_repeats = max(world // kv_heads, 1)
    repeated_kv_heads = kv_heads * kv_head_repeats
    if repeated_kv_heads % world != 0:
        raise ValueError(f"Repeated KV heads ({repeated_kv_heads}) must be divisible by CP world size ({world}).")
    if kv_head_repeats > 1:
        k_und = _repeat_kv_heads(k_und, kv_head_repeats)
        v_und = _repeat_kv_heads(v_und, kv_head_repeats)
        k_gen = _repeat_kv_heads(k_gen, kv_head_repeats)
        v_gen = _repeat_kv_heads(v_gen, kv_head_repeats)

    # all-to-all #1: gather sequence, scatter heads -> [S, H/cp, D]
    q_und = _cp_gather_seq_scatter_heads(q_und, cp_mesh)
    k_und = _cp_gather_seq_scatter_heads(k_und, cp_mesh)
    v_und = _cp_gather_seq_scatter_heads(v_und, cp_mesh)
    q_gen = _cp_gather_seq_scatter_heads(q_gen, cp_mesh)
    k_gen = _cp_gather_seq_scatter_heads(k_gen, cp_mesh)
    v_gen = _cp_gather_seq_scatter_heads(v_gen, cp_mesh)

    # GQA, step 2: locally expand each rank's KV head shard up to its query-head
    # count, so attention runs with equal Q/K/V heads (enable_gqa=False). This lets
    # SDPA dispatch to the flash / memory-efficient kernel (O(S) memory); passing
    # enable_gqa=True instead forces the math fallback, which materializes the full
    # [S, S] score matrix and OOMs on long sequences (CP shards heads across ranks,
    # but each rank still attends over the *full* sequence length). GQA grouping is
    # preserved: local query head i pairs with local KV group i // head_repeats.
    q_local, kv_local = q_und.shape[1], k_und.shape[1]
    if q_local % kv_local != 0:
        raise ValueError(f"Local query heads ({q_local}) must be a multiple of local KV heads ({kv_local}).")
    head_repeats = q_local // kv_local
    if head_repeats > 1:
        k_und = _repeat_kv_heads(k_und, head_repeats)
        v_und = _repeat_kv_heads(v_und, head_repeats)
        k_gen = _repeat_kv_heads(k_gen, head_repeats)
        v_gen = _repeat_kv_heads(v_gen, head_repeats)

    # Understanding pathway: causal self-attention over the full und sequence.
    causal_out = dispatch_attention_fn(
        q_und.unsqueeze(0),
        k_und.unsqueeze(0),
        v_und.unsqueeze(0),
        is_causal=True,
        enable_gqa=False,
        backend=AttentionBackendName.NATIVE,
        parallel_config=None,
    ).squeeze(0)

    # Generation pathway: full attention over cat(und, gen) keys/values.
    all_k = torch.cat([k_und, k_gen], dim=0)
    all_v = torch.cat([v_und, v_gen], dim=0)
    full_out = dispatch_attention_fn(
        q_gen.unsqueeze(0),
        all_k.unsqueeze(0),
        all_v.unsqueeze(0),
        attn_mask=gen_key_mask,
        is_causal=False,
        enable_gqa=False,
        backend=AttentionBackendName.NATIVE,
        parallel_config=None,
    ).squeeze(0)

    # all-to-all #2: gather heads, scatter sequence -> [S/cp, H, D]
    causal_out = _cp_gather_heads_scatter_seq(causal_out, cp_mesh)
    full_out = _cp_gather_heads_scatter_seq(full_out, cp_mesh)
    return causal_out.flatten(-2, -1), full_out.flatten(-2, -1)


class Cosmos3CPAttnProcessor:
    """Cosmos 3 attention processor whose attention core runs Ulysses CP.

    Standalone processor: its ``__call__`` projects + applies RoPE via the shared prologue,
    then brackets the two pathways with all-to-all collectives. The Ulysses mesh is read from
    ``self.cp_mesh``; the per-call generation key mask is read from the attention module
    (stamped each forward by the shard seam).
    """

    def __init__(self, cp_mesh):
        self.cp_mesh = cp_mesh

    def __call__(self, attn, und_seq, gen_seq, rotary_emb):
        q_und, k_und, v_und, q_gen, k_gen, v_gen = _project_qkv_with_rope(attn, und_seq, gen_seq, rotary_emb)
        causal_out, full_out = cosmos3_cp_attention(
            self.cp_mesh,
            q_und,
            k_und,
            v_und,
            q_gen,
            k_gen,
            v_gen,
            gen_key_mask=getattr(attn, "_cp_gen_key_mask", None),
        )
        und_out = attn.to_out(causal_out)
        gen_out = attn.to_add_out(full_out)
        return und_out, gen_out


def enable_cosmos3_context_parallel(transformer, cp_mesh):
    """Wire Ulysses context parallelism onto a ``Cosmos3OmniTransformer`` instance.

    Sets the CP attention processor on every decoder layer and installs the shard/gather
    seams so the model shards each pathway across ``cp_mesh`` and gathers before decode.
    All CP state lives in the closures below, so the model itself stays CP-free.
    """
    processor = Cosmos3CPAttnProcessor(cp_mesh)
    for layer in transformer.layers:
        layer.self_attn.set_processor(processor)

    state = {"meta": None}

    def shard_fn(und_seq, gen_seq, rotary_emb):
        und_seq, gen_seq, rotary_emb, meta = shard_cosmos3_sequence(und_seq, gen_seq, rotary_emb, cp_mesh)
        state["meta"] = meta
        gen_key_mask = build_gen_key_mask(meta, und_seq.dtype, und_seq.device)
        # The processor reads the mask off the attention module each forward.
        for layer in transformer.layers:
            layer.self_attn._cp_gen_key_mask = gen_key_mask
        return und_seq, gen_seq, rotary_emb

    def gather_fn(und_out, gen_out):
        return gather_and_unpad(und_out, gen_out, state["meta"], cp_mesh)

    transformer._cp_shard_fn = shard_fn
    transformer._cp_gather_fn = gather_fn
    return transformer


# =============================================================================
# Dense flash attention (GQA-safe; for TP without CP)
# =============================================================================
class Cosmos3FlashAttnProcessor:
    """Dense attention that expands the GQA KV heads to the query-head count so SDPA
    uses the flash kernel (``enable_gqa=False``) instead of the math fallback, which
    materializes the full ``[S, S]`` score matrix and OOMs on long sequences.

    Standalone processor: no collectives — attention is computed locally over the full
    sequence (the head counts on the attention module are the rank-local values set by
    ``enable_cosmos3_tensor_parallel``). Use this with TP-only; CP installs its own
    processor that already handles the flash dispatch.
    """

    def __call__(self, attn, und_seq, gen_seq, rotary_emb):
        q_und, k_und, v_und, q_gen, k_gen, v_gen = _project_qkv_with_rope(attn, und_seq, gen_seq, rotary_emb)
        repeats = attn.num_attention_heads // attn.num_key_value_heads
        k_und, v_und = _repeat_kv_heads(k_und, repeats), _repeat_kv_heads(v_und, repeats)
        k_gen, v_gen = _repeat_kv_heads(k_gen, repeats), _repeat_kv_heads(v_gen, repeats)

        # Understanding pathway: causal self-attention.
        causal_out = dispatch_attention_fn(
            q_und.unsqueeze(0),
            k_und.unsqueeze(0),
            v_und.unsqueeze(0),
            is_causal=True,
            enable_gqa=False,
            backend=AttentionBackendName.NATIVE,
            parallel_config=None,
        ).squeeze(0)

        # Generation pathway: full attention over cat(und, gen) keys/values.
        all_k = torch.cat([k_und, k_gen], dim=0)
        all_v = torch.cat([v_und, v_gen], dim=0)
        full_out = dispatch_attention_fn(
            q_gen.unsqueeze(0),
            all_k.unsqueeze(0),
            all_v.unsqueeze(0),
            is_causal=False,
            enable_gqa=False,
            backend=AttentionBackendName.NATIVE,
            parallel_config=None,
        ).squeeze(0)
        und_out = attn.to_out(causal_out.flatten(-2, -1))
        gen_out = attn.to_add_out(full_out.flatten(-2, -1))
        return und_out, gen_out


def enable_cosmos3_flash_attention(transformer):
    """Install the dense flash-attention processor on every decoder layer."""
    processor = Cosmos3FlashAttnProcessor()
    for layer in transformer.layers:
        layer.self_attn.set_processor(processor)
    return transformer


# =============================================================================
# Tensor parallelism (shard the attention + MLP weights)
# =============================================================================
def enable_cosmos3_tensor_parallel(transformer, tp_mesh):
    """Shard every decoder layer's attention and MLP projection weights across
    ``tp_mesh`` (Megatron-style tensor parallelism), so a model whose weights don't
    fit on one GPU (e.g. Cosmos3-Super, ~120 GB) can be served across several.

    Layout per layer:
      * column-parallel (output / head dim sharded): ``to_q/to_k/to_v`` and
        ``add_q_proj/add_k_proj/add_v_proj``, plus both MLPs' ``gate_proj/up_proj``;
      * row-parallel (input dim sharded, all-reduce on output): ``to_out`` and
        ``to_add_out``, plus both MLPs' ``down_proj``.

    Each rank then owns ``num_attention_heads / tp`` query heads and
    ``num_key_value_heads / tp`` KV heads, so the per-layer head counts on every
    attention module are rewritten to their local values — the projection + reshape
    code in ``_project_qkv_with_rope`` (used by the installed processor) reads these.
    The embeddings, final norms, lm_head and modality projections stay replicated
    (they're a small fraction of the weights).

    Memory note: weights are loaded to CPU first, then each layer is moved to its GPU
    and sharded in place, so the full model is never materialized on one device.

    Composes with Ulysses CP: this shards weights only and never touches the attention
    processor, so combine it with ``enable_cosmos3_context_parallel`` (CP installs its
    own processor) on a 2-D ``(tp, cp)`` mesh, or with ``enable_cosmos3_flash_attention``
    for TP without CP.
    """
    from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module
    from torch.distributed.tensor.parallel import RowwiseParallel as TorchRowwiseParallel

    # Skip ModelOpt quantizer children to avoid:
    # AttributeError: 'TensorQuantizer' object has no attribute 'weight'
    class RowwiseParallel(TorchRowwiseParallel):
        def _partition_linear_fn(self, name, module, device_mesh):
            if isinstance(module, torch.nn.Linear):
                return super()._partition_linear_fn(name, module, device_mesh)

    tp = tp_mesh.size()
    dev = torch.device("cuda", torch.cuda.current_device())
    plan = {
        "self_attn.to_q": ColwiseParallel(),
        "self_attn.to_k": ColwiseParallel(),
        "self_attn.to_v": ColwiseParallel(),
        "self_attn.to_out": RowwiseParallel(),
        "self_attn.add_q_proj": ColwiseParallel(),
        "self_attn.add_k_proj": ColwiseParallel(),
        "self_attn.add_v_proj": ColwiseParallel(),
        "self_attn.to_add_out": RowwiseParallel(),
        "mlp.gate_proj": ColwiseParallel(),
        "mlp.up_proj": ColwiseParallel(),
        "mlp.down_proj": RowwiseParallel(),
        "mlp_moe_gen.gate_proj": ColwiseParallel(),
        "mlp_moe_gen.up_proj": ColwiseParallel(),
        "mlp_moe_gen.down_proj": RowwiseParallel(),
    }
    for layer in transformer.layers:
        attn = layer.self_attn
        if attn.num_attention_heads % tp != 0 or attn.num_key_value_heads % tp != 0:
            raise ValueError(
                f"TP degree {tp} must divide both the query heads ({attn.num_attention_heads}) "
                f"and the KV heads ({attn.num_key_value_heads})."
            )
        layer.to(dev)  # full layer transiently on this rank's GPU, then sharded in place
        parallelize_module(layer, tp_mesh, plan)
        # Projections now emit only this rank's head shard; the processor reshapes
        # with the local counts.
        attn.num_attention_heads //= tp
        attn.num_key_value_heads //= tp
        attn.num_key_value_groups = attn.num_attention_heads // attn.num_key_value_heads
    return transformer