397 lines
14 KiB
Python
397 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Cosmos3 Edge transformer variant with a Nemotron dense UND backbone."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.distributed import get_tensor_model_parallel_world_size
|
|
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
|
|
from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
|
|
|
|
from .transformer_cosmos3 import (
|
|
Cosmos3VFMTransformer,
|
|
Qwen3VLTextRotaryEmbedding,
|
|
RMSNorm,
|
|
_apply_rotary_pos_emb,
|
|
_nested_get,
|
|
_tf_config_get,
|
|
)
|
|
|
|
COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
|
|
|
|
|
|
def _tf_config_nested_get(config: Any, key: str, default: Any = None) -> Any:
|
|
if config is None:
|
|
return default
|
|
if isinstance(config, dict):
|
|
value = _nested_get(config, key)
|
|
elif hasattr(config, "to_dict"):
|
|
value = _nested_get(config.to_dict(), key)
|
|
elif hasattr(config, "params"):
|
|
value = _nested_get(config.params, key)
|
|
else:
|
|
value = getattr(config, key, None)
|
|
return default if value is None else value
|
|
|
|
|
|
class Cosmos3Relu2MLP(nn.Module):
|
|
"""Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2)."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
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:
|
|
hidden = F.relu(self.up_proj(x))
|
|
return self.down_proj(hidden * hidden)
|
|
|
|
|
|
class Cosmos3EdgeCausalAttention(nn.Module):
|
|
"""Edge UND causal attention.
|
|
|
|
Reasoner self-attention uses raw UND Q/K with RoPE because
|
|
``qk_norm_for_text=false``. The GEN-facing UND key cache is a separate
|
|
normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled;
|
|
no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
hidden_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
head_dim: int,
|
|
rms_norm_eps: float,
|
|
use_und_k_norm_for_gen: bool,
|
|
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.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_und_k_norm_for_gen else None
|
|
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)
|
|
|
|
q_reasoner, k_reasoner = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
|
|
|
|
if self.k_norm_und_for_gen is not None:
|
|
k_gen = F.rms_norm(
|
|
k,
|
|
(self.head_dim,),
|
|
self.k_norm_und_for_gen.weight,
|
|
eps=self.k_norm_und_for_gen.variance_epsilon,
|
|
)
|
|
_, k_gen = _apply_rotary_pos_emb(q, k_gen, freqs_cos, freqs_sin)
|
|
else:
|
|
k_gen = k_reasoner
|
|
|
|
out = self.attn(q_reasoner, k_reasoner, v).reshape(B, S, -1)
|
|
return self.to_out(out), k_gen, v
|
|
|
|
|
|
class Cosmos3EdgeUndDecoderLayer(nn.Module):
|
|
"""Edge UND decoder layer: causal self-attention + ReLU2 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,
|
|
use_und_k_norm_for_gen: bool,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.self_attn = Cosmos3EdgeCausalAttention(
|
|
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,
|
|
use_und_k_norm_for_gen=use_und_k_norm_for_gen,
|
|
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 = Cosmos3Relu2MLP(
|
|
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]:
|
|
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 Cosmos3EdgeLanguageModel(nn.Module):
|
|
"""Nemotron dense UND tower that returns GEN-facing K plus raw V per layer."""
|
|
|
|
_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],
|
|
use_und_k_norm_for_gen: bool,
|
|
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(
|
|
[
|
|
Cosmos3EdgeUndDecoderLayer(
|
|
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,
|
|
use_und_k_norm_for_gen=use_und_k_norm_for_gen,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.layers.{i}",
|
|
)
|
|
for i in range(num_hidden_layers)
|
|
]
|
|
)
|
|
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]]:
|
|
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
|
|
|
|
|
|
class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer):
|
|
"""Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion."""
|
|
|
|
_language_model_cls = Cosmos3EdgeLanguageModel
|
|
_gen_mlp_cls = Cosmos3Relu2MLP
|
|
|
|
@staticmethod
|
|
def _validate_supported_config(model_config: Any) -> None:
|
|
expected_values = {
|
|
"qk_norm_for_text": False,
|
|
"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 Edge transformer config: {key}={actual!r}; expected {expected!r}."
|
|
)
|
|
|
|
backbone_type = _tf_config_get(model_config, "backbone_type", None)
|
|
if backbone_type is None:
|
|
raise ValueError(
|
|
f"Cosmos3 Edge transformer config must declare backbone_type={COSMOS3_EDGE_BACKBONE_TYPE!r}."
|
|
)
|
|
if backbone_type != COSMOS3_EDGE_BACKBONE_TYPE:
|
|
raise ValueError(
|
|
"Unsupported Cosmos3 Edge transformer config: "
|
|
f"backbone_type={backbone_type!r}; expected {COSMOS3_EDGE_BACKBONE_TYPE!r}."
|
|
)
|
|
|
|
required_values = {
|
|
"latent_channel": 48,
|
|
"latent_patch_size": 2,
|
|
"temporal_compression_factor": 4,
|
|
}
|
|
for key, expected in required_values.items():
|
|
actual = _tf_config_get(model_config, key, None)
|
|
if actual is None:
|
|
raise ValueError(f"Cosmos3 Edge transformer config must declare {key}={expected}.")
|
|
if int(actual) != expected:
|
|
raise ValueError(
|
|
f"Unsupported Cosmos3 Edge 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",
|
|
_tf_config_get(model_config, "layer_norm_epsilon", 1e-5),
|
|
)
|
|
)
|
|
|
|
@classmethod
|
|
def _resolve_rope_theta(cls, model_config: Any) -> float:
|
|
return float(_tf_config_get(model_config, "rope_theta", 100_000_000))
|
|
|
|
@classmethod
|
|
def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
|
|
mrope_section = _tf_config_get(model_config, "mrope_section", None)
|
|
if mrope_section is None:
|
|
mrope_section = _tf_config_nested_get(model_config, "mrope_section", [24, 20, 20])
|
|
return list(mrope_section)
|
|
|
|
def _language_model_kwargs(self) -> dict[str, Any]:
|
|
return {"use_und_k_norm_for_gen": self.use_und_k_norm_for_gen}
|
|
|
|
def validate_loaded_weights(self, loaded: set[str]) -> None:
|
|
missing: list[str] = []
|
|
for layer_idx in range(self.num_hidden_layers):
|
|
required_markers = (
|
|
f"language_model.layers.{layer_idx}.mlp.up_proj.",
|
|
f"language_model.layers.{layer_idx}.mlp.down_proj.",
|
|
f"gen_layers.{layer_idx}.mlp.up_proj.",
|
|
f"gen_layers.{layer_idx}.mlp.down_proj.",
|
|
)
|
|
if self.use_und_k_norm_for_gen:
|
|
required_markers = (
|
|
*required_markers,
|
|
f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.",
|
|
)
|
|
missing.extend(
|
|
marker.rstrip(".") for marker in required_markers if not any(marker in name for name in loaded)
|
|
)
|
|
if missing:
|
|
preview = ", ".join(missing[:10])
|
|
suffix = " ..." if len(missing) > 10 else ""
|
|
raise ValueError(
|
|
"Cosmos3 Edge transformer checkpoint is missing required weights for "
|
|
f"{preview}{suffix}. Use a converted Edge checkpoint with ReLU2 UND/GEN MLP weights."
|
|
)
|