451 lines
16 KiB
Python
451 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""MiniCPM4 with PagedAttention + fp32 RoPE/RMSNorm for VoxCPM2.
|
|
|
|
Uses vllm Attention for KV cache, keeps fp32 precision ops from
|
|
minicpm4_hf_compat.py to match native VoxCPM2 numerics.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.config import CacheConfig, VllmConfig
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.attention import Attention
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
from .minicpm4_hf_compat import (
|
|
_apply_rotary_pos_emb,
|
|
_MiniCPMLongRoPE,
|
|
_MiniCPMMLP,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _resolve_lm_cfg(config: Any) -> Any:
|
|
"""Extract lm_config from VoxCPM2Config, converting dict to namespace if needed."""
|
|
lm_cfg = getattr(config, "lm_config", config)
|
|
if isinstance(lm_cfg, dict):
|
|
|
|
class _Cfg:
|
|
pass
|
|
|
|
c = _Cfg()
|
|
for k, v in lm_cfg.items():
|
|
setattr(c, k, v)
|
|
return c
|
|
return lm_cfg
|
|
|
|
|
|
# ===================================================================
|
|
# Attention with vllm PagedAttention backend
|
|
# ===================================================================
|
|
|
|
|
|
class _PagedMiniCPM4Attention(nn.Module):
|
|
"""PagedAttention + fp32 RoPE with separate q/k/v projections."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
kv_channels: int | None,
|
|
layer_idx: int,
|
|
cache_config: CacheConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.hidden_size = hidden_size
|
|
self.num_heads = num_attention_heads
|
|
self.head_dim = kv_channels if kv_channels else hidden_size // num_attention_heads
|
|
self.num_kv_heads = num_key_value_heads
|
|
self.q_size = self.num_heads * self.head_dim
|
|
self.kv_size = self.num_kv_heads * self.head_dim
|
|
|
|
self.q_proj = nn.Linear(hidden_size, self.q_size, bias=False)
|
|
self.k_proj = nn.Linear(hidden_size, self.kv_size, bias=False)
|
|
self.v_proj = nn.Linear(hidden_size, self.kv_size, bias=False)
|
|
self.o_proj = nn.Linear(self.q_size, hidden_size, bias=False)
|
|
self._fused_qkv_weight: torch.Tensor | None = None
|
|
|
|
self.attn = Attention(
|
|
self.num_heads,
|
|
self.head_dim,
|
|
scale=self.head_dim**-0.5,
|
|
num_kv_heads=self.num_kv_heads,
|
|
cache_config=cache_config,
|
|
prefix=f"{prefix}.attn",
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
rope_emb: _MiniCPMLongRoPE | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward: fused QKV → fp32 RoPE → PagedAttention → o_proj."""
|
|
if self._fused_qkv_weight is None:
|
|
self._fused_qkv_weight = torch.cat(
|
|
[
|
|
self.q_proj.weight,
|
|
self.k_proj.weight,
|
|
self.v_proj.weight,
|
|
],
|
|
dim=0,
|
|
).detach()
|
|
qkv = nn.functional.linear(hidden_states, self._fused_qkv_weight)
|
|
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
|
|
|
if rope_emb is not None:
|
|
cos, sin = rope_emb(positions)
|
|
bsz = q.shape[0]
|
|
q_r = q.view(bsz, self.num_heads, self.head_dim)
|
|
k_r = k.view(bsz, self.num_kv_heads, self.head_dim)
|
|
q_r = q_r.unsqueeze(0).transpose(1, 2) # [1, heads, n_tokens, dim]
|
|
k_r = k_r.unsqueeze(0).transpose(1, 2) # [1, kv_heads, n_tokens, dim]
|
|
q_r, k_r = _apply_rotary_pos_emb(q_r, k_r, cos, sin)
|
|
q = q_r.transpose(1, 2).squeeze(0).reshape(bsz, -1) # [n_tokens, q_size]
|
|
k = k_r.transpose(1, 2).squeeze(0).reshape(bsz, -1) # [n_tokens, kv_size]
|
|
|
|
attn_output = self.attn(q, k, v)
|
|
|
|
output = self.o_proj(attn_output)
|
|
return output
|
|
|
|
|
|
# ===================================================================
|
|
# Decoder Layer
|
|
# ===================================================================
|
|
|
|
|
|
class _PagedMiniCPM4DecoderLayer(nn.Module):
|
|
"""Decoder layer: PagedAttention + fp32 RMSNorm + muP scale_depth."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
kv_channels: int | None,
|
|
rms_norm_eps: float,
|
|
layer_idx: int,
|
|
num_hidden_layers: int,
|
|
use_mup: bool,
|
|
scale_depth: float,
|
|
cache_config: CacheConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.self_attn = _PagedMiniCPM4Attention(
|
|
hidden_size=hidden_size,
|
|
num_attention_heads=num_attention_heads,
|
|
num_key_value_heads=num_key_value_heads,
|
|
kv_channels=kv_channels,
|
|
layer_idx=layer_idx,
|
|
cache_config=cache_config,
|
|
prefix=f"{prefix}.self_attn",
|
|
)
|
|
self.mlp = _MiniCPMMLP(hidden_size, intermediate_size)
|
|
self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
|
|
|
self.use_mup = use_mup
|
|
self.scale_depth = scale_depth
|
|
self.num_hidden_layers = num_hidden_layers
|
|
|
|
def _residual_scale(self) -> float:
|
|
if self.use_mup:
|
|
return self.scale_depth / math.sqrt(self.num_hidden_layers)
|
|
return 1.0
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
residual: torch.Tensor | None,
|
|
rope_emb: _MiniCPMLongRoPE | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
# Pre-norm + attention
|
|
residual = hidden_states
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
hidden_states = self.self_attn(positions, hidden_states, rope_emb)
|
|
|
|
scale = self._residual_scale()
|
|
if scale != 1.0:
|
|
hidden_states = residual + hidden_states * scale
|
|
else:
|
|
hidden_states = residual + hidden_states
|
|
|
|
# Pre-norm + FFN
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|
hidden_states = self.mlp(hidden_states)
|
|
|
|
if scale != 1.0:
|
|
hidden_states = residual + hidden_states * scale
|
|
else:
|
|
hidden_states = residual + hidden_states
|
|
|
|
return hidden_states, None
|
|
|
|
|
|
# ===================================================================
|
|
# Full Model
|
|
# ===================================================================
|
|
|
|
|
|
class MiniCPM4PagedForVoxCPM2(nn.Module):
|
|
"""PagedAttention base_lm (28 layers) for VoxCPM2 scaffold."""
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
|
super().__init__()
|
|
config = vllm_config.model_config.hf_config
|
|
cache_config = vllm_config.cache_config
|
|
self.config = config
|
|
|
|
lm_cfg = _resolve_lm_cfg(config)
|
|
|
|
hidden_size = lm_cfg.hidden_size
|
|
num_hidden_layers = lm_cfg.num_hidden_layers
|
|
kv_channels = getattr(lm_cfg, "kv_channels", None)
|
|
|
|
self.vocab_size = lm_cfg.vocab_size
|
|
self.embed_tokens = nn.Embedding(self.vocab_size, hidden_size)
|
|
|
|
rope_scaling = getattr(lm_cfg, "rope_scaling", None)
|
|
if isinstance(rope_scaling, dict):
|
|
rope_scaling_dict = rope_scaling
|
|
elif hasattr(rope_scaling, "__dict__"):
|
|
rope_scaling_dict = {
|
|
"short_factor": rope_scaling.short_factor,
|
|
"long_factor": rope_scaling.long_factor,
|
|
"original_max_position_embeddings": rope_scaling.original_max_position_embeddings,
|
|
}
|
|
else:
|
|
rope_scaling_dict = {}
|
|
|
|
no_rope = getattr(lm_cfg, "no_rope", False)
|
|
if not no_rope:
|
|
self.rope_emb = _MiniCPMLongRoPE(
|
|
hidden_size=hidden_size,
|
|
num_attention_heads=lm_cfg.num_attention_heads,
|
|
kv_channels=kv_channels,
|
|
rope_theta=getattr(lm_cfg, "rope_theta", 10000.0),
|
|
max_position_embeddings=getattr(lm_cfg, "max_position_embeddings", 32768),
|
|
rope_scaling=rope_scaling_dict,
|
|
)
|
|
else:
|
|
self.rope_emb = None
|
|
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
_PagedMiniCPM4DecoderLayer(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=lm_cfg.intermediate_size,
|
|
num_attention_heads=lm_cfg.num_attention_heads,
|
|
num_key_value_heads=lm_cfg.num_key_value_heads,
|
|
kv_channels=kv_channels,
|
|
rms_norm_eps=lm_cfg.rms_norm_eps,
|
|
layer_idx=i,
|
|
num_hidden_layers=num_hidden_layers,
|
|
use_mup=getattr(lm_cfg, "use_mup", False),
|
|
scale_depth=getattr(lm_cfg, "scale_depth", 1.0),
|
|
cache_config=cache_config,
|
|
prefix=f"{prefix}.layers.{i}",
|
|
)
|
|
for i in range(num_hidden_layers)
|
|
]
|
|
)
|
|
|
|
self.norm = RMSNorm(hidden_size, eps=lm_cfg.rms_norm_eps)
|
|
|
|
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
|
["hidden_states", "residual"], hidden_size
|
|
)
|
|
|
|
use_mup = getattr(lm_cfg, "use_mup", False)
|
|
self._scale_emb = getattr(lm_cfg, "scale_emb", 1.0) if use_mup else 1.0
|
|
self._compiled_layers: set[int] = set()
|
|
|
|
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
|
return self.embed_tokens(input_ids) * self._scale_emb
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor | None,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
**kwargs: Any,
|
|
) -> torch.Tensor | IntermediateTensors:
|
|
if inputs_embeds is not None:
|
|
hidden_states = inputs_embeds
|
|
else:
|
|
hidden_states = self.embed_input_ids(input_ids)
|
|
|
|
residual = None
|
|
for layer in self.layers:
|
|
hidden_states, residual = layer(
|
|
positions,
|
|
hidden_states,
|
|
residual,
|
|
self.rope_emb,
|
|
)
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
return hidden_states
|
|
|
|
def precompute_fused_qkv(self) -> None:
|
|
"""Materialize fused QKV weights before CUDA Graph capture."""
|
|
for layer in self.layers:
|
|
attn = layer.self_attn
|
|
if attn._fused_qkv_weight is None:
|
|
attn._fused_qkv_weight = torch.cat(
|
|
[attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight],
|
|
dim=0,
|
|
).detach()
|
|
|
|
def compile_selective(self) -> list[str]:
|
|
"""Compile the full model ``forward`` as one graph.
|
|
|
|
``torch.compile`` is applied at the model level so Dynamo unrolls the
|
|
per-layer Python loop inside the graph. Graph breaks at PagedAttention
|
|
produce sub-graphs, but Dynamo memoises the whole trace once, so the
|
|
per-decode-step dispatch drops to just a few instead of one per layer.
|
|
"""
|
|
if self._compiled_layers:
|
|
return []
|
|
# Null the fused-qkv caches so the compile sees the real weight layout.
|
|
for layer in self.layers:
|
|
layer.self_attn._fused_qkv_weight = None
|
|
self.forward = torch.compile(self.forward, mode="default", fullgraph=False)
|
|
# Mark every layer as compiled so idempotent callers don't double-wrap.
|
|
self._compiled_layers.update(range(len(self.layers)))
|
|
return ["forward (whole model)"]
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
"""Load weights from native checkpoint (base_lm. prefix pre-stripped)."""
|
|
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
|
loaded: set[str] = set()
|
|
|
|
for name, loaded_weight in weights:
|
|
if "rotary_emb.inv_freq" in name:
|
|
continue
|
|
param = params_dict.get(name)
|
|
if param is None:
|
|
continue
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
weight_loader(param, loaded_weight)
|
|
loaded.add(name)
|
|
|
|
return loaded
|
|
|
|
|
|
# ===================================================================
|
|
# Residual LM with PagedAttention (no RoPE, 8 layers)
|
|
# ===================================================================
|
|
|
|
|
|
class MiniCPM4PagedResidualLM(nn.Module):
|
|
"""PagedAttention residual LM (8 layers, no RoPE) for VoxCPM2."""
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
|
super().__init__()
|
|
config = vllm_config.model_config.hf_config
|
|
cache_config = vllm_config.cache_config
|
|
self.config = config
|
|
|
|
lm_cfg = _resolve_lm_cfg(config)
|
|
|
|
hidden_size = lm_cfg.hidden_size
|
|
num_hidden_layers = getattr(config, "residual_lm_num_layers", 8)
|
|
kv_channels = getattr(lm_cfg, "kv_channels", None)
|
|
|
|
self.rope_emb = None
|
|
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
_PagedMiniCPM4DecoderLayer(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=lm_cfg.intermediate_size,
|
|
num_attention_heads=lm_cfg.num_attention_heads,
|
|
num_key_value_heads=lm_cfg.num_key_value_heads,
|
|
kv_channels=kv_channels,
|
|
rms_norm_eps=lm_cfg.rms_norm_eps,
|
|
layer_idx=i,
|
|
num_hidden_layers=num_hidden_layers,
|
|
use_mup=getattr(lm_cfg, "use_mup", False),
|
|
scale_depth=getattr(lm_cfg, "scale_depth", 1.0),
|
|
cache_config=cache_config,
|
|
prefix=f"{prefix}.layers.{i}",
|
|
)
|
|
for i in range(num_hidden_layers)
|
|
]
|
|
)
|
|
|
|
self.norm = RMSNorm(hidden_size, eps=lm_cfg.rms_norm_eps)
|
|
self._compiled_layers: set[int] = set()
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
inputs_embeds: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
hidden_states = inputs_embeds
|
|
residual = None
|
|
for layer in self.layers:
|
|
hidden_states, residual = layer(
|
|
positions,
|
|
hidden_states,
|
|
residual,
|
|
self.rope_emb,
|
|
)
|
|
hidden_states = self.norm(hidden_states)
|
|
return hidden_states
|
|
|
|
def precompute_fused_qkv(self) -> None:
|
|
"""Materialize fused QKV weights before CUDA Graph capture."""
|
|
for layer in self.layers:
|
|
attn = layer.self_attn
|
|
if attn._fused_qkv_weight is None:
|
|
attn._fused_qkv_weight = torch.cat(
|
|
[attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight],
|
|
dim=0,
|
|
).detach()
|
|
|
|
def compile_selective(self) -> list[str]:
|
|
"""Compile the full residual model forward as one graph (same strategy as base_lm)."""
|
|
if self._compiled_layers:
|
|
return []
|
|
for layer in self.layers:
|
|
layer.self_attn._fused_qkv_weight = None
|
|
self.forward = torch.compile(self.forward, mode="default", fullgraph=False)
|
|
self._compiled_layers.update(range(len(self.layers)))
|
|
return ["forward (whole residual)"]
|
|
|
|
def load_weights_from_native(self, native_residual_lm: nn.Module) -> int:
|
|
"""Load weights from native residual_lm. Returns param count."""
|
|
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
|
loaded = 0
|
|
for name, param in native_residual_lm.named_parameters():
|
|
if "rotary_emb" in name:
|
|
continue
|
|
target = params_dict.get(name)
|
|
if target is None:
|
|
continue
|
|
weight_loader = getattr(target, "weight_loader", default_weight_loader)
|
|
weight_loader(target, param.data)
|
|
loaded += 1
|
|
return loaded
|