vendoring mlx-lm components

.github/workflows/tests.yml

@@ -48,6 +48,24 @@ jobs:
       - name: Run core tests
         run: pytest -s mlx_audio/tests/
 
+  # Differential tests: vendored mlx_audio.lm vs upstream mlx-lm.
+  parity:
+    runs-on: macos-14
+    needs: style
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: '3.10'
+
+      - name: Install with parity extra
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[all,dev,parity]"
+
+      - name: Run parity tests
+        run: pytest -s tests/vendor_parity/
+
   # Modular installation tests - validates issue #287
   # Only verifies imports work with minimal deps installed.
   # Full tests run separately with all deps (test files import models that need extra deps).

docs/contributing/adding-a-model.md

@@ -46,6 +46,26 @@ TTS models use the base classes from `mlx_audio/tts/models/base.py`:
 - **`GenerationResult`** -- Dataclass returned by `generate()`. Contains `audio`, `sample_rate`, `token_count`, timing information, and streaming flags.
 - **`BatchGenerationResult`** -- Dataclass for batch generation results.
 
+### Transformer Components
+
+Many speech models wrap a language-model backbone. Import that machinery from
+`mlx_audio.lm`, **not** from `mlx-lm` -- mlx-audio vendors it so the package has
+no runtime dependency on mlx-lm, and a test fails the build if `mlx_lm` is
+imported anywhere outside `sts/voice_pipeline.py`.
+
+| Need | Import from |
+|---|---|
+| KV caches (`KVCache`, `RotatingKVCache`, `BatchKVCache`, `make_prompt_cache`) | `mlx_audio.lm.models.cache` |
+| Attention masks, `scaled_dot_product_attention` | `mlx_audio.lm.models.base` |
+| Samplers, logits processors | `mlx_audio.lm.sample_utils` |
+| `generate_step`, `stream_generate` | `mlx_audio.lm.generate` |
+| Backbones (llama, qwen2, qwen3, gpt2, granite, lfm2, bailing_moe, gemma3) | `mlx_audio.lm.models.<name>` |
+| Quantize / save helpers | `mlx_audio.lm.convert` |
+
+If your model needs a backbone that is not vendored yet, copy it from mlx-lm
+into `mlx_audio/lm/models/` verbatim and add the provenance header used by the
+other files there (upstream path, version, commit).
+
 ### Model Configuration
 
 Create a dataclass for your model's config that extends `BaseModelArgs`:

mlx_audio/codec/models/mimi/modules/__init__.py

@@ -4,7 +4,7 @@
 # flake8: noqa
 """Modules used for building the models."""
 
