mlx_technical_summary.md

# Porting **MobileLLM-R1-950M** to MLX and mlx-lm: Architectural Challenges and Solutions

I spent a some time pairing with Gemini 2.5 Pro and later OpenAI Codex to drag the brand-new facebook/MobileLLM-R1-950M weights onto Apple Silicon.
This write-up is the “why it wasn’t copy-paste” story, plus the gotchas that bit us until the model finally spoke clean English and quantized without drama.

### Goal

Enable **facebook/MobileLLM-R1-950M** to run natively on Apple Silicon using MLX, then create quantized versions compatible with the mlx-lm ecosystem.

---

## 1. Why a Direct "Llama-4 Drop-In" Failed

Although the Hugging Face repo presents MobileLLM-R1-950M as a Llama-4-style dense model, its **config and weights don't align cleanly** with a stock Llama block. The deviations aren't quirks of MLX—they reflect this model's specific architecture:

* **MLP ambiguity**  
  Config advertises both `intermediate_size` and `intermediate_size_mlp`, suggesting a dual-branch feed-forward.  
  Actual weights contain only a SwiGLU branch (`gate_proj`, `up_proj`, `down_proj`).  
  → Solution: **auto-detect MLP variant from weight names** at load time.

* **Grouped-Query Attention (GQA)**  
  `num_attention_heads=24`, `num_key_value_heads=6`.  
  K/V tensors must be **repeated to full head count** for attention shapes to align correctly.

* **QK-norm and scaling**  
  Config includes `use_qk_norm=True` and `attn_scale=0.1`.  
  We add the **RMSNorm on Q/K** as specified, but drop the extra `0.1` multiplier—applying it in MLX's `scaled_dot_product_attention` collapses logits into gibberish.

* **RoPE gating**  
  Config lists all layers under `no_rope_layers`.  
  Disabling RoPE everywhere would eliminate positional encoding entirely.  
  → Treat "all layers disabled" as a config artifact and **apply RoPE everywhere**.

---

## 2. Prompt-Level Deviations

Even after weights loaded correctly, default inference was disrupted by tokenizer settings:

* **Chat template**  
  Default system prompt: *"Please reason step-by-step and put your final answer within \boxed{}."*  
  Without overrides, the model produces verbose "reasoning" outputs.  
  → Added CLI controls: `--system`, `--disable-chat-template`, `--final-only`.

* **Double BOS**  
  Both tokenizer and template inserted BOS tokens.  
  → Fixed with `add_special_tokens=False`.

* **Premature EOS**  
  Template headers (`<|eot_id|>`) were treated as stop tokens.  
  → Limited stopping criteria to true EOS token only.

---

## 3. Sampling Stability

Sampling issues stemmed from API mismatches rather than model problems:

* **Top-p on probabilities** then feeding `mx.random.categorical` produced repetition loops.  
* **Solution:** Apply penalties → scale logits → top-p mask (with `float('-inf')`) → `categorical(logits)`.  
* Added controls for **temperature, repetition penalty, frequency penalty**.

---

## 4. Quantization in mlx-lm: Why Custom Metadata Was Required

mlx-lm provides quantization hooks, but MobileLLM's architecture exposed several challenges:

1. **Frozen gradients during sensitivity analysis** → empty sensitivity lists.  
   → Avoid freezing weights during gradient computation.

2. **Re-quantizing quantized layers** → type errors on second pass.  
   → Skip `QuantizedLinear` layers if already quantized.

3. **Embedding/norm dtype crashes**  
   Standard quantization re-quantized everything, but embeddings must remain float.  
   → Introduced **metadata-driven approach**: config.json records *per-layer bit-widths*. Only specified layers are instantiated as `QuantizedLinear`.

This metadata contract allows **4-bit mixed-precision MobileLLM** to be loaded cleanly by our **metadata-aware `custom_loader.py`**, making it compatible with the mlx-lm ecosystem.

---

## 5. End State

* **MLX path:**  
  Structural fixes (GQA, MLP detection), numerical fixes (QK-norm, RoPE, attn_scale), and prompt controls together yield fluent, stable inference.

* **mlx-lm path:**  
  Custom quantization pipeline produces FP16 and 4-bit models. These can be loaded with our **metadata-aware `custom_loader.py`** and used for inference with our provided scripts.  
  Performance: measurable speedup and reduced VRAM usage on Apple Silicon, with minimal quality degradation.

---

### Takeaway

The MobileLLM-R1-950M port required systematically addressing architectural mismatches (MLP variant detection, GQA handling, QK-norm implementation, RoPE configuration) and developing a metadata-driven quantization approach. Once these were resolved, the model became fully functional in MLX with both float and quantized inference paths.