cosmos_transfer1/diffusion/training/modules/attention.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

from contextlib import nullcontext
from typing import List, Optional, Union

try:
    from megatron.core import parallel_state

    USE_MEGATRON = True
except ImportError:
    USE_MEGATRON = False

import numpy as np
import torch
import torch.nn.functional as F
import transformer_engine as te
from einops import rearrange
from packaging import version
from torch import nn
from torch.nn.attention import SDPBackend
from torch.utils.checkpoint import checkpoint
from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import DotProductAttention
from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb

from cosmos_transfer1.utils import log

# ---------------------- Feed Forward Network -----------------------


class FeedForward(nn.Module):
    """
    Transformer FFN with optional gating

    Parameters:
        d_model (int): Dimensionality of input features.
        d_ff (int): Dimensionality of the hidden layer.
        dropout (float, optional): Dropout rate applied after the activation function. Defaults to 0.1.
        activation (callable, optional): The activation function applied after the first linear layer.
                                         Defaults to nn.ReLU().
        is_gated (bool, optional): If set to True, incorporates gating mechanism to the feed-forward layer.
                                   Defaults to False.
        bias (bool, optional): If set to True, adds a bias to the linear layers. Defaults to True.

    Example:
        >>> ff = FeedForward(d_model=512, d_ff=2048)
        >>> x = torch.randn(64, 10, 512)  # Example input tensor
        >>> output = ff(x)
        >>> print(output.shape)  # Expected shape: (64, 10, 512)
    """

    def __init__(
        self,
        d_model: int,
        d_ff: int,
        dropout: float = 0.1,
        activation=nn.ReLU(),
        is_gated: bool = False,
        bias: bool = False,
    ) -> None:
        super().__init__()

        tp_group = parallel_state.get_tensor_model_parallel_group(check_initialized=False)
        sequence_parallel = getattr(parallel_state, "sequence_parallel", False)
        if tp_group is None:
            tp_size = 1  # TP is not initialized.
        else:
            tp_size = parallel_state.get_tensor_model_parallel_world_size()

        if tp_size == 1:
            self.layer1 = nn.Linear(d_model, d_ff, bias=bias)
            self.layer2 = nn.Linear(d_ff, d_model, bias=bias)
        else:
            assert is_gated is False, "is_gated needs to be False to support Tensor Parallelism"
            assert dropout == 0.0, "dropout needs to be 0.0 to support Tensor Parallelism"
            self.layer1 = te.pytorch.Linear(
                d_model,
                d_ff,
                bias=bias,
                tp_size=tp_size,
                tp_group=tp_group,
                parallel_mode="column",
                sequence_parallel=sequence_parallel,
            )
            self.layer2 = te.pytorch.Linear(
                d_ff,
                d_model,
                bias=bias,
                tp_size=tp_size,
                tp_group=tp_group,
                parallel_mode="row",
                sequence_parallel=sequence_parallel,
            )

        self.dropout = nn.Dropout(dropout)
        self.activation = activation
        self.is_gated = is_gated
        if is_gated:
            self.linear_gate = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x: torch.Tensor):
        g = self.activation(self.layer1(x))
        if self.is_gated:
            x = g * self.linear_gate(x)
        else:
            x = g

        assert self.dropout.p == 0.0, "skipping dropout to save memory"
        return self.layer2(x)


class GPT2FeedForward(FeedForward):
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1, bias: bool = False):
        super().__init__(
            d_model=d_model,
            d_ff=d_ff,
            dropout=dropout,
            activation=nn.GELU(),
            is_gated=False,
            bias=bias,
        )

    def forward(self, x: torch.Tensor):
        assert self.dropout.p == 0.0, "we skip dropout"

        x = self.layer1(x)

        def activation_layer2_forward(x):
            x = self.activation(x)
            x = self.layer2(x)
            return x

        x = checkpoint(activation_layer2_forward, x, use_reentrant=False)
        return x


# ---------------------- Normalization Layer -----------------------