-from mlx_lm.models.cache import KVCache, RotatingKVCache
+from mlx_audio.lm.models.cache import KVCache, RotatingKVCache
 
 from .conv import (
     Conv1d,

mlx_audio/codec/models/mimi/modules/transformer.py

@@ -8,7 +8,8 @@
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx_lm.models.cache import KVCache, RotatingKVCache
+
+from mlx_audio.lm.models.cache import KVCache, RotatingKVCache
 
 
 @dataclass

mlx_audio/convert.py

@@ -486,7 +486,7 @@ def base_requirements(path: str, module) -> bool:
     if not quant_predicate_name:
         return base_requirements
 
-    from mlx_lm.convert import mixed_quant_predicate_builder
+    from mlx_audio.lm.convert import mixed_quant_predicate_builder
 
     mixed_predicate = mixed_quant_predicate_builder(quant_predicate_name, model)
     return lambda p, m: base_requirements(p, m) and mixed_predicate(p, m)
@@ -590,7 +590,12 @@ def convert(
         q_mode: Quantization mode (affine, mxfp4, nvfp4, mxfp8).
         model_domain: Force model domain ("tts", "stt", or "sts"). Auto-detected if None.
     """
-    from mlx_lm.utils import dequantize_model, quantize_model, save_config, save_model
+    from mlx_audio.lm.convert import (
+        dequantize_model,
+        quantize_model,
+        save_config,
+        save_model,
+    )
 
     if quantize and dequantize:
         raise ValueError("Choose either quantize or dequantize, not both.")

mlx_audio/lm/__init__.py

@@ -0,0 +1 @@
+"""Transformer machinery vendored from mlx-lm. Import submodules directly."""

mlx_audio/lm/convert.py

@@ -0,0 +1,184 @@
+# Copyright © 2023-2024 Apple Inc.
+# Derived from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd).
+
+import copy
+import json
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_map, tree_unflatten
+
+MAX_FILE_SIZE_GB = 5
+
+
+def mixed_quant_predicate_builder(recipe: str, model: nn.Module, group_size: int = 64):
+    recipes = {
+        "mixed_2_6": (2, 6),
+        "mixed_3_4": (3, 4),
+        "mixed_3_6": (3, 6),
+        "mixed_4_6": (4, 6),
+    }
+    if recipe not in recipes:
+        raise ValueError(f"Invalid quant recipe {recipe}")
+    low_bits, high_bits = recipes[recipe]
+    down_keys = [name for name, _ in model.named_modules() if "down_proj" in name]
+    if not down_keys:
+        raise ValueError("Model does not have expected keys for mixed quant.")
+    layer_location = next(
+        index for index, key in enumerate(down_keys[0].split(".")) if key.isdigit()
+    )
+    num_layers = len(model.layers)
+
+    def predicate(path: str, module: nn.Module) -> Union[bool, dict]:
+        del module
+        index = (
+            int(path.split(".")[layer_location])
+            if len(path.split(".")) > layer_location
+            else 0
+        )
+        high_precision = (
+            index < num_layers // 8
+            or index >= 7 * num_layers // 8
+            or (index - num_layers // 8) % 3 == 2
+        )
+        wide = (
+            "v_proj" in path
+            or "v_a_proj" in path
+            or "v_b_proj" in path
+            or "down_proj" in path
+        )
+        # lm_head takes high bits regardless of depth, as upstream does.
+        bits = high_bits if (wide and high_precision) or "lm_head" in path else low_bits
+        return {"group_size": group_size, "bits": bits, "mode": "affine"}
+
+    return predicate
+
+
+def quantize_model(
+    model: nn.Module,
+    config: dict,
+    group_size: Optional[int],
+    bits: Optional[int],
+    mode: str = "affine",
+    quant_predicate: Optional[Callable] = None,
+):
+    defaults = {"affine": (64, 4), "mxfp4": (32, 4), "nvfp4": (16, 4), "mxfp8": (32, 8)}
+    group_size, bits = group_size or defaults[mode][0], bits or defaults[mode][1]
+    config = copy.deepcopy(config)
+    quant_predicate = quant_predicate or getattr(model, "quant_predicate", None)
+    params = {"group_size": group_size, "bits": bits, "mode": mode}
+
+    # An existing "quantization" key means the model is already partially
+    # quantized, so record parameters per layer rather than globally.
+    fine_grained = "quantization" in config
+    if not fine_grained:
+        config["quantization"] = params
+
+    def predicate(path, module):
+        if not hasattr(module, "to_quantized") or module.weight.shape[-1] % group_size:
+            return False
+        result = quant_predicate(path, module) if quant_predicate else True
+        if isinstance(result, dict):
+            config["quantization"][path] = result
+        elif fine_grained and result:
+            config["quantization"][path] = params
+        return result
+
+    nn.quantize(model, group_size, bits, mode=mode, class_predicate=predicate)
+    config["quantization_config"] = config["quantization"]
+    return model, config
+
+
+def dequantize_model(model: nn.Module) -> nn.Module:
+    replacements = []
+    for name, module in model.named_modules():
+        if isinstance(module, nn.QuantizedLinear):
+            layer = nn.Linear(*module.weight.shape[::-1], bias="bias" in module)
+        elif isinstance(module, nn.QuantizedEmbedding):
+            layer = nn.Embedding(*module.weight.shape)
+        else:
+            continue
+        layer.weight = mx.dequantize(
+            module.weight,
+            module.scales,
+            module.biases,
+            module.group_size,
+            module.bits,
+            module.mode,
+        )
+        if "bias" in module:
+            layer.bias = module.bias
+        replacements.append((name, layer))
+    if replacements:
+        model.update_modules(tree_unflatten(replacements))
+    return model
+
+
+def save_config(config: dict, config_path: Union[str, Path]) -> None:
+    config = copy.deepcopy(config)
+    config.pop("_name_or_path", None)
+    config.pop("vision_config", None)
+    if "quantization" in config:
+        config["quantization_config"] = config["quantization"]
+    with open(config_path, "w") as handle:
+        json.dump(dict(sorted(config.items())), handle, indent=4)
+
+
+def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list:
+    max_file_size_bytes = max_file_size_gb << 30
+    shards = []
+    shard, shard_size = {}, 0
+    for name, weight in weights.items():
+        if shard_size + weight.nbytes > max_file_size_bytes:
+            shards.append(shard)
+            shard, shard_size = {}, 0
+        shard[name] = weight
+        shard_size += weight.nbytes
+    shards.append(shard)
+    return shards
+
+
+def save_model(
+    save_path: Union[str, Path], model: nn.Module, *, donate_model: bool = False
+) -> None:
+    save_path = Path(save_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    weights = dict(tree_flatten(model.parameters()))
+    total_size = sum(value.nbytes for value in weights.values())
+    shards = make_shards(weights)
+    name_format = (
+        "model-{:05d}-of-{:05d}.safetensors" if len(shards) > 1 else "model.safetensors"
+    )
+    weight_map = {
+        name: name_format.format(index + 1, len(shards))
+        for index, shard in enumerate(shards)
+        for name in shard
+    }
+
+    # Release the model's references before serializing so each shard can be
+    # freed as it is written, rather than holding every weight twice.
+    if donate_model:
+        model.update(tree_map(lambda _: mx.array([]), model.parameters()))
+    weights.clear()
+
+    for index, shard in enumerate(shards):
+        shards[index] = None
+        mx.save_safetensors(
+            str(save_path / name_format.format(index + 1, len(shards))),
+            shard,
+            metadata={"format": "mlx"},
+        )
+        del shard
+
+    with open(save_path / "model.safetensors.index.json", "w") as handle:
+        json.dump(
+            {
+                "metadata": {"total_size": total_size},
+                "weight_map": {k: weight_map[k] for k in sorted(weight_map)},
+            },
+            handle,
+            indent=4,
+        )