Add an MLX backend so MOSS-Music runs on Apple Silicon

mlx/README.md

@@ -0,0 +1,131 @@
+# MOSS-Music on MLX (Apple Silicon)
+
+A native [MLX](https://github.com/ml-explore/mlx) inference backend for **MOSS-Music-8B**,
+so the model runs efficiently on Apple Silicon Macs. The recommended 8-bit recipe fits
+in ~10 GB; its next-token logits match the bf16 model to cosine ≥ 0.9999 with identical
+prefill argmax on the clips tested (see [Parity & quality](#parity--quality)).
+
+**Motivation.** Under PyTorch + Apple's MPS backend several audio-encoder ops fall back
+to CPU, and local generation is effectively unusable (<0.3 tok/s on an M-series laptop,
+often stalling). This backend reimplements the model on MLX and reuses the repository's
+own processor, so the audio→caption pipeline is unchanged; only the model runtime differs.
+
+## Layout
+
+```
+mlx/
+├── moss_music_mlx/
+│   ├── audio_encoder.py   # chunked conv stem + Whisper encoder layers (+ DeepStack taps)
+│   ├── adapters.py        # SwiGLU audio adapter + DeepStack mergers
+│   ├── model.py           # audio↔Qwen3 fusion (masked scatter + DeepStack injection)
+│   ├── convert.py         # HF safetensors → MLX + quantization
+│   └── generate.py        # KV-cached audio-grounded generation + CLI
+└── tests/test_parity.py   # MLX↔PyTorch prefill parity + MLX decode-path self-consistency
+```
+
+The Qwen3 decoder is reused from [`mlx-lm`](https://github.com/ml-explore/mlx-examples)
+(attention, RoPE, q/k-norm, RMSNorm); only the audio encoder, the SwiGLU adapters, and
+the fusion / DeepStack logic are implemented here.
+
+## Install
+
+```bash
+pip install -e ".[mlx]"     # from the repo root: mlx, mlx-lm + CPU torch/transformers for the processor
+```
+
+This makes `moss_music_mlx` importable and `python -m moss_music_mlx.…` runnable from
+anywhere. (Numbers below were produced with `mlx==0.31.2`, `mlx-lm==0.29.1`.)
+
+## Convert
+
+```bash
+# 8-bit (recommended, validated): ~10 GB
+python -m moss_music_mlx.convert <hf_model_dir> ./moss-music-8b-mlx-8bit
+
+# bf16 (no quantization): ~18 GB
+python -m moss_music_mlx.convert <hf_model_dir> ./moss-music-8b-mlx-bf16 --bits 0
+```
+
+The converter prints a per-layer quantization coverage line. The audio encoder is always
+kept at bf16 to preserve audio fidelity; quantization is applied to the Qwen3 layers,
+token embeddings and `lm_head`. `--bits 4` / `--bits 6` also run, but only the 8-bit
+recipe has been parity-checked — treat lower bit-widths as experimental.
+
+> Conversion materializes the full bf16 weights before quantizing — keep **~18–20 GB RAM
+> free** (or convert on a larger machine). Loading a *converted* model needs only its own
+> footprint (~10 GB for 8-bit).
+
+## Generate
+
+```bash
+python -m moss_music_mlx.generate \
+  --model ./moss-music-8b-mlx-8bit \
+  --audio song.mp3 \
+  --prompt "Analyze this track: genre, instruments, key, BPM, structure."
+```
+
+Python API (run with the repo importable, e.g. after `pip install -e .[mlx]`):
+
+```python
+from moss_music_mlx import load_pretrained, generate
+from src.processing_moss_music import MossMusicProcessor
+
+model = load_pretrained("./moss-music-8b-mlx-8bit")
+proc = MossMusicProcessor.from_pretrained("<hf_model_dir>", trust_remote_code=True, enable_time_marker=True)
+print(generate(model, proc, "Describe this music.", audio_path="song.mp3"))
+```
+
+Text streams through an incremental, UTF-8-safe detokenizer (Chinese / multi-byte output
+is not split into replacement characters; control tokens are stripped, `<think>` kept).
+Stop ids are read from the model's `generation_config.json`. `--temp 0` is greedy (the
+parity-validated default); `--temp > 0` uses nucleus + top-k sampling (`--top-p 0.8`,
+`--top-k 50`, matching the reference). Inference is single-audio, batch-1.
+
+## Benchmarks (M4, 24 GB; 75 s clip, greedy)
+
+| Backend | Disk | Load | Prefill (~1.1k tok) | Decode |
+|---|---|---|---|---|
+| PyTorch / MPS (bf16) | 18 GB | 17 s | stalls (>13 min) | <0.3 tok/s |
+| **MLX 8-bit** | **10.2 GB** | **1.5 s** | **3.2 s** | **~23 tok/s** |
+
+A full ~750-token analysis is ≈34 s/song on the 8-bit model. MLX bf16 also runs, but its
+18 GB working set thrashes on a 24 GB Mac, so 8-bit is the recommended local recipe.
+
+> Indicative single-run numbers (one M4, one 75 s clip), wall-clock with `mx.eval`
+> barriers, not a controlled benchmark; the PyTorch/MPS row is what was observed on this
+> machine and will vary by torch/MPS version. Long audio is encoded in `conv_chunksize`
+> (64) windows, so encoder peak memory is bounded regardless of clip length.
+
+## Parity & quality
+
+`tests/test_parity.py` performs two checks:
+
+1. **Prefill parity vs the PyTorch reference** (needs torch + the HF checkpoint): the
+   next-token argmax and logit cosine must match.
+2. **Decode-path self-consistency** (MLX only): a cached step-by-step decode must equal a
+   single full-context forward *of the same MLX model* at every position. This validates
+   the KV-cache / RoPE-offset bookkeeping; the model math itself is guarded by (1).
+
+```bash
+python mlx/tests/test_parity.py --model <hf_dir> --audio clip.wav --mlx-model ./moss-music-8b-mlx-8bit
+# reproduce the bf16-vs-8bit table below:
+python mlx/tests/eval_quant.py --model <hf_dir> --mlx-8bit ./moss-music-8b-mlx-8bit --clips '<clips>/*.wav' --n 5
+```
+
+Observed (greedy; logit-level — 8-bit and bf16 can still diverge after a near-tie token
+over a long generation, as expected for quantization):
+
+| Comparison | Result |
+|---|---|
+| 8-bit vs **fp32 reference** — prefill next token (1 clip) | argmax identical, logit cos **0.999982** |
+| 8-bit vs **bf16** — prefill next token (5 clips, mixed genres¹), via `eval_quant.py` | argmax **5/5**, mean logit cos **0.99998** |
+| 8-bit decode path — cached vs full forward (24 steps) | 23/24 argmax agree, min cos **0.995**² |
+
+¹ Mandopop, Chinese indie, electronic, EDM-pop, metalcore.
+² The single mismatch is a near-tie token under bf16/8-bit non-associativity (cosine stays
+high and there is no cascade), not a cache/RoPE bug.
+
+## Credits
+
+Reuses [`mlx-lm`](https://github.com/ml-explore/mlx-examples)'s Qwen3 implementation
+(MIT) for the language decoder.

mlx/moss_music_mlx/__init__.py

@@ -0,0 +1,32 @@
+"""MLX inference pipeline for MOSS-Music-8B (audio music understanding)."""
+from .model import MossMusicModel
+from .convert import (
+    load_model,
+    load_pretrained,
+    convert,
+    quantize_model,
+    hf_to_mlx_weights,
+    load_config,
+)
+
+__all__ = [
+    "MossMusicModel",
+    "load_model",
+    "load_pretrained",
+    "convert",
+    "quantize_model",
+    "hf_to_mlx_weights",
+    "load_config",
+    "generate",
+    "stream_generate",
+    "generate_ids",
+]
+
+
+def __getattr__(name):
+    # Lazy so that `import moss_music_mlx` does not pull in the torch-based
+    # processor; only generate/stream_generate/generate_ids need it.
+    if name in ("generate", "stream_generate", "generate_ids"):
+        from . import generate as _gen
+        return getattr(_gen, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

mlx/moss_music_mlx/adapters.py

@@ -0,0 +1,22 @@
+"""SwiGLU adapters that project audio-encoder features into the LLM embedding space.
+
+One ``GatedMLP`` adapts the final encoder output (for the masked-scatter into audio
+token positions); ``deepstack_num_inject_layers`` more adapt the tapped encoder
+hidden states for DeepStack injection into the first LLM layers. Mirrors
+``GatedMLP`` in ``src/modeling_moss_music.py`` (gate/up/down, SiLU, no bias).
+"""
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class GatedMLP(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, output_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(input_size, hidden_size, bias=False)
+        self.up_proj = nn.Linear(input_size, hidden_size, bias=False)
+        self.down_proj = nn.Linear(hidden_size, output_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))

mlx/moss_music_mlx/audio_encoder.py

@@ -0,0 +1,188 @@
+"""MLX port of the MOSS-Music audio encoder.
+
+Mirrors ``MossMusicEncoder`` in ``src/modeling_moss_music.py``: a Conv2d stem that
+downsamples the 128-bin log-mel spectrogram 8x in both axes, followed by a stack of
+Whisper-style transformer encoder layers. Hidden states at the configured
+``deepstack_encoder_layer_indexes`` are tapped and returned for DeepStack injection.
+
+Layout note: PyTorch Conv2d is NCHW with weight [out, in, kH, kW]; MLX Conv2d is
+NHWC with weight [out, kH, kW, in]. The converter handles the weight transpose; here
+we just feed NHWC activations and undo the layout to match the reference flatten order
+(channel-outer, freq-inner) before the stem projection.
+"""
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def sinusoid_position_embedding(seq_len: int, dim: int) -> mx.array:
+    """Recompute Whisper-style sinusoid positions (matches SinusoidsPositionEmbedding).
+
+    Recomputed rather than loaded: the reference registers ``inv_timescales`` as a
+    non-persistent buffer, so we reproduce it deterministically instead of relying on
+    the checkpoint.
+    """
+    half = dim // 2
+    log_timescale_increment = math.log(10000.0) / (half - 1)
+    inv_timescales = mx.exp(-log_timescale_increment * mx.arange(half))
+    scaled_time = mx.arange(seq_len)[:, None] * inv_timescales[None, :]
+    pos = mx.concatenate([mx.sin(scaled_time), mx.cos(scaled_time)], axis=1)
+    return pos[None]  # [1, seq_len, dim]
+
+
+class WhisperEncoderLayer(nn.Module):
+    """Whisper encoder layer: LN -> self-attention -> residual, then LN -> FFN ->
+    residual (matches transformers' WhisperEncoderLayer; k_proj has no bias)."""
+
+    def __init__(self, d_model: int, n_heads: int, ffn_dim: int, eps: float):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.self_attn_layer_norm = nn.LayerNorm(d_model, eps=eps)
+        self.q_proj = nn.Linear(d_model, d_model, bias=True)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)  # Whisper k_proj has no bias
+        self.v_proj = nn.Linear(d_model, d_model, bias=True)
+        self.out_proj = nn.Linear(d_model, d_model, bias=True)
+
+        self.final_layer_norm = nn.LayerNorm(d_model, eps=eps)
+        self.fc1 = nn.Linear(d_model, ffn_dim, bias=True)
+        self.fc2 = nn.Linear(ffn_dim, d_model, bias=True)
+
+    def _attn(self, x: mx.array, mask: Optional[mx.array]) -> mx.array:
+        B, T, _ = x.shape
+        q = self.q_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
+        out = out.transpose(0, 2, 1, 3).reshape(B, T, -1)
+        return self.out_proj(out)
+
+    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        x = residual + self._attn(x, mask)
+
+        residual = x
+        x = self.final_layer_norm(x)
+        x = residual + self.fc2(nn.gelu(self.fc1(x)))
+        return x
+
+
+def _conv3_downsample_len(length: int) -> int:
+    """Length after 3 stride-2, pad-1, kernel-3 convolutions."""
+    for _ in range(3):
+        length = (length - 1) // 2 + 1
+    return length
+
+
+class MossMusicAudioEncoder(nn.Module):
+    """Chunked (windowed) audio encoder.
+
+    The reference processes the mel spectrogram in independent windows of
+    ``n_window * 2`` frames: each window goes through the conv stem and the full
+    transformer stack on its own (attention is block-diagonal across windows). Per
+    the checkpoint, ``n_window`` defaults to 200 -> 400-frame windows -> 50
+    downsampled tokens each; the tail window is shorter and padding-masked.
+    """
+
+    def __init__(self, cfg: dict):
+        super().__init__()
+        d_model = cfg["d_model"]
+        dh = cfg["downsample_hidden_size"]
+        # configured tap order is authoritative (mergers pair with it positionally)
+        self.deepstack_indexes = list(cfg.get("deepstack_encoder_layer_indexes") or [])
+        self.deepstack_set = set(self.deepstack_indexes)
+        self.chunk_frames = int(cfg.get("n_window", 200)) * 2
+        # windows are processed in sub-batches to bound peak memory on long audio
+        # (mirrors the reference's conv_chunksize OOM fix); numerically identical.
+        self.conv_chunksize = int(cfg.get("conv_chunksize", 64))
+
+        # NHWC Conv2d stem: 1 -> dh -> dh -> dh, stride 2, pad 1.
+        self.conv1 = nn.Conv2d(1, dh, kernel_size=3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(dh, dh, kernel_size=3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(dh, dh, kernel_size=3, stride=2, padding=1)
+        # 128 mel bins / 8 = 16 freq rows after the stem.
+        self.stem_proj = nn.Linear(dh * (cfg["num_mel_bins"] // cfg["downsample_rate"]), d_model)
+
+        self.layers = [
+            WhisperEncoderLayer(d_model, cfg["encoder_attention_heads"],
+                                cfg["encoder_ffn_dim"], cfg["layer_norm_eps"])
+            for _ in range(cfg["encoder_layers"])
+        ]
+        self.layer_norm = nn.LayerNorm(d_model, eps=cfg["layer_norm_eps"])
+        # reference applies out_proj to last_hidden + deepstack taps only when
+        # output_dim != d_model; Identity (no params) otherwise.
+        out_dim = cfg.get("output_dim", d_model)
+        self.out_proj = nn.Linear(d_model, out_dim, bias=False) if out_dim != d_model else None
+        self.d_model = d_model
+
+    def _stem(self, batch: mx.array) -> mx.array:
+        """[n_chunks, n_mels, chunk_frames] -> [n_chunks, T_ds, d_model]."""
+        x = batch[..., None]  # NHWC: [n_chunks, n_mels, chunk_frames, 1]
+        x = nn.gelu(self.conv1(x))
+        x = nn.gelu(self.conv2(x))
+        x = nn.gelu(self.conv3(x))  # [n_chunks, F=16, T_ds, C=dh]
+        B, Fr, Tds, C = x.shape
+        # reference permute(0,3,1,2).flatten(2): channel outer, freq inner
+        x = x.transpose(0, 2, 3, 1).reshape(B, Tds, C * Fr)
+        x = self.stem_proj(x)
+        return x + sinusoid_position_embedding(x.shape[1], self.d_model).astype(x.dtype)
+
+    def __call__(self, mel: mx.array, feature_len: Optional[int] = None
+                 ) -> Tuple[mx.array, List[mx.array]]:
+        """mel: [1, n_mels, T] or [n_mels, T] -> ([1, N_valid, d], deepstack list)."""
+        if mel.ndim == 3:
+            mel = mel[0]
+        n_mels, total = mel.shape
+        if feature_len is None:
+            feature_len = total
+        cf = self.chunk_frames
+
+        # Split [n_mels, T] into chunk_frames windows; pad the tail to chunk_frames.
+        lengths, off = [], 0
+        while off < feature_len:
+            lengths.append(min(cf, feature_len - off))
+            off += cf
+        chunks, off = [], 0
+        for L in lengths:
+            c = mel[:, off:off + L]
+            if L < cf:
+                c = mx.pad(c, [(0, 0), (0, cf - L)])
+            chunks.append(c)
+            off += L
+        dlens = [_conv3_downsample_len(L) for L in lengths]
+
+        # process windows in sub-batches of conv_chunksize; attention is block-diagonal
+        # per window so this is identical to one big batch, with bounded peak memory.
+        last_parts: List[mx.array] = []
+        deep_parts = {idx: [] for idx in self.deepstack_indexes}
+        for s in range(0, len(chunks), self.conv_chunksize):
+            sub_dlens = dlens[s:s + self.conv_chunksize]
+            x = self._stem(mx.stack(chunks[s:s + self.conv_chunksize], axis=0))
+            col = mx.arange(x.shape[1])[None, :]
+            valid = col < mx.array(sub_dlens)[:, None]
+            mask = mx.where(valid[:, None, None, :], 0.0, -1e9).astype(x.dtype)
+
+            captured = {}
+            for i, layer in enumerate(self.layers):
+                x = layer(x, mask=mask)
+                if i in self.deepstack_set:
+                    captured[i] = x
+            x = self.layer_norm(x)
+            for j, dl in enumerate(sub_dlens):
+                last_parts.append(x[j, :dl])
+            for idx in self.deepstack_indexes:
+                for j, dl in enumerate(sub_dlens):
+                    deep_parts[idx].append(captured[idx][j, :dl])
+
+        def finish(parts: List[mx.array]) -> mx.array:
+            g = mx.concatenate(parts, axis=0)[None]       # [1, N_valid, d]
+            return self.out_proj(g) if self.out_proj is not None else g
+
+        return finish(last_parts), [finish(deep_parts[idx]) for idx in self.deepstack_indexes]