Fix Qwen3.6 converted RMSNorm double shift - #1623

Rebase the native Qwen MTP implementation onto upstream main after ml-explore#1623, retaining the merged RMSNorm conversion fix and stronger upstream regression coverage.

Signed-off-by: Philip John Basile <[email protected]>

… mlx-lm 0.31 loads (#2210)

mlx-lm 0.31.3's Qwen3.5 / Qwen3.6 `TextModel.sanitize` decides whether
to add `+1` to the RMSNorm weights with `has_mtp_weights or
has_unsanitized_conv1d`. A checkpoint that was already converted to MLX
layout but still carries its MTP head (`mtp.*` or `language_model.mtp.*`
tensors, as in MLX conversions such as
`Jundot/Qwen3.6-35B-A3B-oQ4e-mtp`) therefore gets the shift a second
time. Every norm loads at about twice its value, and generation is
garbage from the first token.

mlx-lm fixed this in ml-explore/mlx-lm#1623 (released in 0.32.0) by
deciding on the raw Conv1D layout alone; see also
ml-explore/mlx-lm#1197. unsloth-zoo still accepts `mlx-lm>=0.31.3`, so
installs on 0.31.x keep the bug.

## Which Unsloth paths hit it

Measured with real cached checkpoints plus an injected `mtp.*` shard, on
mlx-lm 0.31.3:

| load path | Qwen3.5-0.8B-bf16 | Qwen3.6-35B-A3B-8bit |
|---|---|---|
| `mlx_lm.load` | norms +1.000, garbage | norms +1.000, garbage |
| `FastMLXModel.from_pretrained` (`text_only` True or False) | correct
(routed to mlx-vlm, which strips `mtp.*` first) | correct |
| `unsloth_zoo.mlx.distill.load_teacher` (GKD teacher, plain
`mlx_lm.load`) | **norms +1.000, garbage** | **norms +1.000, garbage** |

So the user-facing exposure is distillation with such a teacher, plus
any `FastMLXModel` load that reaches the mlx-lm text branch.

## Change

`_ensure_qwen3_5_mtp_norm_sanitize()` wraps
`mlx_lm.models.qwen3_5.TextModel.sanitize` so that `mtp.*` keys are
dropped before the original runs. The original drops them anyway, so the
only behavior that changes is the shift decision, which then matches
mlx-lm 0.32.0. `qwen3_5_moe` delegates to the same `TextModel.sanitize`,
so it is covered too.

- The patch is installed from `load_teacher` and from `FastMLXModel`'s
mlx-lm branch, next to `_ensure_safe_text_wrapper_sanitize`.
- It applies only while the installed sanitize still keys on MTP
weights; on mlx-lm 0.32.x it is a no-op. It is also idempotent.

## Validation

```text
pytest tests/test_mlx_distributed_loader.py tests/test_mlx_model_file_consent.py   # 27 passed
```

- Tests cover the behavior through `load_teacher` for both converted and
raw Conv1D layouts (raw checkpoints still get the shift), using
stand-ins carrying the 0.31.3 and 0.32.0 sanitize bodies. The
`FastMLXModel` text-load test asserts the install.
- A mutation check caught 5 of 5 deliberate breakages: filter removed,
source gate removed, idempotency removed, and each call site removed.
- Real runtime on mlx-lm 0.31.3: with the patch, `load_teacher` loads
norms equal to the checkpoint and answers "Paris" on both models above,
while plain `mlx_lm.load` without Unsloth still shows +1.000 and
garbage.
- On mlx-lm 0.32.2 the gate sees the fixed sanitize and leaves it
untouched.

---------

Co-authored-by: Daniel Han <[email protected]>