def normalize(x: torch.Tensor, dim: Optional[List[int]] = None, eps: float = 0) -> torch.Tensor:
    """
    Normalizes the input tensor along specified dimensions such that the average square norm of elements is adjusted.

    Args:
        x (torch.Tensor): The input tensor to normalize.
        dim (list, optional): The dimensions over which to normalize. If None, normalizes over all dimensions except the first.
        eps (float, optional): A small constant to ensure numerical stability during division.

    Returns:
        torch.Tensor: The normalized tensor.
    """
    if dim is None:
        dim = list(range(1, x.ndim))
    norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32)
    norm = torch.add(eps, norm, alpha=np.sqrt(norm.numel() / x.numel()))
    return x / norm.to(x.dtype)


def get_normalization(name: str, channels: int):
    if name == "I":
        return nn.Identity()
    elif name == "R":
        return te.pytorch.RMSNorm(channels, eps=1e-6)
    else:
        raise ValueError(f"Normalization {name} not found")


# ---------------------- Attention Op -----------------------
# A list of attention ops
if version.parse(torch.__version__) >= version.parse("2.3.0"):
    from torch.nn.attention import SDPBackend, sdpa_kernel

    sdpa_context = sdpa_kernel
    USE_SDPA = True
elif version.parse(torch.__version__) >= version.parse("2.0.0"):
    from torch.backends.cuda import SDPBackend, sdp_kernel

    sdpa_context = sdp_kernel
    USE_SDPA = False
else:
    sdpa_context = nullcontext
    USE_SDPA = False
    print(
        f"No SDP backend available, likely because you are running in pytorch versions < 2.0. In fact, "
        f"you are using PyTorch {torch.__version__}. You might want to consider upgrading."
    )


class BaseAttentionOp(nn.Module):
    def __init__(self):
        super().__init__()


class FusedAttentionOp(BaseAttentionOp):
    def __init__(self):
        super().__init__()


