diff --git a/docs/gelu-call-sites.md b/docs/gelu-call-sites.md index 0ae27aba5..1376e4474 100644 --- a/docs/gelu-call-sites.md +++ b/docs/gelu-call-sites.md @@ -56,7 +56,7 @@ The logit-trace gate cannot tell erf from tanh (see section 0), so these are con | T2 | Phixtral | erf | **tanh (f32, compiled)** | `gelu_new`; mlx-lm `SwitchMLP` default `approx="precise"`. No local checkpoint | | T3 | GPT-2 | erf | **tanh (f32, compiled)** | `activation_function = gelu_new` (`GPT2Config` default); mlx-lm `nn.gelu_approx`. No local checkpoint | | T4 | GPT-NeoX / Pythia | erf | erf, kept | HF `hidden_act = gelu` (`GPTNeoXConfig` default); mlx-lm uses tanh (departs from HF) | -| T5 | GPT-BigCode | erf | erf, unchanged (out of #2548 scope) | HF says `gelu_pytorch_tanh`, mlx-lm erf: the same situation as T6, left for a follow-up | +| T5 | GPT-BigCode | erf | **from `activation_function`** (#2794): `gelu_pytorch_tanh` / `gelu_new` / `gelu_fast` and an absent key -> tanh (f32, compiled), `gelu` -> erf, `quick_gelu` -> fast, any other name is a load error | HF builds the MLP as `ACT2FN[config.activation_function]` and `GPTBigCodeConfig` defaults to `gelu_pytorch_tanh` (santacoder declares it); mlx-lm `nn.gelu` (erf) departs from HF. Eager only (no XLA GPT-BigCode graph); no local checkpoint, smoke not run | | T6 | StarCoder2 (eager) | erf | **tanh (f32, compiled)** | config `hidden_act = gelu_pytorch_tanh` (`Starcoder2Config` default); the OpenXLA StarCoder2 graph, frozen after it ran token-exact against the HF fp32 oracle, was already tanh. mlx-lm `nn.gelu` departs from HF. The issue framed the XLA tanh as the bug; rule 1 makes eager the side that moves | | T8 | Nemotron-NAS | `gelu_new` / `gelu_fast` -> SiLU, unknown -> SiLU | GELU names through `GeluActivation::from_hf_name` (tanh names -> tanh f32 compiled), `silu`, `relu`; any other name is a load error | wrong activation, not an erf/tanh choice | | T10, T11 | ModernBERT MLP and classifier head | erf | erf, kept | HF `hidden_activation = gelu`, `classifier_activation = gelu`; mlx-vlm tanh departs | @@ -79,14 +79,17 @@ The logit-trace gate cannot tell erf from tanh (see section 0), so these are con | V33 | Qwen3-VL-MoE MLP / merger | MLP tanh, merger erf | kept | HF `Qwen3VLMoeVisionMLP` uses `ACT2FN[gelu_pytorch_tanh]`, the patch merger `nn.GELU()`; mlx-vlm has them swapped | | V34 | Qwen3-Omni vision MLP | tanh | kept | `thinker_config.vision_config.hidden_act = gelu_pytorch_tanh`; mlx-vlm `nn.GELU()` departs | | S3 | Aya Vision SigLIP | erf (key absent) | **tanh (f32, compiled)** | `SiglipVisionConfig` default `gelu_pytorch_tanh`; mlx-vlm `GELUMLP` | +| S6 | Idefics2 SigLIP | fast (forced by `use_fast_gelu = true`) | **tanh (f32, compiled)** (#2794) | HF `Idefics2VisionTransformer` uses `ACT2FN[hidden_act]` and `Idefics2VisionConfig` defaults to `gelu_pytorch_tanh`; `idefics2-8b-4bit` omits the key and its tower `model_type` is `idefics2`, so the loader resolves it with `default_missing_hidden_act_to_siglip`. mlx-vlm `FastGELUMLP` departs from HF | +| S7 | Granite 4 Vision SigLIP | fast (forced by `use_fast_gelu = true`) | **tanh (f32, compiled)** (#2794) | `granite-4.0-3b-vision-4bit` declares `gelu_pytorch_tanh`, and its remote `modeling.py` builds the tower with `AutoModel.from_config(vision_config)` (HF `SiglipVisionModel`); mlx-vlm `FastGELUMLP` departs from HF. The Q-Former FFN (V59, erf) is unchanged | | S8 | Granite Vision 3.x SigLIP | erf (key absent) | **tanh (f32, compiled)** | `SiglipVisionConfig` default; mlx-vlm `FastGELUMLP` departs from HF | | S9 | LLaVA 1.5 / LLaVA-NeXT CLIP | erf (key absent) | **fast**, eager and the XLA llava vision path | `CLIPVisionConfig` default `quick_gelu`; mlx-vlm `FastGELUMLP` | | S10 | LLaVA-Interleave SigLIP | tanh (declared) | tanh, kept | `vision_config.hidden_act = gelu_pytorch_tanh`; mlx-vlm fast departs | | S11 | Bunny | erf (hard-coded) | **tanh (f32, compiled)** | the tower is `google/siglip-so400m-patch14-384` (not CLIP), whose remote config declares `gelu_pytorch_tanh`; mlx-vlm `FastGELUMLP` departs | | S12 | Phi-3-V / Phi-3.5-V CLIP | erf (hard-coded) | **fast** | remote code builds its CLIP with `hidden_act="quick_gelu"`; mlx-vlm fast | | S14 | Llama-Nemotron-VL-Embed SigLIP | tanh (declared) | tanh, kept | `vision_config.hidden_act = gelu_pytorch_tanh`; mlx-vlm (granite_vision `FastGELUMLP`) departs | +| S-cfg | Shared `VisionConfig` / `VisionHiddenActivation` (every SigLIP / CLIP tower and the SigLIP text tower) | `gelu_new`, `gelu_fast` and non-GELU names (`relu`, `silu`) parsed as erf; a declared `gelu_pytorch_tanh` ran the eager ~11-pass f32 helper (`TanhF32`) | (#2794) declared names map through `GeluActivation::from_hf_name` (`gelu_new` / `gelu_fast` -> tanh, `gelu_python` -> erf); any other name is a load error, also in Bunny's sparse-config fallback; a declared tanh runs the compiled kernel (`TanhF32Compiled`), the same as an omitted key resolved to the SigLIP default | `ACT2FN` names; the OpenXLA `LlavaVisionConfig` already rejected unknown names and now accepts the same GELU set. The declared and omitted tanh paths were split only to pin the declared one bit-for-bit to the private copy #2423 replaced; the two agree up to f32 rounding | -Shared resolution: `VisionConfig::default_missing_hidden_act_by_model_type` resolves an omitted key by `vision_config.model_type` (SigLIP / SigLIP2 -> tanh, CLIP -> `quick_gelu`, anything else unchanged); the serde default stays erf for the other towers. A declared `quick_gelu` now parses as fast in both the eager `VisionConfig` and the OpenXLA llava vision config (it parsed as erf, or was rejected by XLA, before). The fused single-kernel tanh GELU raised in the #2548 thread is a separate follow-up; nothing here changes `compiled_gelu_tanh_f32`. +Shared resolution: `VisionConfig::default_missing_hidden_act_by_model_type` resolves an omitted key by `vision_config.model_type` (SigLIP / SigLIP2 -> tanh, CLIP -> `quick_gelu`, anything else unchanged); `default_missing_hidden_act_to_siglip` resolves it to tanh for loaders whose tower `model_type` is not a SigLIP one (Gemma 3, PaliGemma, Idefics2); the serde default stays erf for the other towers. A declared `quick_gelu` now parses as fast in both the eager `VisionConfig` and the OpenXLA llava vision config (it parsed as erf, or was rejected by XLA, before). Since #2794 every declared name goes through `GeluActivation::from_hf_name` and a non-GELU name is a load error in both, and every tanh resolution, declared or defaulted, runs the compiled f32 kernel. No production loader passes `use_fast_gelu = true` any more. The fused single-kernel tanh GELU raised in the #2548 thread is a separate follow-up; nothing here changes `compiled_gelu_tanh_f32`. Sections 1 to 4 below are the #2423 inventory at its own revision; where a row's Notes cell ends with a **#2548** tag, the table above is the current state. @@ -96,7 +99,7 @@ Sections 1 to 4 below are the #2423 inventory at its own revision; where a row's |---|---|---|---| | `mlxcel_core::gelu` (FFI) | `cpp/mlx_cxx_bridge.cpp:3286` | erf | f32 constants, result cast back to input dtype | | `mlxcel_core::gelu_approx` (FFI) and `mlxcel_core::utils::gelu_approx` | `mlx_cxx_bridge.cpp:3303`; wrapper `src/lib/mlxcel-core/src/utils.rs:1049` | **erf** (despite the name) | Comment at :3304-3315 claims erf "matches Python nn.GELU(approx="precise")"; that is false, precise = tanh. Doc at `utils.rs:1046` and `lib.rs:1056` also say "Approximate GELU (erf-based)". | -| `compiled_gelu` | `mlx_cxx_bridge.cpp:1871-1893` | erf | Only Rust caller: StarCoder2 | +| `compiled_gelu` | `mlx_cxx_bridge.cpp:1871-1893` | erf | Only Rust caller: StarCoder2. **#2794:** removed (dead after StarCoder2 moved to tanh in #2548) | | `compiled_gelu_approx` | `mlx_cxx_bridge.cpp:1895-1920`, FFI `lib.rs:775` | **erf** | No non-test Rust caller ("Used by: legacy/tests") | | `compiled_geglu_activation` (`get_compiled_geglu`) | `mlx_cxx_bridge.cpp:1922-1951` | erf GeGLU | Not called from Rust directly, but `compiled_gelu_mlp_forward_fp16` uses the same compiled lambda (see below) | | `compiled_geglu_approx_activation` | `mlx_cxx_bridge.cpp:1953-1975` | tanh GeGLU (`gelu_tanh_approx`) | Gemma family fallback / non-fused path | @@ -155,7 +158,7 @@ Legend for formula columns: **erf**, **tanh**, **fast**. "cfg" means the formula | T2 | `src/models/phixtral.rs:609`, `:616` | `SwitchMLP::forward` (sort / no-sort) | Phixtral | **erf** | tanh (`SwitchMLP` default `nn.GELU(approx="precise")`) | [mlx_lm/models/switch_layers.py:229](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/switch_layers.py#L229), [mlx_lm/models/phixtral.py:93](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/phixtral.py#L93) | **MISMATCH** | No local ckpt **#2548:** now tanh (f32, compiled). | | T3 | `src/models/gpt2.rs:579` | `MLP::forward` | GPT-2 | **erf** | tanh | [mlx_lm/models/gpt2.py:80](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt2.py#L80) | **MISMATCH** | HF `activation_function=gelu_new` (tanh) too; module doc :28-33 acknowledges it. No local ckpt **#2548:** now tanh (f32, compiled). | | T4 | `src/models/gpt_neox.rs:883` | `MLP::forward` | GPT-NeoX / Pythia | **erf** | tanh | [mlx_lm/models/gpt_neox.py:100](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt_neox.py#L100) | **MISMATCH** | Deliberate: HF Pythia `hidden_act=gelu` (erf); comment :866-882. No local ckpt **#2548:** erf kept (HF `gelu`). | -| T5 | `src/models/gpt_bigcode.rs:569` | `MLP::forward` | GPT-BigCode | erf | erf (`nn.gelu`) | [mlx_lm/models/gpt_bigcode.py:99](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt_bigcode.py#L99) | MATCH | The issue's known-correct case; HF says `gelu_pytorch_tanh` **#2548:** unchanged, follow-up. | +| T5 | `src/models/gpt_bigcode.rs:569` | `MLP::forward` | GPT-BigCode | erf | erf (`nn.gelu`) | [mlx_lm/models/gpt_bigcode.py:99](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt_bigcode.py#L99) | MATCH | The issue's known-correct case; HF says `gelu_pytorch_tanh` **#2548:** unchanged, follow-up. **#2794:** now from `activation_function` (tanh f32 compiled for santacoder). | | T6 | `src/models/starcoder2.rs:216` | MLP (`compiled_gelu`) | StarCoder2 | erf | erf | [mlx_lm/models/starcoder2.py:82](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/starcoder2.py#L82) | MATCH | HF `hidden_act=gelu_pytorch_tanh`. Local: `starcoder2-3b-4bit`. (XLA emitter uses tanh here, see Appendix) **#2548:** eager now tanh (f32, compiled); XLA unchanged. | | T7 | `src/models/phi3small.rs:475` -> `utils.rs:1060` `gegelu` | GeGELU | Phi-3-small | fast | fast (`a * mx.sigmoid(1.702 a)`) | [mlx_lm/models/phi3small.py:49](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/phi3small.py#L49) | MATCH | Local: `phi-3-small-8k-instruct-aq4_64` | | T8 | `src/models/nemotron_nas.rs:364` | `NASMLP::forward` (cfg `act_fn`) | Nemotron-NAS | cfg: `gelu` -> erf; `gelu_new`/`gelu_fast` -> **silu** (falls to `_`) | `gelu` -> erf, `gelu_new`/`gelu_fast` -> tanh | [mlx_lm/models/nemotron-nas.py:95-97](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/nemotron-nas.py#L95-L97) ([mlx_vlm/models/nemotron_nas/language.py:27-32](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/nemotron_nas/language.py#L27-L32)) | MATCH for `gelu`; **MISMATCH** (silu) for `gelu_new`/`gelu_fast` | Wrong activation entirely for those strings, not an erf/tanh question. Local `nemotron-nas-30b-4bit` is actually `model_type=nemotron_h` **#2548:** GELU names via `from_hf_name`, unknown names rejected. | @@ -256,8 +259,8 @@ Since #2423 the Gemma 3 VLM and PaliGemma loaders resolve a missing key with `Vi | S3 | `src/loading/vlm_siglip.rs:219` | Aya Vision | **erf** (key absent) | tanh (`GELUMLP`) | [mlx_vlm/models/aya_vision/vision.py:10](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/aya_vision/vision.py#L10) | **MISMATCH** | `aya-vision-8b`: absent **#2548:** now tanh (f32, compiled). | | S4 | `src/loading/vlm_smolvlm.rs:309` | SmolVLM / Idefics3 | tanh | tanh (`GELUMLP` via idefics3) | [mlx_vlm/models/idefics3/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/idefics3/vision.py#L8) ([mlx_vlm/models/smolvlm/__init__.py:4](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/smolvlm/__init__.py#L4)) | MATCH | `smolvlm-instruct-bf16`, `idefics3-8b-llama3-4bit`: `gelu_pytorch_tanh` | | S5 | `src/models/colidefics3.rs:162` | ColIdefics3 | cfg (tanh when declared) | tanh | [mlx_vlm/models/colidefics3/__init__.py:2](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/colidefics3/__init__.py#L2) -> idefics3 | MATCH if declared, else MISMATCH | no local ckpt | -| S6 | `src/loading/vlm_idefics2.rs:201` (`use_fast_gelu = true`) | Idefics2 | fast | fast (`FastGELUMLP`) | [mlx_vlm/models/idefics2/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/idefics2/vision.py#L8) | MATCH | `idefics2-8b-4bit`: absent (forced fast) | -| S7 | `src/loading/vlm_granite4_vision.rs:141` (`use_fast_gelu = true`) | Granite 4 Vision | fast | fast | [mlx_vlm/models/granite4_vision/vision.py:77](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/granite4_vision/vision.py#L77) | MATCH | `granite-4.0-3b-vision-4bit`: `gelu_pytorch_tanh` (HF says tanh; both mlx-vlm and mlxcel use fast) | +| S6 | `src/loading/vlm_idefics2.rs:201` (`use_fast_gelu = true`) | Idefics2 | fast | fast (`FastGELUMLP`) | [mlx_vlm/models/idefics2/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/idefics2/vision.py#L8) | MATCH | `idefics2-8b-4bit`: absent (forced fast) **#2794:** now tanh (f32, compiled), follows HF over mlx-vlm. | +| S7 | `src/loading/vlm_granite4_vision.rs:141` (`use_fast_gelu = true`) | Granite 4 Vision | fast | fast | [mlx_vlm/models/granite4_vision/vision.py:77](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/granite4_vision/vision.py#L77) | MATCH | `granite-4.0-3b-vision-4bit`: `gelu_pytorch_tanh` (HF says tanh; both mlx-vlm and mlxcel use fast) **#2794:** now tanh (f32, compiled), follows HF over mlx-vlm. | | S8 | `src/loading/vlm_granite_vision.rs:118` | Granite Vision 3.x | **erf** (key absent) | fast (`FastGELUMLP`) | [mlx_vlm/models/granite_vision/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/granite_vision/vision.py#L8) | **MISMATCH** (erf vs fast) | `granite-vision-3.2-2b-4bit`: absent **#2548:** now tanh (f32, compiled), follows HF over mlx-vlm. | | S9 | `src/loading/vlm_llava.rs:698` (and host preprocessor `:444`) | LLaVA 1.5 / LLaVA-NeXT (CLIP) | **erf** (key absent; HF `CLIPVisionConfig().hidden_act == "quick_gelu"`, verified in transformers 5.20.1) | fast (`FastGELUMLP`) | [mlx_vlm/models/llava/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava/vision.py#L8), [mlx_vlm/models/llava_next/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava_next/vision.py#L8) | **MISMATCH** (erf vs fast) | `llava-1.5-7b-4bit`, `llava-next-mistral-7b-4bit`: absent **#2548:** now fast, eager and XLA. | | S10 | same, SigLIP-towered `llava` | LLaVA-Interleave (Qwen, SigLIP) | tanh | fast | [mlx_vlm/models/llava/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava/vision.py#L8) | **MISMATCH** (tanh vs fast) | `llava-interleave-qwen-0.5b-bf16`: `gelu_pytorch_tanh` **#2548:** tanh kept (declared). | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index c4a1dced9..d9a9cac28 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -1906,30 +1906,6 @@ std::unique_ptr compiled_gpt_oss_swiglu_activation( return std::make_unique(std::move(result[0])); } -// Compiled GELU: x * 0.5 * (1 + erf(x / sqrt(2))) -// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548) -namespace { - static std::function(const std::vector&)> get_compiled_gelu() { - auto fn = [](const std::vector& inputs) -> std::vector { - const auto& x = inputs[0]; - auto sqrt2 = array(std::sqrt(2.0f)); - auto half = array(0.5f); - auto one = array(1.0f); - auto erf_val = mlx::core::erf(mlx::core::divide(x, sqrt2)); - auto scale = mlx::core::multiply(half, mlx::core::add(one, erf_val)); - auto result = mlx::core::multiply(x, scale); - return {mlx::core::astype(result, x.dtype())}; - }; - return compile_shapeless_audited("compiled_gelu", fn); - } -} - -std::unique_ptr compiled_gelu(const MlxArray& x) { - static auto compiled_fn = get_compiled_gelu(); - auto result = compiled_fn({x.inner}); - return std::make_unique(std::move(result[0])); -} - // Compiled `gelu_approx`: despite the name this is the exact erf GELU, // x * 0.5 * (1 + erf(x / sqrt(2))), i.e. MLX `nn.gelu` / `nn.GELU()`. It is // NOT MLX `nn.gelu_approx` (the tanh form); for that use `gelu_tanh` or @@ -3354,8 +3330,8 @@ std::unique_ptr gelu(const MlxArray& a) { // gelu(x) = x * 0.5 * (1 + erf(x / sqrt(2))) // // The f32 constants promote a half-precision input, so the result is cast - // back to the input dtype before it is returned. `compiled_gelu` has always - // done this; leaving it out here let the f32 activation reach the next + // back to the input dtype before it is returned, as the compiled GELU + // helpers do; leaving it out here let the f32 activation reach the next // projection, where MLX promotes the f16 weight to f32 to match it. See // `gelu_approx` for the measurement. auto sqrt2 = array(std::sqrt(2.0f)); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 4307bc122..a72762f41 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -578,10 +578,6 @@ std::unique_ptr qwen4_hyper_inject(const MlxArray& branch, const MlxAr std::unique_ptr compiled_scaled_rms_norm(const MlxArray& x, const MlxArray& scale, float eps); std::unique_ptr compiled_silu(const MlxArray& x); -// Compiled gelu: x * 0.5 * (1 + erf(x / sqrt(2))) — single fused kernel -// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548) -std::unique_ptr compiled_gelu(const MlxArray& x); - // Compiled gelu_approx: despite the name, exact erf GELU // (x * 0.5 * (1 + erf(x / sqrt(2)))), MLX `nn.gelu`, not MLX's tanh // `nn.gelu_approx`; use `gelu_tanh` for that (#2423). diff --git a/src/lib/mlxcel-core/src/ffi_tests.rs b/src/lib/mlxcel-core/src/ffi_tests.rs index 909611fe1..d4abb7a33 100644 --- a/src/lib/mlxcel-core/src/ffi_tests.rs +++ b/src/lib/mlxcel-core/src/ffi_tests.rs @@ -1952,22 +1952,6 @@ fn test_clip_residual_f16_widens_and_returns_f16() { assert_eq!(array_dtype(&out), dtype::FLOAT16); } -#[test] -fn test_compiled_gelu() { - let x = from_slice_f32(&[0.0, 1.0, -1.0, 2.0], &[1, 4]); - - let out = compiled_gelu(&x); - eval(&out); - - assert_eq!(array_shape(&out), vec![1, 4]); - - // gelu(0) = 0, gelu(x) > 0 for x > 0, gelu(x) < 0 for x < 0 (slightly) - let total = sum_all(&out); - eval(&total); - // sum of gelu([0, 1, -1, 2]) should be positive (1 and 2 dominate) - assert!(item_f32(&total) > 0.0); -} - #[test] fn test_compiled_gelu_approx() { let x = from_slice_f32(&[0.0, 1.0, -1.0, 2.0], &[1, 4]); @@ -1983,18 +1967,6 @@ fn test_compiled_gelu_approx() { assert!(item_f32(&total) > 0.0); } -#[test] -fn test_compiled_gelu_preserves_bf16_and_f16_dtype() { - for dtype in [dtype::BFLOAT16, dtype::FLOAT16] { - let x = astype(&from_slice_f32(&[0.0, 1.0, -1.0, 2.0], &[1, 4]), dtype); - let out = compiled_gelu(&x); - eval(&out); - - assert_eq!(array_shape(&out), vec![1, 4]); - assert_eq!(array_dtype(&out), dtype); - } -} - #[test] fn test_compiled_gelu_approx_preserves_bf16_and_f16_dtype() { for dtype in [dtype::BFLOAT16, dtype::FLOAT16] { @@ -2154,29 +2126,6 @@ fn test_gelu_approx_bf16_negative_values() { ); } -#[test] -fn test_compiled_gelu_matches_gelu() { - // Verify compiled_gelu gives same result as non-compiled gelu - let x = from_slice_f32(&[0.5, -0.5, 1.5, -1.5], &[1, 4]); - - let compiled_out = compiled_gelu(&x); - let regular_out = gelu(&x); - - eval(&compiled_out); - eval(®ular_out); - - let compiled_sum = sum_all(&compiled_out); - let regular_sum = sum_all(®ular_out); - eval(&compiled_sum); - eval(®ular_sum); - - let diff = (item_f32(&compiled_sum) - item_f32(®ular_sum)).abs(); - assert!( - diff < 1e-4, - "compiled_gelu and gelu should give same result, diff={diff}" - ); -} - // --- compiled_gelu_approx_mlp_forward: 8-bit decode gate + 4-bit legacy // path coverage (issue #680) ------------------------------------------- // @@ -2185,7 +2134,7 @@ fn test_compiled_gelu_matches_gelu() { // checkpoints. The reference below reproduces the function's own // op-at-a-time fallback (`quantized_matmul` for gate/up/down plus // `gelu_tanh_approx`) using only primitive FFI ops, since none of the -// exposed `gelu`/`gelu_approx`/`compiled_gelu*` helpers implement the +// exposed `gelu`/`gelu_approx`/`compiled_gelu_approx` helpers implement the // tanh approximation (they are all erf-based; see the comments on // `gelu_approx` and `compiled_gelu_approx` in mlx_cxx_bridge.cpp). @@ -2717,7 +2666,6 @@ fn shapeless_compile_audit_harness() { eval(&compiled_relu_squared(&x)); eval(&compiled_silu(&x)); eval(&compiled_gpt_oss_swiglu_activation(&gpt_linear, &gpt_glu)); - eval(&compiled_gelu(&x)); eval(&compiled_gelu_approx(&x)); eval(&compiled_gelu_tanh_f32(&x)); eval(&compiled_geglu_activation(&gate, &x)); @@ -2924,7 +2872,6 @@ fn shapeless_compile_audit_harness() { "compiled_relu_squared", "compiled_silu", "compiled_gpt_oss_swiglu_activation", - "compiled_gelu", "compiled_gelu_approx", "compiled_gelu_tanh_f32", "compiled_geglu_activation", diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 0526745b2..59a04b0ff 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -789,10 +789,6 @@ mod ffi { eps: f32, ) -> UniquePtr; - /// Compiled gelu: x * 0.5 * (1 + erf(x / sqrt(2))) — single fused kernel - /// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548) - fn compiled_gelu(x: &MlxArray) -> UniquePtr; - /// Compiled `gelu_approx`: despite the name, the exact erf GELU /// `x * 0.5 * (1 + erf(x / sqrt(2)))` (MLX `nn.gelu`), not MLX's tanh /// `nn.gelu_approx`. For the tanh form use `gelu_tanh` (#2423). diff --git a/src/lib/mlxcel-core/src/utils.rs b/src/lib/mlxcel-core/src/utils.rs index 9540c0efc..8789e1365 100644 --- a/src/lib/mlxcel-core/src/utils.rs +++ b/src/lib/mlxcel-core/src/utils.rs @@ -1126,7 +1126,8 @@ pub fn compiled_gelu_tanh_f32(x: &MlxArray) -> UniquePtr { /// StarCoder2, Moondream2/3 and RecurrentGemma, the Nemotron-NAS MLP, the /// Nemotron VoiceChat TTS backbone and codec, the Gemma3n dead `gelu_topk` /// branch, and the Phi-4 SigLIP, DeepSeek-VL2, Youtu-VL, Jina-VLM, LFM2-VL, -/// Moondream, PaddleOCR-VL, Qwen2-VL and Molmo vision MLPs (#2548). Not used +/// Moondream, PaddleOCR-VL, Qwen2-VL and Molmo vision MLPs (#2548); the +/// GPT-BigCode MLP and the shared vision `hidden_act` parser (#2794). Not used /// by the remaining `gelu_approx` call sites, which compute erf. Regenerate /// with `grep -rln 'GeluActivation' src --include='*.rs'`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] diff --git a/src/lib/mlxcel-xla/src/emitter/vision_config.rs b/src/lib/mlxcel-xla/src/emitter/vision_config.rs index f8a35239f..9318cc7ed 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision_config.rs @@ -210,6 +210,21 @@ mod tests { assert_eq!(erf.activation, VisionActivation::ExactGelu); } + #[test] + fn declared_hidden_act_accepts_every_act2fn_gelu_name() { + // #2794: the same names the eager `VisionHiddenActivation` accepts. + let mut root: Value = serde_json::from_str(&config_json()).unwrap(); + for (name, expected) in [ + ("gelu_new", VisionActivation::GeluPytorchTanh), + ("gelu_fast", VisionActivation::GeluPytorchTanh), + ("gelu_python", VisionActivation::ExactGelu), + ] { + root["vision_config"]["hidden_act"] = name.into(); + let parsed = LlavaVisionConfig::from_json_str(&root.to_string()).unwrap(); + assert_eq!(parsed.activation, expected, "{name}"); + } + } + #[test] fn pinned_siglip_contract_has_exact_schema() { let path = Path::new("/tmp/mlxcel-llava-hf-1090956d"); @@ -422,7 +437,11 @@ impl LlavaVisionConfig { return Err("vision layer_norm_eps must be finite and positive".to_string()); } let activation = match vision.get("hidden_act").and_then(Value::as_str) { - Some("gelu_pytorch_tanh") => VisionActivation::GeluPytorchTanh, + // The tanh and erf names `transformers` `ACT2FN` knows, the set + // the eager `VisionHiddenActivation` accepts (#2794). + Some("gelu_pytorch_tanh" | "gelu_new" | "gelu_fast") => { + VisionActivation::GeluPytorchTanh + } Some("quick_gelu") => VisionActivation::QuickGelu, // An omitted key is the Hugging Face class default of the tower, // resolved as the eager loaders do @@ -433,7 +452,7 @@ impl LlavaVisionConfig { VisionActivation::GeluPytorchTanh } None if vision_model_type == "clip_vision_model" => VisionActivation::QuickGelu, - None | Some("gelu") => VisionActivation::ExactGelu, + None | Some("gelu" | "gelu_python") => VisionActivation::ExactGelu, other => { return Err(format!( "unsupported vision hidden_act {other:?}; refusing an approximate substitution" diff --git a/src/loading/vlm_granite4_vision.rs b/src/loading/vlm_granite4_vision.rs index e4eb05b9f..162e7774b 100644 --- a/src/loading/vlm_granite4_vision.rs +++ b/src/loading/vlm_granite4_vision.rs @@ -101,6 +101,29 @@ fn parse_i32_list(v: Option<&serde_json::Value>) -> Vec { .unwrap_or_default() } +// Granite 4 Vision runs the tower MLP GELU its `hidden_act` names, not +// mlx-vlm's forced sigmoid `GELU(approx="fast")` (#2794). +const GRANITE4_FORCE_FAST_GELU: bool = false; + +/// Parse the Granite 4 Vision SigLIP tower config (`vision_config`). +/// +/// The checkpoint's remote `modeling.py` builds the tower with +/// `AutoModel.from_config(config.vision_config)`, i.e. Hugging Face +/// `SiglipVisionModel`, whose MLP is `ACT2FN[hidden_act]`; the local +/// `granite-4.0-3b-vision-4bit` declares `gelu_pytorch_tanh`. An omitted key +/// resolves to the `SiglipVisionConfig` class default. mlx-vlm's +/// `FastGELUMLP` departs from Hugging Face here and mlxcel follows Hugging +/// Face (#2794, quick_gelu before). +fn granite4_vision_config(full_config: &serde_json::Value) -> Result { + let raw = full_config + .get("vision_config") + .ok_or_else(|| anyhow::anyhow!("Missing vision_config in config.json"))?; + let mut vision_config: vision::config::VisionConfig = serde_json::from_value(raw.clone()) + .map_err(|e| anyhow::anyhow!("Failed to parse Granite 4 vision config: {}", e))?; + vision_config.default_missing_hidden_act_by_model_type(raw); + Ok(vision_config) +} + /// Load a Granite 4 Vision VLM. pub(crate) fn load_granite4_vision_vlm(model_path: &Path) -> Result { let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; @@ -119,13 +142,7 @@ pub(crate) fn load_granite4_vision_vlm(model_path: &Path) -> Result let text_args: models::granitemoehybrid::ModelArgs = serde_json::from_value(text_config_val) .map_err(|e| anyhow::anyhow!("Failed to parse Granite 4 text config: {}", e))?; - let vision_config: vision::config::VisionConfig = serde_json::from_value( - full_config - .get("vision_config") - .cloned() - .ok_or_else(|| anyhow::anyhow!("Missing vision_config in config.json"))?, - ) - .map_err(|e| anyhow::anyhow!("Failed to parse Granite 4 vision config: {}", e))?; + let vision_config = granite4_vision_config(&full_config)?; let group_size = text_args.group_size(); let bits = text_args.bits(); @@ -134,9 +151,8 @@ pub(crate) fn load_granite4_vision_vlm(model_path: &Path) -> Result let weights = strip_language_model_prefix(load_vlm_weights_common(model_path, None)?); // Vision tower and 8 window-QFormer projectors read by ref before the text - // backbone consumes the weight map. The granite4 SigLIP tower uses the - // sigmoid `GELU(approx="fast")` in its MLP (unlike Granite Vision #539, which - // uses the tanh variant), matching the reference vision encoder. + // backbone consumes the weight map. The tower MLP runs the GELU its config + // selects (see `granite4_vision_config`). let vision_tower = vision::encoders::siglip::SigLipVisionModel::from_weights_with_quant_and_gelu( &weights, @@ -144,7 +160,7 @@ pub(crate) fn load_granite4_vision_vlm(model_path: &Path) -> Result "vision_tower.vision_model", group_size, bits, - true, // use_fast_gelu (sigmoid GELU) + GRANITE4_FORCE_FAST_GELU, ) .map_err(|e| anyhow::anyhow!("Failed to load Granite 4 vision tower: {}", e))?; @@ -258,3 +274,48 @@ pub(crate) fn load_granite4_vision_vlm(model_path: &Path) -> Result Ok(LoadedModel::Granite4VisionVLM(vlm)) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::vision::config::VisionHiddenActivation; + use crate::vision::encoders::siglip::{VisionMlpActivation, select_mlp_activation}; + use serde_json::json; + + fn tower(hidden_act: Option<&str>) -> (VisionHiddenActivation, VisionMlpActivation) { + let mut vision = json!({ + "model_type": "siglip_vision_model", + "num_hidden_layers": 1, + "hidden_size": 8, + "intermediate_size": 16, + "num_attention_heads": 2, + "patch_size": 16 + }); + if let Some(act) = hidden_act { + vision["hidden_act"] = json!(act); + } + let config = granite4_vision_config(&json!({ "vision_config": vision })).unwrap(); + ( + config.hidden_act, + select_mlp_activation(config.hidden_act, GRANITE4_FORCE_FAST_GELU), + ) + } + + #[test] + fn vision_tower_runs_the_hugging_face_tanh_gelu() { + // The local granite-4.0-3b-vision-4bit declares gelu_pytorch_tanh; + // before #2794 the loader forced mlx-vlm's quick_gelu over it. + let expected = ( + VisionHiddenActivation::GeluPytorchTanh, + VisionMlpActivation::PytorchTanhCompiled, + ); + assert_eq!(tower(Some("gelu_pytorch_tanh")), expected); + // An omitted key takes the SiglipVisionConfig default. + assert_eq!(tower(None), expected); + } + + #[test] + fn missing_vision_config_is_an_error() { + assert!(granite4_vision_config(&json!({})).is_err()); + } +} diff --git a/src/loading/vlm_idefics2.rs b/src/loading/vlm_idefics2.rs index 60336624a..1bf925408 100644 --- a/src/loading/vlm_idefics2.rs +++ b/src/loading/vlm_idefics2.rs @@ -149,6 +149,31 @@ fn resolve_perceiver_heads(full_config: &Value) -> Idefics2PerceiverHeads { } } +/// Idefics2 runs the tower MLP GELU its `hidden_act` names, not mlx-vlm's +/// forced sigmoid `GELU(approx="fast")` (#2794). +const IDEFICS2_FORCE_FAST_GELU: bool = false; + +/// Parse the Idefics2 SigLIP tower config (`vision_config`). +/// +/// Hugging Face `Idefics2VisionTransformer` builds its MLP as +/// `ACT2FN[config.hidden_act]`, and `Idefics2VisionConfig().hidden_act` is +/// `gelu_pytorch_tanh`. The tower's `model_type` is `idefics2`, which +/// [`VisionConfig::default_missing_hidden_act_by_model_type`] does not cover, +/// so an omitted key (the local `idefics2-8b-4bit` omits it) is resolved to +/// the SigLIP tanh default explicitly. mlx-vlm's `FastGELUMLP` departs from +/// Hugging Face here and mlxcel follows Hugging Face (#2794, quick_gelu +/// before). +/// +/// [`VisionConfig::default_missing_hidden_act_by_model_type`]: vision::config::VisionConfig::default_missing_hidden_act_by_model_type +fn idefics2_vision_config(full_config: &Value) -> Result { + let mut vision_config: vision::config::VisionConfig = + parse_required_vlm_subconfig(full_config, "vision_config", "Idefics2 vision config")?; + if let Some(raw) = full_config.get("vision_config") { + vision_config.default_missing_hidden_act_to_siglip(raw); + } + Ok(vision_config) +} + /// Load an Idefics2 (`idefics2`) VLM (SigLIP + perceiver-resampler connector + /// Mistral language model). pub(crate) fn load_idefics2_vlm(model_path: &Path) -> Result { @@ -158,9 +183,7 @@ pub(crate) fn load_idefics2_vlm(model_path: &Path) -> Result { let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; - // Vision config (`vision_config`) is a SigLIP tower. - let vision_config: vision::config::VisionConfig = - parse_required_vlm_subconfig(&full_config, "vision_config", "Idefics2 vision config")?; + let vision_config = idefics2_vision_config(&full_config)?; // Text config (`text_config`) is Mistral; normalize the stripped 4-bit // config into full Mistral-7B args. @@ -195,16 +218,15 @@ pub(crate) fn load_idefics2_vlm(model_path: &Path) -> Result { let text_model = models::Llama3Model::from_weights(&text_weights, &text_args) .map_err(|e| anyhow::anyhow!("Failed to load Idefics2 text model: {}", e))?; - // SigLIP vision tower (`vision_model.*`). Idefics2's vision MLP uses the - // sigmoid `GELU(approx="fast")` (not the tanh GELU the other SigLIP consumers - // use), so request the fast-GELU encoder variant. + // SigLIP vision tower (`vision_model.*`), with the MLP GELU its config + // selects (see `idefics2_vision_config`). let vision_model = SigLipVisionModel::from_weights_with_quant_and_gelu( &weights, &vision_config, "vision_model", group_size, bits, - true, + IDEFICS2_FORCE_FAST_GELU, ) .map_err(|e| anyhow::anyhow!("Failed to load Idefics2 SigLIP vision tower: {}", e))?; @@ -336,6 +358,47 @@ mod tests { assert!(args.quantization.is_some()); } + #[test] + fn vision_tower_runs_the_hugging_face_tanh_gelu() { + use crate::vision::config::VisionHiddenActivation; + use crate::vision::encoders::siglip::{VisionMlpActivation, select_mlp_activation}; + + let tower = |hidden_act: Option<&str>| { + let mut vision = json!({ + "model_type": "idefics2", + "num_hidden_layers": 1, + "hidden_size": 8, + "intermediate_size": 16, + "num_attention_heads": 2, + "patch_size": 14 + }); + if let Some(act) = hidden_act { + vision["hidden_act"] = json!(act); + } + let config = idefics2_vision_config(&json!({ "vision_config": vision })).unwrap(); + ( + config.hidden_act, + select_mlp_activation(config.hidden_act, IDEFICS2_FORCE_FAST_GELU), + ) + }; + + // The local idefics2-8b-4bit omits the key: Idefics2VisionConfig + // default `gelu_pytorch_tanh`, not mlx-vlm's quick_gelu. + assert_eq!( + tower(None), + ( + VisionHiddenActivation::GeluPytorchTanh, + VisionMlpActivation::PytorchTanhCompiled + ) + ); + assert_eq!( + tower(Some("gelu_pytorch_tanh")).1, + VisionMlpActivation::PytorchTanhCompiled + ); + // A declared key still decides. + assert_eq!(tower(Some("gelu")).1, VisionMlpActivation::Exact); + } + #[test] fn perceiver_heads_default_when_config_stripped() { let full = json!({ "perceiver_config": { "model_type": "idefics2" } }); diff --git a/src/loading/vlm_llava.rs b/src/loading/vlm_llava.rs index cdb7421cc..aaed614cf 100644 --- a/src/loading/vlm_llava.rs +++ b/src/loading/vlm_llava.rs @@ -646,15 +646,30 @@ fn resolve_llava_vision_hidden_act(vlm_config: &mut VLMConfig, full_config: &Val } } +/// Parse Bunny's `vision_config`, falling back to the +/// `google/siglip-so400m-patch14-384` tower shape when the block is too sparse +/// to describe the tower. +/// +/// A declared `hidden_act` that is not a GELU variant is an error even when +/// the rest of the block is sparse, so the fallback cannot swallow it (#2794). fn parse_bunny_vision_config( full_config: &Value, vision_config_val: &Value, -) -> vision::config::VisionConfig { +) -> Result { + use vision::config::VisionHiddenActivation; + + let declared_act = match vision_config_val.get("hidden_act") { + Some(Value::String(name)) => Some( + VisionHiddenActivation::try_from(name.clone()) + .map_err(|e| anyhow::anyhow!("Bunny vision_config: {e}"))?, + ), + _ => None, + }; let parsed: Option = serde_json::from_value(vision_config_val.clone()).ok(); if let Some(mut config) = parsed { config.default_missing_hidden_act_by_model_type(vision_config_val); - config + Ok(config) } else { let mm_hidden = full_config .get("mm_hidden_size") @@ -664,7 +679,7 @@ fn parse_bunny_vision_config( .get("intermediate_size") .and_then(|v| v.as_u64()) .unwrap_or(4304) as usize; - vision::config::VisionConfig { + Ok(vision::config::VisionConfig { model_type: "siglip_vision_model".to_string(), num_hidden_layers: 27, hidden_size: mm_hidden, @@ -676,9 +691,10 @@ fn parse_bunny_vision_config( layer_norm_eps: 1e-6, // Bunny's `mm_vision_tower` is `google/siglip-so400m-patch14-384`, // whose remote config declares `gelu_pytorch_tanh` (#2548; erf - // before). mlx-vlm's `FastGELUMLP` departs from it. - hidden_act: vision::config::VisionHiddenActivation::SiglipDefaultTanh, - } + // before). mlx-vlm's `FastGELUMLP` departs from it. A key the + // sparse block does declare still decides. + hidden_act: declared_act.unwrap_or(VisionHiddenActivation::GeluPytorchTanh), + }) } } @@ -830,7 +846,7 @@ pub(crate) fn load_llava_bunny_vlm(model_path: &Path) -> Result { .get("vision_config") .cloned() .unwrap_or_default(); - let vision_config = parse_bunny_vision_config(&full_config, &vision_config_val); + let vision_config = parse_bunny_vision_config(&full_config, &vision_config_val)?; let vision_feature_layer = full_config .get("vision_feature_layer") diff --git a/src/loading/vlm_llava_tests.rs b/src/loading/vlm_llava_tests.rs index 80f479ad6..92129d0a6 100644 --- a/src/loading/vlm_llava_tests.rs +++ b/src/loading/vlm_llava_tests.rs @@ -128,7 +128,8 @@ fn parse_bunny_vision_config_uses_defaults_for_sparse_configs() { &json!({ "intermediate_size": 5120 }), - ); + ) + .unwrap(); assert_eq!(config.model_type, "siglip_vision_model"); assert_eq!(config.hidden_size, 1408); @@ -137,6 +138,37 @@ fn parse_bunny_vision_config_uses_defaults_for_sparse_configs() { assert_eq!(config.image_size, 384); } +#[test] +fn parse_bunny_vision_config_rejects_a_non_gelu_hidden_act() { + // Before #2794 the parse error was swallowed by the sparse-config fallback + // and `relu` silently ran as the tanh GELU. + for vision_config in [ + json!({ "hidden_act": "relu" }), + json!({ + "model_type": "siglip_vision_model", + "num_hidden_layers": 2, + "hidden_size": 768, + "intermediate_size": 1536, + "num_attention_heads": 12, + "patch_size": 16, + "hidden_act": "relu" + }), + ] { + let err = parse_bunny_vision_config(&json!({}), &vision_config) + .expect_err("relu is not a SigLIP activation") + .to_string(); + assert!(err.contains("unsupported vision hidden_act"), "{err}"); + } + + // A sparse block still takes the fallback, with the SigLIP tanh default. + let sparse = parse_bunny_vision_config(&json!({}), &json!({})).unwrap(); + assert_eq!(sparse.hidden_size, 1152); + assert_eq!( + sparse.hidden_act, + crate::vision::config::VisionHiddenActivation::GeluPytorchTanh + ); +} + #[test] fn inherit_text_quantization_if_missing_rejects_non_object_configs() { let mut text_config = json!("bad"); @@ -165,7 +197,8 @@ fn parse_bunny_vision_config_preserves_full_configs() { "num_channels": 3, "layer_norm_eps": 1e-5 }), - ); + ) + .unwrap(); assert_eq!(config.hidden_size, 768); assert_eq!(config.intermediate_size, 1536); diff --git a/src/models/gpt_bigcode.rs b/src/models/gpt_bigcode.rs index 5237570c3..f37ebda82 100644 --- a/src/models/gpt_bigcode.rs +++ b/src/models/gpt_bigcode.rs @@ -57,6 +57,7 @@ use mlxcel_core::batched_decode::{SplitAttention, SplitAttentionShape, forward_split_attention}; use mlxcel_core::generate::{BatchedDecodeKind, LanguageModel, route_batched_decode}; use mlxcel_core::layers::{KVCache, LayerNorm, UnifiedEmbedding, UnifiedLinear}; +use mlxcel_core::utils::GeluActivation; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; use serde::Deserialize; @@ -128,6 +129,12 @@ pub struct ModelArgs { #[serde(default)] pub quantization: Option, + + /// The MLP GELU, a Hugging Face `ACT2FN` name. `null` or absent means + /// `gelu_pytorch_tanh`, the `GPTBigCodeConfig` default. Resolved by + /// [`ModelArgs::activation`]. + #[serde(default)] + pub activation_function: Option, } /// A `config.json` token-id field, which may be a single int or a list of ints. @@ -265,9 +272,32 @@ impl ModelArgs { "GPT-BigCode n_inner ({intermediate}) must be between 1 and {MAX_N_INNER}" )); } + self.activation()?; Ok(()) } + /// The MLP activation the checkpoint declares in `activation_function`. + /// + /// Hugging Face builds the MLP as `ACT2FN[config.activation_function]`, + /// so the config decides; an omitted key is the `GPTBigCodeConfig` + /// default `gelu_pytorch_tanh`. mlx-lm departs from Hugging Face here by + /// always applying `nn.gelu` (erf), which mlxcel ran before #2794. The + /// tanh names run as the compiled f32 kernel. A non-GELU name is an error. + pub fn activation(&self) -> Result { + let name = self + .activation_function + .as_deref() + .unwrap_or("gelu_pytorch_tanh"); + GeluActivation::from_hf_name(name) + .map(GeluActivation::in_f32_compiled) + .ok_or_else(|| { + format!( + "GPT-BigCode: unsupported activation_function {name:?}; expected a GELU \ + variant (gelu, gelu_new, gelu_fast, gelu_pytorch_tanh, quick_gelu)" + ) + }) + } + pub fn head_dim(&self) -> usize { self.n_embd / self.n_head } @@ -575,18 +605,19 @@ impl SplitAttention for Attention { pub struct MLP { pub c_fc: UnifiedLinear, pub c_proj: UnifiedLinear, + /// Resolved once at load from `activation_function`. + pub activation: GeluActivation, } impl MLP { pub fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.c_fc.forward(x); - // Upstream applies `nn.gelu`, the erf-based exact GELU, which is what - // `mlxcel_core::utils::gelu_approx` evaluates in this tree (see its own - // doc comment). The checkpoint config names `gelu_pytorch_tanh`, whose - // tanh approximation differs from the erf form by under 1e-3; matching - // upstream keeps this port comparable against the mlx-lm reference, - // which is the same choice `src/models/gpt2.rs` makes. - let h = mlxcel_core::utils::gelu_approx(&h); + // Hugging Face applies `ACT2FN[config.activation_function]`, so the + // checkpoint config decides (`gelu_pytorch_tanh` for every published + // export and the class default). mlx-lm always applies `nn.gelu` + // (erf) and departs from Hugging Face; mlxcel follows Hugging Face + // (#2794). See `ModelArgs::activation`. + let h = self.activation.apply(&h); self.c_proj.forward(&h) } @@ -617,7 +648,13 @@ impl MLP { bits, )?; - Ok(Self { c_fc, c_proj }) + let activation = args.activation()?; + + Ok(Self { + c_fc, + c_proj, + activation, + }) } } diff --git a/src/models/gpt_bigcode_tests.rs b/src/models/gpt_bigcode_tests.rs index c7b143f19..2e4459bb7 100644 --- a/src/models/gpt_bigcode_tests.rs +++ b/src/models/gpt_bigcode_tests.rs @@ -21,8 +21,9 @@ //! HuggingFace export. The forward tests build tiny single-layer models and run //! them on the default device. -use super::{GptBigCodeModel, ModelArgs, TokenIdField, detect_prefix, load_linear}; +use super::{GptBigCodeModel, MLP, ModelArgs, TokenIdField, detect_prefix, load_linear}; use mlxcel_core::layers::UnifiedLinear; +use mlxcel_core::utils::{GeluActivation, array_to_vec_f32}; use mlxcel_core::weights::WeightMap; // Config surface. @@ -78,6 +79,45 @@ fn parses_the_real_santacoder_config() { assert!(args.quantization.is_none()); assert_eq!(args.eos_token_ids(), vec![49152]); assert!(args.validate().is_ok()); + // `activation_function: gelu_pytorch_tanh` is honored (#2794; erf before). + assert_eq!(args.activation(), Ok(GeluActivation::TanhF32Compiled)); +} + +#[test] +fn activation_function_follows_the_hugging_face_act2fn_names() { + let with = |name: Option<&str>| -> ModelArgs { + let mut config = serde_json::json!({ "model_type": "gpt_bigcode" }); + if let Some(name) = name { + config["activation_function"] = name.into(); + } + serde_json::from_value(config).expect("parses") + }; + + // Absent or null: the GPTBigCodeConfig default `gelu_pytorch_tanh`. + assert_eq!(with(None).activation(), Ok(GeluActivation::TanhF32Compiled)); + let null: ModelArgs = + serde_json::from_str(r#"{"model_type": "gpt_bigcode", "activation_function": null}"#) + .expect("parses"); + assert_eq!(null.activation(), Ok(GeluActivation::TanhF32Compiled)); + for name in ["gelu_pytorch_tanh", "gelu_new", "gelu_fast"] { + assert_eq!( + with(Some(name)).activation(), + Ok(GeluActivation::TanhF32Compiled), + "{name}" + ); + } + assert_eq!(with(Some("gelu")).activation(), Ok(GeluActivation::Erf)); + assert_eq!( + with(Some("quick_gelu")).activation(), + Ok(GeluActivation::Fast) + ); + + let relu = with(Some("relu")); + let err = relu + .validate() + .expect_err("relu is not a GPT-BigCode MLP activation"); + assert!(err.contains("unsupported activation_function"), "{err}"); + assert!(err.contains("\"relu\""), "{err}"); } #[test] @@ -333,6 +373,47 @@ fn hf_weights(args: &ModelArgs) -> WeightMap { w } +/// `MLP::forward` evaluates `c_proj(gelu_tanh(c_fc(x)))` for the default +/// config, not the erf GELU mlx-lm applies (#2794). +#[test] +fn mlp_forward_applies_the_configured_tanh_gelu() { + let args = tiny_args(); + let weights = hf_weights(&args); + let mlp = MLP::from_weights(&weights, &args, "transformer.h.0.mlp").expect("MLP loads"); + assert_eq!(mlp.activation, GeluActivation::TanhF32Compiled); + + // Inputs in [-3, 3] so the pre-activations reach the range where the two + // GELU forms differ by ~1e-3. + let data: Vec = (0..24).map(|i| ((i % 13) as f32 - 6.0) * 0.5).collect(); + let x = mlxcel_core::from_slice_f32(&data, &[1, 3, 8]); + + let out = array_to_vec_f32(&mlp.forward(&x)); + let h = mlp.c_fc.forward(&x); + let tanh_ref = array_to_vec_f32( + &mlp.c_proj + .forward(&GeluActivation::TanhF32Compiled.apply(&h)), + ); + let erf_ref = array_to_vec_f32(&mlp.c_proj.forward(&GeluActivation::Erf.apply(&h))); + + let max_diff = |a: &[f32], b: &[f32]| { + a.iter() + .zip(b) + .map(|(x, y)| (x - y).abs()) + .fold(0.0_f32, f32::max) + }; + assert_eq!(out.len(), tanh_ref.len()); + let to_tanh = max_diff(&out, &tanh_ref); + let to_erf = max_diff(&out, &erf_ref); + assert!( + to_tanh <= 1e-5, + "MLP output is {to_tanh} from the tanh GELU" + ); + assert!( + to_erf > 1e-4, + "MLP output is only {to_erf} from the erf GELU" + ); +} + // Key layout. #[test] diff --git a/src/models/siglip_text_tests.rs b/src/models/siglip_text_tests.rs index 6e184f43d..c4d76880a 100644 --- a/src/models/siglip_text_tests.rs +++ b/src/models/siglip_text_tests.rs @@ -284,7 +284,7 @@ fn text_config_defaults_match_the_reference_config() { assert_eq!(args.projection_size(), 768, "defaults to hidden_size"); assert_eq!( args.activation(), - VisionMlpActivation::PytorchTanh, + VisionMlpActivation::PytorchTanhCompiled, "the text tower defaults to gelu_pytorch_tanh, not the vision exact-erf default" ); } @@ -317,7 +317,28 @@ fn text_config_overrides_are_read_including_projection_and_activation() { // vision-side exact-erf default. let nulled = SigLipTextArgs::from_config(&json!({ "text_config": { "hidden_act": null } })).unwrap(); - assert_eq!(nulled.activation(), VisionMlpActivation::PytorchTanh); + assert_eq!( + nulled.activation(), + VisionMlpActivation::PytorchTanhCompiled + ); + + // A declared tanh GELU runs the same compiled kernel (#2794). + let declared = SigLipTextArgs::from_config(&json!({ + "text_config": { "hidden_act": "gelu_pytorch_tanh" } + })) + .unwrap(); + assert_eq!( + declared.activation(), + VisionMlpActivation::PytorchTanhCompiled + ); + + // A non-GELU name is a load error, not a silent erf GELU (#2794). + let err = SigLipTextArgs::from_config(&json!({ "text_config": { "hidden_act": "relu" } })) + .expect_err("relu is not a SigLIP text activation"); + assert!( + format!("{err:#}").contains("unsupported vision hidden_act"), + "{err:#}" + ); // A checkpoint that declares the text fields at the top level (no // `text_config` wrapper) is read as well. diff --git a/src/vision/config.rs b/src/vision/config.rs index 7ed401a7e..72747d1be 100644 --- a/src/vision/config.rs +++ b/src/vision/config.rs @@ -16,6 +16,7 @@ //! //! Deserialization for vision_config and projector config from model config.json +use mlxcel_core::utils::GeluActivation; use serde::{Deserialize, Deserializer}; /// Deserialize null as default value (serde #[serde(default)] only handles missing fields) @@ -29,36 +30,53 @@ where /// Activation selected by a SigLIP/CLIP checkpoint. /// -/// Existing and unknown values retain mlxcel's exact-erf GELU behavior. The -/// explicit `gelu_pytorch_tanh` variant follows the Hugging Face polynomial. -/// A missing key also parses as `ExactGelu`; every loader of a SigLIP or CLIP +/// A declared `hidden_act` maps through [`GeluActivation::from_hf_name`], so +/// every GELU name `transformers` `ACT2FN` knows runs the formula it computes +/// there: `gelu` / `gelu_python` erf, `gelu_pytorch_tanh` / `gelu_new` / +/// `gelu_fast` tanh, `quick_gelu` sigmoid. Any other name, such as `relu` or +/// `silu`, is a deserialization error rather than a silent erf GELU (#2794); +/// the OpenXLA `LlavaVisionConfig` rejects the same names. +/// +/// A missing key parses as `ExactGelu`; every loader of a SigLIP or CLIP /// tower resolves that case to the Hugging Face class default with -/// [`VisionConfig::default_missing_hidden_act_by_model_type`] (or, for Gemma 3 -/// and PaliGemma, [`VisionConfig::default_missing_hidden_act_to_siglip`]). +/// [`VisionConfig::default_missing_hidden_act_by_model_type`] (or, for Gemma 3, +/// PaliGemma and Idefics2, [`VisionConfig::default_missing_hidden_act_to_siglip`]). #[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)] +#[serde(try_from = "String")] pub enum VisionHiddenActivation { - #[serde(rename = "gelu_pytorch_tanh")] + /// The tanh GELU, declared under any of its `ACT2FN` names or resolved + /// from an omitted key on a SigLIP tower. Runs as one compiled f32 kernel + /// (#2423, #2794). GeluPytorchTanh, - /// `gelu_pytorch_tanh` resolved from an omitted key by - /// [`VisionConfig::default_missing_hidden_act_to_siglip`]; never parsed. - /// Same formula as [`GeluPytorchTanh`](Self::GeluPytorchTanh), evaluated as - /// one compiled kernel so the erf-to-tanh switch does not add prefill time. - /// Kept apart because the declared variant's eager f32 path is pinned - /// bit-for-bit to the private copy it replaced (#2423). - #[serde(skip)] - SiglipDefaultTanh, /// `quick_gelu` (`x * sigmoid(1.702 x)`), declared or resolved from an /// omitted key on a CLIP tower (the Hugging Face `CLIPVisionConfig` /// default) by [`VisionConfig::default_missing_hidden_act_by_model_type`]. /// Also set directly by loaders whose CLIP tower config is built in code /// (Phi-3-V). Parsed as `ExactGelu` before #2548. - #[serde(rename = "quick_gelu")] QuickGelu, + /// The erf GELU (`gelu`, `gelu_python`), and the serde default for an + /// absent or null key. #[default] - #[serde(other)] ExactGelu, } +impl TryFrom for VisionHiddenActivation { + type Error = String; + + fn try_from(name: String) -> Result { + match GeluActivation::from_hf_name(&name) { + Some(GeluActivation::Erf) => Ok(Self::ExactGelu), + Some(GeluActivation::Fast) => Ok(Self::QuickGelu), + Some(GeluActivation::Tanh) => Ok(Self::GeluPytorchTanh), + // `from_hf_name` returns only the three variants above. + Some(_) | None => Err(format!( + "unsupported vision hidden_act {name:?}; expected gelu, gelu_python, \ + gelu_pytorch_tanh, gelu_new, gelu_fast or quick_gelu" + )), + } + } +} + /// Vision encoder configuration (SigLIP) #[derive(Debug, Clone, Deserialize)] pub struct VisionConfig { @@ -88,10 +106,12 @@ impl VisionConfig { /// serde default keeps exact erf for every other loader; this is opt-in per /// loader because the CLIP towers' class default is `quick_gelu`, not tanh. /// - /// Used by: the Gemma 3 VLM and PaliGemma loaders (#2423). + /// Used by: the Gemma 3 VLM and PaliGemma loaders (#2423), and Idefics2, + /// whose `Idefics2VisionConfig` default is also `gelu_pytorch_tanh` + /// (#2794). pub(crate) fn default_missing_hidden_act_to_siglip(&mut self, raw: &serde_json::Value) { if !hidden_act_declared(raw) { - self.hidden_act = VisionHiddenActivation::SiglipDefaultTanh; + self.hidden_act = VisionHiddenActivation::GeluPytorchTanh; } } @@ -104,15 +124,15 @@ impl VisionConfig { /// The serde default stays exact erf because the shared struct also /// parses towers whose class default is neither (#2548). /// - /// Used by: the LLaVA (and its host preprocessor), Bunny, Aya Vision and - /// Granite Vision 3.x loaders (#2548). + /// Used by: the LLaVA (and its host preprocessor), Bunny, Aya Vision, + /// Granite Vision 3.x (#2548) and Granite 4 Vision (#2794) loaders. pub(crate) fn default_missing_hidden_act_by_model_type(&mut self, raw: &serde_json::Value) { if hidden_act_declared(raw) { return; } match self.model_type.as_str() { "siglip_vision_model" | "siglip2_vision_model" => { - self.hidden_act = VisionHiddenActivation::SiglipDefaultTanh; + self.hidden_act = VisionHiddenActivation::GeluPytorchTanh; } "clip_vision_model" => { self.hidden_act = VisionHiddenActivation::QuickGelu; @@ -217,22 +237,43 @@ mod tests { } #[test] - fn vision_hidden_act_selects_only_the_explicit_pytorch_tanh_variant() { - let tanh: VisionConfig = serde_json::from_value(config(Some("gelu_pytorch_tanh"))).unwrap(); - assert_eq!(tanh.hidden_act, VisionHiddenActivation::GeluPytorchTanh); - let quick: VisionConfig = serde_json::from_value(config(Some("quick_gelu"))).unwrap(); - assert_eq!(quick.hidden_act, VisionHiddenActivation::QuickGelu); + fn vision_hidden_act_follows_the_hugging_face_act2fn_names() { + let cases = [ + ("gelu_pytorch_tanh", VisionHiddenActivation::GeluPytorchTanh), + ("gelu_new", VisionHiddenActivation::GeluPytorchTanh), + ("gelu_fast", VisionHiddenActivation::GeluPytorchTanh), + ("gelu", VisionHiddenActivation::ExactGelu), + ("gelu_python", VisionHiddenActivation::ExactGelu), + ("quick_gelu", VisionHiddenActivation::QuickGelu), + ]; + for (name, expected) in cases { + let parsed: VisionConfig = serde_json::from_value(config(Some(name))).unwrap(); + assert_eq!(parsed.hidden_act, expected, "{name}"); + } - for value in [config(None), config(Some("gelu")), config(Some("relu"))] { - let parsed: VisionConfig = serde_json::from_value(value).unwrap(); - assert_eq!(parsed.hidden_act, VisionHiddenActivation::ExactGelu); + let absent: VisionConfig = serde_json::from_value(config(None)).unwrap(); + assert_eq!(absent.hidden_act, VisionHiddenActivation::ExactGelu); + let mut null_act = config(None); + null_act["hidden_act"] = serde_json::Value::Null; + let null: VisionConfig = serde_json::from_value(null_act).unwrap(); + assert_eq!(null.hidden_act, VisionHiddenActivation::ExactGelu); + + // Non-GELU names were silently parsed as erf before #2794. + for name in ["relu", "silu"] { + let err = serde_json::from_value::(config(Some(name))) + .expect_err(name) + .to_string(); + assert!( + err.contains("unsupported vision hidden_act"), + "{name}: {err}" + ); } } #[test] fn missing_hidden_act_resolves_to_the_siglip_default_only_when_absent() { let cases = [ - (config(None), VisionHiddenActivation::SiglipDefaultTanh), + (config(None), VisionHiddenActivation::GeluPytorchTanh), (config(Some("gelu")), VisionHiddenActivation::ExactGelu), ( config(Some("gelu_pytorch_tanh")), @@ -248,7 +289,7 @@ mod tests { null_act["hidden_act"] = serde_json::Value::Null; let mut parsed: VisionConfig = serde_json::from_value(null_act.clone()).unwrap(); parsed.default_missing_hidden_act_to_siglip(&null_act); - assert_eq!(parsed.hidden_act, VisionHiddenActivation::SiglipDefaultTanh); + assert_eq!(parsed.hidden_act, VisionHiddenActivation::GeluPytorchTanh); } fn with_model_type(model_type: &str, hidden_act: Option<&str>) -> serde_json::Value { @@ -264,12 +305,12 @@ mod tests { ( "siglip_vision_model", None, - VisionHiddenActivation::SiglipDefaultTanh, + VisionHiddenActivation::GeluPytorchTanh, ), ( "siglip2_vision_model", None, - VisionHiddenActivation::SiglipDefaultTanh, + VisionHiddenActivation::GeluPytorchTanh, ), // CLIPVisionConfig().hidden_act == "quick_gelu". ("clip_vision_model", None, VisionHiddenActivation::QuickGelu), diff --git a/src/vision/encoders/siglip.rs b/src/vision/encoders/siglip.rs index 8b121b7ec..379ec858c 100644 --- a/src/vision/encoders/siglip.rs +++ b/src/vision/encoders/siglip.rs @@ -52,10 +52,10 @@ use crate::vision::config::{VisionConfig, VisionHiddenActivation}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum VisionMlpActivation { Exact, - PytorchTanh, - /// The tanh GELU of a SigLIP tower whose config omits `hidden_act` (Gemma 3, - /// PaliGemma, Aya Vision, Granite Vision 3.x, Bunny), evaluated by the - /// compiled f32 kernel (#2423, #2548). + /// The tanh GELU, declared by the checkpoint or resolved from an omitted + /// key on a SigLIP tower, evaluated by the compiled f32 kernel (#2423, + /// #2548). Before #2794 a declared `gelu_pytorch_tanh` ran the eager f32 + /// helper instead; the two agree up to f32 rounding. PytorchTanhCompiled, FastSigmoid, } @@ -73,8 +73,7 @@ pub(crate) fn select_mlp_activation( VisionMlpActivation::FastSigmoid } else { match hidden_act { - VisionHiddenActivation::GeluPytorchTanh => VisionMlpActivation::PytorchTanh, - VisionHiddenActivation::SiglipDefaultTanh => VisionMlpActivation::PytorchTanhCompiled, + VisionHiddenActivation::GeluPytorchTanh => VisionMlpActivation::PytorchTanhCompiled, VisionHiddenActivation::QuickGelu => VisionMlpActivation::FastSigmoid, VisionHiddenActivation::ExactGelu => VisionMlpActivation::Exact, } @@ -94,13 +93,11 @@ pub(crate) struct VisionMLP { impl VisionMlpActivation { /// The shared `mlxcel_core` formula this activation computes. /// - /// `PytorchTanh` keeps the f32 evaluation the private helper this replaced - /// used (#2423); `PytorchTanhCompiled` is the same formula as one compiled - /// kernel; `Exact` is erf. + /// `PytorchTanhCompiled` is the tanh form as one compiled f32 kernel; + /// `Exact` is erf. fn gelu(self) -> GeluActivation { match self { VisionMlpActivation::Exact => GeluActivation::Erf, - VisionMlpActivation::PytorchTanh => GeluActivation::TanhF32, VisionMlpActivation::PytorchTanhCompiled => GeluActivation::TanhF32Compiled, VisionMlpActivation::FastSigmoid => GeluActivation::Fast, } @@ -581,10 +578,12 @@ impl SigLipVisionModel { } /// Like [`Self::from_weights_with_quant`] but selects the encoder MLP GELU - /// variant. `use_fast_gelu = true` uses the sigmoid `GELU(approx="fast")` - /// that the Idefics2 vision tower uses. Otherwise `hidden_act` selects the - /// explicit PyTorch tanh approximation, with exact-erf GELU as the - /// compatibility default for missing and unknown values. + /// variant. `use_fast_gelu = true` forces the sigmoid `GELU(approx="fast")` + /// regardless of `hidden_act`. No caller passes `true` since #2794 moved + /// Idefics2 and Granite 4 Vision to the Hugging Face tanh GELU; the + /// parameter is kept because the entry point is shared with LLM-jp-VL and + /// its test support. Otherwise `hidden_act` selects the formula (see + /// [`VisionHiddenActivation`]). pub fn from_weights_with_quant_and_gelu( weights: &WeightMap, config: &VisionConfig, @@ -881,18 +880,16 @@ mod tests { select_mlp_activation(VisionHiddenActivation::ExactGelu, false), VisionMlpActivation::Exact ); + // A declared tanh runs the compiled kernel, same as an omitted key + // resolved to the SigLIP default (#2794). assert_eq!( select_mlp_activation(VisionHiddenActivation::GeluPytorchTanh, false), - VisionMlpActivation::PytorchTanh + VisionMlpActivation::PytorchTanhCompiled ); assert_eq!( select_mlp_activation(VisionHiddenActivation::GeluPytorchTanh, true), VisionMlpActivation::FastSigmoid ); - assert_eq!( - select_mlp_activation(VisionHiddenActivation::SiglipDefaultTanh, false), - VisionMlpActivation::PytorchTanhCompiled - ); // CLIP's `quick_gelu` (#2548), declared or the class default. assert_eq!( select_mlp_activation(VisionHiddenActivation::QuickGelu, false), @@ -908,7 +905,9 @@ mod tests { // process inside `cudaStreamEndCapture` in roughly one run in four. let _guard = crate::models::siglip_text::test_guard::lock(); let input = mlxcel_core::from_slice_f32(&[-3.0, -1.0, 0.0, 1.0, 3.0], &[5]); - let output = VisionMlpActivation::PytorchTanh.gelu().apply(&input); + let output = VisionMlpActivation::PytorchTanhCompiled + .gelu() + .apply(&input); mlxcel_core::eval(&output); let expected = [-0.003_637_433, -0.158_808, 0.0, 0.841_192, 2.996_362_7]; for (index, expected) in expected.into_iter().enumerate() { diff --git a/src/vision/encoders/siglip_block_tests.rs b/src/vision/encoders/siglip_block_tests.rs index 25cb1a75d..659351740 100644 --- a/src/vision/encoders/siglip_block_tests.rs +++ b/src/vision/encoders/siglip_block_tests.rs @@ -139,7 +139,7 @@ fn fixture() -> (EncoderLayer, UniquePtr) { &block_config(), 64, 4, - VisionMlpActivation::PytorchTanh, + VisionMlpActivation::PytorchTanhCompiled, ) .unwrap(); let x = mlxcel_core::from_slice_f32(&rng.values(32, 1.0), &[1, 4, 8]);