fredchu/MOSS-Audio-8B-Instruct-MLX

1"""MLX-native MossAudioEncoder.2 3Direct port of src/modeling_moss_audio.py:36-155 (MossAudioEncoder).4Adapted from ml-explore/mlx-examples/whisper/mlx_whisper/whisper.py with:5  - 3× Conv2d stride-2 stem (instead of Whisper's 2× Conv1d)6  - Pre-existing HF Whisper attribute names (q_proj/k_proj/v_proj/out_proj, fc1/fc2,7    self_attn_layer_norm/final_layer_norm) so weight remap is near-identity8  - DeepStack taps: capture hidden state AFTER layers in deepstack_layer_indexes9  - feature_lens-based padding mask10"""11from __future__ import annotations12 13import math14from dataclasses import dataclass, field15from typing import List, Optional, Tuple16 17import mlx.core as mx18import mlx.nn as nn19 20 21# ---- helpers ----------------------------------------------------------22 23 24def sinusoids(length: int, channels: int, max_timescale: float = 10000.0) -> mx.array:25    """Whisper-style sinusoidal position embeddings. Matches mlx-examples whisper."""26    assert channels % 2 == 027    log_timescale_increment = math.log(max_timescale) / (channels // 2 - 1)28    inv_timescales = mx.exp(-log_timescale_increment * mx.arange(channels // 2))29    scaled_time = mx.arange(length)[:, None] * inv_timescales[None, :]30    return mx.concatenate([mx.sin(scaled_time), mx.cos(scaled_time)], axis=1)31 32 33# ---- attention ------------------------------------------------------34 35 36class WhisperAttention(nn.Module):37    """HF-Whisper-style self-attention. Layer-scaling convention (`1/sqrt(head_dim)`38    applied to Q, not split between Q and K like mlx-examples does).39 40    Attribute names match HF so weight remap is identity: q_proj/k_proj/v_proj/out_proj.41    """42 43    def __init__(self, d_model: int, n_heads: int):44        super().__init__()45        self.n_heads = n_heads46        self.head_dim = d_model // n_heads47        assert d_model == self.head_dim * n_heads48        # HF Whisper: q/v/out have bias; k does not49        self.q_proj = nn.Linear(d_model, d_model, bias=True)50        self.k_proj = nn.Linear(d_model, d_model, bias=False)51        self.v_proj = nn.Linear(d_model, d_model, bias=True)52        self.out_proj = nn.Linear(d_model, d_model, bias=True)53 54    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:55        B, T, D = x.shape56        q = self.q_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)57        k = self.k_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)58        v = self.v_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)59        scale = self.head_dim ** -0.560        attn = (q * scale) @ k.transpose(0, 1, 3, 2)  # (B, H, T, T)61        if mask is not None:62            attn = attn + mask63        w = mx.softmax(attn, axis=-1, precise=True)64        out = (w @ v).transpose(0, 2, 1, 3).reshape(B, T, D)65        return self.out_proj(out)66 67 68# ---- encoder layer --------------------------------------------------69 70 71class WhisperEncoderBlock(nn.Module):72    """Pre-LN Whisper encoder block. Matches transformers.WhisperEncoderLayer."""73 74    def __init__(self, d_model: int, n_heads: int, ffn_dim: int):75        super().__init__()76        self.self_attn = WhisperAttention(d_model, n_heads)77        self.self_attn_layer_norm = nn.LayerNorm(d_model)78        self.fc1 = nn.Linear(d_model, ffn_dim)79        self.fc2 = nn.Linear(ffn_dim, d_model)80        self.final_layer_norm = nn.LayerNorm(d_model)81 82    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:83        h = self.self_attn_layer_norm(x)84        x = x + self.self_attn(h, mask=mask)85        h = self.final_layer_norm(x)86        x = x + self.fc2(nn.gelu(self.fc1(h)))87        return x88 89 90# ---- encoder --------------------------------------------------------91 92 93@dataclass94class EncoderConfig:95    num_mel_bins: int = 12896    downsample_hidden_size: int = 48097    d_model: int = 128098    n_heads: int = 2099    ffn_dim: int = 5120100    n_layers: int = 32101    max_source_positions: int = 1500102    layer_norm_eps: float = 1e-5103    output_dim: int = 1280104    deepstack_layer_indexes: List[int] = field(default_factory=lambda: [8, 16, 24])105 106 107class MossAudioEncoderMLX(nn.Module):108    def __init__(self, cfg: EncoderConfig):109        super().__init__()110        self.cfg = cfg111        # Conv2d stem: 1 → 480 → 480 → 480, each stride-2112        # MLX Conv2d expects NHWC, weight shape (OC, kH, kW, IC)113        self.conv1 = nn.Conv2d(1, cfg.downsample_hidden_size, kernel_size=3, stride=2, padding=1)114        self.conv2 = nn.Conv2d(cfg.downsample_hidden_size, cfg.downsample_hidden_size, kernel_size=3, stride=2, padding=1)115        self.conv3 = nn.Conv2d(cfg.downsample_hidden_size, cfg.downsample_hidden_size, kernel_size=3, stride=2, padding=1)116        # After 3× stride-2 on mel-axis (128→64→32→16): flat dim = 480*16 = 7680117        self.stem_proj = nn.Linear(cfg.downsample_hidden_size * 16, cfg.d_model)118        # Precomputed sinusoids, will be sliced119        self._positions = sinusoids(cfg.max_source_positions, cfg.d_model)120        self.layers = [121            WhisperEncoderBlock(cfg.d_model, cfg.n_heads, cfg.ffn_dim)122            for _ in range(cfg.n_layers)123        ]124        self.layer_norm = nn.LayerNorm(cfg.d_model, eps=cfg.layer_norm_eps)125        # MOSS has optional out_proj; for 4B output_dim==d_model, so it's Identity in PyTorch126        # We skip it entirely (equivalent).127        assert cfg.output_dim == cfg.d_model, "non-identity out_proj not yet implemented"128        self._deepstack_set = set(cfg.deepstack_layer_indexes)129 130    def _compute_downsampled_length(self, L: int) -> int:131        """3× stride-2 conv output length: ceil((((L-1)//2+1)-1)//2+1 ... )"""132        def step(n): return (n - 1) // 2 + 1133        return step(step(step(L)))134 135    def __call__(136        self,137        input_features: mx.array,      # (B, n_mels, T) bf16 mel spectrogram138        feature_lens: Optional[mx.array] = None,139        return_deepstack: bool = True,140    ) -> Tuple[mx.array, Optional[List[mx.array]]]:141        if input_features.ndim == 2:142            input_features = input_features[None]143        B, n_mels, T = input_features.shape144        if feature_lens is None:145            feature_lens = mx.full((B,), T, dtype=mx.int32)146 147        # (B, n_mels, T) → (B, n_mels, T, 1) [NHWC with channels-last = 1 input channel]148        # But MLX Conv2d expects input shape (B, H, W, C_in). We map:149        #   H = n_mels (128), W = T (frames), C_in = 1150        x = input_features[..., None]  # (B, n_mels, T, 1)151        x = nn.gelu(self.conv1(x))     # (B, 64, T/2, 480)152        x = nn.gelu(self.conv2(x))     # (B, 32, T/4, 480)153        x = nn.gelu(self.conv3(x))     # (B, 16, T/8, 480)154        # PyTorch reference: (B, C, F, T) → permute(0,3,1,2) → (B, T, C, F) → flatten → (B, T, C*F)155        # MLX is (B, F, T, C) post-conv. Need transpose to (B, T, C, F) to match PT's flatten order.156        B_, H_, W_, C_ = x.shape  # H_=F, W_=T, C_=C157        x = x.transpose(0, 2, 3, 1).reshape(B_, W_, C_ * H_)  # (B, T, C*F)158        x = self.stem_proj(x)  # (B, T', d_model)159 160        # Trim to actual downsampled length (in case input was padded)161        max_len = self._compute_downsampled_length(int(feature_lens.max().item()))162        if x.shape[1] > max_len:163            x = x[:, :max_len, :]164 165        # Add sinusoidal positions166        seq_len = x.shape[1]167        pos = self._positions[:seq_len].astype(x.dtype)168        x = x + pos169 170        # Build attention mask: (B, 1, 1, seq_len) additive171        # padding_mask[b, t] = True if t >= downsampled_len[b]  (this is where we mask out)172        dsl = mx.stack([173            mx.array(self._compute_downsampled_length(int(feature_lens[b].item())), dtype=mx.int32)174            for b in range(B)175        ])  # (B,)176        ar = mx.arange(seq_len, dtype=mx.int32)177        padding = ar[None, :] >= dsl[:, None]  # (B, seq_len) bool178        neg_inf = mx.array(-1e9, dtype=x.dtype)179        mask = mx.where(padding, neg_inf, mx.array(0.0, dtype=x.dtype))180        mask = mask[:, None, None, :]  # (B, 1, 1, seq_len)181 182        deepstack: List[mx.array] = []183        for layer_idx, layer in enumerate(self.layers):184            x = layer(x, mask=mask)185            if return_deepstack and layer_idx in self._deepstack_set:186                # Apply the final layer_norm snapshot at this point, per MOSS's output_deepstack_hidden_states187                # Actually, MOSS captures x BEFORE the final layer_norm — matches what PyTorch does.188                deepstack.append(x)189 190        x = self.layer_norm(x)191        return x, (deepstack if return_deepstack else None)192 193 194# ---- GatedMLP (for audio_adapter + deepstack_audio_merger_list) ----195 196 197class GatedMLP(nn.Module):198    """MOSS's GatedMLP: down(silu(gate(x)) * up(x)). SwiGLU convention.199 200    Matches MOSS/src/modeling_moss_audio.py:155-164.201    All linears are bias=False.202    """203 204    def __init__(self, input_size: int, hidden_size: int, output_size: int):205        super().__init__()206        self.gate_proj = nn.Linear(input_size, hidden_size, bias=False)207        self.up_proj = nn.Linear(input_size, hidden_size, bias=False)208        self.down_proj = nn.Linear(hidden_size, output_size, bias=False)209 210    def __call__(self, x: mx.array) -> mx.array:211        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))212 213 214__all__ = ["sinusoids", "WhisperAttention", "WhisperEncoderBlock",215           "EncoderConfig", "MossAudioEncoderMLX", "GatedMLP"]216