sgl-project--sglang/python/sglang/srt/models/midashenglm.py

import collections
import collections.abc
import logging
from collections.abc import Callable, Sequence
from typing import Iterable, List, Optional, Tuple, TypeAlias, cast

import torch
import torch.nn as nn
import torchaudio.functional as F
from transformers import PretrainedConfig

from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.conv import Conv2dLayer
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.managers.mm_utils import (
    MultiModalityDataPaddingPatternMultimodalTokens,
    general_mm_embed_routine,
)
from sglang.srt.managers.schedule_batch import (
    Modality,
    MultimodalDataItem,
    MultimodalInputs,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.qwen2 import Qwen2ForCausalLM
from sglang.srt.utils import add_prefix

logger = logging.getLogger(__name__)
_Tuple2: TypeAlias = int | tuple[int, int] | Sequence[int]


def _resolve_tuple2(x: _Tuple2) -> tuple[int, int]:
    if isinstance(x, collections.abc.Sequence):
        assert (
            len(x) == 2
        ), f"Expected a sequence of length 2, got {x} with length {len(x)}"
        return cast(tuple[int, int], tuple(x))
    return (x, x)


def calculate_mel_frames_dasheng(
    audio_length_samples: int,
    n_fft: int = 512,
    hop_size: int = 160,
    dasheng_subsampling: int = 4,
    center=True,
    model_subsampling: int = 5,
) -> int:
    """Calculate the number of Mel-spectrogram frames."""
    if center:
        audio_length_samples = audio_length_samples + n_fft

    return (
        int(1 + ((audio_length_samples - n_fft) / hop_size))
        // dasheng_subsampling
        // model_subsampling
    )


class AudioPatchEmbed(nn.Module):
    def __init__(
        self,
        input_size: _Tuple2 = 64,
        patch_size: _Tuple2 = 16,
        patch_stride: _Tuple2 = 16,
        in_chans: int = 1,
        embed_dim: int = 768,
        norm_layer: Callable | None = None,
        flatten: bool = False,
    ):
        super().__init__()
        self.input_size = _resolve_tuple2(input_size)
        self.patch_size = _resolve_tuple2(patch_size)
        self.patch_stride = _resolve_tuple2(patch_stride)
        self.grid_size = (
            self.input_size[0] // self.patch_stride[0],
            self.input_size[1] // self.patch_stride[1],
        )
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten
        self.proj = Conv2dLayer(
            in_chans,
            embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_stride,
        )
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj(x)
        if self.flatten:
            x = torch.permute(torch.flatten(x, 2, 3), (0, 2, 1))
        x = self.norm(x)
        return x


class LayerScale(nn.Module):
    def __init__(self, dim, init_values=1e-5, inplace=False):
        super().__init__()
        self.inplace = inplace
        self.gamma = nn.Parameter(init_values * torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.mul_(self.gamma) if self.inplace else x * self.gamma


class DashengMlp(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_features: int | None = None,
        out_features: int | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = ColumnParallelLinear(
            input_size=in_features,
            output_size=hidden_features,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("fc1", prefix),
        )
        self.act = nn.GELU()
        self.fc2 = RowParallelLinear(
            input_size=hidden_features,
            output_size=out_features,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("fc2", prefix),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.fc1(x)
        x = self.act(x)
        x, _ = self.fc2(x)
        return x


class DashengAttention(nn.Module):
    """Audio encoder attention using VisionAttention for compatibility."""

    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        qkv_bias: bool = False,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        assert dim % num_heads == 0, "dim should be divisible by num_heads"
        self.embed_dim = dim
        self.num_heads = num_heads
        self.head_dim = self.embed_dim // self.num_heads
        self.scale = self.head_dim**-0.5

        self.attn = VisionAttention(
            embed_dim=dim,
            num_heads=num_heads,
            projection_size=dim,
            use_qkv_parallel=True,
            proj_bias=True,
            qkv_bias=qkv_bias,
            qkv_backend="sdpa",
            softmax_in_single_precision=False,
            flatten_batch=False,
            quant_config=quant_config,
            prefix=prefix,
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        """
        Args:
            x: [B, N, C] tensor
            mask: [B, N] boolean mask
        """
        attn_mask = None
        if mask is not None:
            attn_mask = mask.unsqueeze(1).unsqueeze(2)  # [B, 1, 1, N]
            attn_mask = attn_mask.float()
            attn_mask = (1.0 - attn_mask) * -10000.0

        x = self.attn(x, attn_mask=attn_mask)
        return x


class DashengBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = False,
        init_values: float | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.attn = DashengAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            quant_config=quant_config,
            prefix=add_prefix("attn", prefix),
        )
        self.ls1 = (
            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        )
        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
        self.mlp = DashengMlp(
            in_features=dim,
            hidden_features=int(dim * mlp_ratio),
            quant_config=quant_config,
            prefix=add_prefix("mlp", prefix),
        )
        self.ls2 = (
            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        )

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x = x + self.ls1(self.attn(self.norm1(x), mask))
        x = x + self.ls2(self.mlp(self.norm2(x)))
        return x


class DashengFrontend(nn.Module):
    """Audio frontend that converts waveforms to log mel-spectrograms."""

    def __init__(self, config: PretrainedConfig):
        super().__init__()
        self.n_fft = config.n_fft
        self.hop_length = config.hop_length
        self.win_length = config.win_length
        self.center = config.center
        spectrogram_window = torch.hann_window(config.win_length)
        self.register_buffer(
            "spectrogram_window",
            spectrogram_window,
            persistent=False,
        )
        self.spectrogram_window: torch.Tensor
        melscale_fbanks = F.melscale_fbanks(
            n_freqs=config.n_fft // 2 + 1,
            f_min=config.f_min,
            f_max=config.f_max,
            n_mels=config.n_mels,
            sample_rate=config.sample_rate,
        )
        self.register_buffer("melscale_fbanks", melscale_fbanks, persistent=False)
        self.melscale_fbanks: torch.Tensor

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Convert waveform to log mel-spectrogram.

        Args:
            waveform: [B, T] tensor of audio samples

        Returns:
            log_mel_spectrogram: [B, n_mels, time] tensor
        """
        spectrogram = F.spectrogram(
            waveform=waveform.to(torch.float32),
            pad=0,
            window=self.spectrogram_window,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            power=2,
            normalized=False,
            center=self.center,
        )
        mel_spectrogram = (spectrogram.mT @ self.melscale_fbanks.to(torch.float32)).mT
        log_mel_spectrogram = F.amplitude_to_DB(
            mel_spectrogram.unsqueeze(1),
            multiplier=10,
            amin=1e-10,
            db_multiplier=0,
            top_db=120,
        ).squeeze(1)
        return log_mel_spectrogram.to(waveform.dtype)


class DashengAudioTransformer(nn.Module):
    """Audio encoder transformer."""

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.target_length = config.target_length
        self.hop_length = config.hop_length
        self.front_end = DashengFrontend(config)
        self.init_bn = nn.BatchNorm2d(config.n_mels, momentum=0.01)
        self.patch_embed = AudioPatchEmbed(
            input_size=(config.n_mels, config.target_length),
            embed_dim=config.embed_dim,
            in_chans=config.input_channels,
            patch_size=config.patch_size,
            flatten=False,
            patch_stride=config.patch_stride,
        )
        self.time_pos_embed = nn.Parameter(
            torch.empty(1, config.embed_dim, 1, self.patch_embed.grid_size[1])
        )
        self.freq_pos_embed = nn.Parameter(
            torch.empty(1, config.embed_dim, self.patch_embed.grid_size[0], 1)
        )
        self.blocks = nn.ModuleList(
            DashengBlock(
                dim=config.embed_dim,
                num_heads=config.num_heads,
                mlp_ratio=config.mlp_ratio,
                qkv_bias=config.qkv_bias,
                init_values=config.init_values,
                quant_config=quant_config,
                prefix=add_prefix(f"blocks.{i}", prefix),
            )
            for i in range(config.depth)
        )
        self.norm = nn.LayerNorm(config.embed_dim, eps=1e-6)

    def forward_features(
        self,
        x: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        t = x.shape[-1]
        x = x + self.time_pos_embed[:, :, :, :t]
        x = x + self.freq_pos_embed[:, :, :, :]
        x = torch.permute(torch.flatten(x, 2, 3), (0, 2, 1))
        for block in self.blocks:
            x = block(x, mask)
        x = self.norm(x)
        return x

    def _to_mask(self, lengths: torch.Tensor, max_length: int) -> torch.Tensor:
        batch_size = len(lengths)
        idx = torch.arange(max_length, device=lengths.device)
        idx = idx.repeat(batch_size).view(batch_size, max_length)
        mask = (idx < lengths.unsqueeze(-1)).bool()
        return mask

    def forward(
        self,
        x: torch.Tensor,
        x_length: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Args:
            x: [B, T] audio waveform tensor
            x_length: [B] tensor of audio lengths

        Returns:
            x: [B, seq_len, embed_dim] encoded features
            mask: [B, seq_len] mask tensor
        """
        x = self.front_end(x)
        x = x.to(self.time_pos_embed.dtype)
        target_length_in_patches = self.target_length // 4
        x = x.unsqueeze(1)
        x = torch.permute(x, (0, 2, 1, 3))
        x = self.init_bn(x)
        x = torch.permute(x, (0, 2, 1, 3))
        x = self.patch_embed(x)
        t = x.shape[-1]
        input_splits = x.split(target_length_in_patches, dim=-1)
        if x_length is not None:
            assert len(x_length) == len(
                x
            ), "batchsizes of input x and x_length need to be same"
            assert x_length.ndim == 1, "Lengths are of size (B,)"
            scaled_lengths = (x_length / (self.hop_length * 4)).long()
            mask = self._to_mask(max_length=t, lengths=scaled_lengths)
            split_masks = mask.split(target_length_in_patches, dim=-1)
        else:
            mask = None
            split_masks = [None] * len(input_splits)
        outputs = []
        for split_x, split_mask in zip(input_splits, split_masks):
            forward_kwargs = {}
            forward_kwargs["mask"] = split_mask
            split_x = self.forward_features(split_x, **forward_kwargs)
            outputs.append(split_x)
        x = torch.cat(outputs, dim=1)
        return x, mask


class AudioProjectorSubsample(nn.Module):
    """Audio projector with subsampling."""

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        downsample_rate=5,
        dtype: torch.dtype | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.k = downsample_rate
        self.fc1 = ColumnParallelLinear(
            input_size=in_dim * self.k,
            output_size=out_dim,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("net.0", prefix),
        )
        self.act = nn.GELU()
        self.fc2 = RowParallelLinear(
            input_size=out_dim,
            output_size=out_dim,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("net.2", prefix),
        )

    def forward(self, x, mask=None):
        batch_size, seq_len, dim = x.shape
        num_frames_to_discard = seq_len % self.k
        if num_frames_to_discard > 0:
            x = x[:, :-num_frames_to_discard, :]
            if mask is not None:
                mask = mask[:, :-num_frames_to_discard]
        if mask is None:
            mask = torch.ones(x.shape[:-1], dtype=torch.long, device=x.device)
        x = x.reshape(batch_size, -1, self.k * dim)
        x, _ = self.fc1(x)
        x = self.act(x)
        x, _ = self.fc2(x)
        mask = mask.reshape(batch_size, -1, self.k)
        mask = mask.any(dim=-1).long()
        return x, mask


class MiDashengLMModel(nn.Module):
    """MiDashengLM model for audio-language processing."""

    default_bitsandbytes_target_modules = [
        ".fc1.",
        ".fc2.",
        ".gate_up_proj.",
        ".down_proj.",
        ".q_proj.",
        ".k_proj.",
        ".v_proj.",
        ".o_proj.",
    ]

    bitsandbytes_stacked_params_mapping = {
        "q_proj": ("qkv_proj", 0),
        "k_proj": ("qkv_proj", 1),
        "v_proj": ("qkv_proj", 2),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        rope_scaling = config.text_config.rope_parameters
        if rope_scaling:
            if "mrope_section" in rope_scaling:
                # Remove mrope_section from rope_parameters so downstream
                # code treats this as standard rotary embedding.
                del rope_scaling["mrope_section"]
        self.audio_encoder = DashengAudioTransformer(
            config.audio_encoder_config,
            quant_config=quant_config,
            prefix=add_prefix("audio_encoder", prefix),
        )
        self.audio_projector = AudioProjectorSubsample(
            in_dim=config.audio_encoder_config.embed_dim,
            out_dim=config.text_config.hidden_size,
            downsample_rate=config.subsample_factor,
            quant_config=quant_config,
            prefix=add_prefix("audio_projector", prefix),
        )
        self.language_model = Qwen2ForCausalLM(
            config.text_config,
            quant_config=quant_config,
            prefix=add_prefix("decoder", prefix),
        )
        self.logits_processor = self.language_model.logits_processor
        self.quant_config = quant_config

    def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
        """Pad input IDs with multimodal tokens."""
        pattern = MultiModalityDataPaddingPatternMultimodalTokens()
        return pattern.pad_input_tokens(input_ids, mm_inputs)

    def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
        """Process audio inputs and return embeddings.

        Args:
            items: List of multimodal data items containing audio features

        Returns:
            audio_embeddings: Concatenated audio embeddings
        """
        logger.debug("=" * 80)
        logger.debug(f"get_audio_feature called with {len(items)} items")
        logger.debug("=" * 80)
        for i, item in enumerate(items):
            logger.debug(f"Item {i} feature shape: {item.feature.shape}")
            logger.debug(
                f"Item {i} audio_length: {getattr(item, 'audio_length', 'NOT SET')}"
            )
            logger.debug(f"Item {i} pad_value: {getattr(item, 'pad_value', 'NOT SET')}")
            logger.debug(f"Item {i} hash: {getattr(item, 'hash', 'NOT SET')}")
        input_values = torch.cat([item.feature for item in items], dim=0)
        logger.debug(f"Concatenated input_values shape: {input_values.shape}")
        audio_lengths = []
        for item in items:
            if hasattr(item, "audio_length") and item.audio_length is not None:
                audio_lengths.append(item.audio_length)
            else:
                audio_lengths.append(item.feature.shape[-1])
        audio_length = torch.tensor(audio_lengths, device=input_values.device)
        logger.debug(f"audio_length: {audio_length}")
        encoder_out, encoder_atts = self.audio_encoder(input_values, audio_length)
        logger.debug(f"Encoder output shape: {encoder_out.shape}")
        audio_embeddings, _ = self.audio_projector(encoder_out, encoder_atts)
        audio_embeddings = audio_embeddings.to(input_values.dtype)
        logger.debug(f"Projector output shape: {audio_embeddings.shape}")
        batch_size, max_audio_tokens, embed_dim = audio_embeddings.shape
        logger.debug(f"Using all {max_audio_tokens} audio tokens from projector output")
        masked_audio_features = audio_embeddings.reshape(-1, embed_dim)
        logger.debug(f"Final output shape: {masked_audio_features.shape}")
        logger.debug(
            f"Stats: min={masked_audio_features.min().item():.4f}, max={masked_audio_features.max().item():.4f}"
        )
        logger.debug(
            f"Audio embeddings dtype: {masked_audio_features.dtype}, device: {masked_audio_features.device}"
        )
        logger.debug(
            f"First 5 values of first audio token: {masked_audio_features[0, :5].tolist()}"
        )
        logger.debug("=" * 80)
        return masked_audio_features

    def get_input_embeddings(self):
        return self.language_model.model.embed_tokens

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        **kwargs,
    ):
        """Run forward pass for MiDashengLM.

        Args:
            input_ids: Flattened (concatenated) input_ids corresponding to a batch.
            positions: Flattened (concatenated) position ids corresponding to a batch.
            forward_batch: Forward batch information including multimodal data.
        """
        if forward_batch.contains_mm_inputs():
            logger.debug("=" * 80)
            logger.debug(f"input_ids shape: {input_ids.shape}")
            logger.debug(f"input_ids first 20: {input_ids[:20].tolist()}")
            logger.debug(
                f"input_ids unique values count: {len(torch.unique(input_ids))}"
            )
            if forward_batch.mm_inputs and len(forward_batch.mm_inputs) > 0:
                mm_input = forward_batch.mm_inputs[0]
                if mm_input and len(mm_input.mm_items) > 0:
                    pad_value = mm_input.mm_items[0].pad_value
                    logger.debug(f"Expected pad_value: {pad_value}")
                    logger.debug(
                        f"Count of pad_value in input_ids: {(input_ids == pad_value).sum().item()}"
                    )
                    if hasattr(mm_input, "audio_token_id") and mm_input.audio_token_id:
                        logger.debug(f"audio_token_id: {mm_input.audio_token_id}")
                        logger.debug(
                            f"Count of audio_token_id in input_ids: {(input_ids == mm_input.audio_token_id).sum().item()}"
                        )
            logger.debug("=" * 80)

        return general_mm_embed_routine(
            input_ids=input_ids,
            forward_batch=forward_batch,
            language_model=self.language_model,
            positions=positions,
            data_embedding_funcs={Modality.AUDIO: self.get_audio_feature},
        )

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        """Load model weights."""
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        buffers_dict = dict(self.named_buffers())
        audio_encoder_loaded = []
        audio_projector_loaded = []
        skipped_weights = []
        decoder_weights = []
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
                continue
            if name.startswith("decoder"):
                decoder_weights.append((name, loaded_weight))
                continue
            original_name = name
            if "audio_encoder.front_end" in name:
                if ".mel_scale.fb" in name:
                    name = name.replace(".mel_scale.fb", ".melscale_fbanks")
                elif ".spectrogram.window" in name:
                    name = name.replace(".spectrogram.window", ".spectrogram_window")
            if "audio_encoder" in name and ".attn.qkv." in name:
                name = name.replace(".attn.qkv.", ".attn.attn.qkv_proj.")
            if "audio_encoder" in name and ".attn.proj." in name:
                name = name.replace(".attn.proj.", ".attn.attn.proj.")
            if "audio_projector" in name:
                name = name.replace(".net.0.", ".fc1.")
                name = name.replace(".net.2.", ".fc2.")
            if (
                name.endswith(".bias")
                and name not in params_dict
                and name not in buffers_dict
            ):
                skipped_weights.append(f"{original_name} (bias not in params/buffers)")
                continue
            if name in params_dict:
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
            elif name in buffers_dict:
                buffers_dict[name].copy_(loaded_weight)
            else:
                if "audio_projector" in original_name:
                    skipped_weights.append(f"{original_name} -> {name} (NOT IN MODEL)")
                else:
                    skipped_weights.append(f"{original_name} (not in model)")
                continue

            if "audio_encoder" in original_name:
                audio_encoder_loaded.append(original_name)
            elif "audio_projector" in original_name:
                audio_projector_loaded.append(original_name)
        if decoder_weights:
            logger.debug(
                f"Passing {len(decoder_weights)} decoder weights to language_model.load_weights()"
            )
            decoder_weights_stripped = [
                (name.replace("decoder.", "", 1), weight)
                for name, weight in decoder_weights
            ]
            self.language_model.load_weights(decoder_weights_stripped)
        logger.debug("=" * 80)
        logger.debug(f"Audio encoder weights loaded: {len(audio_encoder_loaded)}")
        logger.debug(f"Audio projector weights loaded: {len(audio_projector_loaded)}")
        logger.debug(
            f"Decoder weights passed to language_model: {len(decoder_weights)}"
        )
        logger.debug(f"Skipped weights: {len(skipped_weights)}")
        encoder_skipped = [s for s in skipped_weights if "audio_encoder" in s]
        projector_skipped = [s for s in skipped_weights if "audio_projector" in s]
        if projector_skipped:
            logger.debug("Skipped audio_projector weights:")
            for s in projector_skipped:
                logger.debug(f"  {s}")
        if encoder_skipped:
            logger.debug(f"Skipped audio_encoder weights: {len(encoder_skipped)}")
            non_bias_skipped = [s for s in encoder_skipped if "bias" not in s]
            if non_bias_skipped:
                logger.debug("  First 10 non-bias skipped:")
                for s in non_bias_skipped[:10]:
                    logger.debug(f"    {s}")
        logger.debug("=" * 80)

    def get_embed_and_head(self):
        return (
            self.language_model.model.embed_tokens.weight,
            self.language_model.lm_head.weight,
        )


EntryClass = [MiDashengLMModel]