From d87d73aaa640a6d31fe85726562496998879f042 Mon Sep 17 00:00:00 2001 From: Ben Wu <4549761+dthinkr@users.noreply.github.com> Date: Sat, 27 Jun 2026 20:10:00 +0100 Subject: [PATCH] feat(mlx): add MLX inference backend for Apple Silicon Port MOSS-Music-8B to MLX so it runs efficiently on Apple Silicon, where the PyTorch/MPS path falls back to CPU and is effectively unusable for local generation. The audio encoder (chunked conv stem, Whisper layers, DeepStack taps), the SwiGLU adapters and the audio-to-Qwen3 fusion are reimplemented in MLX; the Qwen3 decoder is reused from mlx-lm. The repository's own MossMusicProcessor is reused unchanged for mel features and audio-token placeholders. Includes: - HF to MLX conversion with 8-bit quantization (audio encoder kept bf16) - KV-cached generation with a UTF-8-safe incremental detokenizer and a CLI - prefill parity vs the PyTorch reference, decode-path self-consistency tests, artifact-free unit tests, and a bf16-vs-8bit eval harness - an [mlx] optional-dependency group; mlx/ mirrors the sglang/ backend layout Validation: 8-bit vs fp32 reference prefill argmax identical, logit cos 0.99999; 8-bit vs bf16 over 5 mixed-genre clips argmax 5/5, mean cos 0.99998; ~10.2 GB, ~23 tok/s on an M4 (vs PyTorch/MPS stalling at <0.3 tok/s). --- mlx/README.md | 131 +++++++++++++++++++ mlx/moss_music_mlx/__init__.py | 32 +++++ mlx/moss_music_mlx/adapters.py | 22 ++++ mlx/moss_music_mlx/audio_encoder.py | 188 ++++++++++++++++++++++++++++ mlx/moss_music_mlx/convert.py | 170 +++++++++++++++++++++++++ mlx/moss_music_mlx/generate.py | 141 +++++++++++++++++++++ mlx/moss_music_mlx/model.py | 111 ++++++++++++++++ mlx/tests/eval_quant.py | 74 +++++++++++ mlx/tests/test_parity.py | 153 ++++++++++++++++++++++ mlx/tests/test_units.py | 80 ++++++++++++ pyproject.toml | 17 ++- 11 files changed, 1118 insertions(+), 1 deletion(-) create mode 100644 mlx/README.md create mode 100644 mlx/moss_music_mlx/__init__.py create mode 100644 mlx/moss_music_mlx/adapters.py create mode 100644 mlx/moss_music_mlx/audio_encoder.py create mode 100644 mlx/moss_music_mlx/convert.py create mode 100644 mlx/moss_music_mlx/generate.py create mode 100644 mlx/moss_music_mlx/model.py create mode 100644 mlx/tests/eval_quant.py create mode 100644 mlx/tests/test_parity.py create mode 100644 mlx/tests/test_units.py diff --git a/mlx/README.md b/mlx/README.md new file mode 100644 index 0000000..67c25c8 --- /dev/null +++ b/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 ./moss-music-8b-mlx-8bit + +# bf16 (no quantization): ~18 GB +python -m moss_music_mlx.convert ./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("", 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, `` 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 --audio clip.wav --mlx-model ./moss-music-8b-mlx-8bit +# reproduce the bf16-vs-8bit table below: +python mlx/tests/eval_quant.py --model --mlx-8bit ./moss-music-8b-mlx-8bit --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. diff --git a/mlx/moss_music_mlx/__init__.py b/mlx/moss_music_mlx/__init__.py new file mode 100644 index 0000000..679e24c --- /dev/null +++ b/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}") diff --git a/mlx/moss_music_mlx/adapters.py b/mlx/moss_music_mlx/adapters.py new file mode 100644 index 0000000..0a47b78 --- /dev/null +++ b/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)) diff --git a/mlx/moss_music_mlx/audio_encoder.py b/mlx/moss_music_mlx/audio_encoder.py new file mode 100644 index 0000000..1a0f95c --- /dev/null +++ b/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] diff --git a/mlx/moss_music_mlx/convert.py b/mlx/moss_music_mlx/convert.py new file mode 100644 index 0000000..d8de363 --- /dev/null +++ b/mlx/moss_music_mlx/convert.py @@ -0,0 +1,170 @@ +"""HF (PyTorch safetensors) -> MLX conversion and 8-bit quantization for MOSS-Music. + +Usage: + python -m moss_music_mlx.convert [--bits 8] [--group-size 64] + +The language model, adapters, lm_head and token embeddings are quantized; the audio +encoder is kept at high precision (bf16) to preserve audio fidelity — this is the +"near-lossless" 8-bit recipe. Quantization metadata is written to the output config so +the model reloads with the matching layer structure. +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +from typing import Dict, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten + +from .model import MossMusicModel + +_CONV = {"conv1", "conv2", "conv3"} +# tokenizer / processor assets copied verbatim so the output dir is self-contained +_ASSETS = ( + "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", + "added_tokens.json", "special_tokens_map.json", "chat_template.jinja", + "preprocessor_config.json", "generation_config.json", +) + + +def load_config(model_dir: str) -> dict: + with open(os.path.join(model_dir, "config.json")) as f: + return json.load(f) + + +def _map_key(name: str) -> str: + if name.startswith("audio_encoder."): + return name.replace(".self_attn.", ".") # flat WhisperEncoderLayer + return name + + +def hf_to_mlx_weights(model_dir: str, dtype=mx.bfloat16) -> Dict[str, mx.array]: + """Read the sharded HF checkpoint into MLX arrays with the layout fixes applied.""" + with open(os.path.join(model_dir, "model.safetensors.index.json")) as f: + index = json.load(f)["weight_map"] + shards = sorted(set(index.values())) + out: Dict[str, mx.array] = {} + for shard in shards: + for name, arr in mx.load(os.path.join(model_dir, shard)).items(): + if "embed_positions.inv_timescales" in name: + continue # recomputed deterministically + parts = name.split(".") + if (name.startswith("audio_encoder.") and len(parts) > 1 and parts[1] in _CONV + and name.endswith(".weight") and arr.ndim == 4): + arr = arr.transpose(0, 2, 3, 1) # [out,in,kH,kW] -> NHWC [out,kH,kW,in] + out[_map_key(name)] = arr.astype(dtype) + return out + + +def _quant_predicate(group_size: int): + def pred(path: str, module: nn.Module): + if path.startswith("audio_encoder"): + return False # keep audio encoder high precision + if not isinstance(module, (nn.Linear, nn.Embedding)): + return False + return module.weight.shape[-1] % group_size == 0 + return pred + + +def quantize_model(model: MossMusicModel, bits: int = 8, group_size: int = 64) -> MossMusicModel: + nn.quantize(model, group_size=group_size, bits=bits, class_predicate=_quant_predicate(group_size)) + return model + + +def load_model(model_dir: str, dtype=mx.bfloat16) -> MossMusicModel: + """Build an un-quantized model straight from an HF checkpoint (for parity checks).""" + model = MossMusicModel(load_config(model_dir)) + model.load_weights(list(hf_to_mlx_weights(model_dir, dtype=dtype).items()), strict=True) + mx.eval(model.parameters()) + return model + + +def _resolve_eos_ids(model_dir: str, config: dict) -> list: + """Collect stop-token ids from generation_config.json (preferred) and config, + normalizing the int-or-list eos_token_id shapes Qwen3 checkpoints use.""" + ids = set() + sources = [config, config.get("language_config", {}) or {}] + gen_path = os.path.join(model_dir, "generation_config.json") + if os.path.exists(gen_path): + with open(gen_path) as f: + sources.insert(0, json.load(f)) + for src in sources: + eos = src.get("eos_token_id") + if isinstance(eos, int): + ids.add(eos) + elif isinstance(eos, (list, tuple)): + ids.update(int(e) for e in eos) + return sorted(ids) + + +def load_pretrained(out_dir: str) -> MossMusicModel: + """Load a converted (possibly quantized) MLX model directory.""" + with open(os.path.join(out_dir, "config.json")) as f: + config = json.load(f) + model = MossMusicModel(config) + q = config.get("quantization") + if q: + nn.quantize(model, group_size=q["group_size"], bits=q["bits"], + class_predicate=_quant_predicate(q["group_size"])) + model.load_weights(os.path.join(out_dir, "model.safetensors")) + model.eos_token_ids = _resolve_eos_ids(out_dir, config) + mx.eval(model.parameters()) + return model + + +def convert(model_dir: str, out_dir: str, bits: Optional[int] = 8, group_size: int = 64): + config = load_config(model_dir) + model = MossMusicModel(config) + model.load_weights(list(hf_to_mlx_weights(model_dir, dtype=mx.bfloat16).items()), strict=True) + + quantization = None + if bits: + quantize_model(model, bits=bits, group_size=group_size) + quantization = {"group_size": group_size, "bits": bits} + + os.makedirs(out_dir, exist_ok=True) + weights = dict(tree_flatten(model.parameters())) + mx.eval(weights) + mx.save_safetensors(os.path.join(out_dir, "model.safetensors"), weights) + + if quantization: + config["quantization"] = quantization + with open(os.path.join(out_dir, "config.json"), "w") as f: + json.dump(config, f, indent=2) + for a in _ASSETS: + src = os.path.join(model_dir, a) + if os.path.exists(src): + shutil.copy2(src, os.path.join(out_dir, a)) + + nbytes = sum(v.nbytes for v in weights.values()) + print(f"saved {len(weights)} tensors, {nbytes/1e9:.2f} GB -> {out_dir}" + + (f" ({bits}-bit, group {group_size})" if bits else " (bf16)")) + if bits: + # quantized layers carry a `.scales` tensor; report coverage so the size/ + # quality claims are auditable and a silently-skipped large layer is visible. + quant = sorted(k[:-len(".scales")] for k in weights if k.endswith(".scales")) + print(f" quantized {len(quant)} layers (audio_encoder kept bf16)") + skipped = [k[:-len(".weight")] for k in weights + if k.endswith(".weight") and not k.startswith("audio_encoder") + and k[:-len(".weight")] + ".scales" not in weights + and any(t in k for t in ("proj", "lm_head", "embed_tokens"))] + if skipped: + print(f" WARNING: {len(skipped)} projection/embedding layers left bf16: {skipped[:6]}") + + +def main(): + p = argparse.ArgumentParser(description="Convert MOSS-Music HF checkpoint to MLX") + p.add_argument("model_dir") + p.add_argument("out_dir") + p.add_argument("--bits", type=int, default=8, help="quantization bits (0 = bf16, no quant)") + p.add_argument("--group-size", type=int, default=64) + args = p.parse_args() + convert(args.model_dir, args.out_dir, bits=(args.bits or None), group_size=args.group_size) + + +if __name__ == "__main__": + main() diff --git a/mlx/moss_music_mlx/generate.py b/mlx/moss_music_mlx/generate.py new file mode 100644 index 0000000..b44711e --- /dev/null +++ b/mlx/moss_music_mlx/generate.py @@ -0,0 +1,141 @@ +"""Audio-grounded generation for the MLX MOSS-Music model. + +Reuses the repository's own ``MossMusicProcessor`` (mel features + audio-token +placeholders) and ``load_audio``; runs an incremental KV-cached decode loop over the +MLX model. DeepStack/audio fusion happens once, on the prefill step. + +Text is emitted with an incremental detokenizer that decodes the running id list and +yields only newly-completed characters, holding back partial multi-byte UTF-8 (so +Chinese / emoji stream correctly). Special control tokens are stripped to match the +reference (``src/hf_inference.py``); ```` reasoning is preserved. +""" +from __future__ import annotations + +import os +import sys +import time +from typing import Callable, Iterable, Iterator, Optional + +import mlx.core as mx +from mlx_lm.models.cache import KVCache +from mlx_lm.sample_utils import make_sampler + +try: # prefer the installed package; fall back to a source checkout + from src.processing_moss_music import MossMusicProcessor + from src.audio_io import load_audio +except ImportError: + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + from src.processing_moss_music import MossMusicProcessor # noqa: E402 + from src.audio_io import load_audio # noqa: E402 + +_FALLBACK_EOS = 151645 # <|im_end|> + + +def resolve_stop_ids(model, processor) -> set: + """Stop token ids: prefer the model's (read from config at load), then the + tokenizer's eos, then the <|im_end|> fallback.""" + ids = set(getattr(model, "eos_token_ids", None) or []) + eos = getattr(getattr(processor, "tokenizer", None), "eos_token_id", None) + if isinstance(eos, int): + ids.add(eos) + ids.add(_FALLBACK_EOS) + return ids + + +def generate_ids(model, processor, prompt: str, audio_path: Optional[str] = None, + max_new_tokens: int = 768, temp: float = 0.0, top_p: float = 0.8, + top_k: int = 50, stop_ids: Optional[Iterable[int]] = None) -> Iterator[int]: + """Core decode loop: yields generated token ids (excludes the stop token). + + temp=0 is greedy (top_p/top_k ignored); temp>0 uses nucleus + top-k sampling + matching the reference's defaults (top_p=0.8, top_k=50).""" + if audio_path is not None: + raw = load_audio(audio_path, sample_rate=processor.config.mel_sr) + inputs = processor(text=prompt, audios=[raw], return_tensors="pt") + mel = mx.array(inputs["audio_data"].float().numpy()).astype(mx.bfloat16) + feature_len = int(inputs["audio_data_seqlens"][0]) + audio_mask = mx.array((inputs["input_ids"] == processor.audio_token_id).numpy()) + else: + inputs = processor(text=prompt, return_tensors="pt") + mel = audio_mask = feature_len = None + + stops = set(stop_ids) if stop_ids is not None else resolve_stop_ids(model, processor) + sampler = make_sampler(temp=temp, top_p=top_p, top_k=top_k) + input_ids = mx.array(inputs["input_ids"].numpy()) + cache = [KVCache() for _ in model.language_model.layers] + + logits = model(input_ids, mel=mel, audio_mask=audio_mask, feature_len=feature_len, cache=cache) + y = sampler(mx.softmax(logits[:, -1, :], axis=-1).log()) + mx.eval(y) + for _ in range(max_new_tokens): + tok = int(y.item()) + if tok in stops: + break + yield tok + logits = model(y[None], cache=cache) + y = sampler(mx.softmax(logits[:, -1, :], axis=-1).log()) + mx.eval(y) + + +def stream_generate(model, processor, prompt: str, audio_path: Optional[str] = None, + max_new_tokens: int = 768, temp: float = 0.0, top_p: float = 0.8, + top_k: int = 50, stop_ids: Optional[Iterable[int]] = None, + on_token: Optional[Callable[[int], None]] = None) -> Iterator[str]: + """Yield decoded text deltas, UTF-8-safe (holds back partial multi-byte chars).""" + def decode(ids): + return processor.decode(ids, skip_special_tokens=True) + + ids, emitted = [], "" + for tok in generate_ids(model, processor, prompt, audio_path, max_new_tokens, + temp, top_p, top_k, stop_ids): + if on_token is not None: + on_token(tok) + ids.append(tok) + text = decode(ids) + if text.endswith("�"): + continue # mid multi-byte char; wait for the completing token + if len(text) > len(emitted): + yield text[len(emitted):] + emitted = text + text = decode(ids) # flush any remainder + if len(text) > len(emitted): + yield text[len(emitted):] + + +def generate(model, processor, prompt: str, audio_path: Optional[str] = None, + max_new_tokens: int = 768, temp: float = 0.0, top_p: float = 0.8, + top_k: int = 50, stop_ids: Optional[Iterable[int]] = None) -> str: + return "".join(stream_generate(model, processor, prompt, audio_path, + max_new_tokens, temp, top_p, top_k, stop_ids)) + + +def main(): + import argparse + from .convert import load_pretrained + + p = argparse.ArgumentParser(description="Generate with the MLX MOSS-Music model") + p.add_argument("--model", required=True, help="converted MLX model dir") + p.add_argument("--audio", help="audio file (wav/mp3/...)") + p.add_argument("--prompt", default="Please describe this music.") + p.add_argument("--max-new-tokens", type=int, default=768) + p.add_argument("--temp", type=float, default=0.0, help="0 = greedy") + p.add_argument("--top-p", type=float, default=0.8, help="nucleus (used when temp>0)") + p.add_argument("--top-k", type=int, default=50, help="top-k (used when temp>0)") + args = p.parse_args() + + t = time.time() + model = load_pretrained(args.model) + processor = MossMusicProcessor.from_pretrained(args.model, trust_remote_code=True, enable_time_marker=True) + print(f"[loaded in {time.time()-t:.1f}s]\n", flush=True) + + t, n = time.time(), [0] + for piece in stream_generate(model, processor, args.prompt, args.audio, + args.max_new_tokens, args.temp, args.top_p, args.top_k, + on_token=lambda _t: n.__setitem__(0, n[0] + 1)): + print(piece, end="", flush=True) + dt = time.time() - t + print(f"\n\n[{n[0]} tokens in {dt:.1f}s = {n[0]/max(dt,1e-9):.1f} tok/s]", flush=True) + + +if __name__ == "__main__": + main() diff --git a/mlx/moss_music_mlx/model.py b/mlx/moss_music_mlx/model.py new file mode 100644 index 0000000..f7ae88f --- /dev/null +++ b/mlx/moss_music_mlx/model.py @@ -0,0 +1,111 @@ +"""MLX MossMusicModel: audio encoder + adapters fused into a Qwen3 LLM. + +Audio features replace the ``<|AUDIO|>`` placeholder embeddings (masked scatter), +and DeepStack features are added to the hidden states at those positions after each +of the first ``len(deepstack mergers)`` decoder layers. The Qwen3 stack itself is +reused verbatim from ``mlx_lm`` (attention, RoPE, q/k-norm, RMSNorm). +""" +from __future__ import annotations + +from typing import List, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.qwen3 import ModelArgs as Qwen3Args, Qwen3Model +from mlx_lm.models.base import create_attention_mask + +from .audio_encoder import MossMusicAudioEncoder +from .adapters import GatedMLP + + +def _qwen3_args(lc: dict) -> Qwen3Args: + return Qwen3Args( + model_type="qwen3", + hidden_size=lc["hidden_size"], + num_hidden_layers=lc["num_hidden_layers"], + intermediate_size=lc["intermediate_size"], + num_attention_heads=lc["num_attention_heads"], + rms_norm_eps=lc["rms_norm_eps"], + vocab_size=lc["vocab_size"], + num_key_value_heads=lc["num_key_value_heads"], + max_position_embeddings=lc["max_position_embeddings"], + rope_theta=lc["rope_theta"], + head_dim=lc["head_dim"], + tie_word_embeddings=lc.get("tie_word_embeddings", False), + rope_scaling=lc.get("rope_scaling"), + ) + + +class MossMusicModel(nn.Module): + def __init__(self, config: dict): + super().__init__() + ac = config["audio_config"] + lc = config["language_config"] + self.audio_token_id = int(config.get("audio_token_id", 151654)) + + self.audio_encoder = MossMusicAudioEncoder(ac) + self.audio_adapter = GatedMLP(ac["output_dim"], config["adapter_hidden_size"], lc["hidden_size"]) + + k = len(ac.get("deepstack_encoder_layer_indexes") or []) + if config.get("deepstack_num_inject_layers") is not None: + k = min(k, int(config["deepstack_num_inject_layers"])) + self.deepstack_audio_merger_list = [ + GatedMLP(ac["output_dim"], config["adapter_hidden_size"], lc["hidden_size"]) + for _ in range(k) + ] + + self.language_model = Qwen3Model(_qwen3_args(lc)) + self.lm_head = nn.Linear(lc["hidden_size"], lc["vocab_size"], bias=False) + + @staticmethod + def _scatter(h: mx.array, src: mx.array, mask_row: mx.array, idx: mx.array, add: bool) -> mx.array: + """Place/add src[0, k] into h[0, j] for the k-th audio position j.""" + gathered = src[0][idx] # [T, D] + if add: + return h + mx.where(mask_row[:, None], gathered, 0.0)[None] + return mx.where(mask_row[:, None], gathered, h[0])[None] + + def audio_embeds(self, mel: mx.array, feature_len: Optional[int] = None): + """Run encoder + adapters -> (fused audio embeds, list of deepstack embeds).""" + enc_last, deepstack = self.audio_encoder(mel, feature_len) + adapted = self.audio_adapter(enc_last) + merged = [self.deepstack_audio_merger_list[i](deepstack[i]) + for i in range(len(self.deepstack_audio_merger_list))] + return adapted, merged + + def __call__(self, input_ids: mx.array, mel: Optional[mx.array] = None, + audio_mask: Optional[mx.array] = None, feature_len: Optional[int] = None, + cache=None) -> mx.array: + if input_ids.shape[0] != 1: + raise ValueError("MLX MossMusicModel supports batch size 1 (single audio) only.") + h = self.language_model.embed_tokens(input_ids) + + deepstack = [] + if mel is not None: + if audio_mask is None: + raise ValueError("audio_mask is required when mel is provided.") + adapted, deepstack = self.audio_embeds(mel, feature_len) + mask_row = audio_mask[0] + # fail loudly on a placeholder/encoder mismatch instead of mis-gathering + n_audio = int(mask_row.sum().item()) + if n_audio != adapted.shape[1]: + raise ValueError( + f"audio token count ({n_audio}) != audio embeds ({adapted.shape[1]}); " + "check the processor's <|AUDIO|> placeholders vs the encoder output length.") + for j, ds in enumerate(deepstack): + if ds.shape[1] != n_audio: + raise ValueError(f"deepstack[{j}] length {ds.shape[1]} != audio tokens {n_audio}") + idx = mx.maximum(mx.cumsum(mask_row.astype(mx.int32)) - 1, 0) + h = self._scatter(h, adapted, mask_row, idx, add=False) + + layers = self.language_model.layers + if cache is None: + cache = [None] * len(layers) + attn_mask = create_attention_mask(h, cache[0]) + + for i, (layer, c) in enumerate(zip(layers, cache)): + h = layer(h, attn_mask, c) + if i < len(deepstack): + h = self._scatter(h, deepstack[i], mask_row, idx, add=True) + + return self.lm_head(self.language_model.norm(h)) diff --git a/mlx/tests/eval_quant.py b/mlx/tests/eval_quant.py new file mode 100644 index 0000000..b0ef0d6 --- /dev/null +++ b/mlx/tests/eval_quant.py @@ -0,0 +1,74 @@ +"""Reproduce the README's 8-bit-vs-bf16 prefill parity numbers. + +Loads the bf16 MLX model (the port baseline) and the converted 8-bit model one at a +time (both don't fit in 24 GB together) and, for each clip, compares the next-token +prefill logits — cosine + argmax. The bf16-vs-8bit delta is the quantization cost. + + python mlx/tests/eval_quant.py --model --mlx-8bit <8bit_dir> \ + --clips '/path/to/clips/*.wav' --n 5 +""" +import argparse +import glob +import gc +import os +import sys + +import numpy as np + +_REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, _REPO) +sys.path.insert(0, os.path.join(_REPO, "mlx")) + +PROMPT = "Analyze this music track. Output CAPTION, KEY, BPM, STRUCTURE." + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True, help="HF checkpoint dir (for bf16 baseline + processor)") + ap.add_argument("--mlx-8bit", required=True, help="converted 8-bit MLX dir") + ap.add_argument("--clips", required=True, help="glob for audio clips") + ap.add_argument("--n", type=int, default=5) + args = ap.parse_args() + + import mlx.core as mx + from moss_music_mlx.convert import load_model, load_pretrained + from src.processing_moss_music import MossMusicProcessor + from src.audio_io import load_audio + + clips = sorted(glob.glob(args.clips))[: args.n] + assert clips, f"no clips matched {args.clips}" + proc = MossMusicProcessor.from_pretrained(args.model, trust_remote_code=True, enable_time_marker=True) + + pre = [] + for c in clips: + inp = proc(text=PROMPT, audios=[load_audio(c, sample_rate=proc.config.mel_sr)], return_tensors="pt") + pre.append((mx.array(inp["input_ids"].numpy()), + mx.array(inp["audio_data"].float().numpy()).astype(mx.bfloat16), + mx.array((inp["input_ids"] == proc.audio_token_id).numpy()), + int(inp["audio_data_seqlens"][0]))) + + def prefill(model): + out = [] + for ids, mel, am, fl in pre: + logits = model(ids, mel=mel, audio_mask=am, feature_len=fl) + mx.eval(logits) + out.append(np.asarray(logits[0, -1].astype(mx.float32))) + return out + + print(f"{len(clips)} clips. loading bf16 baseline...", flush=True) + m = load_model(args.model, dtype=mx.bfloat16); bf16 = prefill(m) + del m; gc.collect(); mx.clear_cache() + print("loading 8-bit...", flush=True) + m = load_pretrained(args.mlx_8bit); q8 = prefill(m); del m; gc.collect() + + cosall, match = [], 0 + print("\n=== bf16 vs 8-bit (prefill next-token logits) ===") + for c, a, b in zip(clips, bf16, q8): + cos = float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9)) + am = int(a.argmax() == b.argmax()); match += am; cosall.append(cos) + print(f" {os.path.basename(c)[:40]:40} cos={cos:.5f} argmax {'=' if am else 'x'}") + print(f"\nmean cos = {np.mean(cosall):.5f} argmax match = {match}/{len(clips)}") + + +if __name__ == "__main__": + main() diff --git a/mlx/tests/test_parity.py b/mlx/tests/test_parity.py new file mode 100644 index 0000000..847b776 --- /dev/null +++ b/mlx/tests/test_parity.py @@ -0,0 +1,153 @@ +"""Parity / self-consistency checks for the MLX MOSS-Music port. + +Two checks: + 1. Prefill parity vs the PyTorch reference (needs torch + the HF checkpoint): + one prefill on a real clip; the next-token argmax and logit cosine must match. + 2. Decode-path (KV-cache) self-consistency (MLX only, fast): a cached step-by-step + decode must match a single full-context forward of the same MLX model at every + position. This validates the KVCache / RoPE-offset bookkeeping (which prefill-only + parity never exercises); it does NOT independently re-check the math against torch + at each decoded position — a bug shared by both MLX paths would cancel out. The + prefill check (1) is what guards the model math. + + python mlx/tests/test_parity.py --model --audio clip.wav [--mlx-model ] + [--decode-steps 24] [--skip-torch] +""" +import argparse +import os +import sys + +import numpy as np + +_REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, _REPO) +sys.path.insert(0, os.path.join(_REPO, "mlx")) + +PROMPT = "Please analyze this music track: genre, instruments, key, BPM and structure." + + +def _cos(a, b): + return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9)) + + +def _processor(model_dir): + from src.processing_moss_music import MossMusicProcessor + return MossMusicProcessor.from_pretrained(model_dir, trust_remote_code=True, enable_time_marker=True) + + +def _mlx_inputs(inp, proc): + import mlx.core as mx + return dict( + input_ids=mx.array(inp["input_ids"].numpy()), + mel=mx.array(inp["audio_data"].float().numpy()).astype(mx.bfloat16), + audio_mask=mx.array((inp["input_ids"] == proc.audio_token_id).numpy()), + feature_len=int(inp["audio_data_seqlens"][0]), + ) + + +def _load_mlx(model_dir, mlx_model_dir): + import mlx.core as mx + from moss_music_mlx.convert import load_model, load_pretrained + return load_pretrained(mlx_model_dir) if mlx_model_dir else load_model(model_dir, dtype=mx.bfloat16) + + +def torch_reference(model_dir, audio_path, proc): + import torch + from src.audio_io import load_audio + from src.modeling_moss_music import MossMusicModel + + model = MossMusicModel.from_pretrained(model_dir, trust_remote_code=True, dtype=torch.float32).eval() + raw = load_audio(audio_path, sample_rate=proc.config.mel_sr) + inp = proc(text=PROMPT, audios=[raw], return_tensors="pt") + audio_mask = inp["input_ids"] == proc.audio_token_id + with torch.no_grad(): + logits = model(input_ids=inp["input_ids"], audio_data=inp["audio_data"].float(), + audio_data_seqlens=inp["audio_data_seqlens"], audio_input_mask=audio_mask, + use_cache=False, return_dict=True).logits + return inp, np.asarray(logits[0, -1].float()) + + +def prefill_parity(model, proc, inp, ref): + import mlx.core as mx + mi = _mlx_inputs(inp, proc) + logits = model(mi["input_ids"], mel=mi["mel"], audio_mask=mi["audio_mask"], feature_len=mi["feature_len"]) + mx.eval(logits) + mlx = np.asarray(logits[0, -1].astype(mx.float32)) + cos = _cos(ref, mlx) + print(f"[prefill] cos={cos:.6f} argmax ref={int(ref.argmax())} mlx={int(mlx.argmax())}") + assert int(ref.argmax()) == int(mlx.argmax()), "prefill next-token argmax mismatch" + assert cos > 0.999, f"prefill logit cosine too low: {cos}" + + +def decode_consistency(model, proc, inp, steps): + """Cached decode logits must equal a full-context forward at every position.""" + import mlx.core as mx + from mlx_lm.models.cache import KVCache + mi = _mlx_inputs(inp, proc) + P = mi["input_ids"].shape[1] + + cache = [KVCache() for _ in model.language_model.layers] + logits = model(mi["input_ids"], mel=mi["mel"], audio_mask=mi["audio_mask"], + feature_len=mi["feature_len"], cache=cache) + cached = [np.asarray(logits[0, -1].astype(mx.float32))] + toks = [] + y = mx.argmax(logits[:, -1, :], axis=-1) + for _ in range(steps): + toks.append(int(y.item())) + logits = model(y[None], cache=cache) + cached.append(np.asarray(logits[0, -1].astype(mx.float32))) + y = mx.argmax(logits[:, -1, :], axis=-1) + + forced = toks[:-1] + full_ids = mx.concatenate([mi["input_ids"], mx.array([forced])], axis=1) if forced else mi["input_ids"] + full_mask = (mx.concatenate([mi["audio_mask"], mx.zeros((1, len(forced)), dtype=mi["audio_mask"].dtype)], axis=1) + if forced else mi["audio_mask"]) + full = model(full_ids, mel=mi["mel"], audio_mask=full_mask, feature_len=mi["feature_len"], cache=None) + mx.eval(full) + + agree, mincos = 0, 1.0 + for i in range(steps): + fl = np.asarray(full[0, P - 1 + i].astype(mx.float32)) + cos = _cos(fl, cached[i]) + mincos = min(mincos, cos) + agree += int(fl.argmax() == cached[i].argmax()) + print(f"[decode] {agree}/{steps} steps argmax-agree, min cos={mincos:.6f}") + # A cache/RoPE-offset bug tanks the cosine and cascades; a lone argmax flip at + # high cosine is just a near-tie token under bf16/8-bit non-associativity. + assert mincos > 0.99, f"decode-path logit cosine too low (cache/RoPE bug?): {mincos}" + assert agree >= int(0.9 * steps), f"cached decode diverges systematically ({agree}/{steps})" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True, help="HF checkpoint dir") + ap.add_argument("--audio", required=True) + ap.add_argument("--mlx-model", default=None, help="converted MLX dir (else build from HF)") + ap.add_argument("--decode-steps", type=int, default=48) + ap.add_argument("--skip-torch", action="store_true", help="run only the MLX decode-path check") + args = ap.parse_args() + + from moss_music_mlx.convert import load_config + proc = _processor(args.model) + model = _load_mlx(args.model, args.mlx_model) + if args.mlx_model: + q = load_config(args.mlx_model).get("quantization") + print(f"[artifact] {args.mlx_model} — {q['bits']}-bit (group {q['group_size']})" if q + else f"[artifact] {args.mlx_model} — bf16 (no quantization)") + else: + print("[artifact] un-quantized bf16 MLX built from the HF checkpoint") + + if not args.skip_torch: + inp, ref = torch_reference(args.model, args.audio, proc) + prefill_parity(model, proc, inp, ref) + else: + from src.audio_io import load_audio + inp = proc(text=PROMPT, audios=[load_audio(args.audio, sample_rate=proc.config.mel_sr)], + return_tensors="pt") + + decode_consistency(model, proc, inp, args.decode_steps) + print("PARITY OK ✅") + + +if __name__ == "__main__": + main() diff --git a/mlx/tests/test_units.py b/mlx/tests/test_units.py new file mode 100644 index 0000000..146b357 --- /dev/null +++ b/mlx/tests/test_units.py @@ -0,0 +1,80 @@ +"""Artifact-free unit tests for the load-bearing pure logic (no model / torch needed). + +Run with: pytest mlx/tests/test_units.py +""" +import os +import sys + +import numpy as np +import mlx.core as mx + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.dirname( + os.path.abspath(__file__)))), "mlx")) + +from moss_music_mlx.audio_encoder import _conv3_downsample_len +from moss_music_mlx.convert import _map_key, _resolve_eos_ids + + +def _chunk_lengths(feature_len, cf=400): + lengths, off = [], 0 + while off < feature_len: + lengths.append(min(cf, feature_len - off)) + off += cf + return lengths + + +def test_audio_token_count_invariant(): + """Per-window downsampled tokens must sum to the processor's full-length count, + so audio_mask placeholders always match the encoder output length.""" + for fl in [1, 8, 50, 399, 400, 401, 799, 800, 801, 1234, 7500, 12000, 36000]: + per_window = sum(_conv3_downsample_len(L) for L in _chunk_lengths(fl)) + assert per_window == _conv3_downsample_len(fl), fl + + +def test_map_key_remap(): + f = "audio_encoder.layers.3.self_attn.q_proj.weight" + assert _map_key(f) == "audio_encoder.layers.3.q_proj.weight" + # must NOT corrupt the similarly-named layer norm + assert _map_key("audio_encoder.layers.3.self_attn_layer_norm.weight") == \ + "audio_encoder.layers.3.self_attn_layer_norm.weight" + # language model keys are untouched + assert _map_key("language_model.layers.0.self_attn.q_proj.weight") == \ + "language_model.layers.0.self_attn.q_proj.weight" + + +def test_resolve_eos_ids(): + nodir = "/nonexistent-dir-xyz" + assert _resolve_eos_ids(nodir, {"eos_token_id": 151645}) == [151645] + assert _resolve_eos_ids(nodir, {"eos_token_id": [151645, 151643]}) == [151643, 151645] + # also pulls from language_config + assert _resolve_eos_ids(nodir, {"language_config": {"eos_token_id": 7}}) == [7] + + +def _mlx_scatter(h, src, mask): + """The exact cumsum-scatter used in model.py.""" + mask_row = mask[0] + idx = mx.maximum(mx.cumsum(mask_row.astype(mx.int32)) - 1, 0) + gathered = src[0][idx] + return mx.where(mask_row[:, None], gathered, h[0])[None] + + +def test_masked_scatter_equivalence(): + """cumsum-scatter must equal PyTorch-style masked_scatter_ (row-major fill).""" + rng = np.random.default_rng(0) + D = 4 + patterns = [ + [0, 1, 1, 1, 0], # contiguous middle + [1, 1, 0, 0, 0], # leading audio + [0, 0, 0, 1, 1], # trailing audio + [1, 0, 1, 0, 1], # interleaved + [1, 1, 1, 1, 1], # all audio + ] + for p in patterns: + mask = np.array([p], dtype=bool) + n = int(mask.sum()) + h = rng.standard_normal((1, len(p), D)).astype(np.float32) + src = rng.standard_normal((1, n, D)).astype(np.float32) + ref = h.copy() + ref[0, mask[0]] = src[0] # numpy row-major masked scatter + got = np.asarray(_mlx_scatter(mx.array(h), mx.array(src), mx.array(mask))) + assert np.allclose(got, ref, atol=1e-6), p diff --git a/pyproject.toml b/pyproject.toml index 4d81bfa..5896507 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,11 +34,26 @@ torch-runtime = [ "transformers==4.57.1", ] +# Apple Silicon (MLX) inference backend — see mlx/README.md. +# torch/torchaudio/transformers are CPU-only here (used by the shared processor for +# mel features); the model itself runs on MLX. Numbers in mlx/README.md were produced +# with mlx==0.31.2 and mlx-lm==0.29.1. +mlx = [ + "mlx>=0.31", + "mlx-lm>=0.29", + "torch>=2.4", + "torchaudio>=2.4", + "transformers>=4.57", +] + [project.urls] Homepage = "https://github.com/OpenMOSS/MOSS-Music" Repository = "https://github.com/OpenMOSS/MOSS-Music" [tool.setuptools] +# Discover both the core `src` package and the MLX backend package under `mlx/`, +# so `pip install -e .[mlx]` makes `moss_music_mlx` importable / `-m`-runnable. [tool.setuptools.packages.find] -include = ["src", "src.*"] +where = [".", "mlx"] +include = ["src*", "moss_music_mlx*"]