custom_mlx_lm/README.md

Custom MLX-LM Conversion, Quantization, and Inference

Overview

• Scripts here convert the HF safetensors model to MLX format, optionally apply mixed-precision dynamic quantization, and run inference with prompt formatting consistent with inference.py.
• Quant layout is persisted in config.json so the loader can re-materialize only Linear layers as QuantizedLinear while keeping embeddings and norms in float.

Key scripts

• custom_convert_2.py
  • Convert and optionally quantize.
  • Mixed precision uses calibration data and a sensitivity-driven split between 4-bit and 8-bit Linear layers.
  • Saves weights to weights.npz and writes quantization metadata to config.json.
• custom_loader.py
  • Loads the model with the correct module types (QuantizedLinear vs float) based on config metadata, then applies saved weights.
  • Leaves embeddings and layernorms in float.
• inference_mlx_lm.py (CLI: mobilellm-infer)
  • Runs generation. Uses chat_template.jinja when present, else prepends BOS, matching inference.py behavior.
• quant_summary.py
  • Prints a summary of per-layer bit-widths and checks quantized tensors exist in weights.npz.

Quickstart

• Mixed-precision dynamic quantization
  • uv run python custom_mlx_lm/custom_convert_2.py --hf-path . --mlx-path MobileLLM-R1-950M-mixed-4bit-mlx --dynamic-quant --target-bpw 4.5 --report-ppl
  • Group size defaults to 64 when not provided.
• Uniform quantization
  • uv run python custom_mlx_lm/custom_convert_2.py --hf-path . --mlx-path MobileLLM-R1-950M-4bit-mlx --quantize --bits 4 --report-ppl
• Summarize quant layout
  • uv run python custom_mlx_lm/quant_summary.py --model-path MobileLLM-R1-950M-mixed-4bit-mlx --show 8
• Inference
  • mobilellm-infer --model-path MobileLLM-R1-950M-mixed-4bit-mlx --prompt "What is the nearest prime to 9^2?"

Notes and defaults

• Calibration: load_data uses WikiText-like data; dynamic quant computes sensitivities once and chooses 4/8-bit per Linear layer to target the requested bits-per-weight. Reported PPL is from the same set.
• Group size: defaults to 64 when quantizing if not provided.
• Prompt formatting: by default uses chat_template.jinja if present; otherwise prepends BOS for stable behavior across float and quant models.

Troubleshooting

• Empty sensitivities (ValueError: min() arg is empty)
  • Fixed: ensure Linear weights are not frozen during sensitivity estimation; grads must exist.
• Unable to quantize model of type QuantizedLinear
  • Fixed: second quantization pass now targets only remaining float Linear layers.
• [dequantize] The matrix should be given as a uint32
  • Fixed: loader does not blanket-quantize; it re-materializes only Linear layers from per-layer bits map before loading weights, leaving embeddings in float.

Rationale and behavior

• Persist per-layer bits: enables deterministic, loader-driven reconstruction of quant modules and prevents accidental quantization of unsupported modules.
• Keep embeddings float: avoids dtype mismatch and preserves quality.
• Match inference.py formatting: improves output consistency between float and quant variants.