746 lines
30 KiB
Python
746 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
# ruff: noqa: E501
|
|
"""Dots-path speech encoder for inference only (single GPU).
|
|
|
|
Ported from cybertron_alm ``dots_audio_encoder/modeling_whisper.py``.
|
|
Upstream ``WhisperEncoder`` is exposed as :class:`DotsSpeechEncoder`.
|
|
"""
|
|
|
|
import math
|
|
from typing import Any, ClassVar
|
|
|
|
import torch
|
|
from torch import nn
|
|
from transformers.activations import ACT2FN
|
|
from transformers.modeling_outputs import BaseModelOutput
|
|
from transformers.modeling_utils import PreTrainedModel
|
|
from transformers.models.whisper.configuration_whisper import WhisperConfig
|
|
from transformers.utils import logging
|
|
|
|
from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported
|
|
|
|
_FLASH_ATTN_VERSION = 3 if is_fa_version_supported(3) else 2
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
__all__ = ["DotsSpeechEncoder"]
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim: int, eps: float = 1e-6):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
var = x.pow(2).mean(-1, keepdim=True)
|
|
x = x * torch.rsqrt(var + self.eps)
|
|
return self.weight * x
|
|
|
|
|
|
def swiglu(x: torch.Tensor) -> torch.Tensor:
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.nn.functional.silu(x1) * x2
|
|
|
|
|
|
class RotaryEmbedding(nn.Module):
|
|
def __init__(
|
|
self,
|
|
head_dim: int,
|
|
rope_parameters: dict,
|
|
base_seq_len: int = 0,
|
|
):
|
|
super().__init__()
|
|
self.partial_rotary_factor = float(
|
|
rope_parameters.get("partial_rotary_factor", 1.0)
|
|
)
|
|
self.rope_theta = float(rope_parameters.get("rope_theta", 10000.0))
|
|
self.rope_type = rope_parameters.get("rope_type", "default")
|
|
rotary_dim = int(head_dim * self.partial_rotary_factor)
|
|
# keep even rotary dimension to align cos/sin pairs
|
|
self.rotary_dim = (rotary_dim // 2) * 2
|
|
self.attention_scaling = 1.0
|
|
if self.rope_type != "default":
|
|
logger.warning(
|
|
"RoPE type %s is not implemented in WhisperLv3; using default",
|
|
self.rope_type,
|
|
)
|
|
if self.rotary_dim > 0:
|
|
inv_freq = 1.0 / (
|
|
self.rope_theta
|
|
** (
|
|
torch.arange(0, self.rotary_dim, 2, dtype=torch.float)
|
|
/ max(self.rotary_dim, 1)
|
|
)
|
|
)
|
|
else:
|
|
inv_freq = torch.tensor([])
|
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
self.base_seq_len = max(0, int(base_seq_len))
|
|
self._cache: (
|
|
tuple[int, torch.dtype, torch.device, torch.Tensor, torch.Tensor] | None
|
|
) = None
|
|
|
|
@torch.no_grad()
|
|
def get_cos_sin(
|
|
self, position_ids: torch.Tensor, dtype: torch.dtype, device: torch.device
|
|
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
|
if self.rotary_dim == 0:
|
|
return None, None
|
|
seq_len = position_ids.shape[-1]
|
|
if position_ids.shape[0] == 1 and self._cache is not None:
|
|
cached_seq_len, cached_dtype, cached_device, cached_cos, cached_sin = (
|
|
self._cache
|
|
)
|
|
if (
|
|
cached_seq_len >= seq_len
|
|
and cached_dtype == dtype
|
|
and cached_device == device
|
|
):
|
|
return cached_cos[:, :seq_len, :], cached_sin[:, :seq_len, :]
|
|
|
|
if position_ids.shape[0] == 1:
|
|
cache_seq_len = max(seq_len, self.base_seq_len)
|
|
position_ids = torch.arange(cache_seq_len, device=device)[None, :]
|
|
# [1, D/2, 1]
|
|
if self.inv_freq.dtype != torch.float32:
|
|
inv_freq = 1.0 / (
|
|
self.rope_theta
|
|
** (
|
|
torch.arange(
|
|
0, self.rotary_dim, 2, dtype=torch.float, device=device
|
|
)
|
|
/ max(self.rotary_dim, 1)
|
|
)
|
|
)
|
|
else:
|
|
inv_freq = self.inv_freq.to(device=device)
|
|
inv_freq_expanded = inv_freq[None, :, None]
|
|
# [B, 1, T]
|
|
position_ids_expanded = position_ids[:, None, :].float()
|
|
# Force float32 since bfloat16 loses precision on long contexts
|
|
# See https://github.com/huggingface/transformers/pull/29285
|
|
device_type = device.type if isinstance(device, torch.device) 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() * self.attention_scaling
|
|
sin = emb.sin() * self.attention_scaling
|
|
cos = cos.to(dtype)
|
|
sin = sin.to(dtype)
|
|
if position_ids.shape[0] == 1:
|
|
self._cache = (position_ids.shape[-1], dtype, device, cos, sin)
|
|
return cos[:, :seq_len, :], sin[:, :seq_len, :]
|
|
return cos, sin
|
|
|
|
|
|
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 | None,
|
|
sin: torch.Tensor | None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if cos is None or sin is None:
|
|
return q, k
|
|
|
|
if q.dim() == 4:
|
|
# [B, T, H, D]
|
|
rotary_dim = cos.shape[-1]
|
|
cos = cos.unsqueeze(2) # [B, T, 1, rotary_dim]
|
|
sin = sin.unsqueeze(2)
|
|
q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
|
|
k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
|
|
q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
|
|
k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
|
|
return torch.cat([q_embed, q_pass], dim=-1), torch.cat(
|
|
[k_embed, k_pass], dim=-1
|
|
)
|
|
|
|
if q.dim() == 3:
|
|
# [tokens, H, D] (varlen flattened)
|
|
rotary_dim = cos.shape[-1]
|
|
cos = cos.unsqueeze(1) # [tokens, 1, rotary_dim]
|
|
sin = sin.unsqueeze(1)
|
|
q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
|
|
k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
|
|
q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
|
|
k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
|
|
return torch.cat([q_embed, q_pass], dim=-1), torch.cat(
|
|
[k_embed, k_pass], dim=-1
|
|
)
|
|
|
|
return q, k
|
|
|
|
|
|
class WhisperPositionalEmbedding(nn.Embedding):
|
|
def __init__(
|
|
self, num_positions: int, embedding_dim: int, padding_idx: int | None = None
|
|
):
|
|
super().__init__(num_positions, embedding_dim)
|
|
|
|
def forward(self, input_ids, past_key_values_length=0):
|
|
return self.weight[
|
|
past_key_values_length : past_key_values_length + input_ids.shape[1]
|
|
]
|
|
|
|
|
|
class WhisperAttention(nn.Module):
|
|
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
|
|
|
def __init__(
|
|
self,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
dropout: float = 0.0,
|
|
is_decoder: bool = False,
|
|
bias: bool = True,
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.num_heads = num_heads
|
|
self.dropout = dropout
|
|
self.head_dim = embed_dim // num_heads
|
|
|
|
if (self.head_dim * num_heads) != self.embed_dim:
|
|
raise ValueError(
|
|
f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
|
|
f" and `num_heads`: {num_heads})."
|
|
)
|
|
self.scaling = self.head_dim**-0.5
|
|
self.is_decoder = is_decoder
|
|
|
|
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
|
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
|
|
def forward_flash_attn(
|
|
self,
|
|
hidden_states,
|
|
cu_seqlens_q=None,
|
|
cu_seqlens_kv=None,
|
|
max_seqlen_q=None,
|
|
max_seqlen_kv=None,
|
|
output_attentions=False,
|
|
rotary_cos=None,
|
|
rotary_sin=None,
|
|
):
|
|
"""Dense eager attention with SGLang FA3 for packed variable-length input."""
|
|
query_states = self.q_proj(hidden_states)
|
|
key_states = self.k_proj(hidden_states)
|
|
value_states = self.v_proj(hidden_states)
|
|
if cu_seqlens_q is None:
|
|
bsz, tgt_len, _ = hidden_states.size()
|
|
query_states = query_states.view(
|
|
bsz, tgt_len, self.num_heads, self.head_dim
|
|
)
|
|
key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
|
|
value_states = value_states.view(
|
|
bsz, tgt_len, self.num_heads, self.head_dim
|
|
)
|
|
if rotary_cos is not None and rotary_sin is not None:
|
|
cos = rotary_cos[:, :tgt_len, :]
|
|
sin = rotary_sin[:, :tgt_len, :]
|
|
query_states, key_states = apply_rotary_pos_emb(
|
|
query_states, key_states, cos, sin
|
|
)
|
|
attn_output, attn_probs = self._eager_attention(
|
|
query_states, key_states, value_states, output_attentions
|
|
)
|
|
attn_output = attn_output.view(bsz, tgt_len, self.embed_dim)
|
|
else:
|
|
query_states = query_states.view(-1, self.num_heads, self.head_dim)
|
|
key_states = key_states.view(-1, self.num_heads, self.head_dim)
|
|
value_states = value_states.view(-1, self.num_heads, self.head_dim)
|
|
if rotary_cos is not None and rotary_sin is not None:
|
|
query_states, key_states = apply_rotary_pos_emb(
|
|
query_states, key_states, rotary_cos, rotary_sin
|
|
)
|
|
result = flash_attn_varlen_func(
|
|
query_states,
|
|
key_states,
|
|
value_states,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_k=cu_seqlens_kv,
|
|
max_seqlen_q=max_seqlen_q,
|
|
max_seqlen_k=max_seqlen_kv,
|
|
softmax_scale=self.scaling,
|
|
causal=self.is_decoder,
|
|
return_softmax_lse=output_attentions,
|
|
fa_version=_FLASH_ATTN_VERSION,
|
|
)
|
|
if isinstance(result, tuple):
|
|
attn_output = result[0]
|
|
attn_probs = result[1] if output_attentions else None
|
|
else:
|
|
attn_output = result
|
|
attn_probs = None
|
|
attn_output = attn_output.view(-1, self.embed_dim)
|
|
|
|
return self.out_proj(attn_output), attn_probs
|
|
|
|
def _eager_attention(self, query, key, value, output_attentions):
|
|
query = query.transpose(1, 2)
|
|
key = key.transpose(1, 2)
|
|
value = value.transpose(1, 2)
|
|
scores = torch.matmul(query, key.transpose(-2, -1)) * self.scaling
|
|
if self.is_decoder:
|
|
seq_len = scores.shape[-1]
|
|
causal_mask = torch.ones(
|
|
seq_len, seq_len, dtype=torch.bool, device=scores.device
|
|
).triu(1)
|
|
scores.masked_fill_(causal_mask, torch.finfo(scores.dtype).min)
|
|
probabilities = torch.softmax(scores, dim=-1, dtype=torch.float32).to(
|
|
query.dtype
|
|
)
|
|
output = torch.matmul(probabilities, value).transpose(1, 2)
|
|
return output, probabilities if output_attentions else None
|
|
|
|
|
|
# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Whisper
|
|
class WhisperEncoderLayer(nn.Module):
|
|
def __init__(self, config: WhisperConfig):
|
|
super().__init__()
|
|
self.embed_dim = config.d_model
|
|
use_causal = getattr(config, "use_causal", False)
|
|
self.self_attn = WhisperAttention(
|
|
embed_dim=self.embed_dim,
|
|
num_heads=config.encoder_attention_heads,
|
|
dropout=config.attention_dropout,
|
|
is_decoder=use_causal, # enable causal attention when use_causal=True
|
|
)
|
|
self.use_causal = use_causal
|
|
norm_cls = RMSNorm if getattr(config, "use_rms_norm", False) else nn.LayerNorm
|
|
self.self_attn_layer_norm = norm_cls(self.embed_dim)
|
|
self.dropout = config.dropout
|
|
self.activation_fn = (
|
|
ACT2FN[config.activation_function]
|
|
if config.activation_function != "swiglu"
|
|
else swiglu
|
|
)
|
|
self.use_swiglu = config.activation_function == "swiglu"
|
|
self.activation_dropout = config.activation_dropout
|
|
ffn_dim = config.encoder_ffn_dim
|
|
fc1_out = ffn_dim * 2 if self.use_swiglu else ffn_dim
|
|
self.fc1 = nn.Linear(self.embed_dim, fc1_out)
|
|
self.fc2 = nn.Linear(ffn_dim, self.embed_dim)
|
|
self.final_layer_norm = norm_cls(self.embed_dim)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
cu_seqlens_q: torch.Tensor = None,
|
|
cu_seqlens_kv: torch.Tensor = None,
|
|
max_seqlen_q: int | None = None,
|
|
max_seqlen_kv: int | None = None,
|
|
output_attentions: bool = False,
|
|
rotary_cos: torch.Tensor | None = None,
|
|
rotary_sin: torch.Tensor | None = None,
|
|
) -> tuple[Any, ...]:
|
|
"""
|
|
Args:
|
|
hidden_states: Input to the layer of shape `(batch, seq_len,
|
|
embed_dim)`, or `(total_tokens, embed_dim)` when `cu_seqlens_q`
|
|
is given.
|
|
cu_seqlens_q: Cumulative query sequence lengths for packed
|
|
variable-length input. If `None`, the input is treated as a
|
|
padded batch.
|
|
cu_seqlens_kv: Cumulative key/value sequence lengths for packed
|
|
variable-length input.
|
|
max_seqlen_q: Longest query sequence in the packed batch.
|
|
max_seqlen_kv: Longest key/value sequence in the packed batch.
|
|
output_attentions: Whether to also return the attention weights.
|
|
rotary_cos: Cosine component of the rotary position embedding.
|
|
rotary_sin: Sine component of the rotary position embedding.
|
|
|
|
Returns:
|
|
A tuple of the output hidden states, followed by the attention
|
|
weights if `output_attentions` is `True`.
|
|
"""
|
|
residual = hidden_states
|
|
hidden_states = self.self_attn_layer_norm(hidden_states)
|
|
|
|
hidden_states, attn_weights = self.self_attn.forward_flash_attn(
|
|
hidden_states=hidden_states,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_kv=cu_seqlens_kv,
|
|
max_seqlen_q=max_seqlen_q,
|
|
max_seqlen_kv=max_seqlen_kv,
|
|
output_attentions=output_attentions,
|
|
rotary_cos=rotary_cos,
|
|
rotary_sin=rotary_sin,
|
|
)
|
|
hidden_states = nn.functional.dropout(
|
|
hidden_states, p=self.dropout, training=self.training
|
|
)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.final_layer_norm(hidden_states)
|
|
if self.use_swiglu:
|
|
hidden_states = swiglu(self.fc1(hidden_states))
|
|
else:
|
|
hidden_states = self.activation_fn(self.fc1(hidden_states))
|
|
hidden_states = nn.functional.dropout(
|
|
hidden_states, p=self.activation_dropout, training=self.training
|
|
)
|
|
hidden_states = self.fc2(hidden_states)
|
|
hidden_states = nn.functional.dropout(
|
|
hidden_states, p=self.dropout, training=self.training
|
|
)
|
|
hidden_states = residual + hidden_states
|
|
|
|
if hidden_states.dtype == torch.float16 and (
|
|
torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
|
|
):
|
|
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
|
hidden_states = torch.clamp(
|
|
hidden_states, min=-clamp_value, max=clamp_value
|
|
)
|
|
|
|
outputs: tuple[Any, ...] = (hidden_states,)
|
|
|
|
if output_attentions:
|
|
outputs += (attn_weights,)
|
|
|
|
return outputs
|
|
|
|
|
|
class DotsSpeechPreTrainedModel(PreTrainedModel):
|
|
config_class = WhisperConfig
|
|
base_model_prefix = "model"
|
|
main_input_name = "input_features"
|
|
supports_gradient_checkpointing = False
|
|
_no_split_modules: ClassVar[list[str]] = ["WhisperEncoderLayer"]
|
|
|
|
def _init_weights(self, module):
|
|
std = self.config.init_std
|
|
if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d)):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.bias is not None:
|
|
module.bias.data.zero_()
|
|
elif isinstance(module, nn.Embedding):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.padding_idx is not None:
|
|
module.weight.data[module.padding_idx].zero_()
|
|
|
|
|
|
class DotsSpeechEncoder(DotsSpeechPreTrainedModel):
|
|
"""
|
|
Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
|
|
[`WhisperEncoderLayer`].
|
|
|
|
Args:
|
|
config: WhisperConfig
|
|
"""
|
|
|
|
def __init__(self, config: WhisperConfig):
|
|
super().__init__(config)
|
|
self.dropout = config.dropout
|
|
|
|
embed_dim = config.d_model
|
|
self.use_latent_input = getattr(config, "use_latent_input", False)
|
|
self.use_causal = getattr(config, "use_causal", False)
|
|
self.use_conv2d_stem = getattr(config, "use_conv2d_stem", False)
|
|
latent_dim = getattr(config, "latent_dim", None)
|
|
if self.use_latent_input and latent_dim is None:
|
|
raise ValueError(
|
|
"DotsSpeechEncoder: use_latent_input=True requires config.latent_dim"
|
|
)
|
|
if self.use_conv2d_stem and self.use_latent_input:
|
|
raise ValueError(
|
|
"DotsSpeechEncoder: use_conv2d_stem and use_latent_input are mutually exclusive"
|
|
)
|
|
self.num_mel_bins = (
|
|
latent_dim if self.use_latent_input and latent_dim is not None else 128
|
|
)
|
|
self.padding_idx = config.pad_token_id
|
|
self.max_source_positions = config.max_source_positions
|
|
self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
|
|
|
|
if self.use_conv2d_stem:
|
|
# Conv2D stem: 3 layers of stride-2 for 8x downsampling
|
|
dhs = getattr(config, "downsample_hidden_size", 480)
|
|
# Causal: keep freq padding, remove time padding (handled by pad(14,0) in forward)
|
|
conv_padding = (1, 0) if self.use_causal else 1
|
|
self.conv2d1 = nn.Conv2d(
|
|
1, dhs, kernel_size=3, stride=2, padding=conv_padding
|
|
)
|
|
self.conv2d2 = nn.Conv2d(
|
|
dhs, dhs, kernel_size=3, stride=2, padding=conv_padding
|
|
)
|
|
self.conv2d3 = nn.Conv2d(
|
|
dhs, dhs, kernel_size=3, stride=2, padding=conv_padding
|
|
)
|
|
# After 3x stride-2 on freq=128: 128→64→32→16; linear projects dhs*16 → embed_dim
|
|
freq_after = self.num_mel_bins
|
|
for _ in range(3):
|
|
freq_after = (freq_after + 1) // 2
|
|
self.conv_out = nn.Linear(dhs * freq_after, embed_dim, bias=False)
|
|
self.conv1 = None
|
|
self.conv2 = None
|
|
elif self.use_latent_input:
|
|
# Latent path keeps length (stride=1) and applies GLU after conv1
|
|
self.conv1 = nn.Conv1d(
|
|
self.num_mel_bins, embed_dim * 2, kernel_size=3, stride=1, padding=1
|
|
)
|
|
self.conv2 = nn.Conv1d(
|
|
embed_dim, embed_dim, kernel_size=3, stride=1, padding=1
|
|
)
|
|
else:
|
|
self.conv1 = nn.Conv1d(
|
|
self.num_mel_bins, embed_dim, kernel_size=3, padding=1
|
|
)
|
|
self.conv2 = nn.Conv1d(
|
|
embed_dim, embed_dim, kernel_size=3, stride=2, padding=1
|
|
)
|
|
|
|
self.use_rope = getattr(config, "use_rope", False)
|
|
rope_parameters = getattr(config, "rope_parameters", {}) or {}
|
|
self.rope_parameters = rope_parameters
|
|
if self.use_rope:
|
|
head_dim = embed_dim // config.encoder_attention_heads
|
|
self.rotary_embedding = RotaryEmbedding(
|
|
head_dim,
|
|
rope_parameters,
|
|
base_seq_len=self.max_source_positions,
|
|
)
|
|
self.embed_positions = None
|
|
else:
|
|
self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
|
|
|
|
self.layers = nn.ModuleList(
|
|
[WhisperEncoderLayer(config) for _ in range(config.encoder_layers)]
|
|
)
|
|
norm_cls = RMSNorm if getattr(config, "use_rms_norm", False) else nn.LayerNorm
|
|
self.layer_norm = norm_cls(config.d_model)
|
|
|
|
self.post_init()
|
|
|
|
def get_input_embeddings(self) -> nn.Module:
|
|
if self.use_conv2d_stem:
|
|
return self.conv2d1
|
|
return self.conv1
|
|
|
|
def set_input_embeddings(self, value: nn.Module):
|
|
self.conv1 = value
|
|
|
|
@staticmethod
|
|
def _temporal_mask(feat, valid_lens):
|
|
"""Zero out temporal positions >= valid_lens. feat: [B,C,F,T], valid_lens: [B] tensor."""
|
|
T = feat.shape[-1]
|
|
mask = torch.arange(T, device=feat.device)[None, :] < valid_lens[:, None]
|
|
return mask[:, None, None, :]
|
|
|
|
def _conv2d_stem_one_chunk(self, chunk, chunk_valid_mel_lens=None):
|
|
"""Run 3x Conv2d(stride=2) + GELU with per-layer masking."""
|
|
# Causal: left-pad time by 14 = 2 + 4 + 8 (receptive field of the 3-layer
|
|
# stride-2 stack mapped back to input time). Combined with conv padding=(1,0),
|
|
# this makes the whole conv2d stem strictly causal.
|
|
if self.use_causal:
|
|
chunk = nn.functional.pad(chunk, (14, 0))
|
|
if chunk_valid_mel_lens is not None:
|
|
chunk_valid_mel_lens = chunk_valid_mel_lens + 14
|
|
# Step 0: mask mel — silence_mel(-1.5) → 0
|
|
if chunk_valid_mel_lens is not None:
|
|
chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
|
|
chunk = nn.functional.gelu(self.conv2d1(chunk))
|
|
# Step 1: mask conv1 output — gelu(conv(0)+bias) → 0
|
|
if chunk_valid_mel_lens is not None:
|
|
chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
|
|
chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
|
|
chunk = nn.functional.gelu(self.conv2d2(chunk))
|
|
# Step 2: mask conv2 output — gelu(bias) → 0
|
|
if chunk_valid_mel_lens is not None:
|
|
chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
|
|
chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
|
|
chunk = nn.functional.gelu(self.conv2d3(chunk))
|
|
# Step 3: mask conv3 output — gelu(bias) → 0
|
|
if chunk_valid_mel_lens is not None:
|
|
chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
|
|
chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
|
|
return chunk
|
|
|
|
def _forward_conv2d_stem(
|
|
self, input_features, input_seq_lens=None, audio_sample_lens=None
|
|
):
|
|
"""Conv2D stem: 3x Conv2d(stride=2) → 8x downsample. Returns [B, T/8, embed_dim]."""
|
|
x = input_features.unsqueeze(1) # [B, 1, 128, T]
|
|
|
|
if audio_sample_lens is not None:
|
|
hop_length = 160
|
|
if not isinstance(audio_sample_lens, torch.Tensor):
|
|
audio_sample_lens = torch.tensor(audio_sample_lens, device=x.device)
|
|
valid_mel_lens = audio_sample_lens.to(x.device) // hop_length
|
|
else:
|
|
valid_mel_lens = None
|
|
x = self._conv2d_stem_one_chunk(x, valid_mel_lens)
|
|
# [B, dhs, F_out, T/8] → [B, T/8, dhs*freq_after] → [B, T/8, embed_dim]
|
|
batch, channels, frequency, time = x.shape
|
|
x = x.permute(0, 3, 1, 2).reshape(batch, time, channels * frequency)
|
|
inputs_embeds = self.conv_out(x)
|
|
|
|
return inputs_embeds
|
|
|
|
def _forward_conv1d_stem(self, input_features):
|
|
"""Standard Conv1d stem: conv1 + conv2(stride=2), 2x downsample. Returns [B, T/2, embed_dim]."""
|
|
assert self.conv1 is not None and self.conv2 is not None
|
|
if self.use_causal:
|
|
x = nn.functional.pad(input_features, (2, 0))
|
|
x = nn.functional.gelu(self.conv1(x))
|
|
x = nn.functional.pad(x, (2, 0))
|
|
x = nn.functional.gelu(self.conv2(x))
|
|
else:
|
|
x = nn.functional.gelu(self.conv1(input_features))
|
|
x = nn.functional.gelu(self.conv2(x))
|
|
return x.permute(0, 2, 1)
|
|
|
|
def _forward_latent_stem(self, input_features):
|
|
"""Latent input stem: conv1+GLU + conv2, stride=1, no downsample. Returns [B, T, embed_dim]."""
|
|
assert self.conv1 is not None and self.conv2 is not None
|
|
if self.use_causal:
|
|
x = nn.functional.pad(input_features, (2, 0))
|
|
x = nn.functional.glu(self.conv1(x), dim=1)
|
|
x = nn.functional.pad(x, (2, 0))
|
|
x = nn.functional.gelu(self.conv2(x))
|
|
else:
|
|
x = nn.functional.glu(self.conv1(input_features), dim=1)
|
|
x = nn.functional.gelu(self.conv2(x))
|
|
return x.permute(0, 2, 1)
|
|
|
|
def forward(
|
|
self,
|
|
input_features,
|
|
input_seq_lens=None,
|
|
audio_sample_lens=None,
|
|
output_attentions=None,
|
|
output_hidden_states=None,
|
|
return_dict=None,
|
|
):
|
|
"""Mel `[B, n_mels, T]` → encoder hidden states. Inference-only."""
|
|
output_attentions = (
|
|
output_attentions
|
|
if output_attentions is not None
|
|
else self.config.output_attentions
|
|
)
|
|
output_hidden_states = (
|
|
output_hidden_states
|
|
if output_hidden_states is not None
|
|
else self.config.output_hidden_states
|
|
)
|
|
return_dict = (
|
|
return_dict if return_dict is not None else self.config.use_return_dict
|
|
)
|
|
if self.use_conv2d_stem:
|
|
inputs_embeds = self._forward_conv2d_stem(
|
|
input_features, input_seq_lens, audio_sample_lens
|
|
)
|
|
elif self.use_latent_input:
|
|
inputs_embeds = self._forward_latent_stem(input_features)
|
|
else:
|
|
inputs_embeds = self._forward_conv1d_stem(input_features)
|
|
rotary_cos = None
|
|
rotary_sin = None
|
|
position_ids = torch.arange(
|
|
inputs_embeds.shape[1], device=inputs_embeds.device
|
|
)[None, :]
|
|
if self.use_rope:
|
|
rotary_cos, rotary_sin = self.rotary_embedding.get_cos_sin(
|
|
position_ids, inputs_embeds.dtype, inputs_embeds.device
|
|
)
|
|
hidden_states = inputs_embeds
|
|
else:
|
|
assert self.embed_positions is not None
|
|
embed_pos = self.embed_positions.weight[: inputs_embeds.shape[1]]
|
|
hidden_states = inputs_embeds + embed_pos
|
|
|
|
encoder_states: tuple[torch.Tensor, ...] | None = (
|
|
() if output_hidden_states else None
|
|
)
|
|
all_attentions: tuple[torch.Tensor, ...] | None = (
|
|
() if output_attentions else None
|
|
)
|
|
|
|
if input_seq_lens is not None:
|
|
# Build varlen metadata and pack valid tokens without per-sample cat loops.
|
|
# Read max on whatever device the caller provided: when it is a CPU
|
|
# tensor this avoids a device->host sync (one per forward).
|
|
max_seqlen_q = int(input_seq_lens.max().item())
|
|
input_seq_lens = input_seq_lens.to(
|
|
device=hidden_states.device, dtype=torch.long
|
|
)
|
|
B, S, D = hidden_states.shape
|
|
max_seqlen_kv = max_seqlen_q
|
|
cu_seqlens_q = torch.nn.functional.pad(
|
|
input_seq_lens.cumsum(0, dtype=torch.int32), (1, 0)
|
|
)
|
|
cu_seqlens_kv = cu_seqlens_q
|
|
|
|
token_positions = torch.arange(S, device=hidden_states.device)[None, :]
|
|
valid_token_mask = token_positions < input_seq_lens[:, None]
|
|
hidden_states = hidden_states[valid_token_mask]
|
|
if rotary_cos is not None and rotary_sin is not None:
|
|
packed_positions = token_positions.expand(B, S)[valid_token_mask]
|
|
rotary_cos = rotary_cos.squeeze(0).index_select(0, packed_positions)
|
|
rotary_sin = rotary_sin.squeeze(0).index_select(0, packed_positions)
|
|
else:
|
|
cu_seqlens_q = None
|
|
cu_seqlens_kv = None
|
|
max_seqlen_q = None
|
|
max_seqlen_kv = None
|
|
|
|
for encoder_layer in self.layers:
|
|
if output_hidden_states:
|
|
assert encoder_states is not None
|
|
encoder_states = encoder_states + (hidden_states,)
|
|
layer_outputs = encoder_layer(
|
|
hidden_states,
|
|
cu_seqlens_q,
|
|
cu_seqlens_kv,
|
|
max_seqlen_q,
|
|
max_seqlen_kv,
|
|
output_attentions=output_attentions,
|
|
rotary_cos=rotary_cos,
|
|
rotary_sin=rotary_sin,
|
|
)
|
|
hidden_states = layer_outputs[0]
|
|
if output_attentions:
|
|
assert all_attentions is not None
|
|
all_attentions = all_attentions + (layer_outputs[1],)
|
|
|
|
if input_seq_lens is not None:
|
|
# Recover packed varlen output directly on device.
|
|
recover_positions = torch.arange(max_seqlen_q, device=hidden_states.device)[
|
|
None, :
|
|
]
|
|
recover_mask = recover_positions < input_seq_lens[:, None]
|
|
recovered = hidden_states.new_zeros(B, max_seqlen_q, D)
|
|
recovered[recover_mask] = hidden_states
|
|
hidden_states = recovered
|
|
|
|
hidden_states = self.layer_norm(hidden_states)
|
|
if output_hidden_states:
|
|
assert encoder_states is not None
|
|
encoder_states = encoder_states + (hidden_states,)
|
|
|
|
if not return_dict:
|
|
return tuple(
|
|
v
|
|
for v in [hidden_states, encoder_states, all_attentions]
|
|
if v is not None
|
|
)
|
|
return BaseModelOutput(
|
|
last_hidden_state=hidden_states,
|
|
hidden_states=encoder_states,
|
|
attentions=all_attentions,
|
|
)
|