汪志鹏andDidan Deng
0059ec878c
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
1625 lines
60 KiB
Python
1625 lines
60 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Copyright (c) 2026 SandAI. All Rights Reserved.
|
|
# Ported from daVinci-MagiHuman inference/model/dit/dit_module.py
|
|
# Adaptations: removed Ulysses context-parallelism, inlined Modality/VarlenHandler.
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, IntEnum
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import rearrange, repeat
|
|
from torch.nn import Parameter
|
|
from torch.nn import functional as F
|
|
from vllm.distributed import (
|
|
get_tensor_model_parallel_world_size,
|
|
)
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.vllm_flash_attn import flash_attn_varlen_func as _vllm_fa_varlen
|
|
|
|
try:
|
|
from magi_compiler.api import magi_register_custom_op
|
|
from magi_compiler.config import CompileConfig
|
|
except Exception:
|
|
|
|
class CompileConfig: # type: ignore[no-redef]
|
|
pass
|
|
|
|
def magi_register_custom_op(*args, **kwargs): # type: ignore[no-redef]
|
|
def decorator(func):
|
|
return func
|
|
|
|
return decorator
|
|
|
|
|
|
def magi_compile(*args, **kwargs):
|
|
"""No-op stub — vllm-omni handles execution; magi compilation is skipped."""
|
|
|
|
def decorator(cls_or_fn):
|
|
return cls_or_fn
|
|
|
|
return decorator
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Inlined from inference/common/sequence_schema.py
|
|
# ---------------------------------------------------------------------------
|
|
class Modality(IntEnum):
|
|
VIDEO = 0
|
|
AUDIO = 1
|
|
TEXT = 2
|
|
|
|
|
|
@dataclass
|
|
class VarlenHandler:
|
|
cu_seqlens_q: torch.Tensor
|
|
cu_seqlens_k: torch.Tensor
|
|
max_seqlen_q: int
|
|
max_seqlen_k: int
|
|
|
|
|
|
def _is_hopper_arch() -> bool:
|
|
if not torch.cuda.is_available():
|
|
return False
|
|
return torch.cuda.get_device_capability()[0] == 9
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FFA handler for local / flex attention
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class FFAHandler:
|
|
q_ranges: torch.Tensor
|
|
k_ranges: torch.Tensor
|
|
max_seqlen_q: int
|
|
max_seqlen_k: int
|
|
attn_type_map: torch.Tensor
|
|
softmax_scale: float
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Activation helpers
|
|
# ---------------------------------------------------------------------------
|
|
class MLPActivationType(Enum):
|
|
SWIGLU7 = "swiglu7"
|
|
GELU7 = "gelu7"
|
|
|
|
|
|
def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None):
|
|
out_dtype = x.dtype if out_dtype is None else out_dtype
|
|
x = x.to(torch.float32)
|
|
x_glu, x_linear = x[..., ::2], x[..., 1::2]
|
|
x_glu = x_glu.clamp(min=None, max=limit)
|
|
x_linear = x_linear.clamp(min=-limit, max=limit)
|
|
out_glu = x_glu * torch.sigmoid(alpha * x_glu)
|
|
return (out_glu * (x_linear + 1)).to(out_dtype)
|
|
|
|
|
|
def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None):
|
|
out_dtype = x.dtype if out_dtype is None else out_dtype
|
|
x = x.to(torch.float32)
|
|
x_glu = x.clamp(min=None, max=limit)
|
|
out_glu = x_glu * torch.sigmoid(alpha * x_glu)
|
|
return out_glu.to(out_dtype)
|
|
|
|
|
|
def create_activation_func(activation_type: MLPActivationType) -> Callable:
|
|
match activation_type:
|
|
case MLPActivationType.SWIGLU7:
|
|
return swiglu7
|
|
case MLPActivationType.GELU7:
|
|
return gelu7
|
|
case _:
|
|
raise ValueError(f"Unknown activation type: {activation_type}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Modality dispatcher (permutation helper)
|
|
# ---------------------------------------------------------------------------
|
|
class ModalityDispatcher:
|
|
permuted_modality_mapping: torch.Tensor
|
|
group_size: torch.Tensor
|
|
group_size_cpu: list[int]
|
|
num_modalities: int
|
|
|
|
def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
|
|
self.modality_mapping = modality_mapping
|
|
self.num_modalities = num_modalities
|
|
self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping)
|
|
self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32)
|
|
self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()]
|
|
|
|
def _precompute_permute_mapping(self, modality_mapping):
|
|
self.permute_mapping = torch.argsort(modality_mapping)
|
|
self.inv_permute_mapping = torch.argsort(self.permute_mapping)
|
|
return modality_mapping[self.permute_mapping]
|
|
|
|
def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
|
|
return list(torch.split(x, self.group_size_cpu, dim=0))
|
|
|
|
def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor:
|
|
return torch.cat(processed_groups, dim=0)
|
|
|
|
@staticmethod
|
|
def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
|
|
return x[permute_mapping]
|
|
|
|
@staticmethod
|
|
def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
|
|
return x[inv_permute_mapping]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Positional / rotary embedding helpers
|
|
# ---------------------------------------------------------------------------
|
|
def freq_bands(
|
|
num_bands: int, temperature: float = 10000.0, step: int = 2, device: torch.device | None = None
|
|
) -> torch.Tensor:
|
|
exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands
|
|
return 1.0 / (temperature**exp)
|
|
|
|
|
|
def rotate_half(x, interleaved=False):
|
|
if not interleaved:
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
else:
|
|
x1, x2 = x[..., ::2], x[..., 1::2]
|
|
return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
|
|
|
|
|
|
def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
|
|
ro_dim = cos.shape[-1] * 2
|
|
assert ro_dim <= x.shape[-1]
|
|
cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
|
|
sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
|
|
return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], dim=-1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fourier positional embedding
|
|
# ---------------------------------------------------------------------------
|
|
class ElementWiseFourierEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
max_res: int = 224,
|
|
temperature: float = 10000.0,
|
|
in_pixels: bool = True,
|
|
linear_bands: bool = False,
|
|
learnable: bool = False,
|
|
device: torch.device = torch.device("cpu"),
|
|
dtype: torch.dtype = torch.float32,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.in_pixels = in_pixels
|
|
self.learnable = learnable
|
|
self.temperature = temperature
|
|
self.max_res = max_res
|
|
self.linear_bands = linear_bands
|
|
self.device = device
|
|
self.dtype = dtype
|
|
bands = self.get_default_bands()
|
|
self.bands = nn.Parameter(bands, requires_grad=self.learnable)
|
|
|
|
def forward(self, coords: torch.Tensor) -> torch.Tensor:
|
|
coords_xyz = coords[:, :3]
|
|
sizes = coords[:, 3:6]
|
|
refs = coords[:, 6:9]
|
|
|
|
scales = (refs - 1) / (sizes - 1)
|
|
scales[(refs == 1) & (sizes == 1)] = 1
|
|
assert not scales.isnan().any(), "scales has nan"
|
|
assert not scales.isinf().any(), "scales has inf"
|
|
|
|
centers = (sizes - 1) / 2
|
|
centers[:, 0] = 0
|
|
coords_xyz = coords_xyz - centers
|
|
|
|
bands = self.bands.to(coords.device, coords.dtype)
|
|
proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * bands
|
|
sin_proj = proj.sin()
|
|
cos_proj = proj.cos()
|
|
return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)
|
|
|
|
def reset_parameters(self):
|
|
self.bands.copy_(self.get_default_bands())
|
|
|
|
def get_default_bands(self):
|
|
if self.in_pixels:
|
|
raise NotImplementedError("in_pixels are not implemented yet")
|
|
return freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Multi-modality RMSNorm
|
|
# ---------------------------------------------------------------------------
|
|
class MultiModalityRMSNorm(nn.Module):
|
|
__constants__ = ["dim", "eps", "num_modality"]
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.eps = eps
|
|
self.num_modality = num_modality
|
|
self.weight = nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32))
|
|
if num_modality > 1:
|
|
self.forward = self.forward_multi_experts
|
|
else:
|
|
self.forward = self.forward_single_expert
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self):
|
|
nn.init.zeros_(self.weight)
|
|
|
|
def rms(self, x: torch.Tensor) -> torch.Tensor:
|
|
t = x.float()
|
|
return t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps)
|
|
|
|
def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor:
|
|
original_dtype = x.dtype
|
|
t = self.rms(x)
|
|
weight_chunked = self.weight.chunk(self.num_modality, dim=0)
|
|
t_list = modality_dispatcher.dispatch(t)
|
|
for i in range(self.num_modality):
|
|
t_list[i] = t_list[i] * (weight_chunked[i] + 1)
|
|
t = modality_dispatcher.undispatch(*t_list)
|
|
return t.to(original_dtype)
|
|
|
|
def forward_single_expert(
|
|
self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher | None = None
|
|
) -> torch.Tensor:
|
|
t, original_dtype = x.float(), x.dtype
|
|
t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps)
|
|
return (t * (self.weight + 1)).to(original_dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Linear layers with bf16 compute and MoE dispatch
|
|
# ---------------------------------------------------------------------------
|
|
class _BF16ComputeLinear(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(
|
|
ctx,
|
|
input: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.Tensor | None,
|
|
output_dtype: torch.dtype | None,
|
|
compute_dtype: torch.dtype = torch.bfloat16,
|
|
):
|
|
input_cast = input.to(compute_dtype)
|
|
weight_cast = weight.to(compute_dtype)
|
|
output = torch.matmul(input_cast, weight_cast.t())
|
|
if bias is not None:
|
|
output = output + bias.to(compute_dtype)
|
|
return output.to(output_dtype)
|
|
|
|
|
|
class BaseLinear(nn.Module):
|
|
__constants__ = ["in_features", "out_features", "num_layers", "num_experts"]
|
|
|
|
def __init__(
|
|
self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None
|
|
):
|
|
super().__init__()
|
|
factory_kwargs = {"device": device, "dtype": torch.bfloat16}
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.num_layers_for_initialization = num_layers_for_initialization
|
|
self.num_experts = num_experts
|
|
self.use_bias = bias
|
|
self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs))
|
|
if bias:
|
|
self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs))
|
|
else:
|
|
self.register_parameter("bias", None)
|
|
|
|
def forward(
|
|
self,
|
|
input: torch.Tensor,
|
|
output_dtype: torch.dtype | None = None,
|
|
modality_dispatcher: ModalityDispatcher | None = None,
|
|
) -> torch.Tensor:
|
|
output_dtype = input.dtype if output_dtype is None else output_dtype
|
|
return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16)
|
|
|
|
|
|
class NativeMoELinear(BaseLinear):
|
|
def forward(
|
|
self,
|
|
input: torch.Tensor,
|
|
output_dtype: torch.dtype | None = None,
|
|
modality_dispatcher: ModalityDispatcher | None = None,
|
|
) -> torch.Tensor:
|
|
output_dtype = input.dtype if output_dtype is None else output_dtype
|
|
input_list = modality_dispatcher.dispatch(input) # type: ignore
|
|
weight_chunked = self.weight.chunk(self.num_experts, dim=0)
|
|
if self.bias is not None:
|
|
bias_chunked = self.bias.chunk(self.num_experts, dim=0)
|
|
for i in range(self.num_experts):
|
|
input_list[i] = _BF16ComputeLinear.apply(
|
|
input_list[i],
|
|
weight_chunked[i],
|
|
bias_chunked[i] if self.bias is not None else None,
|
|
output_dtype,
|
|
torch.bfloat16,
|
|
)
|
|
return modality_dispatcher.undispatch(*input_list) # type: ignore
|
|
|
|
|
|
def create_linear(
|
|
in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None
|
|
) -> BaseLinear | NativeMoELinear:
|
|
if num_experts == 1:
|
|
return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype)
|
|
else:
|
|
return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MoE TP parallel linear wrappers: per-expert vLLM parallel layers
|
|
# ---------------------------------------------------------------------------
|
|
class MoEQKVParallelLinear(nn.Module):
|
|
"""Per-expert QKVParallelLinear with modality dispatch.
|
|
|
|
Wraps ``num_experts`` independent QKVParallelLinear instances.
|
|
Forward: dispatch tokens by modality → per-expert QKV matmul (TP-sharded)
|
|
→ undispatch.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
head_size: int,
|
|
total_num_heads: int,
|
|
total_num_kv_heads: int,
|
|
num_experts: int,
|
|
bias: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.experts = nn.ModuleList(
|
|
[
|
|
QKVParallelLinear(
|
|
hidden_size=hidden_size,
|
|
head_size=head_size,
|
|
total_num_heads=total_num_heads,
|
|
total_num_kv_heads=total_num_kv_heads,
|
|
bias=bias,
|
|
return_bias=False,
|
|
)
|
|
for _ in range(num_experts)
|
|
]
|
|
)
|
|
# Expose per-rank head info from the first expert (all are identical).
|
|
self.num_heads = self.experts[0].num_heads
|
|
self.num_kv_heads = self.experts[0].num_kv_heads
|
|
self.head_size = head_size
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
modality_dispatcher: ModalityDispatcher,
|
|
) -> torch.Tensor:
|
|
x_list = modality_dispatcher.dispatch(x)
|
|
out_list: list[torch.Tensor] = []
|
|
for i in range(self.num_experts):
|
|
out = self.experts[i](x_list[i])
|
|
out_list.append(out)
|
|
return modality_dispatcher.undispatch(*out_list)
|
|
|
|
|
|
class MoEColumnParallelLinear(nn.Module):
|
|
"""Per-expert ColumnParallelLinear with modality dispatch.
|
|
|
|
Forward: dispatch → per-expert column-parallel matmul → undispatch.
|
|
Output stays TP-local (no gather).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
num_experts: int,
|
|
bias: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.experts = nn.ModuleList(
|
|
[
|
|
ColumnParallelLinear(
|
|
input_size=input_size,
|
|
output_size=output_size,
|
|
bias=bias,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
)
|
|
for _ in range(num_experts)
|
|
]
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
modality_dispatcher: ModalityDispatcher,
|
|
) -> torch.Tensor:
|
|
x_list = modality_dispatcher.dispatch(x)
|
|
out_list: list[torch.Tensor] = []
|
|
for i in range(self.num_experts):
|
|
out = self.experts[i](x_list[i])
|
|
out_list.append(out)
|
|
return modality_dispatcher.undispatch(*out_list)
|
|
|
|
|
|
class MoERowParallelLinear(nn.Module):
|
|
"""Per-expert RowParallelLinear with modality dispatch.
|
|
|
|
Forward: dispatch → per-expert row-parallel matmul (includes all-reduce)
|
|
→ undispatch.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
num_experts: int,
|
|
bias: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.experts = nn.ModuleList(
|
|
[
|
|
RowParallelLinear(
|
|
input_size=input_size,
|
|
output_size=output_size,
|
|
bias=bias,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
)
|
|
for _ in range(num_experts)
|
|
]
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
modality_dispatcher: ModalityDispatcher,
|
|
) -> torch.Tensor:
|
|
x_list = modality_dispatcher.dispatch(x)
|
|
out_list: list[torch.Tensor] = []
|
|
for i in range(self.num_experts):
|
|
out = self.experts[i](x_list[i])
|
|
out_list.append(out)
|
|
return modality_dispatcher.undispatch(*out_list)
|
|
|
|
|
|
def validate_magi_human_tp_constraints(
|
|
*,
|
|
hidden_size: int,
|
|
num_heads_q: int,
|
|
num_heads_kv: int,
|
|
tensor_parallel_size: int,
|
|
) -> None:
|
|
"""Validate MagiHuman TP divisibility constraints.
|
|
|
|
Both shared layers (num_modality == 1) and MoE layers (num_modality == 3)
|
|
support TP via vLLM's parallel linear layers (QKVParallelLinear /
|
|
ColumnParallelLinear / RowParallelLinear). MoE layers use per-expert
|
|
parallel layers with modality dispatch.
|
|
|
|
Supported tp_sizes given default config (hidden=5120, heads_q=40, kv=8): 1, 2, 4.
|
|
"""
|
|
tp = tensor_parallel_size
|
|
if tp <= 1:
|
|
return
|
|
errors: list[str] = []
|
|
if num_heads_q % tp != 0:
|
|
errors.append(f"num_heads_q ({num_heads_q}) must be divisible by tensor_parallel_size ({tp})")
|
|
if num_heads_kv % tp != 0:
|
|
errors.append(f"num_heads_kv ({num_heads_kv}) must be divisible by tensor_parallel_size ({tp})")
|
|
# SWIGLU layers use intermediate = int(hidden * 8/3) // 4 * 4
|
|
intermediate_swiglu = int(hidden_size * 4 * 2 / 3) // 4 * 4
|
|
if intermediate_swiglu % tp != 0:
|
|
errors.append(
|
|
f"swiglu intermediate_size ({intermediate_swiglu}) must be divisible by "
|
|
f"tensor_parallel_size ({tp}). Supported tp values: 1, 2, 4"
|
|
)
|
|
# GELU7 MoE layers use intermediate = hidden * 4
|
|
intermediate_gelu = hidden_size * 4
|
|
if intermediate_gelu % tp != 0:
|
|
errors.append(f"gelu intermediate_size ({intermediate_gelu}) must be divisible by tensor_parallel_size ({tp})")
|
|
if errors:
|
|
raise ValueError("MagiHuman TP constraint violations:\n" + "\n".join(f" - {e}" for e in errors))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Flash attention (no context-parallelism) — uses vllm's flash attention
|
|
# ---------------------------------------------------------------------------
|
|
|
|
HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None
|
|
|
|
|
|
def _fa_varlen_simple(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
had_batch = query.ndim == 4
|
|
if had_batch:
|
|
query = query.squeeze(0)
|
|
key = key.squeeze(0)
|
|
value = value.squeeze(0)
|
|
seq_len = query.shape[0]
|
|
cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=query.device)
|
|
out = _vllm_fa_varlen(
|
|
q=query,
|
|
k=key,
|
|
v=value,
|
|
cu_seqlens_q=cu_seqlens,
|
|
cu_seqlens_k=cu_seqlens,
|
|
max_seqlen_q=seq_len,
|
|
max_seqlen_k=seq_len,
|
|
)
|
|
if had_batch:
|
|
out = out.unsqueeze(0)
|
|
return out
|
|
|
|
|
|
@magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True)
|
|
def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
|
|
return _fa_varlen_simple(query, key, value)
|
|
|
|
|
|
def _split_q_range_with_no_overlap(
|
|
q_ranges: torch.Tensor, k_ranges: torch.Tensor
|
|
) -> tuple[list[list[int]], list[list[list[int]]]]:
|
|
range_boundary = torch.unique(q_ranges, sorted=True).tolist()
|
|
candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])]
|
|
q_ranges = q_ranges.tolist()
|
|
k_ranges = k_ranges.tolist()
|
|
for q_range, k_range in zip(q_ranges, k_ranges):
|
|
q_start, q_end = q_range
|
|
for q_range_cand in candidates:
|
|
if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end:
|
|
q_range_cand[2].append(k_range)
|
|
q_ranges_out = []
|
|
k_ranges_out = []
|
|
for q_range_cand in candidates:
|
|
if len(q_range_cand[2]) > 0:
|
|
q_ranges_out.append(q_range_cand[0:2])
|
|
k_ranges_out.append(q_range_cand[2])
|
|
return q_ranges_out, k_ranges_out
|
|
|
|
|
|
def _flash_attn_with_correction(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
q_ranges: list[list[int]],
|
|
k_range_list: list[list[list[int]]],
|
|
):
|
|
output = torch.zeros_like(query)
|
|
output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device)
|
|
|
|
for q_range, k_ranges in zip(q_ranges, k_range_list):
|
|
q_start, q_end = q_range
|
|
q_chunk = query[q_start:q_end]
|
|
q_len = q_chunk.shape[0]
|
|
|
|
# Concatenate all k_ranges into a single key/value block, then run one
|
|
# flash-attention call. This avoids the need to merge per-chunk LSEs.
|
|
k_parts = [key[ks:ke] for ks, ke in k_ranges]
|
|
v_parts = [value[ks:ke] for ks, ke in k_ranges]
|
|
k_combined = torch.cat(k_parts, dim=0) if len(k_parts) > 1 else k_parts[0]
|
|
v_combined = torch.cat(v_parts, dim=0) if len(v_parts) > 1 else v_parts[0]
|
|
k_len = k_combined.shape[0]
|
|
|
|
cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=query.device)
|
|
cu_k = torch.tensor([0, k_len], dtype=torch.int32, device=query.device)
|
|
qo_out = _vllm_fa_varlen(
|
|
q=q_chunk,
|
|
k=k_combined,
|
|
v=v_combined,
|
|
cu_seqlens_q=cu_q,
|
|
cu_seqlens_k=cu_k,
|
|
max_seqlen_q=q_len,
|
|
max_seqlen_k=k_len,
|
|
)
|
|
output[q_start:q_end] = qo_out
|
|
return output, output_lse
|
|
|
|
|
|
def _flex_flash_attn_func_infer_output_meta(
|
|
query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
output = torch.empty_like(query)
|
|
output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device)
|
|
return output, output_lse
|
|
|
|
|
|
@magi_register_custom_op(
|
|
name="infra::flex_flash_attn_func",
|
|
mutates_args=(),
|
|
infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta,
|
|
is_subgraph_boundary=True,
|
|
)
|
|
def flex_flash_attn_func(
|
|
query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if HAS_MAGI_ATTENTION and _is_hopper_arch():
|
|
from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func
|
|
|
|
return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges)
|
|
else:
|
|
q_ranges_split, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges)
|
|
return _flash_attn_with_correction(query, key, value, q_ranges_split, k_range_list)
|
|
|
|
|
|
def flash_attn_no_cp(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
|
|
q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
|
|
return flash_attn_func(q, k, v).squeeze(0)
|
|
|
|
|
|
def flex_flash_attn_no_cp(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
q_ranges: torch.Tensor,
|
|
k_ranges: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
q, k, v = q.to(torch.bfloat16).squeeze(0), k.to(torch.bfloat16).squeeze(0), v.to(torch.bfloat16).squeeze(0)
|
|
out, _ = flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges)
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Attention module (no context-parallelism)
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class AttentionConfig:
|
|
hidden_size: int
|
|
num_heads_q: int
|
|
num_heads_kv: int
|
|
head_dim: int
|
|
params_dtype: torch.dtype
|
|
checkpoint_qk_layernorm_rope: bool
|
|
num_modality: int
|
|
num_layers: int
|
|
use_local_attn: bool = False
|
|
enable_attn_gating: bool = False
|
|
|
|
|
|
class Attention(torch.nn.Module):
|
|
config: AttentionConfig
|
|
|
|
def __init__(self, config: AttentionConfig):
|
|
super().__init__()
|
|
self.config = config
|
|
self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality)
|
|
self.gating_size = config.num_heads_q if config.enable_attn_gating else 0
|
|
|
|
# Both shared blocks (num_modality == 1) and MoE blocks (num_modality > 1)
|
|
# use vLLM's parallel linear layers for TP support.
|
|
# MoE blocks wrap per-expert parallel layers with modality dispatch.
|
|
if config.num_modality == 1:
|
|
# QKVParallelLinear handles GQA head-sharding for any tp_size.
|
|
# The combined checkpoint weight [Q, K, V, G] is split during
|
|
# load_weights: Q+K+V → linear_qkv, G → linear_gating.
|
|
self.linear_qkv = QKVParallelLinear(
|
|
hidden_size=config.hidden_size,
|
|
head_size=config.head_dim,
|
|
total_num_heads=config.num_heads_q,
|
|
total_num_kv_heads=config.num_heads_kv,
|
|
bias=False,
|
|
return_bias=False,
|
|
)
|
|
self.linear_proj = RowParallelLinear(
|
|
input_size=config.num_heads_q * config.head_dim,
|
|
output_size=config.hidden_size,
|
|
bias=False,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
)
|
|
if config.enable_attn_gating:
|
|
self.linear_gating = ColumnParallelLinear(
|
|
input_size=config.hidden_size,
|
|
output_size=config.num_heads_q,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
)
|
|
else:
|
|
self.linear_gating = None
|
|
else:
|
|
# MoE blocks: per-expert TP-sharded parallel layers.
|
|
self.linear_qkv = MoEQKVParallelLinear(
|
|
hidden_size=config.hidden_size,
|
|
head_size=config.head_dim,
|
|
total_num_heads=config.num_heads_q,
|
|
total_num_kv_heads=config.num_heads_kv,
|
|
num_experts=config.num_modality,
|
|
bias=False,
|
|
)
|
|
self.linear_proj = MoERowParallelLinear(
|
|
input_size=config.num_heads_q * config.head_dim,
|
|
output_size=config.hidden_size,
|
|
num_experts=config.num_modality,
|
|
bias=False,
|
|
)
|
|
if config.enable_attn_gating:
|
|
self.linear_gating = MoEColumnParallelLinear(
|
|
input_size=config.hidden_size,
|
|
output_size=config.num_heads_q,
|
|
num_experts=config.num_modality,
|
|
bias=False,
|
|
)
|
|
else:
|
|
self.linear_gating = None
|
|
|
|
self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality)
|
|
self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality)
|
|
|
|
# q_size / kv_size reflect the per-rank head count when tp > 1.
|
|
# Both shared and MoE QKV layers expose .num_heads / .num_kv_heads.
|
|
if config.num_modality == 1:
|
|
self.q_size = self.linear_qkv.num_heads * config.head_dim
|
|
self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim
|
|
self._local_heads_q = self.linear_qkv.num_heads
|
|
self._local_heads_kv = self.linear_qkv.num_kv_heads
|
|
else:
|
|
self.q_size = self.linear_qkv.num_heads * config.head_dim
|
|
self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim
|
|
self._local_heads_q = self.linear_qkv.num_heads
|
|
self._local_heads_kv = self.linear_qkv.num_kv_heads
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
rope: torch.Tensor,
|
|
permute_mapping: torch.Tensor,
|
|
inv_permute_mapping: torch.Tensor,
|
|
varlen_handler: VarlenHandler,
|
|
local_attn_handler: FFAHandler | None,
|
|
modality_dispatcher: ModalityDispatcher,
|
|
) -> torch.Tensor:
|
|
hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
|
|
|
|
if self.config.num_modality == 1:
|
|
# vLLM parallel layers with return_bias=False return a single tensor.
|
|
qkv = self.linear_qkv(hidden_states).to(torch.float32)
|
|
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
|
if self.linear_gating is not None:
|
|
g = self.linear_gating(hidden_states).to(torch.float32)
|
|
else:
|
|
g = hidden_states.new_empty(hidden_states.shape[0], 0)
|
|
else:
|
|
# MoE TP path: per-expert QKV parallel layers.
|
|
qkv = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32)
|
|
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
|
if self.linear_gating is not None:
|
|
g = self.linear_gating(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32)
|
|
else:
|
|
g = hidden_states.new_empty(hidden_states.shape[0], 0)
|
|
|
|
q = q.view(-1, self._local_heads_q, self.config.head_dim)
|
|
k = k.view(-1, self._local_heads_kv, self.config.head_dim)
|
|
v = v.view(-1, self._local_heads_kv, self.config.head_dim)
|
|
g = g.view(k.shape[0], self._local_heads_q, -1)
|
|
|
|
q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
|
|
k = self.k_norm(k, modality_dispatcher=modality_dispatcher)
|
|
|
|
q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0)
|
|
k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0)
|
|
v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0)
|
|
|
|
sin_emb, cos_emb = rope.tensor_split(2, -1)
|
|
q = apply_rotary_emb_torch(q, cos_emb, sin_emb)
|
|
k = apply_rotary_emb_torch(k, cos_emb, sin_emb)
|
|
|
|
if self.config.use_local_attn and local_attn_handler is not None:
|
|
self_attn_out = flex_flash_attn_no_cp(q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges)
|
|
else:
|
|
self_attn_out = flash_attn_no_cp(q, k, v)
|
|
self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping)
|
|
|
|
if self.config.enable_attn_gating:
|
|
self_attn_out = self_attn_out * torch.sigmoid(g)
|
|
|
|
self_attn_out = self_attn_out.view(-1, self._local_heads_q * self.config.head_dim).to(torch.bfloat16)
|
|
if self.config.num_modality == 1:
|
|
return self.linear_proj(self_attn_out)
|
|
return self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MLP module
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class MLPConfig:
|
|
hidden_size: int
|
|
intermediate_size: int
|
|
activation_type: MLPActivationType
|
|
params_dtype: torch.dtype
|
|
num_modality: int = 1
|
|
num_layers: int = 1
|
|
gated_act: bool = False
|
|
|
|
|
|
class MLP(torch.nn.Module):
|
|
config: MLPConfig
|
|
|
|
def __init__(self, config: MLPConfig):
|
|
super().__init__()
|
|
num_experts = config.num_modality
|
|
self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality)
|
|
intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size
|
|
|
|
# Both shared blocks (num_experts == 1) and MoE blocks (num_experts > 1)
|
|
# use vLLM's parallel linear layers for TP support.
|
|
if num_experts == 1:
|
|
# ColumnParallelLinear shards the output dim uniformly. For
|
|
# SWIGLU7 the interleaved [up0, gate0, up1, gate1, ...] format
|
|
# is preserved within each rank's contiguous slice, so swiglu7
|
|
# (which uses x[..., ::2] / x[..., 1::2]) still works correctly.
|
|
self.up_gate_proj = ColumnParallelLinear(
|
|
input_size=config.hidden_size,
|
|
output_size=intermediate_size_up,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
)
|
|
self.down_proj = RowParallelLinear(
|
|
input_size=config.intermediate_size,
|
|
output_size=config.hidden_size,
|
|
bias=False,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
)
|
|
else:
|
|
# MoE blocks: per-expert TP-sharded parallel layers.
|
|
self.up_gate_proj = MoEColumnParallelLinear(
|
|
input_size=config.hidden_size,
|
|
output_size=intermediate_size_up,
|
|
num_experts=num_experts,
|
|
bias=False,
|
|
)
|
|
self.down_proj = MoERowParallelLinear(
|
|
input_size=config.intermediate_size,
|
|
output_size=config.hidden_size,
|
|
num_experts=num_experts,
|
|
bias=False,
|
|
)
|
|
self.activation_func = create_activation_func(config.activation_type)
|
|
|
|
def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor:
|
|
x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
|
|
if isinstance(self.up_gate_proj, ColumnParallelLinear):
|
|
x = self.up_gate_proj(x).to(torch.float32)
|
|
x = self.activation_func(x).to(torch.bfloat16)
|
|
return self.down_proj(x).to(torch.float32)
|
|
# MoE TP path: per-expert column/row parallel layers.
|
|
x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32)
|
|
x = self.activation_func(x).to(torch.bfloat16)
|
|
x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32)
|
|
return x
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Adapter (per-modality embedders + RoPE)
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class AdapterConfig:
|
|
hidden_size: int
|
|
num_attention_heads: int
|
|
text_in_channels: int
|
|
video_in_channels: int
|
|
audio_in_channels: int
|
|
params_dtype: torch.dtype
|
|
|
|
|
|
class Adapter(torch.nn.Module):
|
|
config: AdapterConfig
|
|
|
|
def __init__(self, config: AdapterConfig):
|
|
super().__init__()
|
|
self.config = config
|
|
self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
|
|
self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
|
|
self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
|
|
self.rope = ElementWiseFourierEmbed(
|
|
config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False
|
|
)
|
|
|
|
def forward(self, x, coords_mapping, video_mask, audio_mask, text_mask):
|
|
rope = self.rope(coords_mapping)
|
|
|
|
text_input = x[text_mask, : self.config.text_in_channels]
|
|
audio_input = x[audio_mask, : self.config.audio_in_channels]
|
|
video_input = x[video_mask, : self.config.video_in_channels]
|
|
|
|
text_out = self.text_embedder(text_input)
|
|
audio_out = self.audio_embedder(audio_input)
|
|
video_out = self.video_embedder(video_input)
|
|
|
|
output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=x.dtype)
|
|
output_x[text_mask] = text_out
|
|
output_x[audio_mask] = audio_out
|
|
output_x[video_mask] = video_out
|
|
return output_x, rope
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Transformer layer (no CP)
|
|
# ---------------------------------------------------------------------------
|
|
class TransFormerLayer(torch.nn.Module):
|
|
def __init__(self, config: Any, layer_idx: int):
|
|
super().__init__()
|
|
num_modality = 3 if layer_idx in config.mm_layers else 1
|
|
use_local_attn = layer_idx in config.local_attn_layers
|
|
self.post_norm = layer_idx in config.post_norm_layers
|
|
attention_config = AttentionConfig(
|
|
hidden_size=config.hidden_size,
|
|
num_heads_q=config.num_heads_q,
|
|
num_heads_kv=config.num_heads_kv,
|
|
head_dim=config.head_dim,
|
|
params_dtype=config.params_dtype,
|
|
checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope,
|
|
num_modality=num_modality,
|
|
num_layers=config.num_layers,
|
|
use_local_attn=use_local_attn,
|
|
enable_attn_gating=config.enable_attn_gating,
|
|
)
|
|
self.attention: Attention = Attention(attention_config)
|
|
|
|
activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7
|
|
if activation_type == MLPActivationType.SWIGLU7:
|
|
gated_act = True
|
|
intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4
|
|
else:
|
|
gated_act = False
|
|
intermediate_size = config.hidden_size * 4
|
|
mlp_config = MLPConfig(
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
activation_type=activation_type,
|
|
params_dtype=config.params_dtype,
|
|
num_modality=num_modality,
|
|
num_layers=config.num_layers,
|
|
gated_act=gated_act,
|
|
)
|
|
self.mlp: MLP = MLP(mlp_config)
|
|
if self.post_norm:
|
|
self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
|
|
self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
rope: torch.Tensor,
|
|
permute_mapping: torch.Tensor,
|
|
inv_permute_mapping: torch.Tensor,
|
|
varlen_handler: VarlenHandler,
|
|
local_attn_handler: FFAHandler | None,
|
|
modality_dispatcher: ModalityDispatcher,
|
|
) -> torch.Tensor:
|
|
attn_out = self.attention(
|
|
hidden_states,
|
|
rope,
|
|
permute_mapping,
|
|
inv_permute_mapping,
|
|
varlen_handler,
|
|
local_attn_handler,
|
|
modality_dispatcher,
|
|
)
|
|
if self.post_norm:
|
|
attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
|
|
hidden_states = hidden_states + attn_out
|
|
|
|
mlp_out = self.mlp(hidden_states, modality_dispatcher)
|
|
if self.post_norm:
|
|
mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher)
|
|
hidden_states = hidden_states + mlp_out
|
|
return hidden_states
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TransformerBlock with magi_compile
|
|
# ---------------------------------------------------------------------------
|
|
is_base_model = True
|
|
|
|
|
|
def config_patch(compile_config: CompileConfig) -> CompileConfig:
|
|
global is_base_model
|
|
if is_base_model:
|
|
is_base_model = False
|
|
else:
|
|
compile_config.offload_config.gpu_resident_weight_ratio = 0.0
|
|
return compile_config
|
|
|
|
|
|
@magi_compile(
|
|
config_patch=config_patch, dynamic_arg_dims={"x": 0, "rope": 0, "permute_mapping": 0, "inv_permute_mapping": 0}
|
|
)
|
|
class TransformerBlock(torch.nn.Module):
|
|
def __init__(self, model_config: Any):
|
|
super().__init__()
|
|
self.layers: list[TransFormerLayer] = nn.ModuleList()
|
|
for layer_idx in range(model_config.num_layers):
|
|
self.layers.append(TransFormerLayer(model_config, layer_idx))
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
rope: torch.Tensor,
|
|
permute_mapping: torch.Tensor,
|
|
inv_permute_mapping: torch.Tensor,
|
|
varlen_handler: VarlenHandler,
|
|
local_attn_handler: FFAHandler | None,
|
|
modality_dispatcher: ModalityDispatcher,
|
|
) -> torch.Tensor:
|
|
for layer in self.layers:
|
|
x = layer(
|
|
x, rope, permute_mapping, inv_permute_mapping, varlen_handler, local_attn_handler, modality_dispatcher
|
|
)
|
|
return x
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Internal config for TransformerBlock / DiTModel construction
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class TransformerConfig:
|
|
hidden_size: int
|
|
video_in_channels: int
|
|
audio_in_channels: int
|
|
text_in_channels: int
|
|
params_dtype: torch.dtype
|
|
post_process_dtype: torch.dtype
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DiTModel (no context-parallelism)
|
|
# ---------------------------------------------------------------------------
|
|
class DiTModel(torch.nn.Module):
|
|
config: TransformerConfig
|
|
_layerwise_offload_blocks_attr = "blocks"
|
|
|
|
@property
|
|
def blocks(self) -> nn.ModuleList:
|
|
return self.block.layers
|
|
|
|
def __init__(self, model_config: Any):
|
|
super().__init__()
|
|
validate_magi_human_tp_constraints(
|
|
hidden_size=model_config.hidden_size,
|
|
num_heads_q=model_config.hidden_size // model_config.head_dim,
|
|
num_heads_kv=model_config.num_query_groups,
|
|
tensor_parallel_size=get_tensor_model_parallel_world_size(),
|
|
)
|
|
self.config = TransformerConfig(
|
|
hidden_size=model_config.hidden_size,
|
|
video_in_channels=model_config.video_in_channels,
|
|
audio_in_channels=model_config.audio_in_channels,
|
|
text_in_channels=model_config.text_in_channels,
|
|
params_dtype=model_config.params_dtype,
|
|
post_process_dtype=torch.float32,
|
|
)
|
|
adapter_config = AdapterConfig(
|
|
hidden_size=model_config.hidden_size,
|
|
num_attention_heads=model_config.num_heads_q,
|
|
text_in_channels=model_config.text_in_channels,
|
|
video_in_channels=model_config.video_in_channels,
|
|
audio_in_channels=model_config.audio_in_channels,
|
|
params_dtype=torch.float32,
|
|
)
|
|
self.adapter: Adapter = Adapter(adapter_config)
|
|
self.block: TransformerBlock = TransformerBlock(model_config=model_config)
|
|
self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size)
|
|
self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size)
|
|
self.final_linear_video = nn.Linear(
|
|
self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32
|
|
)
|
|
self.final_linear_audio = nn.Linear(
|
|
self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
coords_mapping: torch.Tensor,
|
|
modality_mapping: torch.Tensor,
|
|
varlen_handler: VarlenHandler,
|
|
local_attn_handler: FFAHandler | None,
|
|
):
|
|
modality_dispatcher = ModalityDispatcher(modality_mapping, 3)
|
|
permute_mapping = modality_dispatcher.permute_mapping
|
|
inv_permute_mapping = modality_dispatcher.inv_permute_mapping
|
|
video_mask = modality_mapping == Modality.VIDEO
|
|
audio_mask = modality_mapping == Modality.AUDIO
|
|
text_mask = modality_mapping == Modality.TEXT
|
|
|
|
x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
|
|
|
|
x = x.to(self.config.params_dtype)
|
|
x = ModalityDispatcher.permute(x, permute_mapping)
|
|
|
|
x = self.block(
|
|
x,
|
|
rope,
|
|
permute_mapping=permute_mapping,
|
|
inv_permute_mapping=inv_permute_mapping,
|
|
varlen_handler=varlen_handler,
|
|
local_attn_handler=local_attn_handler,
|
|
modality_dispatcher=modality_dispatcher,
|
|
)
|
|
|
|
x = ModalityDispatcher.inv_permute(x, inv_permute_mapping)
|
|
|
|
x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
|
|
x_video = self.final_norm_video(x_video)
|
|
x_video = self.final_linear_video(x_video)
|
|
|
|
x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
|
|
x_audio = self.final_norm_audio(x_audio)
|
|
x_audio = self.final_linear_audio(x_audio)
|
|
|
|
x_out = torch.zeros(
|
|
x.shape[0],
|
|
max(self.config.video_in_channels, self.config.audio_in_channels),
|
|
device=x.device,
|
|
dtype=x.dtype,
|
|
)
|
|
x_out[video_mask, : self.config.video_in_channels] = x_video
|
|
x_out[audio_mask, : self.config.audio_in_channels] = x_audio
|
|
|
|
return x_out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public config dataclass for building DiTModel from JSON
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class MagiHumanDiTConfig:
|
|
num_layers: int = 40
|
|
hidden_size: int = 5120
|
|
head_dim: int = 128
|
|
num_query_groups: int = 8
|
|
video_in_channels: int = 48 * 4
|
|
audio_in_channels: int = 64
|
|
text_in_channels: int = 3584
|
|
checkpoint_qk_layernorm_rope: bool = False
|
|
params_dtype: torch.dtype = torch.float32
|
|
mm_layers: list = field(default_factory=lambda: [0, 1, 2, 3, 36, 37, 38, 39])
|
|
local_attn_layers: list = field(default_factory=list)
|
|
enable_attn_gating: bool = True
|
|
gelu7_layers: list = field(default_factory=lambda: [0, 1, 2, 3])
|
|
post_norm_layers: list = field(default_factory=list)
|
|
|
|
def __post_init__(self):
|
|
self.num_heads_q = self.hidden_size // self.head_dim
|
|
self.num_heads_kv = self.num_query_groups
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from .pipeline_magi_human import EvalInput
|
|
|
|
|
|
# ===========================================================================
|
|
# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py)
|
|
# ===========================================================================
|
|
def _unfold_3d(
|
|
x: torch.Tensor,
|
|
kernel_size: tuple[int, int, int],
|
|
stride: tuple[int, int, int],
|
|
) -> torch.Tensor:
|
|
"""Pure-PyTorch 3D unfold matching UnfoldAnd behavior.
|
|
|
|
After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW).
|
|
UnfoldAnd permutes kernel dims next to channel before reshape so that the
|
|
col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics.
|
|
Without this permute, .view() interleaves spatial and kernel positions.
|
|
|
|
Args:
|
|
x: (N, C, D, H, W)
|
|
kernel_size: (kD, kH, kW)
|
|
stride: (sD, sH, sW)
|
|
Returns:
|
|
(N, C*kD*kH*kW, L) where L = product of output spatial dims.
|
|
"""
|
|
ndim = len(kernel_size)
|
|
for d in range(ndim):
|
|
x = x.unfold(d + 2, kernel_size[d], stride[d])
|
|
# x: (N, C, oD, oH, oW, kD, kH, kW)
|
|
# Permute to (N, C, kD, kH, kW, oD, oH, oW) so that view groups correctly
|
|
perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2))
|
|
x = x.permute(*perm).contiguous()
|
|
|
|
batch_size = x.shape[0]
|
|
col_dim = 1
|
|
for i in range(1, ndim + 2):
|
|
col_dim *= x.shape[i]
|
|
spatial = 1
|
|
for i in range(ndim + 2, 2 * ndim + 2):
|
|
spatial *= x.shape[i]
|
|
return x.view(batch_size, col_dim, spatial)
|
|
|
|
|
|
def calc_local_qk_range(
|
|
num_video_tokens,
|
|
num_audio_and_txt_tokens,
|
|
num_frames,
|
|
frame_receptive_field,
|
|
):
|
|
token_per_frame = num_video_tokens // num_frames
|
|
total_tokens = num_video_tokens + num_audio_and_txt_tokens
|
|
|
|
q_range_list = []
|
|
k_range_list = []
|
|
for i in range(num_frames):
|
|
q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame]))
|
|
k_range_list.append(
|
|
torch.tensor(
|
|
[
|
|
(i - frame_receptive_field) * token_per_frame,
|
|
(i + frame_receptive_field + 1) * token_per_frame,
|
|
]
|
|
)
|
|
)
|
|
local_q_range = torch.stack(q_range_list, dim=0)
|
|
local_k_range = torch.stack(k_range_list, dim=0)
|
|
|
|
local_k_range[local_k_range < 0] = 0
|
|
local_k_range[local_k_range > num_video_tokens] = num_video_tokens
|
|
|
|
video_q_range = torch.tensor([[0, num_video_tokens]])
|
|
video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]])
|
|
|
|
at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]])
|
|
at_k_ranges = torch.tensor([[0, total_tokens]])
|
|
|
|
q_ranges = (
|
|
torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
|
|
)
|
|
k_ranges = (
|
|
torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
|
|
)
|
|
return q_ranges, k_ranges
|
|
|
|
|
|
def calc_local_attn_ffa_handler(
|
|
num_video_tokens,
|
|
num_audio_and_txt_tokens,
|
|
num_frames,
|
|
frame_receptive_field,
|
|
):
|
|
q_ranges, k_ranges = calc_local_qk_range(
|
|
num_video_tokens,
|
|
num_audio_and_txt_tokens,
|
|
num_frames,
|
|
frame_receptive_field,
|
|
)
|
|
total = num_video_tokens + num_audio_and_txt_tokens
|
|
return FFAHandler(
|
|
q_ranges=q_ranges,
|
|
k_ranges=k_ranges,
|
|
max_seqlen_q=total,
|
|
max_seqlen_k=total,
|
|
attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32),
|
|
softmax_scale=None,
|
|
)
|
|
|
|
|
|
def get_coords(
|
|
shape: list[int],
|
|
ref_feat_shape: list[int],
|
|
offset_thw: list[int] | None = None,
|
|
device: torch.device = torch.device("cpu"),
|
|
dtype: torch.dtype = torch.float32,
|
|
):
|
|
if offset_thw is None:
|
|
offset_thw = [0, 0, 0]
|
|
ori_t, ori_h, ori_w = shape
|
|
ref_t, ref_h, ref_w = ref_feat_shape
|
|
|
|
offset_t, offset_h, offset_w = offset_thw
|
|
time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t
|
|
height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h
|
|
width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w
|
|
|
|
time_grid, height_grid, width_grid = torch.meshgrid(
|
|
time_rng,
|
|
height_rng,
|
|
width_rng,
|
|
indexing="ij",
|
|
)
|
|
coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3)
|
|
|
|
meta = torch.tensor(
|
|
[ori_t, ori_h, ori_w, ref_t, ref_h, ref_w],
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
meta_expanded = meta.expand(coords_flat.size(0), -1)
|
|
return torch.cat([coords_flat, meta_expanded], dim=-1)
|
|
|
|
|
|
@dataclass
|
|
class SingleData:
|
|
video_x_t: torch.Tensor
|
|
audio_x_t: torch.Tensor
|
|
audio_feat_len: int
|
|
txt_feat: torch.Tensor
|
|
txt_feat_len: int
|
|
t: int
|
|
h: int
|
|
w: int
|
|
patch_size: int
|
|
t_patch_size: int
|
|
spatial_rope_interpolation: Literal["inter", "extra"]
|
|
ref_audio_offset: int
|
|
text_offset: int
|
|
coords_style: Literal["v1", "v2"] = "v1"
|
|
|
|
def __post_init__(self):
|
|
self.video_token_num = self.video_x_t.shape[0]
|
|
self.audio_x_t = self.audio_x_t[: self.audio_feat_len]
|
|
self.txt_feat = self.txt_feat[: self.txt_feat_len]
|
|
self.video_channel = self.video_x_t.shape[-1]
|
|
self.audio_channel = self.audio_x_t.shape[-1]
|
|
self.txt_channel = self.txt_feat.shape[-1]
|
|
|
|
@property
|
|
def device(self):
|
|
return self.video_x_t.device
|
|
|
|
@property
|
|
def default_dtype(self):
|
|
return self.video_x_t.dtype
|
|
|
|
@property
|
|
def total_token_num(self):
|
|
return self.video_token_num + self.audio_feat_len + self.txt_feat_len
|
|
|
|
@property
|
|
def token_sequence(self):
|
|
tensors = [self.video_x_t, self.audio_x_t, self.txt_feat]
|
|
max_channel = max(t.shape[-1] for t in tensors)
|
|
padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors]
|
|
return torch.cat(padded, dim=0)
|
|
|
|
@property
|
|
def modality_mapping(self):
|
|
v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device)
|
|
a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device)
|
|
t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device)
|
|
return torch.cat([v_map, a_map, t_map], dim=0)
|
|
|
|
def default_coords(self, shape, ref_feat_shape, offset_thw=None):
|
|
if offset_thw is None:
|
|
offset_thw = [0, 0, 0]
|
|
return get_coords(
|
|
shape=shape,
|
|
ref_feat_shape=ref_feat_shape,
|
|
offset_thw=offset_thw,
|
|
device=self.device,
|
|
dtype=self.default_dtype,
|
|
)
|
|
|
|
@property
|
|
def coords_mapping(self):
|
|
if self.spatial_rope_interpolation == "inter":
|
|
video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32)
|
|
else:
|
|
video_ref_feat_shape = (
|
|
self.t // self.t_patch_size,
|
|
self.h // self.patch_size,
|
|
self.w // self.patch_size,
|
|
)
|
|
|
|
video_coords = self.default_coords(
|
|
shape=(
|
|
self.t // self.t_patch_size,
|
|
self.h // self.patch_size,
|
|
self.w // self.patch_size,
|
|
),
|
|
ref_feat_shape=video_ref_feat_shape,
|
|
)
|
|
|
|
if self.coords_style == "v1":
|
|
audio_coords = self.default_coords(
|
|
shape=(self.audio_feat_len, 1, 1),
|
|
ref_feat_shape=(self.t // self.t_patch_size, 1, 1),
|
|
)
|
|
text_coords = self.default_coords(
|
|
shape=(self.txt_feat_len, 1, 1),
|
|
ref_feat_shape=(2, 1, 1),
|
|
offset_thw=[self.text_offset, 0, 0],
|
|
)
|
|
elif self.coords_style == "v2":
|
|
magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1
|
|
audio_coords = self.default_coords(
|
|
shape=(self.audio_feat_len, 1, 1),
|
|
ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1),
|
|
)
|
|
text_coords = self.default_coords(
|
|
shape=(self.txt_feat_len, 1, 1),
|
|
ref_feat_shape=(1, 1, 1),
|
|
offset_thw=[-self.txt_feat_len, 0, 0],
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown coords_style: {self.coords_style}")
|
|
|
|
return torch.cat([video_coords, audio_coords, text_coords], dim=0)
|
|
|
|
def depack_token_sequence(self, token_sequence):
|
|
video_x_t = token_sequence[: self.video_token_num, : self.video_channel]
|
|
video_x_t = rearrange(
|
|
video_x_t,
|
|
"(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)",
|
|
H=self.h // self.patch_size,
|
|
W=self.w // self.patch_size,
|
|
pT=self.t_patch_size,
|
|
pH=self.patch_size,
|
|
pW=self.patch_size,
|
|
).contiguous()
|
|
audio_x_t = token_sequence[
|
|
self.video_token_num : self.video_token_num + self.audio_feat_len,
|
|
: self.audio_channel,
|
|
]
|
|
return video_x_t, audio_x_t
|
|
|
|
|
|
@dataclass
|
|
class SimplePackedData:
|
|
items: list[SingleData]
|
|
|
|
@property
|
|
def token_sequence(self):
|
|
return torch.cat([item.token_sequence for item in self.items], dim=0)
|
|
|
|
@property
|
|
def modality_mapping(self):
|
|
return torch.cat([item.modality_mapping for item in self.items], dim=0)
|
|
|
|
@property
|
|
def coords_mapping(self):
|
|
return torch.cat([item.coords_mapping for item in self.items], dim=0)
|
|
|
|
@property
|
|
def total_token_num(self):
|
|
return sum(item.total_token_num for item in self.items)
|
|
|
|
def __getitem__(self, index):
|
|
return self.items[index]
|
|
|
|
@property
|
|
def cu_seqlen(self):
|
|
cu = torch.cumsum(
|
|
torch.tensor([item.total_token_num for item in self.items]),
|
|
dim=0,
|
|
)
|
|
return F.pad(cu, (1, 0))
|
|
|
|
@property
|
|
def max_seqlen(self):
|
|
return torch.tensor(max(item.total_token_num for item in self.items))
|
|
|
|
def depack_token_sequence(self, token_sequence):
|
|
video_list, audio_list = [], []
|
|
parts = torch.split(
|
|
token_sequence,
|
|
[item.total_token_num for item in self.items],
|
|
dim=0,
|
|
)
|
|
for item, part in zip(self.items, parts):
|
|
v, a = item.depack_token_sequence(part)
|
|
video_list.append(v)
|
|
audio_list.append(a)
|
|
return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0)
|
|
|
|
|
|
class MagiDataProxy:
|
|
def __init__(
|
|
self,
|
|
patch_size: int = 2,
|
|
t_patch_size: int = 1,
|
|
frame_receptive_field: int = 11,
|
|
spatial_rope_interpolation: str = "extra",
|
|
ref_audio_offset: int = 1000,
|
|
text_offset: int = 0,
|
|
coords_style: str = "v2",
|
|
):
|
|
self.patch_size = patch_size
|
|
self.t_patch_size = t_patch_size
|
|
self.frame_receptive_field = frame_receptive_field
|
|
self.spatial_rope_interpolation = spatial_rope_interpolation
|
|
self.ref_audio_offset = ref_audio_offset
|
|
self.text_offset = text_offset
|
|
self.coords_style = coords_style
|
|
self._kernel = (t_patch_size, patch_size, patch_size)
|
|
self._stride = (t_patch_size, patch_size, patch_size)
|
|
self._saved_data: dict[str, Any] = {}
|
|
|
|
def saved_for_output(self, **kwargs):
|
|
self._saved_data.update(kwargs)
|
|
|
|
def get_saved_data(self, key: str):
|
|
return self._saved_data[key]
|
|
|
|
def img2tokens(self, x_t: torch.Tensor):
|
|
x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride)
|
|
return rearrange(
|
|
x_t_unfolded,
|
|
"N col_dim num_tokens -> N num_tokens col_dim",
|
|
).contiguous()
|
|
|
|
def process_input(self, transported_data: EvalInput):
|
|
batch_size, _, t, h, w = transported_data.x_t.shape
|
|
x_t = self.img2tokens(transported_data.x_t)
|
|
audio_x_t = transported_data.audio_x_t.contiguous()
|
|
text_in = transported_data.txt_feat.contiguous()
|
|
|
|
simple_packed_data = SimplePackedData(items=[])
|
|
for i in range(batch_size):
|
|
single_data = SingleData(
|
|
video_x_t=x_t[i],
|
|
audio_x_t=audio_x_t[i],
|
|
audio_feat_len=transported_data.audio_feat_len[i],
|
|
txt_feat=text_in[i],
|
|
txt_feat_len=transported_data.txt_feat_len[i],
|
|
t=t,
|
|
h=h,
|
|
w=w,
|
|
patch_size=self.patch_size,
|
|
t_patch_size=self.t_patch_size,
|
|
spatial_rope_interpolation=self.spatial_rope_interpolation,
|
|
ref_audio_offset=self.ref_audio_offset,
|
|
text_offset=self.text_offset,
|
|
coords_style=self.coords_style,
|
|
)
|
|
simple_packed_data.items.append(single_data)
|
|
|
|
if self.frame_receptive_field != -1:
|
|
assert batch_size == 1, "local attention only supports batch size 1"
|
|
local_attn_handler = calc_local_attn_ffa_handler(
|
|
num_video_tokens=simple_packed_data[0].video_token_num,
|
|
num_audio_and_txt_tokens=(simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len),
|
|
num_frames=t,
|
|
frame_receptive_field=self.frame_receptive_field,
|
|
)
|
|
if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor):
|
|
local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item()
|
|
if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor):
|
|
local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item()
|
|
else:
|
|
local_attn_handler = None
|
|
|
|
varlen_handler = VarlenHandler(
|
|
cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(),
|
|
cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(),
|
|
max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(),
|
|
max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(),
|
|
)
|
|
|
|
self.saved_for_output(simple_packed_data=simple_packed_data)
|
|
|
|
x = simple_packed_data.token_sequence
|
|
coords_mapping = simple_packed_data.coords_mapping
|
|
modality_mapping = simple_packed_data.modality_mapping
|
|
return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler)
|
|
|
|
def process_output(self, x: torch.Tensor):
|
|
simple_packed_data: SimplePackedData = self.get_saved_data("simple_packed_data")
|
|
return simple_packed_data.depack_token_sequence(x)
|