1992 lines
82 KiB
Python
1992 lines
82 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Cosmos3 VFM Transformer for vllm-omni.
|
|
|
|
Implements the Mixture-of-Transformers architecture with two pathways:
|
|
- Understanding (UND): causal self-attention on text tokens (Qwen3-VL backbone)
|
|
- Generation (GEN): cross-attention where visual Q attends to [K_und, K_gen]
|
|
|
|
Ported from the TRT-LLM integration (tekit branch user/shreyasm/cosmos3).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Iterator
|
|
from contextlib import AbstractContextManager, contextmanager, nullcontext
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from cache_dit import ForwardPattern
|
|
from vllm.distributed import get_tensor_model_parallel_world_size
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|
QuantizationConfig,
|
|
)
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
|
|
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class RMSNorm(_VllmRMSNorm):
|
|
"""Cosmos3-local RMSNorm that uses the FP32 native implementation."""
|
|
|
|
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.forward_native(x)
|
|
|
|
def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.forward_native(x)
|
|
|
|
def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.forward_native(x)
|
|
|
|
|
|
def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]:
|
|
"""Return (ulysses_size, ulysses_rank, ulysses_pg) from vllm-omni parallel state.
|
|
|
|
Returns (1, 0, None) when sequence parallelism is not active.
|
|
"""
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_sp_group,
|
|
get_ulysses_parallel_rank,
|
|
get_ulysses_parallel_world_size,
|
|
)
|
|
|
|
size = get_ulysses_parallel_world_size()
|
|
if size <= 1:
|
|
return 1, 0, None
|
|
return size, get_ulysses_parallel_rank(), get_sp_group().ulysses_group
|
|
|
|
|
|
def _is_sp_active() -> bool:
|
|
"""Check whether sequence parallelism is active in the current forward context.
|
|
|
|
Follows the Bagel pattern: read ``forward_context.sp_active`` which returns
|
|
True when ``sequence_parallel_size > 1`` even without ``_sp_plan`` hooks.
|
|
"""
|
|
|
|
if not is_forward_context_available():
|
|
return False
|
|
return get_forward_context().sp_active
|
|
|
|
|
|
def _tf_config_get(config: Any, key: str, default: Any) -> Any:
|
|
"""Read a value from TransformerConfig, dict, or simple namespace."""
|
|
if config is None:
|
|
return default
|
|
if hasattr(config, "get"):
|
|
return config.get(key, default)
|
|
return getattr(config, key, default)
|
|
|
|
|
|
def _nested_get(value: Any, key: str) -> Any:
|
|
if isinstance(value, dict):
|
|
if key in value:
|
|
return value[key]
|
|
for child in value.values():
|
|
found = _nested_get(child, key)
|
|
if found is not None:
|
|
return found
|
|
elif isinstance(value, list | tuple):
|
|
for child in value:
|
|
found = _nested_get(child, key)
|
|
if found is not None:
|
|
return found
|
|
return None
|
|
|
|
|
|
def _od_config_get(od_config: Any, key: str, default: Any = None) -> Any:
|
|
"""Read Cosmos3 options from runtime, model, or transformer config."""
|
|
if od_config is None:
|
|
return default
|
|
for attr in ("custom_pipeline_args", "model_config"):
|
|
source = getattr(od_config, attr, None) or {}
|
|
if isinstance(source, dict):
|
|
if key in source:
|
|
return source[key]
|
|
found = _nested_get(source, key)
|
|
if found is not None:
|
|
return found
|
|
tf_model_config = getattr(od_config, "tf_model_config", None)
|
|
if isinstance(tf_model_config, dict):
|
|
if key in tf_model_config:
|
|
return tf_model_config[key]
|
|
found = _nested_get(tf_model_config, key)
|
|
if found is not None:
|
|
return found
|
|
value = _tf_config_get(tf_model_config, key, None)
|
|
return default if value is None else value
|
|
|
|
|
|
def _as_bool(value: Any) -> bool:
|
|
if isinstance(value, str):
|
|
return value.strip().lower() in {"1", "true", "yes", "on"}
|
|
return bool(value)
|
|
|
|
|
|
def resolve_sound_gen(od_config: Any) -> bool:
|
|
"""Capability gate shared by the pipeline and transformer.
|
|
|
|
Explicit ``sound_gen`` flag wins (including an explicit False);
|
|
otherwise infer from the presence of any sound-width key in od_config.
|
|
"""
|
|
sound_gen_value = _od_config_get(od_config, "sound_gen", None)
|
|
if sound_gen_value is not None:
|
|
return _as_bool(sound_gen_value)
|
|
for key in ("sound_dim", "io_channels", "vocoder_input_dim", "latent_ch"):
|
|
if _od_config_get(od_config, key, None) is not None:
|
|
return True
|
|
return False
|
|
|
|
|
|
class DomainAwareLinear(nn.Module):
|
|
"""Linear projection with one weight/bias pair per action embodiment domain."""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
num_domains: int,
|
|
*,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
) -> None:
|
|
super().__init__()
|
|
self.input_size = int(input_size)
|
|
self.output_size = int(output_size)
|
|
self.num_domains = int(num_domains)
|
|
self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size, dtype=dtype)
|
|
self.bias = nn.Embedding(self.num_domains, self.output_size, dtype=dtype)
|
|
nn.init.xavier_uniform_(self.fc.weight)
|
|
nn.init.zeros_(self.bias.weight)
|
|
|
|
def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:
|
|
if domain_id.ndim == 0:
|
|
domain_id = domain_id.unsqueeze(0)
|
|
domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1)
|
|
if x.shape[0] != domain_id.shape[0]:
|
|
raise ValueError(
|
|
"Cosmos3 action domain_id batch size must match action tokens: "
|
|
f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}."
|
|
)
|
|
if torch.any((domain_id < 0) | (domain_id >= self.num_domains)):
|
|
raise ValueError(f"Cosmos3 action domain_id must be in [0, {self.num_domains}), got {domain_id.tolist()}.")
|
|
|
|
weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size)
|
|
bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size)
|
|
if x.ndim == 2:
|
|
return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias
|
|
if x.ndim == 3:
|
|
return torch.bmm(x, weight) + bias.unsqueeze(1)
|
|
raise ValueError(f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}.")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rotary Position Embeddings (mRoPE)
|
|
# ---------------------------------------------------------------------------
|
|
def compute_mrope_position_ids_text(
|
|
num_tokens: int,
|
|
temporal_offset: int,
|
|
) -> tuple[torch.Tensor, int]:
|
|
"""Generate 3D mRoPE position IDs for text tokens.
|
|
|
|
Text tokens: all three axes (t, h, w) share the same monotonically
|
|
increasing position IDs.
|
|
"""
|
|
ids = torch.arange(num_tokens, dtype=torch.long) + temporal_offset
|
|
mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous()
|
|
return mrope_ids, temporal_offset + num_tokens
|
|
|
|
|
|
def compute_mrope_position_ids_vision(
|
|
grid_t: int,
|
|
grid_h: int,
|
|
grid_w: int,
|
|
temporal_offset: int | float,
|
|
fps: float | None = None,
|
|
base_fps: float = 24.0,
|
|
temporal_compression_factor: int = 4,
|
|
base_temporal_compression_factor: int | None = None,
|
|
enable_fps_modulation: bool = True,
|
|
start_frame_offset: int = 0,
|
|
) -> tuple[torch.Tensor, int | float]:
|
|
"""Generate 3D mRoPE position IDs for vision tokens.
|
|
|
|
Creates a (t, h, w) position grid with spatial indices reset per segment
|
|
(Qwen3VL-style). Flattened in t-major order.
|
|
"""
|
|
fps_modulation = enable_fps_modulation and fps is not None
|
|
|
|
if fps_modulation:
|
|
tps = fps / temporal_compression_factor
|
|
effective_base_tcf = (
|
|
base_temporal_compression_factor
|
|
if base_temporal_compression_factor is not None
|
|
else temporal_compression_factor
|
|
)
|
|
base_tps = base_fps / effective_base_tcf
|
|
frame_indices = torch.arange(grid_t, dtype=torch.float32)
|
|
t_index = (
|
|
((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset)
|
|
.view(-1, 1)
|
|
.expand(-1, grid_h * grid_w)
|
|
.flatten()
|
|
)
|
|
else:
|
|
t_index = (
|
|
torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
|
|
+ int(temporal_offset)
|
|
+ start_frame_offset
|
|
)
|
|
|
|
h_index = torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten()
|
|
w_index = torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten()
|
|
|
|
if fps_modulation:
|
|
mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0)
|
|
else:
|
|
mrope_ids = torch.stack([t_index, h_index, w_index], dim=0)
|
|
|
|
next_offset = math.floor(mrope_ids.max().item()) + 1
|
|
return mrope_ids, next_offset
|
|
|
|
|
|
def compute_mrope_position_ids_sound(
|
|
grid_t: int,
|
|
temporal_offset: int | float,
|
|
sound_latent_fps: float,
|
|
base_fps: float = 24.0,
|
|
temporal_compression_factor_sound: int = 1,
|
|
enable_fps_modulation: bool = True,
|
|
) -> tuple[torch.Tensor, int | float]:
|
|
"""Generate mRoPE IDs for sound tokens as a (T, 1, 1) grid."""
|
|
return compute_mrope_position_ids_vision(
|
|
grid_t=grid_t,
|
|
grid_h=1,
|
|
grid_w=1,
|
|
temporal_offset=temporal_offset,
|
|
fps=sound_latent_fps,
|
|
base_fps=base_fps,
|
|
temporal_compression_factor=temporal_compression_factor_sound,
|
|
base_temporal_compression_factor=temporal_compression_factor_sound,
|
|
enable_fps_modulation=enable_fps_modulation,
|
|
)
|
|
|
|
|
|
def compute_mrope_position_ids_action(
|
|
grid_t: int,
|
|
temporal_offset: int | float,
|
|
action_fps: float | None,
|
|
base_fps: float = 24.0,
|
|
base_temporal_compression_factor: int = 4,
|
|
enable_fps_modulation: bool = True,
|
|
start_frame_offset: int = 1,
|
|
) -> tuple[torch.Tensor, int | float]:
|
|
"""Generate mRoPE IDs for action tokens as a frame-rate (T, 1, 1) grid."""
|
|
return compute_mrope_position_ids_vision(
|
|
grid_t=grid_t,
|
|
grid_h=1,
|
|
grid_w=1,
|
|
temporal_offset=temporal_offset,
|
|
fps=action_fps,
|
|
base_fps=base_fps,
|
|
temporal_compression_factor=1,
|
|
base_temporal_compression_factor=base_temporal_compression_factor,
|
|
enable_fps_modulation=enable_fps_modulation,
|
|
start_frame_offset=start_frame_offset,
|
|
)
|
|
|
|
|
|
class Qwen3VLTextRotaryEmbedding(nn.Module):
|
|
"""Multi-dimensional rotary position embedding for Qwen3-VL."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
head_dim: int,
|
|
rope_theta: float,
|
|
mrope_section: list[int],
|
|
) -> None:
|
|
super().__init__()
|
|
self.head_dim = head_dim
|
|
self.rope_theta = rope_theta
|
|
self.mrope_section = mrope_section
|
|
inv_freq = 1.0 / (
|
|
rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float) / head_dim)
|
|
)
|
|
self.attention_scaling = 1.0
|
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
|
|
def apply_interleaved_mrope(self, freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor:
|
|
"""Reorganize from chunked [TTT...HHH...WWW] to interleaved [THTHW...]."""
|
|
freqs_t = freqs[0]
|
|
for dim, offset in enumerate((1, 2), start=1):
|
|
length = mrope_section[dim] * 3
|
|
idx = slice(offset, length, 3)
|
|
freqs_t[..., idx] = freqs[dim, ..., idx]
|
|
return freqs_t
|
|
|
|
@torch.no_grad()
|
|
def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if position_ids.ndim == 2:
|
|
position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
|
|
inv_freq_expanded = (
|
|
self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device)
|
|
)
|
|
position_ids_expanded = position_ids[:, :, None, :].float()
|
|
|
|
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
|
|
freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|
cos = emb.cos() * self.attention_scaling
|
|
sin = emb.sin() * self.attention_scaling
|
|
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RoPE application (Qwen3/Llama style)
|
|
# ---------------------------------------------------------------------------
|
|
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
|
def _apply_rotary_pos_emb(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Qwen3-style RoPE: (x * cos) + (rotate_half(x) * sin).
|
|
|
|
Args:
|
|
q: [B, S, h, D]
|
|
k: [B, S, H_kv, D]
|
|
cos: [1, S, 1, D] or broadcastable
|
|
sin: [1, S, 1, D] or broadcastable
|
|
"""
|
|
q_embed = (q * cos) + (_rotate_half(q) * sin)
|
|
k_embed = (k * cos) + (_rotate_half(k) * sin)
|
|
return q_embed, k_embed
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Timestep Embedder
|
|
# ---------------------------------------------------------------------------
|
|
class TimestepEmbedder(nn.Module):
|
|
"""Embeds scalar timesteps into vector representations."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
frequency_embedding_size: int = 256,
|
|
max_period: int = 10000,
|
|
) -> None:
|
|
super().__init__()
|
|
# Following diffusers naming pattern here for checkpoint compatibility.
|
|
self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True)
|
|
self.act = nn.SiLU()
|
|
self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
|
|
self.frequency_embedding_size = frequency_embedding_size
|
|
self.hidden_size = hidden_size
|
|
|
|
half = frequency_embedding_size // 2
|
|
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
|
|
self.register_buffer("freqs", freqs, persistent=False)
|
|
|
|
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
|
args = t[:, None] * self.freqs[None]
|
|
t_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
return self.linear_2(self.act(self.linear_1(t_freq)))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GatedMLP (replaces TRT-LLM GatedMLP)
|
|
# ---------------------------------------------------------------------------
|
|
class Cosmos3GatedMLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int = 4096,
|
|
intermediate_size: int = 12288,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.gate_proj = ColumnParallelLinear(
|
|
hidden_size,
|
|
intermediate_size,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.gate_proj",
|
|
)
|
|
self.up_proj = ColumnParallelLinear(
|
|
hidden_size,
|
|
intermediate_size,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.up_proj",
|
|
)
|
|
self.down_proj = RowParallelLinear(
|
|
intermediate_size,
|
|
hidden_size,
|
|
bias=False,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.down_proj",
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Attention Modules
|
|
# ---------------------------------------------------------------------------
|
|
class Cosmos3CausalAttention(nn.Module):
|
|
"""Understanding pathway: causal self-attention on text tokens.
|
|
|
|
Returns (output, K, V) where K/V are post-norm, post-RoPE for the
|
|
generation pathway's cross-attention.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
hidden_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
head_dim: int,
|
|
rms_norm_eps: float,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.hidden_size = hidden_size
|
|
self.num_heads = num_attention_heads
|
|
self.num_kv_heads = num_key_value_heads
|
|
self.head_dim = head_dim
|
|
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
self.num_heads_local = self.num_heads // tp_size
|
|
self.num_kv_heads_local = self.num_kv_heads // tp_size
|
|
|
|
self.to_q = ColumnParallelLinear(
|
|
hidden_size,
|
|
self.num_heads * self.head_dim,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_q",
|
|
)
|
|
self.to_k = ColumnParallelLinear(
|
|
hidden_size,
|
|
self.num_kv_heads * self.head_dim,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_k",
|
|
)
|
|
self.to_v = ColumnParallelLinear(
|
|
hidden_size,
|
|
self.num_kv_heads * self.head_dim,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_v",
|
|
)
|
|
self.to_out = RowParallelLinear(
|
|
self.num_heads * self.head_dim,
|
|
hidden_size,
|
|
bias=False,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_out",
|
|
)
|
|
|
|
self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
|
|
self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
|
|
|
|
# skip_sequence_parallel=True because the UND pathway is
|
|
# computed once and replicated across SP ranks.
|
|
# Only the GEN pathway is sequence-sharded.
|
|
self.attn = FrameworkAttention(
|
|
num_heads=self.num_heads,
|
|
head_size=self.head_dim,
|
|
causal=True,
|
|
softmax_scale=1.0 / (self.head_dim**0.5),
|
|
num_kv_heads=self.num_kv_heads,
|
|
skip_sequence_parallel=True,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
freqs_cos: torch.Tensor,
|
|
freqs_sin: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
B, S, _ = hidden_states.shape
|
|
|
|
q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim)
|
|
k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
|
|
v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
|
|
|
|
# Per-head QK norm
|
|
q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
|
|
k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
|
|
|
|
# Qwen3-style RoPE
|
|
q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
|
|
|
|
out = self.attn(q, k, v).reshape(B, S, -1)
|
|
return self.to_out(out), k, v
|
|
|
|
|
|
class Cosmos3CrossAttention(nn.Module):
|
|
"""Generation pathway: full attention where visual Q attends to all K/V.
|
|
|
|
* **Non-SP path**: explicit ``cat([k_und, k_gen])``. Text conditioning is
|
|
always present because K/V are physically concatenated.
|
|
|
|
* **SP path** (Ulysses active): ``k_und``/``v_und`` are passed as
|
|
``joint_key``/``joint_value`` in ``AttentionMetadata``. The Ulysses
|
|
wrapper head-slices the replicated UND K/V and performs all-to-all on the
|
|
sharded GEN Q/K/V so every query sees the full context.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
hidden_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
head_dim: int,
|
|
rms_norm_eps: float,
|
|
quant_config: QuantizationConfig | None = None,
|
|
qk_norm: bool = True,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.hidden_size = hidden_size
|
|
self.num_heads = num_attention_heads
|
|
self.num_kv_heads = num_key_value_heads
|
|
self.head_dim = head_dim
|
|
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
self.num_heads_local = self.num_heads // tp_size
|
|
self.num_kv_heads_local = self.num_kv_heads // tp_size
|
|
|
|
self.to_q = ColumnParallelLinear(
|
|
hidden_size,
|
|
self.num_heads * self.head_dim,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_q",
|
|
)
|
|
self.to_k = ColumnParallelLinear(
|
|
hidden_size,
|
|
self.num_kv_heads * self.head_dim,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_k",
|
|
)
|
|
self.to_v = ColumnParallelLinear(
|
|
hidden_size,
|
|
self.num_kv_heads * self.head_dim,
|
|
bias=False,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_v",
|
|
)
|
|
self.to_out = RowParallelLinear(
|
|
self.num_heads * self.head_dim,
|
|
hidden_size,
|
|
bias=False,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_out",
|
|
)
|
|
|
|
self.qk_norm = qk_norm
|
|
if self.qk_norm:
|
|
self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
|
|
self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
|
|
|
|
self.attn = FrameworkAttention(
|
|
num_heads=self.num_heads,
|
|
head_size=self.head_dim,
|
|
causal=False,
|
|
softmax_scale=1.0 / (self.head_dim**0.5),
|
|
num_kv_heads=self.num_kv_heads,
|
|
)
|
|
# Multi-control attention operates on one full [control_i, target]
|
|
# sequence at a time. Keep those sequences replicated when Ulysses is
|
|
# active; sharding the concatenated [control_1, ..., control_N, target]
|
|
# layout would split the per-control ranges across ranks.
|
|
self.multi_control_attn = FrameworkAttention(
|
|
num_heads=self.num_heads,
|
|
head_size=self.head_dim,
|
|
causal=False,
|
|
softmax_scale=1.0 / (self.head_dim**0.5),
|
|
num_kv_heads=self.num_kv_heads,
|
|
prefix=f"{prefix}.multi_control_attn",
|
|
skip_sequence_parallel=True,
|
|
)
|
|
|
|
# TODO(follow-up): collapse _forward_local and _forward_sp into a single
|
|
# joint-based path when NoParallelAttention can process joint_key/value.
|
|
# Currently the non-SP path must concatenate the replicated UND K/V explicitly,
|
|
# while the SP path passes it as joint_*.
|
|
|
|
# -- Non-SP path: explicit K/V concatenation + framework Attention --------
|
|
|
|
def _forward_local(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
k_und: torch.Tensor,
|
|
v_und: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
B, S_gen = q.shape[:2]
|
|
k_all = torch.cat([k_und, k], dim=1)
|
|
v_all = torch.cat([v_und, v], dim=1)
|
|
|
|
out = self.attn(q, k_all, v_all)
|
|
return out.reshape(B, S_gen, -1)
|
|
|
|
# -- SP path: framework Attention with joint_key/value -------------------
|
|
|
|
def _forward_sp(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
k_und: torch.Tensor,
|
|
v_und: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
B, S_gen = q.shape[:2]
|
|
|
|
# Zero-length joint_query satisfies the Ulysses contract
|
|
# (joint_query, joint_key, joint_value must all be present) without
|
|
# adding text tokens to Q. joint_len=0 keeps post_attention on the
|
|
# standard reverse-all-to-all path (no joint-output splitting).
|
|
joint_q = q.new_empty(B, 0, self.num_heads_local, self.head_dim)
|
|
|
|
attn_metadata = AttentionMetadata(
|
|
joint_query=joint_q,
|
|
joint_key=k_und,
|
|
joint_value=v_und,
|
|
joint_strategy="front",
|
|
)
|
|
out = self.attn(q, k, v, attn_metadata)
|
|
return out.reshape(B, S_gen, -1)
|
|
|
|
def _forward_multi_control(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
k_und: torch.Tensor,
|
|
v_und: torch.Tensor,
|
|
control_token_sizes: tuple[int, ...],
|
|
control_weights: tuple[float, ...],
|
|
) -> torch.Tensor:
|
|
"""Run independent attention per control and weight target outputs.
|
|
|
|
For control ``i``, both its queries and the target queries attend to
|
|
``[text, control_i, target]``. Each control keeps the output from its
|
|
own pass, while target outputs are combined using normalized weights.
|
|
"""
|
|
if len(control_token_sizes) != len(control_weights):
|
|
raise ValueError(
|
|
"Cosmos3 control token sizes and weights must have the same length: "
|
|
f"sizes={len(control_token_sizes)}, weights={len(control_weights)}."
|
|
)
|
|
control_tokens = sum(control_token_sizes)
|
|
if control_tokens <= 0 or control_tokens >= q.shape[1]:
|
|
raise ValueError(
|
|
"Cosmos3 multi-control attention requires non-empty control and target token ranges, "
|
|
f"got control_tokens={control_tokens}, total_tokens={q.shape[1]}."
|
|
)
|
|
|
|
q_target = q[:, control_tokens:]
|
|
k_target = k[:, control_tokens:]
|
|
v_target = v[:, control_tokens:]
|
|
control_outputs: list[torch.Tensor] = []
|
|
target_output: torch.Tensor | None = None
|
|
start = 0
|
|
for size, weight in zip(control_token_sizes, control_weights, strict=True):
|
|
if size <= 0:
|
|
raise ValueError(f"Cosmos3 control token sizes must be positive, got {control_token_sizes}.")
|
|
end = start + size
|
|
q_control = q[:, start:end]
|
|
k_control = k[:, start:end]
|
|
v_control = v[:, start:end]
|
|
q_pair = torch.cat([q_control, q_target], dim=1)
|
|
k_pair = torch.cat([k_und, k_control, k_target], dim=1)
|
|
v_pair = torch.cat([v_und, v_control, v_target], dim=1)
|
|
pair_output = self.multi_control_attn(q_pair, k_pair, v_pair)
|
|
control_outputs.append(pair_output[:, :size])
|
|
weighted_target = pair_output[:, size:] * weight
|
|
target_output = weighted_target if target_output is None else target_output + weighted_target
|
|
start = end
|
|
|
|
if start != control_tokens or target_output is None:
|
|
raise RuntimeError("Cosmos3 multi-control attention failed to cover all control token ranges.")
|
|
output = torch.cat([*control_outputs, target_output], dim=1)
|
|
return output.reshape(q.shape[0], q.shape[1], -1)
|
|
|
|
# -- Public forward: routes to the appropriate path ----------------------
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
k_und: torch.Tensor,
|
|
v_und: torch.Tensor,
|
|
freqs_cos: torch.Tensor,
|
|
freqs_sin: torch.Tensor,
|
|
control_token_sizes: tuple[int, ...] | None = None,
|
|
control_weights: tuple[float, ...] | None = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
hidden_states: [B, S_gen_local, hidden_size] (may be sequence-sharded)
|
|
k_und: [B, S_und, H_kv_local, D] pre-computed UND keys (TP-sharded, post-norm/RoPE)
|
|
v_und: [B, S_und, H_kv_local, D] pre-computed UND values (TP-sharded)
|
|
freqs_cos: [B, S_gen_local, 1, D]
|
|
freqs_sin: [B, S_gen_local, 1, D]
|
|
"""
|
|
B, S_gen, _ = hidden_states.shape
|
|
|
|
q = self.to_q(hidden_states).view(B, S_gen, self.num_heads_local, self.head_dim)
|
|
k = self.to_k(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
|
|
v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
|
|
|
|
# Per-head QK norm
|
|
if self.qk_norm:
|
|
q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
|
|
k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
|
|
|
|
# Qwen3-style RoPE
|
|
q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
|
|
|
|
if control_token_sizes is not None or control_weights is not None:
|
|
if control_token_sizes is None or control_weights is None:
|
|
raise ValueError("Cosmos3 multi-control attention requires both control token sizes and weights.")
|
|
out = self._forward_multi_control(
|
|
q,
|
|
k,
|
|
v,
|
|
k_und,
|
|
v_und,
|
|
control_token_sizes,
|
|
control_weights,
|
|
)
|
|
elif _is_sp_active():
|
|
out = self._forward_sp(q, k, v, k_und, v_und)
|
|
else:
|
|
out = self._forward_local(q, k, v, k_und, v_und)
|
|
|
|
return self.to_out(out)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Decoder Layers
|
|
# ---------------------------------------------------------------------------
|
|
class Cosmos3UndDecoderLayer(nn.Module):
|
|
"""Understanding pathway decoder layer: causal self-attention + MLP."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
head_dim: int,
|
|
rms_norm_eps: float,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.self_attn = Cosmos3CausalAttention(
|
|
hidden_size=hidden_size,
|
|
num_attention_heads=num_attention_heads,
|
|
num_key_value_heads=num_key_value_heads,
|
|
head_dim=head_dim,
|
|
rms_norm_eps=rms_norm_eps,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.self_attn",
|
|
)
|
|
self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
self.mlp = Cosmos3GatedMLP(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.mlp",
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
freqs: tuple[torch.Tensor, torch.Tensor],
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Returns (hidden_states, K, V) where K/V are for GEN cross-attention."""
|
|
residual = hidden_states
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
|
|
cos, sin = freqs
|
|
attn_out, k, v = self.self_attn(hidden_states, cos, sin)
|
|
hidden_states = residual + attn_out
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|
hidden_states = residual + self.mlp(hidden_states)
|
|
|
|
return hidden_states, k, v
|
|
|
|
|
|
class Cosmos3GenDecoderLayer(nn.Module):
|
|
"""Generation pathway decoder layer: cross-attention (to UND K/V) + MLP."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
layer_idx: int | None = None,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
head_dim: int,
|
|
rms_norm_eps: float,
|
|
quant_config: QuantizationConfig | None = None,
|
|
mlp_cls: type[nn.Module] = Cosmos3GatedMLP,
|
|
qk_norm: bool = True,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.layer_idx = layer_idx
|
|
self.cross_attention = Cosmos3CrossAttention(
|
|
hidden_size=hidden_size,
|
|
num_attention_heads=num_attention_heads,
|
|
num_key_value_heads=num_key_value_heads,
|
|
head_dim=head_dim,
|
|
rms_norm_eps=rms_norm_eps,
|
|
quant_config=quant_config,
|
|
qk_norm=qk_norm,
|
|
prefix=f"{prefix}.cross_attention",
|
|
)
|
|
self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
self.mlp = mlp_cls(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.mlp",
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
*,
|
|
k_und: torch.Tensor | None = None,
|
|
v_und: torch.Tensor | None = None,
|
|
freqs_cos: torch.Tensor | None = None,
|
|
freqs_sin: torch.Tensor | None = None,
|
|
cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None,
|
|
freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
control_token_sizes: tuple[int, ...] | None = None,
|
|
control_weights: tuple[float, ...] | None = None,
|
|
) -> torch.Tensor:
|
|
if cached_kv is not None:
|
|
if self.layer_idx is None:
|
|
raise ValueError("Cosmos3 GEN layer requires layer_idx when cached_kv is provided.")
|
|
k_und, v_und = cached_kv[self.layer_idx]
|
|
if freqs_gen is not None:
|
|
freqs_cos, freqs_sin = freqs_gen
|
|
if k_und is None or v_und is None or freqs_cos is None or freqs_sin is None:
|
|
raise ValueError("Cosmos3 GEN layer requires k_und/v_und/freqs_cos/freqs_sin.")
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
|
|
hidden_states = self.cross_attention(
|
|
hidden_states,
|
|
k_und=k_und,
|
|
v_und=v_und,
|
|
freqs_cos=freqs_cos,
|
|
freqs_sin=freqs_sin,
|
|
control_token_sizes=control_token_sizes,
|
|
control_weights=control_weights,
|
|
)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|
hidden_states = residual + self.mlp(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Language Model (Understanding pathway)
|
|
# ---------------------------------------------------------------------------
|
|
class Cosmos3LanguageModel(nn.Module):
|
|
"""Understanding pathway: a standard causal LM that processes text tokens.
|
|
|
|
Returns per-layer K/V tensors for the generation pathway's cross-attention.
|
|
The UND pathway is independent of the denoising step, so its K/V can be
|
|
computed once and reused across all sampling steps.
|
|
"""
|
|
|
|
_layerwise_offload_blocks_attrs = ["layers"]
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
num_hidden_layers: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
head_dim: int,
|
|
vocab_size: int,
|
|
rms_norm_eps: float,
|
|
rope_theta: float,
|
|
mrope_section: list[int],
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
|
|
self.rotary_emb = Qwen3VLTextRotaryEmbedding(
|
|
head_dim=head_dim,
|
|
rope_theta=rope_theta,
|
|
mrope_section=mrope_section,
|
|
)
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
Cosmos3UndDecoderLayer(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
num_attention_heads=num_attention_heads,
|
|
num_key_value_heads=num_key_value_heads,
|
|
head_dim=head_dim,
|
|
rms_norm_eps=rms_norm_eps,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.layers.{i}",
|
|
)
|
|
for i in range(num_hidden_layers)
|
|
]
|
|
)
|
|
# TODO: Not used right now, will be used in the future for prompt upsampler.
|
|
self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
text_ids: torch.Tensor,
|
|
freqs: tuple[torch.Tensor, torch.Tensor],
|
|
) -> list[tuple[torch.Tensor, torch.Tensor]]:
|
|
"""
|
|
Args:
|
|
text_ids: [B, S] token IDs
|
|
freqs: (cos, sin) each [B, S, 1, D]
|
|
|
|
Returns:
|
|
List of (K, V) per layer, each [B, S, H_kv, D].
|
|
|
|
No padding mask is applied: with right-padding + causal self-attention,
|
|
real query positions only attend to real keys, and the caller trims pad
|
|
K/V via ``max_real_len`` before the GEN cross-attention sees them.
|
|
"""
|
|
hidden = self.embed_tokens(text_ids)
|
|
|
|
cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
|
|
for layer in self.layers:
|
|
hidden, k, v = layer(hidden, freqs)
|
|
cached_kv.append((k, v))
|
|
|
|
return cached_kv
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main Transformer
|
|
# ---------------------------------------------------------------------------
|
|
class Cosmos3GenSPPrepare(nn.Module):
|
|
"""Module boundary used by _sp_plan to shard GEN states and RoPE together."""
|
|
|
|
def forward(
|
|
self,
|
|
hidden_gen: torch.Tensor,
|
|
freqs_cos: torch.Tensor,
|
|
freqs_sin: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
return hidden_gen, freqs_cos, freqs_sin
|
|
|
|
|
|
class Cosmos3VFMTransformer(nn.Module):
|
|
"""Cosmos3 VFM Transformer: UND language model + GEN denoising layers.
|
|
|
|
The UND pathway runs once per generation (K/V cached). The GEN pathway
|
|
runs at each denoising step over the target video/image latent stream and
|
|
optional transfer-control, action, and sound latent streams.
|
|
|
|
Layerwise offloading uses ``gen_layers`` as the GEN block container. The
|
|
nested UND language model declares its own ``layers`` container so the two
|
|
pathways are offloaded as independent rings.
|
|
|
|
Model-level CPU offload is Cosmos3-local: the UND ``reasoner`` and GEN
|
|
``generator`` components are swapped at the phase boundaries marked by
|
|
``with self._offload_context(...)``.
|
|
|
|
Sequence parallelism uses ``_sp_plan`` to shard/gather the GEN pathway at
|
|
module boundaries. ``Cosmos3CrossAttention`` checks
|
|
``forward_context.sp_active`` at runtime and routes to the framework
|
|
``Attention`` layer (with Ulysses all-to-all) or plain SDPA accordingly.
|
|
"""
|
|
|
|
_cache_dit_adapter_config = CacheDiTAdapterConfig(
|
|
# Cosmos3 GEN blocks return only hidden_states. Per-layer UND K/V
|
|
# conditioning uses the transformer's cache-dit fallback path.
|
|
block_forward_patterns={
|
|
"gen_layers": ForwardPattern.Pattern_3,
|
|
},
|
|
has_separate_cfg=True,
|
|
check_forward_pattern=False,
|
|
)
|
|
|
|
_repeated_blocks = ["Cosmos3GenDecoderLayer"]
|
|
|
|
_layerwise_offload_blocks_attrs = ["gen_layers"]
|
|
|
|
packed_modules_mapping = {}
|
|
|
|
_language_model_cls = Cosmos3LanguageModel
|
|
_gen_mlp_cls = Cosmos3GatedMLP
|
|
|
|
@staticmethod
|
|
def _is_transformer_block(name: str, module) -> bool:
|
|
return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit()
|
|
|
|
_hsdp_shard_conditions = [_is_transformer_block]
|
|
|
|
# Modules whose parameters must NOT be FSDP-sharded at the root level.
|
|
# time_embedder is cast to fp32 by post_load_weights for precision; if it
|
|
# were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16),
|
|
# the dtype upcast would be silently reverted, causing dtype mismatch in forward.
|
|
_hsdp_ignored_modules = ["time_embedder"]
|
|
|
|
_sp_plan = {
|
|
"gen_sp_prepare": {
|
|
0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),
|
|
1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
|
|
2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
|
|
},
|
|
"gen_sp_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
|
}
|
|
|
|
@staticmethod
|
|
def _validate_supported_config(model_config: Any) -> None:
|
|
"""Fail loudly when a checkpoint requests an unsupported architecture."""
|
|
expected_values = {
|
|
"qk_norm_for_diffusion": True,
|
|
"qk_norm_for_text": True,
|
|
"position_embedding_type": "unified_3d_mrope",
|
|
"unified_3d_mrope_reset_spatial_ids": True,
|
|
"joint_attn_implementation": "two_way",
|
|
}
|
|
for key, expected in expected_values.items():
|
|
actual = _tf_config_get(model_config, key, expected)
|
|
if actual != expected:
|
|
raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.")
|
|
|
|
@classmethod
|
|
def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
|
|
return float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
|
|
|
|
@classmethod
|
|
def _resolve_rope_theta(cls, model_config: Any) -> float:
|
|
return float(_tf_config_get(model_config, "rope_theta", 5_000_000))
|
|
|
|
@classmethod
|
|
def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
|
|
rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
|
|
return list(rope_scaling.get("mrope_section", [24, 20, 20]))
|
|
|
|
def _language_model_kwargs(self) -> dict[str, Any]:
|
|
return {}
|
|
|
|
def validate_loaded_weights(self, loaded: set[str]) -> None:
|
|
del loaded
|
|
|
|
def __init__(
|
|
self,
|
|
od_config: OmniDiffusionConfig,
|
|
*,
|
|
temporal_compression_factor: int | None = None,
|
|
sound_gen: bool = False,
|
|
sound_dim: int | None = None,
|
|
sound_latent_fps: float | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
model_config = od_config.tf_model_config
|
|
self._validate_supported_config(model_config)
|
|
|
|
self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096))
|
|
self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36))
|
|
self.num_attention_heads = int(_tf_config_get(model_config, "num_attention_heads", 32))
|
|
self.num_key_value_heads = int(_tf_config_get(model_config, "num_key_value_heads", 8))
|
|
self.head_dim = int(_tf_config_get(model_config, "head_dim", 128))
|
|
self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288))
|
|
self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936))
|
|
self.rms_norm_eps = self._resolve_rms_norm_eps(model_config)
|
|
self.rope_theta = self._resolve_rope_theta(model_config)
|
|
self.mrope_section = self._resolve_mrope_section(model_config)
|
|
self.qk_norm_for_diffusion = bool(_tf_config_get(model_config, "qk_norm_for_diffusion", True))
|
|
self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2))
|
|
self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48))
|
|
self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001))
|
|
self.base_fps = float(_tf_config_get(model_config, "base_fps", 24.0))
|
|
self.sound_gen = sound_gen
|
|
self.sound_dim = sound_dim
|
|
self.sound_latent_fps = sound_latent_fps
|
|
if self.sound_gen and (sound_dim is None or sound_latent_fps is None):
|
|
raise ValueError(
|
|
"Cosmos3VFMTransformer requires an explicit sound_dim and sound_latent_fps when sound_gen is True; "
|
|
"the pipeline must pass Cosmos3SoundTokenizer.latent_ch so the audio projection "
|
|
"layers are sized from the authoritative AVAE latent width."
|
|
)
|
|
action_gen_value = _od_config_get(od_config, "action_gen", None)
|
|
action_dim_value = _od_config_get(od_config, "action_dim", None)
|
|
if action_dim_value is None:
|
|
action_dim_value = _od_config_get(od_config, "max_action_dim", None)
|
|
self.action_gen = _as_bool(action_gen_value) if action_gen_value is not None else False
|
|
self.action_dim = int(action_dim_value if action_dim_value is not None else 64)
|
|
self.num_embodiment_domains = int(_od_config_get(od_config, "num_embodiment_domains", 32))
|
|
if temporal_compression_factor is None:
|
|
temporal_compression_factor = _tf_config_get(model_config, "temporal_compression_factor", 4)
|
|
self.temporal_compression_factor = int(temporal_compression_factor)
|
|
self.temporal_compression_factor_sound = int(
|
|
_tf_config_get(model_config, "temporal_compression_factor_sound", 1)
|
|
)
|
|
self.temporal_compression_factor_sound = int(
|
|
_tf_config_get(model_config, "temporal_compression_factor_sound", 1)
|
|
)
|
|
self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True))
|
|
self.temporal_modality_margin = int(
|
|
_tf_config_get(
|
|
model_config,
|
|
"unified_3d_mrope_temporal_modality_margin",
|
|
15000,
|
|
)
|
|
)
|
|
self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size
|
|
|
|
self.use_und_k_norm_for_gen = _tf_config_get(model_config, "use_und_k_norm_for_gen", None)
|
|
|
|
dtype = od_config.dtype
|
|
quant_config = getattr(od_config, "quantization_config", None) if od_config else None
|
|
|
|
self.language_model = self._language_model_cls(
|
|
hidden_size=self.hidden_size,
|
|
intermediate_size=self.intermediate_size,
|
|
num_hidden_layers=self.num_hidden_layers,
|
|
num_attention_heads=self.num_attention_heads,
|
|
num_key_value_heads=self.num_key_value_heads,
|
|
head_dim=self.head_dim,
|
|
vocab_size=self.vocab_size,
|
|
rms_norm_eps=self.rms_norm_eps,
|
|
rope_theta=self.rope_theta,
|
|
mrope_section=self.mrope_section,
|
|
quant_config=quant_config,
|
|
prefix="language_model",
|
|
**self._language_model_kwargs(),
|
|
)
|
|
|
|
# Video projection layers are small; not worth quantizing.
|
|
self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size)
|
|
self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim)
|
|
self.time_embedder = TimestepEmbedder(self.hidden_size)
|
|
if self.action_gen:
|
|
self.action_proj_in = DomainAwareLinear(
|
|
self.action_dim,
|
|
self.hidden_size,
|
|
self.num_embodiment_domains,
|
|
dtype=dtype,
|
|
)
|
|
self.action_proj_out = DomainAwareLinear(
|
|
self.hidden_size,
|
|
self.action_dim,
|
|
self.num_embodiment_domains,
|
|
dtype=dtype,
|
|
)
|
|
self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size, dtype=dtype))
|
|
if self.sound_gen:
|
|
self.audio_proj_in = nn.Linear(self.sound_dim, self.hidden_size)
|
|
self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
|
|
self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
|
|
|
|
self.gen_layers = nn.ModuleList(
|
|
[
|
|
Cosmos3GenDecoderLayer(
|
|
layer_idx=i,
|
|
hidden_size=self.hidden_size,
|
|
intermediate_size=self.intermediate_size,
|
|
num_attention_heads=self.num_attention_heads,
|
|
num_key_value_heads=self.num_key_value_heads,
|
|
head_dim=self.head_dim,
|
|
rms_norm_eps=self.rms_norm_eps,
|
|
quant_config=quant_config,
|
|
mlp_cls=self._gen_mlp_cls,
|
|
qk_norm=self.qk_norm_for_diffusion,
|
|
prefix=f"gen_layers.{i}",
|
|
)
|
|
for i in range(self.num_hidden_layers)
|
|
]
|
|
)
|
|
|
|
self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
|
|
self.gen_sp_prepare = Cosmos3GenSPPrepare()
|
|
self.gen_sp_gather = nn.Identity()
|
|
|
|
# Cached state (populated on first forward, reused across denoising steps)
|
|
self.cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
|
|
self.cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None
|
|
|
|
self._model_cpu_offload_enabled = False
|
|
self._model_cpu_offload_device: torch.device | None = None
|
|
self._model_cpu_offload_mover: SequentialOffloadHook | None = None
|
|
self._active_model_cpu_offload_component: str | None = None
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
offload_device = getattr(self, "_model_cpu_offload_device", None)
|
|
if getattr(self, "_model_cpu_offload_enabled", False) and offload_device is not None:
|
|
return offload_device
|
|
return next(self.parameters()).device
|
|
|
|
def _model_cpu_offload_components(self) -> dict[str, list[nn.Module]]:
|
|
"""Cosmos3's mutually-exclusive reasoner/generator component sets."""
|
|
return {
|
|
"reasoner": [self.language_model.layers],
|
|
"generator": [self.gen_layers],
|
|
}
|
|
|
|
def _model_cpu_offload_component_tensor_ids(self) -> set[int]:
|
|
component_tensors: set[int] = set()
|
|
for modules in self._model_cpu_offload_components().values():
|
|
for module in modules:
|
|
component_tensors.update(id(param) for param in module.parameters())
|
|
component_tensors.update(id(buffer) for buffer in module.buffers())
|
|
return component_tensors
|
|
|
|
def _move_model_cpu_offload_residents(self) -> None:
|
|
"""Keep non-reasoner/non-generator weights resident on the target device."""
|
|
component_tensors = self._model_cpu_offload_component_tensor_ids()
|
|
device = self._model_cpu_offload_device
|
|
if device is None:
|
|
return
|
|
with torch.no_grad():
|
|
for param in self.parameters():
|
|
if id(param) not in component_tensors and param.data.device != device:
|
|
param.data = param.data.to(device, non_blocking=False)
|
|
for buffer in self.buffers():
|
|
if id(buffer) not in component_tensors and buffer.device != device:
|
|
buffer.data = buffer.data.to(device, non_blocking=False)
|
|
|
|
def _offload_model_cpu_component(self, name: str) -> None:
|
|
mover = self._model_cpu_offload_mover
|
|
if mover is None:
|
|
raise RuntimeError("Cosmos3 model CPU offload is not enabled")
|
|
for module in self._model_cpu_offload_components()[name]:
|
|
mover._to_cpu(module)
|
|
|
|
def _load_model_cpu_component(self, name: str) -> None:
|
|
mover = self._model_cpu_offload_mover
|
|
if mover is None:
|
|
raise RuntimeError("Cosmos3 model CPU offload is not enabled")
|
|
for module in self._model_cpu_offload_components()[name]:
|
|
mover._to_gpu(module)
|
|
|
|
def enable_model_cpu_offload(
|
|
self,
|
|
*,
|
|
device: torch.device,
|
|
pin_memory: bool = True,
|
|
use_hsdp: bool = False,
|
|
) -> None:
|
|
"""Enable Cosmos3 reasoner/generator CPU swapping inside ``forward``."""
|
|
if getattr(self, "_model_cpu_offload_enabled", False):
|
|
return
|
|
|
|
from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
|
|
|
|
self._model_cpu_offload_device = torch.device(device)
|
|
self._model_cpu_offload_mover = SequentialOffloadHook(
|
|
offload_targets=[],
|
|
device=self._model_cpu_offload_device,
|
|
pin_memory=pin_memory,
|
|
use_hsdp=use_hsdp,
|
|
)
|
|
self._move_model_cpu_offload_residents()
|
|
for name in self._model_cpu_offload_components():
|
|
self._offload_model_cpu_component(name)
|
|
self._model_cpu_offload_enabled = True
|
|
self._active_model_cpu_offload_component = None
|
|
logger.info("Cosmos3 component-level CPU offload enabled on %s", self._model_cpu_offload_device)
|
|
|
|
def disable_model_cpu_offload(self) -> None:
|
|
if not getattr(self, "_model_cpu_offload_enabled", False):
|
|
return
|
|
for name in self._model_cpu_offload_components():
|
|
self._load_model_cpu_component(name)
|
|
if self._model_cpu_offload_device is not None and self._model_cpu_offload_device.type != "cpu":
|
|
current_omni_platform.synchronize()
|
|
self._model_cpu_offload_enabled = False
|
|
self._model_cpu_offload_device = None
|
|
self._model_cpu_offload_mover = None
|
|
self._active_model_cpu_offload_component = None
|
|
|
|
def _activate_model_cpu_offload_component(self, name: str) -> None:
|
|
if not getattr(self, "_model_cpu_offload_enabled", False):
|
|
return
|
|
components = self._model_cpu_offload_components()
|
|
if name not in components:
|
|
raise ValueError(f"Unknown Cosmos3 offload component: {name!r} (known: {list(components)})")
|
|
if self._active_model_cpu_offload_component == name:
|
|
return
|
|
self._active_model_cpu_offload_component = None
|
|
for other in components:
|
|
if other != name:
|
|
self._offload_model_cpu_component(other)
|
|
self._load_model_cpu_component(name)
|
|
self._active_model_cpu_offload_component = name
|
|
|
|
@contextmanager
|
|
def _model_cpu_offload_context(self, name: str) -> Iterator[None]:
|
|
self._activate_model_cpu_offload_component(name)
|
|
yield
|
|
|
|
def _offload_context(self, name: str) -> AbstractContextManager[None]:
|
|
if not getattr(self, "_model_cpu_offload_enabled", False):
|
|
return nullcontext()
|
|
return self._model_cpu_offload_context(name)
|
|
|
|
# -- Patchify / Unpatchify -----------------------------------------------
|
|
|
|
def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]:
|
|
"""Returns (hp, wp, H_padded, W_padded)."""
|
|
p = self.latent_patch_size
|
|
H_padded = ((h + p - 1) // p) * p
|
|
W_padded = ((w + p - 1) // p) * p
|
|
return H_padded // p, W_padded // p, H_padded, W_padded
|
|
|
|
def patchify(self, latents: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
|
|
"""[B, C, t, h, w] -> [B, t*hp*wp, p*p*C], padding h/w if needed."""
|
|
B = latents.shape[0]
|
|
p = self.latent_patch_size
|
|
C = self.latent_channel_size
|
|
hp, wp, H_padded, W_padded = self._pad_to_patch_size(h, w)
|
|
|
|
if H_padded != h or W_padded != w:
|
|
latents = F.pad(latents, (0, W_padded - w, 0, H_padded - h))
|
|
|
|
x = latents.reshape(B, C, t, hp, p, wp, p)
|
|
x = x.permute(0, 2, 3, 5, 4, 6, 1) # [B, t, hp, wp, p, p, C]
|
|
return x.reshape(B, t * hp * wp, p * p * C)
|
|
|
|
def unpatchify(self, tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
|
|
"""[B, t*hp*wp, p*p*C] -> [B, C, t, h, w], cropping padding if needed."""
|
|
B = tokens.shape[0]
|
|
p = self.latent_patch_size
|
|
C = self.latent_channel_size
|
|
hp, wp, H_padded, W_padded = self._pad_to_patch_size(h, w)
|
|
|
|
x = tokens.reshape(B, t, hp, wp, p, p, C)
|
|
x = x.permute(0, 6, 1, 2, 4, 3, 5) # [B, C, t, hp, p, wp, p]
|
|
x = x.reshape(B, C, t, H_padded, W_padded)
|
|
|
|
if H_padded != h or W_padded != w:
|
|
x = x[:, :, :, :h, :w]
|
|
return x
|
|
|
|
def pack_sound(self, sound_latents: torch.Tensor) -> torch.Tensor:
|
|
"""[B, C_sound, T_sound] -> [B, T_sound, C_sound]."""
|
|
if sound_latents.ndim != 3:
|
|
raise ValueError(f"Cosmos3 sound latents must have shape [B, C, T], got {tuple(sound_latents.shape)}.")
|
|
if sound_latents.shape[1] != self.sound_dim:
|
|
raise ValueError(
|
|
f"Cosmos3 sound latent channel mismatch: expected {self.sound_dim}, got {sound_latents.shape[1]}."
|
|
)
|
|
return sound_latents.permute(0, 2, 1).contiguous()
|
|
|
|
@staticmethod
|
|
def unpack_sound(tokens: torch.Tensor) -> torch.Tensor:
|
|
"""[B, T_sound, C_sound] -> [B, C_sound, T_sound]."""
|
|
return tokens.permute(0, 2, 1).contiguous()
|
|
|
|
def pack_action(self, action_latents: torch.Tensor) -> torch.Tensor:
|
|
"""Validate and return action latents as [B, T_action, D_action] tokens."""
|
|
if action_latents.ndim != 3:
|
|
raise ValueError(f"Cosmos3 action latents must have shape [B, T, D], got {tuple(action_latents.shape)}.")
|
|
if action_latents.shape[-1] != self.action_dim:
|
|
raise ValueError(
|
|
f"Cosmos3 action latent dimension mismatch: expected {self.action_dim}, got {action_latents.shape[-1]}."
|
|
)
|
|
return action_latents.contiguous()
|
|
|
|
@staticmethod
|
|
def unpack_action(tokens: torch.Tensor) -> torch.Tensor:
|
|
"""Return [B, T_action, D_action] action predictions."""
|
|
return tokens.contiguous()
|
|
|
|
# -- RoPE computation ----------------------------------------------------
|
|
|
|
def _compute_rope_freqs(
|
|
self,
|
|
text_mask: torch.Tensor,
|
|
t: int,
|
|
hp: int,
|
|
wp: int,
|
|
fps: float | None,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
t_action: int | None = None,
|
|
action_start_frame_offset: int = 1,
|
|
action_fps: float | None = None,
|
|
t_sound: int | None = None,
|
|
num_vision_items: int = 1,
|
|
share_vision_temporal_positions: bool = False,
|
|
) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
|
|
"""Compute mRoPE cos/sin for UND text and GEN media pathways."""
|
|
if num_vision_items <= 0:
|
|
raise ValueError(f"Cosmos3 num_vision_items must be positive, got {num_vision_items}.")
|
|
B = text_mask.shape[0]
|
|
S_text = text_mask.shape[1]
|
|
text_lengths = text_mask.sum(dim=1).long()
|
|
effective_fps = fps if fps is not None and t > 1 else None
|
|
action_frames = int(t_action or 0)
|
|
sound_frames = int(t_sound or 0)
|
|
|
|
text_pos_list = []
|
|
gen_pos_list = []
|
|
for b in range(B):
|
|
real_len = int(text_lengths[b].item())
|
|
t_pos, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0)
|
|
media_temporal_offset = t_offset + self.temporal_modality_margin
|
|
gen_positions = []
|
|
if num_vision_items == 1 or share_vision_temporal_positions:
|
|
v_pos, _ = compute_mrope_position_ids_vision(
|
|
t,
|
|
hp,
|
|
wp,
|
|
temporal_offset=media_temporal_offset,
|
|
fps=effective_fps,
|
|
base_fps=self.base_fps,
|
|
temporal_compression_factor=self.temporal_compression_factor,
|
|
enable_fps_modulation=self.enable_fps_modulation,
|
|
)
|
|
gen_positions.extend([v_pos] * num_vision_items)
|
|
else:
|
|
vision_offset: int | float = media_temporal_offset
|
|
for _ in range(num_vision_items):
|
|
v_pos, vision_offset = compute_mrope_position_ids_vision(
|
|
t,
|
|
hp,
|
|
wp,
|
|
temporal_offset=vision_offset,
|
|
fps=effective_fps,
|
|
base_fps=self.base_fps,
|
|
temporal_compression_factor=self.temporal_compression_factor,
|
|
enable_fps_modulation=self.enable_fps_modulation,
|
|
)
|
|
gen_positions.append(v_pos)
|
|
if action_frames > 0:
|
|
a_pos, _ = compute_mrope_position_ids_action(
|
|
action_frames,
|
|
temporal_offset=media_temporal_offset,
|
|
action_fps=action_fps if action_fps is not None else fps,
|
|
base_fps=self.base_fps,
|
|
base_temporal_compression_factor=self.temporal_compression_factor,
|
|
enable_fps_modulation=self.enable_fps_modulation,
|
|
start_frame_offset=action_start_frame_offset,
|
|
)
|
|
gen_positions.append(a_pos)
|
|
if sound_frames > 0:
|
|
s_pos, _ = compute_mrope_position_ids_sound(
|
|
sound_frames,
|
|
temporal_offset=media_temporal_offset,
|
|
sound_latent_fps=self.sound_latent_fps,
|
|
base_fps=self.base_fps,
|
|
temporal_compression_factor_sound=getattr(self, "temporal_compression_factor_sound", 1),
|
|
enable_fps_modulation=self.enable_fps_modulation,
|
|
)
|
|
gen_positions.append(s_pos)
|
|
pos_dtype = gen_positions[0].dtype
|
|
for pos in gen_positions[1:]:
|
|
pos_dtype = torch.promote_types(pos_dtype, pos.dtype)
|
|
v_pos = torch.cat([pos.to(pos_dtype) for pos in gen_positions], dim=1)
|
|
if real_len < S_text:
|
|
t_pos = torch.cat(
|
|
[t_pos, torch.zeros(3, S_text - real_len, dtype=t_pos.dtype)],
|
|
dim=1,
|
|
)
|
|
text_pos_list.append(t_pos)
|
|
gen_pos_list.append(v_pos)
|
|
|
|
text_pos_ids = torch.stack(text_pos_list, dim=1).to(device) # [3, B, S_text]
|
|
gen_pos_ids = torch.stack(gen_pos_list, dim=1).to(device) # [3, B, S_gen]
|
|
|
|
rotary_emb = self.language_model.rotary_emb
|
|
_dummy = torch.tensor([], dtype=dtype, device=device)
|
|
cos_und, sin_und = rotary_emb(_dummy, position_ids=text_pos_ids)
|
|
cos_gen, sin_gen = rotary_emb(_dummy, position_ids=gen_pos_ids)
|
|
|
|
freqs_und = (cos_und.unsqueeze(2), sin_und.unsqueeze(2)) # (B, S, 1, D)
|
|
freqs_gen = (cos_gen.unsqueeze(2), sin_gen.unsqueeze(2))
|
|
return freqs_und, freqs_gen
|
|
|
|
# -- Cache management ----------------------------------------------------
|
|
|
|
def reset_cache(self) -> None:
|
|
self.cached_kv = None
|
|
self.cached_freqs_gen = None
|
|
|
|
@staticmethod
|
|
def _validate_gen_sequence_parallel(
|
|
*,
|
|
s_gen: int,
|
|
s_video: int,
|
|
s_control: int,
|
|
s_action: int,
|
|
s_sound: int,
|
|
has_action: bool,
|
|
has_sound: bool,
|
|
has_control: bool,
|
|
ulysses_size: int,
|
|
) -> None:
|
|
if ulysses_size <= 1 or s_gen % ulysses_size == 0:
|
|
return
|
|
|
|
detail_parts = []
|
|
if has_control:
|
|
detail_parts.append(f"control tokens {s_control}")
|
|
detail_parts.append(f"video tokens {s_video}")
|
|
if has_action:
|
|
detail_parts.append(f"action tokens {s_action}")
|
|
if has_sound:
|
|
detail_parts.append(f"sound tokens {s_sound}")
|
|
detail = " = " + " + ".join(detail_parts) if len(detail_parts) > 1 else ""
|
|
adjust_detail = (
|
|
"Adjust the spatial resolution, frame count, action chunk size, "
|
|
"sound duration, or sound latent FPS so the combined media sequence is a "
|
|
"multiple of ulysses_degree."
|
|
if has_control or has_action or has_sound
|
|
else (
|
|
"Adjust the spatial resolution so that "
|
|
"t * ceil(h/patch) * ceil(w/patch) is a multiple "
|
|
"of ulysses_degree."
|
|
)
|
|
)
|
|
raise ValueError(
|
|
f"GEN sequence length ({s_gen}{detail}) must be divisible by "
|
|
f"ulysses_degree ({ulysses_size}). {adjust_detail}"
|
|
)
|
|
|
|
def sound_latent_frames_for_sequence_parallel(
|
|
self,
|
|
*,
|
|
video_shape: tuple[int, int, int],
|
|
sound_frames: int,
|
|
num_vision_items: int = 1,
|
|
) -> int:
|
|
# Sound is the only modality the packed GEN sequence pairs with here: action and
|
|
# sound are never generated together (the pipeline rejects action+sound), so the
|
|
# base is just the vision tokens.
|
|
#
|
|
# Note: padded frames go through attention and are only trimmed on decode, so SP
|
|
# output is not bit-exact with non-SP (the extra frame perturbs the kept ones).
|
|
from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_world_size
|
|
|
|
ulysses_size = get_ulysses_parallel_world_size()
|
|
if ulysses_size <= 1 or sound_frames <= 0:
|
|
return sound_frames
|
|
t, h, w = video_shape
|
|
hp, wp, _, _ = self._pad_to_patch_size(h, w)
|
|
base = num_vision_items * t * hp * wp
|
|
pad = (-(base + sound_frames)) % ulysses_size
|
|
return sound_frames + pad
|
|
|
|
# -- Forward -------------------------------------------------------------
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
text_ids: torch.Tensor,
|
|
text_mask: torch.Tensor,
|
|
video_shape: tuple[int, int, int],
|
|
fps: float | None = None,
|
|
action_latents: torch.Tensor | None = None,
|
|
action_domain_ids: torch.Tensor | None = None,
|
|
action_noisy_mask: torch.Tensor | None = None,
|
|
action_start_frame_offset: int = 1,
|
|
action_fps: float | None = None,
|
|
sound_latents: torch.Tensor | None = None,
|
|
noisy_frame_mask: torch.Tensor | None = None,
|
|
control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None,
|
|
control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None,
|
|
transfer_share_vision_temporal_positions: bool = True,
|
|
**kwargs,
|
|
) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
|
"""
|
|
Args:
|
|
hidden_states: [B, C, t, h, w] noisy latents
|
|
timestep: [B] diffusion timestep
|
|
text_ids: [B, S_text] tokenized text
|
|
text_mask: [B, S_text] attention mask (1=real, 0=pad)
|
|
video_shape: (t, h, w) in latent space
|
|
fps: video frame rate for temporal mRoPE modulation
|
|
action_latents: Optional [B, T_action, D_action] noisy action latents.
|
|
action_domain_ids: Optional [B] embodiment domain IDs for action projections.
|
|
action_noisy_mask: Optional [B, T_action, 1] mask where 1=noisy
|
|
action token and 0=clean conditioned token.
|
|
sound_latents: Optional [B, C_sound, T_sound] noisy sound latents.
|
|
noisy_frame_mask: Optional [B, 1, t, 1, 1] mask where 1=noisy (add
|
|
timestep embedding, predict velocity) and 0=conditioned (clean
|
|
context, skip timestep embedding). None means all target vision
|
|
frames are noisy, as in T2I/T2V.
|
|
control_latents: Optional transfer-control latents. Controls are
|
|
clean vision context and are packed before the noisy target.
|
|
control_weights: Optional non-negative relative weights for
|
|
transfer controls. Values are normalized to sum to one.
|
|
|
|
Returns:
|
|
[B, C, t, h, w] velocity prediction, or
|
|
tuple outputs in video, action, sound order when action/sound streams
|
|
are provided. Transfer-control streams condition the video prediction
|
|
and are not returned.
|
|
"""
|
|
if kwargs:
|
|
raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}")
|
|
t, h, w = video_shape
|
|
hp, wp, _, _ = self._pad_to_patch_size(h, w)
|
|
text_lengths = text_mask.sum(dim=1)
|
|
min_real_len = int(text_lengths.min().item())
|
|
max_real_len = int(text_lengths.max().item())
|
|
if min_real_len != max_real_len:
|
|
raise ValueError(
|
|
f"Cosmos3 requires identical real text lengths within a batch "
|
|
f"(got min={min_real_len}, max={max_real_len})."
|
|
)
|
|
has_action = action_latents is not None
|
|
has_sound = sound_latents is not None
|
|
if control_latents is None:
|
|
control_latent_list: list[torch.Tensor] = []
|
|
elif isinstance(control_latents, torch.Tensor):
|
|
control_latent_list = [control_latents]
|
|
else:
|
|
control_latent_list = list(control_latents)
|
|
has_control = len(control_latent_list) > 0
|
|
if control_weights is None:
|
|
normalized_control_weights = (
|
|
[1.0 / len(control_latent_list)] * len(control_latent_list) if has_control else []
|
|
)
|
|
else:
|
|
if isinstance(control_weights, torch.Tensor):
|
|
raw_control_weights = control_weights.detach().cpu().flatten().tolist()
|
|
else:
|
|
raw_control_weights = list(control_weights)
|
|
if len(raw_control_weights) != len(control_latent_list):
|
|
raise ValueError(
|
|
"Cosmos3 transfer control_weights length must match control_latents: "
|
|
f"weights={len(raw_control_weights)}, controls={len(control_latent_list)}."
|
|
)
|
|
normalized_control_weights = [float(weight) for weight in raw_control_weights]
|
|
if any(not math.isfinite(weight) or weight < 0.0 for weight in normalized_control_weights):
|
|
raise ValueError(
|
|
"Cosmos3 transfer control_weights must be finite and non-negative, "
|
|
f"got {normalized_control_weights}."
|
|
)
|
|
control_weight_sum = sum(normalized_control_weights)
|
|
if control_weight_sum <= 0.0:
|
|
raise ValueError("Cosmos3 transfer control_weights must have a positive sum.")
|
|
normalized_control_weights = [weight / control_weight_sum for weight in normalized_control_weights]
|
|
use_multi_control_attention = len(control_latent_list) > 1
|
|
if has_control and (has_action or has_sound):
|
|
raise ValueError("Cosmos3 transfer control latents cannot be combined with action or sound latents.")
|
|
if has_action and not self.action_gen:
|
|
raise ValueError(
|
|
"Cosmos3 action generation was requested, but this transformer "
|
|
"was initialized without action modules. Check that the "
|
|
"transformer config enables action_gen."
|
|
)
|
|
if has_sound and not self.sound_gen:
|
|
raise ValueError(
|
|
"Cosmos3 sound generation was requested, but this transformer "
|
|
"was initialized without sound modules. Check that the "
|
|
"transformer config enables sound_gen or defines sound_dim."
|
|
)
|
|
|
|
# Query Ulysses state at runtime
|
|
ulysses_size, _, _ = _get_ulysses_state()
|
|
if use_multi_control_attention and ulysses_size > 1:
|
|
logger.warning_once(
|
|
"Cosmos3 multi-control attention requires full [control_i | target] "
|
|
"sequences, so the GEN pathway runs replicated on every Ulysses rank. "
|
|
"Sequence parallelism will not reduce per-rank memory or latency for "
|
|
"multi-control transfer requests."
|
|
)
|
|
|
|
# Pack action/sound tokens (no learned weights) up front so the UND
|
|
# cache sizing knows their token lengths. The modality projections are
|
|
# deferred into the generator offload context below, so model-level
|
|
# offload does not stage GEN weights before swapping to the reasoner.
|
|
action_tokens = None
|
|
sound_tokens = None
|
|
s_action = 0
|
|
s_sound = 0
|
|
if action_latents is not None:
|
|
if action_latents.shape[0] != hidden_states.shape[0]:
|
|
raise ValueError(
|
|
"Cosmos3 action and video batch sizes must match: "
|
|
f"video={hidden_states.shape[0]}, action={action_latents.shape[0]}."
|
|
)
|
|
if action_domain_ids is None:
|
|
action_domain_ids = torch.zeros(action_latents.shape[0], dtype=torch.long, device=action_latents.device)
|
|
action_tokens = self.pack_action(action_latents)
|
|
s_action = action_tokens.shape[1]
|
|
if sound_latents is not None:
|
|
if sound_latents.shape[0] != hidden_states.shape[0]:
|
|
raise ValueError(
|
|
"Cosmos3 sound and video batch sizes must match: "
|
|
f"video={hidden_states.shape[0]}, sound={sound_latents.shape[0]}."
|
|
)
|
|
sound_tokens = self.pack_sound(sound_latents)
|
|
s_sound = sound_tokens.shape[1]
|
|
|
|
# Run UND pathway once and cache K/V (replicated across all ranks).
|
|
# freqs_gen (M-RoPE cos/sin) is derived purely from shape/fps and is
|
|
# cached alongside UND K/V by both the bespoke and session-state paths.
|
|
need_kv = self.cached_kv is None
|
|
if need_kv or self.cached_freqs_gen is None:
|
|
freqs_und, freqs_gen = self._compute_rope_freqs(
|
|
text_mask,
|
|
t,
|
|
hp,
|
|
wp,
|
|
fps,
|
|
hidden_states.device,
|
|
hidden_states.dtype,
|
|
t_action=s_action,
|
|
action_start_frame_offset=action_start_frame_offset,
|
|
action_fps=action_fps,
|
|
t_sound=s_sound,
|
|
num_vision_items=len(control_latent_list) + 1,
|
|
share_vision_temporal_positions=transfer_share_vision_temporal_positions,
|
|
)
|
|
self.cached_freqs_gen = freqs_gen
|
|
|
|
if need_kv:
|
|
with self._offload_context("reasoner"):
|
|
cached_kv_full = self.language_model(text_ids, freqs_und)
|
|
# Trim to real text length (remove padding). K/V stay replicated;
|
|
# the framework Attention layer head-slices them via joint_key/value.
|
|
self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full]
|
|
|
|
with self._offload_context("generator"):
|
|
# Patchify latents and project to hidden space after UND cache
|
|
# construction, so model-level offload does not stage GEN weights
|
|
# before immediately swapping to the reasoner.
|
|
hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
|
|
s_video = hidden_video.shape[1]
|
|
s_control = 0
|
|
hidden_controls: list[torch.Tensor] = []
|
|
for idx, control in enumerate(control_latent_list):
|
|
if control.shape != hidden_states.shape:
|
|
raise ValueError(
|
|
"Cosmos3 transfer control latent shape must match target latent shape: "
|
|
f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}."
|
|
)
|
|
hidden_control = self.proj_in(
|
|
self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w)
|
|
)
|
|
hidden_controls.append(hidden_control)
|
|
s_control += hidden_control.shape[1]
|
|
control_token_sizes = tuple(control.shape[1] for control in hidden_controls)
|
|
multi_control_weights = tuple(normalized_control_weights) if use_multi_control_attention else None
|
|
multi_control_token_sizes = control_token_sizes if use_multi_control_attention else None
|
|
hidden_action = None
|
|
hidden_sound = None
|
|
if action_tokens is not None:
|
|
assert action_domain_ids is not None
|
|
hidden_action = self.action_proj_in(action_tokens, action_domain_ids)
|
|
hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype)
|
|
if sound_tokens is not None:
|
|
hidden_sound = self.audio_proj_in(sound_tokens)
|
|
hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype)
|
|
|
|
# Timestep embedding (fp32 for precision).
|
|
# For I2V: only add to noisy tokens, not conditioned ones.
|
|
# Conditioned frames are clean context and should not receive
|
|
# the diffusion timestep signal.
|
|
with torch.autocast(current_omni_platform.device_type, enabled=False):
|
|
time_embed = self.time_embedder((timestep * self.timestep_scale).float())
|
|
time_embed = time_embed.to(hidden_states.dtype)
|
|
|
|
if noisy_frame_mask is not None:
|
|
# Build per-token mask from per-frame mask.
|
|
# noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1]
|
|
token_noisy_mask = (
|
|
noisy_frame_mask[:, 0, :, 0, 0] # [B, t]
|
|
.unsqueeze(-1) # [B, t, 1]
|
|
.expand(-1, -1, hp * wp) # [B, t, hp*wp]
|
|
.reshape(hidden_video.shape[0], -1, 1) # [B, t*hp*wp, 1]
|
|
)
|
|
hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask
|
|
else:
|
|
hidden_video = hidden_video + time_embed.unsqueeze(1)
|
|
|
|
if hidden_action is not None:
|
|
if action_noisy_mask is None:
|
|
hidden_action = hidden_action + time_embed.unsqueeze(1)
|
|
else:
|
|
if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1):
|
|
raise ValueError(
|
|
"Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
|
|
f"got {tuple(action_noisy_mask.shape)}."
|
|
)
|
|
action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device)
|
|
hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
|
|
|
|
if hidden_sound is not None:
|
|
hidden_sound = hidden_sound + time_embed.unsqueeze(1)
|
|
hidden_parts = [*hidden_controls, hidden_video]
|
|
if hidden_action is not None:
|
|
hidden_parts.append(hidden_action)
|
|
if hidden_sound is not None:
|
|
hidden_parts.append(hidden_sound)
|
|
hidden_gen = torch.cat(hidden_parts, dim=1)
|
|
|
|
# Run GEN layers. UND K/V (replicated) is passed to each layer;
|
|
# the Cosmos3CrossAttention forwards them as joint_key/value so the
|
|
# framework Attention handles the Ulysses head-slicing internally.
|
|
if self.cached_kv is None or self.cached_freqs_gen is None:
|
|
raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
|
|
self._validate_gen_sequence_parallel(
|
|
s_gen=hidden_gen.shape[1],
|
|
s_video=s_video,
|
|
s_control=s_control,
|
|
s_action=s_action,
|
|
s_sound=s_sound,
|
|
has_action=has_action,
|
|
has_sound=has_sound,
|
|
has_control=has_control,
|
|
ulysses_size=1 if use_multi_control_attention else ulysses_size,
|
|
)
|
|
freqs_cos, freqs_sin = self.cached_freqs_gen
|
|
if not use_multi_control_attention:
|
|
hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
|
|
freqs_gen = (freqs_cos, freqs_sin)
|
|
|
|
if len(self.gen_layers) == len(self.cached_kv):
|
|
for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
|
|
hidden_gen = layer(
|
|
hidden_gen,
|
|
k_und=k_und,
|
|
v_und=v_und,
|
|
freqs_cos=freqs_cos,
|
|
freqs_sin=freqs_sin,
|
|
control_token_sizes=multi_control_token_sizes,
|
|
control_weights=multi_control_weights,
|
|
)
|
|
# Cache-dit's block wrapper may return a tuple; unwrap it.
|
|
if isinstance(hidden_gen, tuple):
|
|
hidden_gen = hidden_gen[0]
|
|
else:
|
|
# Cache-dit patches gen_layers to a grouped wrapper.
|
|
for layer in self.gen_layers:
|
|
hidden_gen = layer(
|
|
hidden_gen,
|
|
cached_kv=self.cached_kv,
|
|
freqs_gen=freqs_gen,
|
|
control_token_sizes=multi_control_token_sizes,
|
|
control_weights=multi_control_weights,
|
|
)
|
|
if isinstance(hidden_gen, tuple):
|
|
hidden_gen = hidden_gen[0]
|
|
|
|
if not use_multi_control_attention:
|
|
hidden_gen = self.gen_sp_gather(hidden_gen)
|
|
|
|
# Final norm and project back to latent space
|
|
hidden_gen = self.norm_moe_gen(hidden_gen)
|
|
if not has_action and not has_sound and not has_control:
|
|
return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
|
|
|
|
split_sizes = []
|
|
if has_control:
|
|
split_sizes.append(s_control)
|
|
split_sizes.append(s_video)
|
|
if has_action:
|
|
split_sizes.append(s_action)
|
|
if has_sound:
|
|
split_sizes.append(s_sound)
|
|
split_hidden = hidden_gen.split(split_sizes, dim=1)
|
|
split_idx = 0
|
|
if has_control:
|
|
split_idx += 1
|
|
hidden_video = split_hidden[split_idx]
|
|
split_idx += 1
|
|
video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
|
|
if has_control:
|
|
return video_pred
|
|
outputs: list[torch.Tensor] = [video_pred]
|
|
if has_action:
|
|
hidden_action = split_hidden[split_idx]
|
|
split_idx += 1
|
|
assert action_domain_ids is not None
|
|
outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
|
|
if has_sound:
|
|
hidden_sound = split_hidden[split_idx]
|
|
outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
|
|
return tuple(outputs)
|
|
|
|
def post_load_weights(self) -> None:
|
|
"""Post-load processing: ensure correct dtypes."""
|
|
self.time_embedder.to(torch.float32)
|