Update README: mixed-precision (8/4) at 92.2% recommended

README.md

@@ -16,8 +16,8 @@ tags:
 - quantized
 - int4
 - 4bit
+- mixed-precision
 - dwq
-- distillation-aware-quantization
 - on-device
 - apple-silicon
 - qwen3
@@ -28,23 +28,23 @@ base_model: zhifeixie/Mega-ASR
 base_model_relation: quantized
 ---
 
-# Mega-ASR — MLX 4-bit DWQ
+# Mega-ASR — MLX 4-bit
 
-MLX 4-bit (DWQ-distilled) deployment of [zhifeixie/Mega-ASR](https://huggingface.co/zhifeixie/Mega-ASR),
+MLX deployment of [zhifeixie/Mega-ASR](https://huggingface.co/zhifeixie/Mega-ASR),
 the 1.7B-parameter robust multilingual ASR foundation model built on Qwen3-ASR-1.7B.
 
-Quantized via [`mlx_lm.quant.dwq`](https://github.com/ml-explore/mlx-lm) (Distillation-aware
-Weight Quantization) — the 4-bit weights are fine-tuned to match the bf16 teacher's
-logits on a calibration corpus, recovering most of the post-training drift.
+Two LLM variants ship in this repo. The **recommended** one is the mixed-precision
+build — 8-bit attention + 4-bit MLP layers — which closes the quality gap to ONNX
+GPTQ at the smallest viable size.
 
 ## What's in this repo
 
 | File | Size | Role |
 | --- | ---: | --- |
-| `mlx/llm-dwq4/model.safetensors` | **923 MB** | Qwen3 1.7B LLM, MLX 4-bit DWQ, group_size=64 (4.5 bits/weight) |
-| `mlx/llm-dwq4/{config,tokenizer*}.json` | — | mlx-lm load metadata + LLM tokenizer |
-| `onnx/audio_encoder_fp32.onnx` | **1.27 GB** | 24-layer Whisper-style audio encoder (ONNX fp32, run via onnxruntime). MLX port is planned. |
-| `tokenizer/*` | — | Original Qwen3-ASR tokenizer (with audio special tokens). Required because the LLM-only tokenizer drops `<\|audio_pad\|>`, `<\|audio_start\|>`, etc. |
+| `mlx/llm-mixed8_4/` | **1.5 GB** | **Recommended** Qwen3 LLM, 8-bit attention + 4-bit MLP (5.0 bpw avg) |
+| `mlx/llm-dwq4/` | 923 MB | 4-bit DWQ-distilled (smallest, slight quality drop) |
+| `onnx/audio_encoder_fp32.onnx` | 1.27 GB | 24-layer Whisper-style audio encoder (ONNX fp32, run via onnxruntime). MLX port is on the roadmap. |
+| `tokenizer/*` | — | Original Qwen3-ASR tokenizer (with audio special tokens `<\|audio_pad\|>` etc.) |
 | `examples/*.wav` | ~3 MB | 8 noisy benchmark clips from Voices-in-the-Wild-Bench |
 | `inference.py` | — | End-to-end ASR pipeline: ONNX encoder + MLX LLM |
 
@@ -53,17 +53,38 @@ logits on a calibration corpus, recovering most of the post-training drift.
 8-clip [Voices-in-the-Wild-Bench](https://github.com/xzf-thu/Voices-in-the-Wild-Bench)
 agreement (1 − WER), prompt forced to `language English`:
 
-| Variant | Encoder | Decoder | Agreement | Total size |
-| --- | --- | --- | ---: | ---: |
-| PT bf16 (original) | fp16 | fp16 | 95.1% | 7.5 GB |
-| ONNX recommended (GPTQ) | INT8 ONNX | INT4 GPTQ | 92.7% | 2.3 GB |
-| **MLX 4-bit DWQ (this repo)** | **fp32 ONNX** | **MLX 4-bit DWQ** | **89.9%** | **~2.2 GB** |
-| MLX 4-bit (no DWQ) | fp32 ONNX | MLX 4-bit | 89.1% | ~2.2 GB |
+| Variant | Encoder | LLM | Bpw | Agreement | Total size |
+| --- | --- | --- | ---: | ---: | ---: |
+| PT bf16 (original) | fp16 | fp16 | 16 | 95.1% | 7.5 GB |
+| ONNX recommended (GPTQ) | INT8 ONNX | INT4 GPTQ | ~4.5 | 92.7% | 2.3 GB |
+| **MLX recommended (mixed)** | **fp32 ONNX** | **MLX 8/4 mixed** | **5.0** | **92.2%** | **~2.8 GB** |
+| MLX 4-bit DWQ | fp32 ONNX | MLX 4-bit DWQ | 4.5 | 89.9% | ~2.2 GB |
+| MLX 4-bit (no DWQ) | fp32 ONNX | MLX 4-bit | 4.5 | 89.1% | ~2.2 GB |
+
+The mixed variant gets all 6 "easy" samples perfect and improves the 2 hard
+samples (`echo`, `recording`) — only the audio-quality-limited tail remains.
+
+### Why mixed precision
 
-DWQ recovers +0.8% over straight mlx-lm 4-bit quantization. Audio encoder is
-still ONNX fp32 because the `qwen3_asr_audio_encoder` architecture (custom
-chunked Whisper-style with Conv2d front-end) isn't natively supported by
-mlx-lm or mlx-whisper yet — porting it to pure MLX is on the roadmap.
+Pure 4-bit MLX hits a quality wall around 89% because mlx-lm's affine
+quantization is naive groupwise (no calibration, no GPTQ-style error
+redistribution). Attention layers are the most quality-sensitive in Qwen3 —
+keeping them at 8-bit while dropping MLP layers to 4-bit recovers all the
+4-bit quality loss at only ~12% more weight memory than uniform 8-bit.
+
+| Variant | Attention | MLP | Bpw | Agreement |
+| --- | --- | --- | ---: | ---: |
+| pure 4-bit | 4-bit | 4-bit | 4.5 | 89.1% |
+| **mixed 8/4** | **8-bit** | **4-bit** | **5.0** | **92.2%** |
+| mixed 8/6 | 8-bit | 6-bit | 6.5 | 91.4% |
+| 6-bit | 6-bit | 6-bit | 6.5 | 90.7% |
+| 8-bit | 8-bit | 8-bit | 8.5 | 92.2% |
+
+The mixed 8/4 build is Pareto-optimal — same quality as full 8-bit at ~60%
+of its size, and 2.3 percentage points higher agreement than DWQ-distilled
+4-bit. DWQ on plain-text data couldn't bridge the gap because Mega-ASR's
+inference distribution (scattered audio embeddings into a text prompt) is
+out-of-distribution for the bf16 teacher's plain-text calibration corpus.
 
 ## Inference
 
@@ -72,13 +93,13 @@ pip install mlx mlx-lm onnxruntime soundfile transformers librosa numpy
 git clone https://huggingface.co/Reza2kn/mega-asr-mlx
 cd mega-asr-mlx
 python inference.py --encoder-path onnx/audio_encoder_fp32.onnx \
-                    --mlx-llm-path mlx/llm-dwq4 \
+                    --mlx-llm-path mlx/llm-mixed8_4 \
                     --examples-dir examples
 ```
 
-The pipeline runs:
+Pipeline:
 1. mel features (Whisper preprocessor) → ONNX audio encoder (onnxruntime CPU) → audio embeddings (1, F, 2048)
-2. tokenize the Qwen3-ASR chat prompt with `audio_pad_id=151676`, expand the single pad placeholder to F copies
+2. tokenize the Qwen3-ASR chat prompt with `audio_pad_id=151676`, expand the single placeholder to F copies
 3. embed all tokens via `model.model.embed_tokens` (MLX), scatter audio embeddings at the audio_pad positions
 4. greedy decode via MLX with `input_embeddings`
 
@@ -87,24 +108,27 @@ The pipeline runs:
 - LLM extracted from `zhifeixie/Mega-ASR/Qwen3-ASR-1.7B/` by stripping the
   `thinker.model.` prefix from layer weights and dropping the tied `lm_head`
   (relies on `tie_word_embeddings=True`).
-- `mlx_lm.convert --hf-path <extracted> -q --q-bits 4 --q-group-size 64` →
-  base MLX 4-bit (4.501 bits/weight).
-- `mlx_lm.quant.dwq --model <bf16-teacher> --bits 4 --group-size 64
-  --num-samples 64 --max-seq-length 256 --learning-rate 1e-6` →
-  DWQ-distilled weights. 64 steps of distillation on tulu-3-sft-mixture
-  reduced validation KL loss from ~0.18 to ~0.14.
+- **Mixed-precision quant** via `mlx_lm.utils.quantize_model` with a
+  per-layer `quant_predicate`:
+  - q_proj / k_proj / v_proj / o_proj → 8-bit
+  - gate_proj / up_proj / down_proj → 4-bit
+  - group_size=64, mode=affine
+- **DWQ variant** via `mlx_lm.quant.dwq --bits 4 --group-size 64
+  --num-samples 64 --max-seq-length 256 --learning-rate 1e-6`. 64 distillation
+  steps on tulu-3-sft-mixture reduced KL loss from ~0.18 to ~0.14.
 - Audio encoder ONNX is reused unchanged from
   [Reza2kn/mega-asr-onnx](https://huggingface.co/Reza2kn/mega-asr-onnx).
 
 ## Performance
 
-| Hardware | Cold load | Warm (3-4s audio) |
+| Hardware | Cold load | Warm (3-4 s audio) |
 | --- | ---: | ---: |
-| M-series Mac (MLX) | ~3 s | ~1.5 s (LLM @ ~60 tps) |
+| M-series Mac (MLX, mixed8_4) | ~3 s | ~1.5 s (LLM @ ~50 tps) |
+| M-series Mac (MLX, dwq4) | ~3 s | ~1.5 s (LLM @ ~60 tps) |
 
 ## Credits
 
 - Original model: [zhifeixie/Mega-ASR](https://huggingface.co/zhifeixie/Mega-ASR) (1.7B params, Apache-2.0)
-- MLX port + DWQ: this repo
+- MLX port + mixed-precision + DWQ: this repo
 - Benchmark: [Voices-in-the-Wild-Bench](https://github.com/xzf-thu/Voices-in-the-Wild-Bench)
 - DWQ tool: [`mlx_lm.quant.dwq`](https://github.com/ml-explore/mlx-lm) (Apple Inc.)