Skip to content

fix(models): PhiMoE norm epsilon key and SAM antialiased bicubic a - #2507

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-2465-phimoe-sam
Oct 20, 2026
Merged

inureyes merged 2 commits into
mainfrom
fix/issue-2465-phimoe-sam

Conversation

@inureyes

Copy link
Copy Markdown
Member

Closes #2465

Summary

Two numeric-parity fixes in one PR per the issue: PhiMoE reads the checkpoint's norm-epsilon key, and the antialiased bicubic kernel uses ATen's antialiased coefficient.

What changed

  • src/models/phimoe.rs: layer_norm_eps now carries #[serde(alias = "rms_norm_eps")]. Phi-3.5-MoE checkpoints spell it rms_norm_eps: 1e-5 (the key mlx-lm's phimoe.py reads); previously the checkpoint value was ignored and the LayerNorms ran at the 1e-6 default. The layer_norm_eps spelling keeps parsing and 1e-6 stands when neither key exists.
  • src/vision/encoders/resample_matrix.rs: interp_matrix picks a = -0.5 for a cubic kernel whenever antialias is set (ATen's HelperInterpCubic::aa_filter, what mlx-vlm uses) and keeps Keys' -0.75 without it. Non-antialiased output is bit-identical (pinned bytes in test).
  • tests/fixtures/deepseekocr_sam/: torch-generated reference for the AA cubic matrix (16 to 5, 16 to 8, 16 to 24) and the SAM grid resample, from generate_reference.py with torch 2.14.1+cpu and mlx-vlm 0.7.8 (73d366ec), which agree within 9.8e-07.

Deviation from the proposed scale > 1 gate

Torch 2.14.1 uses a = -0.5 for the antialiased kernel in both directions, not only when downsampling: at the 16 to 24 upsample -0.5 matches torch within 6.1e-07 while -0.75 misses by 3.6e-02; at 16 to 5, 1.0e-07 vs 1.1e-02. Implemented as -0.5 whenever a cubic kernel is antialiased; support still stretches only when downsampling. No production caller upsamples with antialias (SAM downsamples 1024 to 64), so no existing behavior changes beyond the AA-downsample path.

Checkpoint comparison

On models/mlx/phi-3.5-moe-4bit with the prompt "The capital of France is" and 16 greedy tokens, mlx-lm (CPU-only mlx 0.32.4 from the oracle venv) produces Paris.\n\nIncorrect grammar: Me and my friend is going to the while mlxcel (CUDA MLX pinned 243bdf09, GPU) produces Paris. It is not only the country's largest; both start Paris. and diverge at token 3. A control run of mlxcel with the checkpoint's rms_norm_eps stripped (the pre-fix 1e-6 default) produces the identical stream, so the epsilon change does not alter generation on this prompt and the divergence is the cross-device numeric path, recorded rather than claimed as token-exact.

Test plan

  • cargo test --profile test-fast phimoe -- --test-threads=1 (15 passed)
  • cargo test --profile test-fast resample_matrix -- --test-threads=1 (9 passed)
  • cargo test --profile test-fast deepseekocr -- --test-threads=1 (22 passed)
  • cargo fmt --all -- --check and cargo clippy --profile test-fast -- -D warnings (clean)

Phi-3.5-MoE checkpoints spell the norm epsilon `rms_norm_eps: 1e-5`, which is the key mlx-lm's phimoe.py reads, but ModelArgs only read `layer_norm_eps`, so the checkpoint value was ignored and the norms ran at the 1e-6 default. Accept both spellings with a serde alias and keep 1e-6 when neither key exists.

The SAM vision encoder's antialiased bicubic resize used Keys' a = -0.75, while ATen's antialiased kernel (HelperInterpCubic::aa_filter) and mlx-vlm use a = -0.5 for PIL compatibility. interp_matrix now picks a = -0.5 for a cubic kernel whenever antialias is set and keeps -0.75 without it, so the non-antialiased output stays bit for bit; MoonViT's edge-replicated grid pins those bytes.

Torch 2.14.1 uses a = -0.5 for the antialiased kernel in both directions, which supersedes the issue's proposed scale > 1 gate: at the 16 -> 24 upsample a = -0.5 matches torch within 6.1e-07 while -0.75 misses by 3.6e-02. The support still stretches only when downsampling, and no production caller resamples upward with antialias set (SAM downsamples 1024 -> 64).

Fixtures in tests/fixtures/deepseekocr_sam/ are generated out of band by generate_reference.py with torch 2.14.1+cpu and mlx-vlm 0.7.8 (73d366ec), which agree within 9.8e-07. A greedy check on the Phi-3.5-MoE checkpoint produces the identical 16-token stream with 1e-5 and with the pre-fix 1e-6 default, so the epsilon change does not alter generation; the divergence from mlx-lm is the cross-device numeric path and is recorded in the PR.

Refs #2465.
@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:low Low priority area:models Model architectures, weights, loading, metadata labels Oct 20, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Oct 20, 2026
@inureyes
inureyes merged commit 49db023 into main Oct 20, 2026
21 of 28 checks passed
@inureyes
inureyes deleted the fix/issue-2465-phimoe-sam branch October 30, 2026 13:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:models Model architectures, weights, loading, metadata priority:low Low priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(models): PhiMoE norm epsilon key and SAM antialiased bicubic a

1 participant