Repository navigation
fix(models): PhiMoE norm epsilon key and SAM antialiased bicubic a - #2507
Merged
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_epsnow carries#[serde(alias = "rms_norm_eps")]. Phi-3.5-MoE checkpoints spell itrms_norm_eps: 1e-5(the key mlx-lm'sphimoe.pyreads); previously the checkpoint value was ignored and the LayerNorms ran at the 1e-6 default. Thelayer_norm_epsspelling keeps parsing and 1e-6 stands when neither key exists.src/vision/encoders/resample_matrix.rs:interp_matrixpicksa = -0.5for a cubic kernel wheneverantialiasis set (ATen'sHelperInterpCubic::aa_filter, what mlx-vlm uses) and keeps Keys'-0.75without 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, fromgenerate_reference.pywith 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.5for the antialiased kernel in both directions, not only when downsampling: at the 16 to 24 upsample-0.5matches torch within 6.1e-07 while-0.75misses by 3.6e-02; at 16 to 5, 1.0e-07 vs 1.1e-02. Implemented as-0.5whenever a cubic kernel is antialiased; support still stretches only when downsampling. No production caller upsamples withantialias(SAM downsamples 1024 to 64), so no existing behavior changes beyond the AA-downsample path.Checkpoint comparison
On
models/mlx/phi-3.5-moe-4bitwith the prompt "The capital of France is" and 16 greedy tokens, mlx-lm (CPU-only mlx 0.32.4 from the oracle venv) producesParis.\n\nIncorrect grammar: Me and my friend is going to thewhile mlxcel (CUDA MLX pinned 243bdf09, GPU) producesParis. It is not only the country's largest; both startParis.and diverge at token 3. A control run of mlxcel with the checkpoint'srms_norm_epsstripped (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