anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Qwen3TTS][Bugfix] Replace vLLM fused layers with HF-compatible numerics in code predictor (#2277)

Browse Source 
Signed-off-by: linyueqian <linyueqian@outlook.com>
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-03-28 10:47:17 +08:00
1 parent a353594e2e
commit f55ea28005
1 file changed
+148 -137
Download Patch File Download Diff File Expand all files Collapse all files
vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py
+148 -137
View File
@@ -8,18 +8,9 @@ import torch.nn.functional as F
from vllm.config import VllmConfig
from vllm.config.vllm import set_current_vllm_config
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
)
from vllm.model_executor.models.utils import is_pp_missing_parameter
from vllm_omni.platforms import current_omni_platform
@@ -29,20 +20,83 @@ logger = init_logger(__name__)
# ===================================================================
# Standalone Code Predictor Layers (no vLLM paged attention)
# HF-numerics-compatible layers for code predictor
# ===================================================================
#
# These replace vLLM's Qwen3DecoderLayer for the code predictor.
# Input is batch-major [B, seq_len, H], attention uses F.scaled_dot_product_attention.
# Weight names match the checkpoint (self_attn.qkv_proj, mlp.gate_up_proj, etc.)
# so load_weights works unchanged.
# These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32,
# rotate_half RoPE) to produce outputs numerically identical to the
# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel,
# get_rope) introduce small precision differences that compound across
# the 15 autoregressive steps of the code predictor, causing severe
# audio quality degradation (UTMOS ~4.26 → ~2.66).
#
# See: https://github.com/vllm-project/vllm-omni/issues/2274
class _RMSNorm(nn.Module):
"""RMSNorm matching HuggingFace's Qwen3TTSRMSNorm exactly.
Computes variance in float32 to avoid bfloat16 precision loss.
"""
def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
"""Rotates half the hidden dims of the input."""
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
class _RotaryEmbedding(nn.Module):
"""RoPE matching HuggingFace's Qwen3TTSRotaryEmbedding exactly.
Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False).
"""
def __init__(self, config: Qwen3TTSTalkerCodePredictorConfig) -> None:
super().__init__()
head_dim = getattr(
config,
"head_dim",
config.hidden_size // config.num_attention_heads,
)
# Standard default RoPE
rope_theta = getattr(config, "rope_theta", 10000.0)
inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
# position_ids: [batch, seq_len]
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
position_ids_expanded = position_ids[:, None, :].float()
# Force float32 (matching HF)
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
with torch.autocast(device_type=device_type, enabled=False):
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
emb = torch.cat((freqs, freqs), dim=-1)
cos = emb.cos()
sin = emb.sin()
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
class _CodePredictorAttention(nn.Module):
"""Standalone multi-head attention for code predictor.
Uses F.scaled_dot_product_attention (SDPA) instead of vLLM's paged Attention.
Supports fused QKV, RoPE, q/k normalization, and native GQA via enable_gqa.
Uses F.scaled_dot_product_attention with HF-compatible RoPE and RMSNorm.
Input: [B, seq_len, hidden_size], output: [B, seq_len, hidden_size].
"""
@@ -61,54 +115,52 @@ class _CodePredictorAttention(nn.Module):
"head_dim",
config.hidden_size // config.num_attention_heads,
)
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = self.head_dim**-0.5
self._use_gqa = self.num_kv_heads != self.num_heads
self.qkv_proj = QKVParallelLinear(
hidden_size=self.hidden_size,
head_size=self.head_dim,
total_num_heads=self.num_heads,
total_num_kv_heads=self.num_kv_heads,
# Separate q/k/v projections matching HF (no fused packing)
self.q_proj = nn.Linear(
self.hidden_size,
self.num_heads * self.head_dim,
bias=getattr(config, "attention_bias", False),
prefix=f"{prefix}.qkv_proj",
disable_tp=True,
)
self.o_proj = RowParallelLinear(
input_size=self.num_heads * self.head_dim,
output_size=self.hidden_size,
self.k_proj = nn.Linear(
self.hidden_size,
self.num_kv_heads * self.head_dim,
bias=getattr(config, "attention_bias", False),
)
self.v_proj = nn.Linear(
self.hidden_size,
self.num_kv_heads * self.head_dim,
bias=getattr(config, "attention_bias", False),
)
self.o_proj = nn.Linear(
self.num_heads * self.head_dim,
self.hidden_size,
bias=False,
prefix=f"{prefix}.o_proj",
disable_tp=True,
)
self.rotary_emb = get_rope(
self.head_dim,
max_position=config.max_position_embeddings,
rope_parameters=getattr(config, "rope_parameters", None),
dual_chunk_attention_config=None,
)
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps)
def forward(
self,
hidden_states: torch.Tensor,
position_ids: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
bsz, seq_len, _ = hidden_states.shape
hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim)
hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim)
qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1))
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2)
k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2)
v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2)
q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
q, k = self.rotary_emb(position_ids, q, k)
q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
cos, sin = position_embeddings
# cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads
cos = cos.unsqueeze(1) # [batch, 1, seq_len, head_dim]
sin = sin.unsqueeze(1)
q = (q * cos) + (_rotate_half(q) * sin)
k = (k * cos) + (_rotate_half(k) * sin)
attn_out = F.scaled_dot_product_attention(
q,
@@ -119,13 +171,13 @@ class _CodePredictorAttention(nn.Module):
enable_gqa=self._use_gqa,
)
attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1)
output, _ = self.o_proj(attn_out)
return output.view(bsz, seq_len, -1)
attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.o_proj(attn_out)
return output
class _CodePredictorMLP(nn.Module):
"""SiLU-gated MLP for code predictor, matching Qwen3MLP structure."""
"""SiLU-gated MLP for code predictor, matching HF's Qwen3TTSTalkerTextMLP."""
def __init__(
self,
@@ -134,27 +186,12 @@ class _CodePredictorMLP(nn.Module):
prefix: str = "",
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
input_size=config.hidden_size,
output_sizes=[config.intermediate_size] * 2,
bias=False,
prefix=f"{prefix}.gate_up_proj",
disable_tp=True,
)
self.down_proj = RowParallelLinear(
input_size=config.intermediate_size,
output_size=config.hidden_size,
bias=False,
prefix=f"{prefix}.down_proj",
disable_tp=True,
)
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate_up, _ = self.gate_up_proj(x)
gate, up = gate_up.chunk(2, dim=-1)
x = F.silu(gate) * up
x, _ = self.down_proj(x)
return x
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class _CodePredictorDecoderLayer(nn.Module):
@@ -169,17 +206,17 @@ class _CodePredictorDecoderLayer(nn.Module):
super().__init__()
self.self_attn = _CodePredictorAttention(config, prefix=f"{prefix}.self_attn")
self.mlp = _CodePredictorMLP(config, prefix=f"{prefix}.mlp")
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.input_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
self,
hidden_states: torch.Tensor,
position_ids: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(hidden_states, position_ids)
hidden_states = self.self_attn(hidden_states, position_embeddings)
hidden_states = residual + hidden_states
residual = hidden_states
@@ -210,7 +247,8 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
self.layers = nn.ModuleList(
[_CodePredictorDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)]
)
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.rotary_emb = _RotaryEmbedding(config)
# Codec embeddings: one per residual group. Stored in talker hidden dim
# (some checkpoints use talker_hidden_size != code_predictor hidden_size).
@@ -228,61 +266,24 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
position_ids: torch.Tensor,
) -> torch.Tensor:
hidden_states = inputs_embeds
position_embeddings = self.rotary_emb(hidden_states, position_ids)
for layer in self.layers:
hidden_states = layer(hidden_states, position_ids)
hidden_states = layer(hidden_states, position_embeddings)
hidden_states = self.norm(hidden_states)
return hidden_states
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
params_dict = dict(self.named_parameters(remove_duplicate=False))
loaded_params: set[str] = set()
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
mapped = name.replace(weight_name, param_name)
if mapped.endswith(".bias") and mapped not in params_dict:
continue
if is_pp_missing_parameter(mapped, self):
continue
if mapped.endswith("scale"):
mapped = maybe_remap_kv_scale_name(mapped, params_dict)
if mapped is None:
continue
param = params_dict.get(mapped)
if param is None:
continue
weight_loader = getattr(param, "weight_loader", default_weight_loader)
if weight_loader == default_weight_loader:
weight_loader(param, loaded_weight)
else:
weight_loader(param, loaded_weight, shard_id)
loaded_params.add(mapped)
break
else:
mapped = maybe_remap_kv_scale_name(name, params_dict)
if mapped is None:
continue
if name.endswith(".bias") and mapped not in params_dict:
continue
if is_pp_missing_parameter(mapped, self):
continue
param = params_dict.get(mapped)
if param is None:
continue
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(mapped)
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_params.add(name)
return loaded_params
@@ -299,17 +300,18 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
~O(T^2) extra attention FLOPs (negligible for T=16, 5 layers) for
zero KV cache management overhead and a simpler execution model.
Optimizations over baseline:
1. torch.compile on model forward -- fuses 60+ small kernel launches per step
into fewer fused kernels (4x speedup on model_fwd, ~75% of total time).
Uses HF-compatible layers (plain nn.Linear, float32 RMSNorm, rotate_half
RoPE) to ensure numerical fidelity with the reference implementation.
Precision matters here because small errors compound across 15 AR steps.
Optimizations preserved:
1. torch.compile on model forward -- fuses small kernel launches.
2. Pre-allocated embedding buffer [B, max_seq, H] -- no torch.cat per step.
3. Projection caching -- each token projected once and cached, avoids O(T^2)
redundant projections.
4. Pre-allocated position_ids [max_seq] -- no torch.arange per step.
3. Projection caching -- each token projected once and cached.
4. Pre-allocated position_ids -- no torch.arange per step.
5. Inline sampling -- no custom op / forward_context overhead.
6. No context managers in forward().
7. Cached module references -- bypass nn.Module.__call__ and ModuleList indexing.
8. Pre-allocated output tensor.
6. Cached module references -- bypass nn.Module.__call__ overhead.
7. CUDA graphs per batch-size bucket.
"""
def __init__(
@@ -409,10 +411,20 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
self._compiled_model_fwd = self.model.forward
return
self._compiled_model_fwd = torch.compile(self.model.forward, mode="default", dynamic=False)
# torch.compile fuses RMSNorm/RoPE in ways that lose float32
# precision, compounding across 15 AR steps. Use torch.compile
# with options that disable the problematic fusions while still
# getting kernel fusion benefits for the linear layers and SDPA.
self._compiled_model_fwd = torch.compile(
self.model.forward,
dynamic=False,
options={
"epilogue_fusion": False,
},
)
self._warmup_buckets()
self._capture_cuda_graphs()
logger.info("code_predictor: torch.compile (mode=default) + CUDA graphs")
logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs")
def _padded_bsz(self, bsz: int) -> int:
for bucket in self._bucket_sizes:
@@ -430,11 +442,11 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
max_seq = self._num_groups + 1
device = next(self.model.parameters()).device
base_pos = torch.arange(max_seq, device=device, dtype=torch.long)
proj_buf = self._proj_buf
for bsz in self._bucket_sizes:
pos_ids = base_pos if bsz == 1 else base_pos.repeat(bsz)
# position_ids: [batch, seq_len] for HF-style RoPE
pos_ids = torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(bsz, -1)
self._bucket_pos_ids[bsz] = pos_ids
for _ in range(3):
self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids)
@@ -517,8 +529,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
full_pos_ids = self._bucket_pos_ids.get(padded_bsz)
if full_pos_ids is None:
base_pos = torch.arange(max_seq, device=device, dtype=torch.long)
full_pos_ids = base_pos if padded_bsz == 1 else base_pos.repeat(padded_bsz)
full_pos_ids = torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1)
# Use captured CUDA graph if available, otherwise call compiled fn.
cuda_graph_entry = self._cuda_graphs.get(padded_bsz)