class TorchAttentionOp(FusedAttentionOp):
    def __init__(self, backend: Optional[Union[List[SDPBackend], SDPBackend]] = SDPBackend.EFFICIENT_ATTENTION):
        super().__init__()
        self.backend = backend
        self.sdpa_context = sdpa_context if self.backend is not None else nullcontext
        if self.backend is not None:
            log.warning(
                "SDPA context manager is not working well with torch.compile, causing graph breaks and "
                "significant slowdowns. If you are using torch.compile you'll most likely want to turn off "
                "this context manager."
            )

    def forward(
        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Computes the scaled dot-product attention over the input tensors using the specified backend.
        B is the batch size, M the sequence length, H the number of heads, and K the embeding size per head

        check F.scaled_dot_product_attention
        Args:
            q (Tensor): The query tensor of shape [B, Mq, H, K] / [B, ..., H, K]
            k (Tensor): The key tensor of shape [B, Mk, H, V] / [B, ..., H, K]
            v (Tensor): The value tensor of shape [B, Mk, H, V] / [B, ..., H, V]

            mask (Optional[Tensor]): An optional mask tensor. Follow scaled_dot_product_attention API, mask should be a boolean tensor with shape [B, H, Mq, Mk]

        Returns:
            Tensor: [B, Mq, H, V] / [B, ..., H, V]
        """
        in_q_shape = q.shape
        in_k_shape = k.shape
        q = rearrange(q, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
        k = rearrange(k, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
        v = rearrange(v, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
        if mask is not None:
            assert mask.dtype == torch.bool, "Mask should be a boolean tensor"
        with self.sdpa_context(self.backend):
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)  # scale is dim_head ** -0.5 per default
        return rearrange(out, "b h ... l -> b ... h l").view(*in_q_shape[:-1], in_k_shape[-1])


class Attention(nn.Module):
    """
    Generalized attention impl. **With TP support**.

    Allowing for both self-attention and cross-attention configurations depending on whether a `context_dim` is provided.
    If `context_dim` is None, self-attention is assumed.

    Parameters:
        query_dim (int): Dimension of each query vector.
        context_dim (int, optional): Dimension of each context vector. If None, self-attention is assumed.
        heads (int, optional): Number of attention heads. Defaults to 8.
        dim_head (int, optional): Dimension of each head. Defaults to 64.
        dropout (float, optional): Dropout rate applied to the output of the attention block. Defaults to 0.0.
        attn_op (BaseAttentionOp, optional): Custom attention operation to be used instead of the default.
        qkv_bias (bool, optional): If True, adds a learnable bias to query, key, and value projections. Defaults to False.
        out_bias (bool, optional): If True, adds a learnable bias to the output projection. Defaults to False.
        qkv_norm (str, optional): A string representing normalization strategies for query, key, and value projections.
                                  Defaults to "SSI".
        norm_args (dict, optional): Arguments to pass to the normalization function. Defaults to an empty dict.

    Examples:
        >>> attn = Attention(query_dim=128, context_dim=256, heads=4, dim_head=32, dropout=0.1)
        >>> query = torch.randn(10, 128)  # Batch size of 10
        >>> context = torch.randn(10, 256)  # Batch size of 10
        >>> output = attn(query, context)  # Perform the attention operation

    Note:
        https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
    """

    def __init__(
        self,
        query_dim: int,
        context_dim=None,
        heads=8,
        dim_head=64,
        dropout=0.0,
        attn_op: Optional[BaseAttentionOp] = None,
        qkv_bias: bool = False,
        out_bias: bool = False,
        qkv_norm: str = "SSI",
        norm_args: dict = {},
        backend: str = "transformer_engine",
        qkv_format: str = "bshd",
    ) -> None:
        super().__init__()
        log.debug(
            f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
            f"{heads} heads with a dimension of {dim_head}. Norm options are {qkv_norm} and norm args are {norm_args}."
        )
        self.is_selfattn = context_dim is None  # self attention

        inner_dim = dim_head * heads
        context_dim = query_dim if context_dim is None else context_dim

        self.heads = heads
        self.dim_head = dim_head
        self.qkv_format = qkv_format
        norm_dim = dim_head
        self.backend = backend
        self.query_dim = query_dim
        self.context_dim = context_dim
        self.inner_dim = inner_dim
        tp_group = parallel_state.get_tensor_model_parallel_group(check_initialized=False) if USE_MEGATRON else None
        if tp_group is None:
            self.tp_size = 1  # TP is not initialized.
        else:
            self.tp_size = parallel_state.get_tensor_model_parallel_world_size()

        if self.backend == "torch":
            assert (
                self.tp_size == 1
            ), f"Attention backend {self.backend} cannot use TP size > 1. Attempted: {self.tp_size}"

        assert self.heads % self.tp_size == 0, "the number of heads should be divisible by TP size"

        if self.tp_size == 1:
            self.to_q = nn.Sequential(
                nn.Linear(query_dim, inner_dim, bias=qkv_bias),
                get_normalization(qkv_norm[0], norm_dim, **norm_args),
            )
            self.to_k = nn.Sequential(
                nn.Linear(context_dim, inner_dim, bias=qkv_bias),
                get_normalization(qkv_norm[1], norm_dim, **norm_args),
            )
            self.to_v = nn.Sequential(
                nn.Linear(context_dim, inner_dim, bias=qkv_bias),
                get_normalization(qkv_norm[2], norm_dim, **norm_args),
            )

            self.to_out = nn.Sequential(
                nn.Linear(inner_dim, query_dim, bias=out_bias),
                nn.Dropout(dropout),
            )
        else:  # TP enabled.
            sequence_parallel = getattr(parallel_state, "sequence_parallel", False)
            if sequence_parallel:
                assert qkv_format == "sbhd", "sequence parallel only supports sbhd format"

            self.to_q = nn.Sequential(
                te.pytorch.Linear(
                    query_dim,
                    inner_dim,
                    bias=qkv_bias,
                    tp_size=self.tp_size,
                    tp_group=tp_group,
                    sequence_parallel=sequence_parallel,
                    parallel_mode="column",
                ),
                get_normalization(qkv_norm[0], norm_dim, **norm_args),
            )
            self.to_k = nn.Sequential(
                te.pytorch.Linear(
                    context_dim,
                    inner_dim,
                    bias=qkv_bias,
                    tp_size=self.tp_size,
                    tp_group=tp_group,
                    sequence_parallel=sequence_parallel,
                    parallel_mode="column",
                ),
                get_normalization(qkv_norm[1], norm_dim, **norm_args),
            )
            self.to_v = nn.Sequential(
                te.pytorch.Linear(
                    context_dim,
                    inner_dim,
                    bias=qkv_bias,
                    tp_size=self.tp_size,
                    tp_group=tp_group,
                    sequence_parallel=sequence_parallel,
                    parallel_mode="column",
                ),
                get_normalization(qkv_norm[2], norm_dim, **norm_args),
            )

            self.to_out = nn.Sequential(
                te.pytorch.Linear(
                    inner_dim,
                    query_dim,
                    bias=out_bias,
                    tp_size=self.tp_size,
                    tp_group=tp_group,
                    parallel_mode="row",
                    sequence_parallel=sequence_parallel,
                ),
                nn.Dropout(dropout),
            )

        if attn_op:  # use what is given
            self.attn_op = attn_op
        elif self.backend == "transformer_engine":
            sequence_parallel = getattr(parallel_state, "sequence_parallel", False) if USE_MEGATRON else False
            self.attn_op: BaseAttentionOp = DotProductAttention(
                self.heads,
                self.dim_head,
                num_gqa_groups=self.heads,
                attention_dropout=0,
                qkv_format=qkv_format,
                attn_mask_type="no_mask",
                tp_size=self.tp_size,
                tp_group=tp_group,
                sequence_parallel=sequence_parallel,
            )
        elif self.backend == "torch":
            self.attn_op = TorchAttentionOp(None)
        else:
            raise ValueError(f"Backend {backend} not found")

    def cal_qkv(
        self, x, context=None, mask=None, rope_emb=None, **kwargs
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        del kwargs

        """
        self.to_q, self.to_k, self.to_v are nn.Sequential with projection + normalization layers.
        Before 07/24/2024, these modules normalize across all heads.
        After 07/24/2024, to support tensor parallelism and follow the common practice in the community,
        we support to normalize per head.
        To keep the checkpoint copatibility with the previous code,
        we keep the nn.Sequential but call the projection and the normalization layers separately.
        """

        q = self.to_q[0](x)
        context = x if context is None else context
        k = self.to_k[0](context)
        v = self.to_v[0](context)
        q, k, v = map(
            lambda t: rearrange(t, "b ... (n c) -> b ... n c", n=self.heads // self.tp_size, c=self.dim_head),
            (q, k, v),
        )

        def apply_norm_and_rotary_pos_emb(q, k, v, rope_emb):
            q = self.to_q[1](q)
            k = self.to_k[1](k)
            v = self.to_v[1](v)
            if self.is_selfattn and rope_emb is not None:  # only apply to self-attention!
                q = apply_rotary_pos_emb(q, rope_emb, tensor_format=self.qkv_format, fused=True)
                k = apply_rotary_pos_emb(k, rope_emb, tensor_format=self.qkv_format, fused=True)
            return q, k, v

        q, k, v = checkpoint(apply_norm_and_rotary_pos_emb, q, k, v, rope_emb, use_reentrant=False)

        return q, k, v

    def cal_attn(self, q, k, v, mask=None):
        if self.backend == "transformer_engine":
            seq_dim = self.qkv_format.index("s")
            assert (
                q.shape[seq_dim] > 1 and k.shape[seq_dim] > 1
            ), "Seqlen must be larger than 1 for TE Attention starting with 1.8 TE version."
            out = self.attn_op(q, k, v, core_attention_bias_type="no_bias", core_attention_bias=None)  # [B, Mq, H, V]
            return self.to_out(out)
        elif self.backend == "torch":
            out = self.attn_op(q, k, v, mask=mask)  # [B, Mq, H, V]
            return self.to_out(rearrange(out, " b ... n c -> b ... (n c)"))
        else:
            raise ValueError(f"Backend {self.backend} not found")

    def forward(
        self,
        x,
        context=None,
        mask=None,
        rope_emb=None,
        **kwargs,
    ):
        """
        Args:
            x (Tensor): The query tensor of shape [B, Mq, K]
            context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
        """
        q, k, v = self.cal_qkv(x, context, mask, rope_emb=rope_emb, **kwargs)
        return self.cal_attn(q, k, v, mask)