Repository navigation
investigate: gelu_approx computes erf GELU where upstream models use the tanh approximation #2423
Description
Problem / Background
mlxcel_core::gelu_approx (src/lib/mlxcel-core/src/utils.rs:1049, C++ gelu_approx at src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:3287 and compiled_gelu_approx at :1891) computes the exact erf GELU, x * 0.5 * (1 + erf(x / sqrt(2))). In MLX, nn.gelu_approx is the tanh approximation 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3))), nn.gelu_fast_approx is x * sigmoid(1.702 * x), and nn.GELU(approx="precise") and approx="tanh" both select the tanh form, while only approx="none" (and nn.gelu) is erf. Source: https://github.com/ml-explore/mlx/blob/main/python/mlx/nn/layers/activations.py (gelu, gelu_approx, gelu_fast_approx, class GELU).
The erf choice is deliberate. The comment in gelu_approx() says:
The tanh approximation ... produced NaN for negative bf16 inputs in the SigLIP vision encoder MLP because power(x, 3.0f) uses exp(3log(x)) which is undefined for x<0. Even with xx*x the MLX lazy graph fuses operations that can overflow bf16 intermediates within the large SigLIP computation graph (4096 patches). The erf-based GELU ... matches Python nn.GELU(approx="precise") output within floating-point tolerance.
The last sentence is wrong: approx="precise" is the tanh form. Meanwhile the tree already runs tanh GELU in production: the shared C++ helper gelu_tanh_approx (src/lib/mlxcel-core/cpp/mlx_cxx_internal.h:91, uses x*x*x, no power) backs every fused GeGLU path (Gemma 4 and others), and at least five private Rust copies exist (src/models/bert.rs:303, src/models/kokoro/ops.rs:223, src/vision/encoders/molmo2.rs:90, src/vision/encoders/siglip.rs:92, src/vision/encoders/kimi_vl.rs:158).
Sites known to diverge
- InternViT MLP,
src/vision/encoders/internvl.rs:208-213: upstreamnn.GELU(approx="precise")(https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/internvl_chat/vision.py); the local comment claimsgelu_approxis the "precise" GELU. - MiniCPM-V 4.6 mergers,
src/vision/encoders/minicpmv4_6.rs:343and:393: upstreamnn.GELU(approx="precise")(https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/minicpmv4_6/minicpmv4_6.py). - MiniCPM-o vision MLP (
MiniCPMOMlp),src/vision/encoders/minicpmo.rs:237: upstreamnn.GELU(approx="precise")(https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/minicpmo/vision.py).
Not every call site is wrong: src/models/gpt_bigcode.rs:561-569 ports upstream nn.gelu (erf), so erf is correct there and only the function name misleads. There are 83 non-test gelu_approx references under src/, mostly in src/models/ and src/vision/encoders/. The new ports #2393 (InternVL HF layout) and #2398 (MiniCPM-V 4.7) already require the tanh form for their own paths.
Acceptance Criteria
- Inventory every
gelu_approx,compiled_gelu_approx,gelu, and private tanh-helper call site, each mapped to the upstream mlx-lm / mlx-vlm activation it ports (erf, tanh, or fast/sigmoid) with a web URL, posted as a table in an issue comment or committed with the PR. - Re-test the bf16 stability rationale: run tanh GELU computed with
x*x*x(asgelu_tanh_approxdoes), both in input dtype and in f32 then cast back, on the SigLIP encoder (4096 patches) with a bf16 checkpoint, and report whether any NaN or overflow reproduces. - For each family where upstream uses tanh, measure erf against tanh with the teacher-forced logit trace in
docs/benchmarks.md(examples/logit_trace.rsplusscripts/compare_logit_traces.py) against the mlx-vlm / mlx-lm reference, gating on disagreement at decided positions, including the reverted arm that must fail. (PR fix(vision): add GELU selector, use tanh where checkpoints declare it #2552: measured for InternVL, MiniCPM-V 4.6, Gemma 3, PaliGemma, Gemma 4 vision, Gemma3n sparse; remaining families in fix(models): decide the GELU formula for the families #2423 left open #2548) - Decide per family. If behavior changes, add an explicit activation selector (erf / tanh / fast) in
mlxcel-corereusinggelu_tanh_approxrather than silently changinggelu_approx's semantics, fold the private Rust copies into it, and update the// Used by:comments perdocs/code-guidelines.md. (PR fix(vision): add GELU selector, use tanh where checkpoints declare it #2552: selector added and private copies folded; decided for the measured families; remaining in fix(models): decide the GELU formula for the families #2423 left open #2548) - Correct the misleading comments:
gelu_approx()andcompiled_gelu_approxinmlx_cxx_bridge.cpp, the doc comment atutils.rs:1046, andinternvl.rs:210-211. - Any behavior change is reached through the real CLI and server load path and smoke-tested on a real checkpoint per family changed. (PR fix(vision): add GELU selector, use tanh where checkpoints declare it #2552: smoke-tested for every family it changed; remaining in fix(models): decide the GELU formula for the families #2423 left open #2548)
Verification
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --check
Activity
fix(vision): use tanh GELU where the checkpoint declares it
docs: record GELU smoke tests and the Gemma 4 reference A/B
inureyes commentedon Oct 22, 2026
Investigation results, implemented in #2552 (uses
Refs; this issue stays open).Full inventory, about 110 call sites each mapped to its upstream formula with pinned links (mlx-lm c39139d70, mlx-vlm v0.7.8, mlx-audio v0.5.9) and the Hugging Face activation where it differs: docs/gelu-call-sites.md. Measurements: docs/benchmark_results/gelu-erf-vs-tanh-2026-10-23.md.
bf16 stability. tanh GELU with
x*x*x, in input dtype and in f32, on gemma-3-4b-it-4bit's bf16 SigLIP tower at 4096 patches (27 x [1, 4096, 4304], max |input| 161 to 177), MiniCPM-V 4.6 bf16 and PaliGemma2 f16: zero NaN and zero Inf. The original NaN came frompower(x, 3.0f)(exp(3 log x)), not from tanh.Gate. For every measured family the reverted (erf) arm does not fail: swapping erf and tanh moves logits by at most the bf16 quantum and changes no decided token, inside mlxcel and inside the float32 mlx-vlm reference alike. The decided-position gate cannot distinguish the two, so each family was decided by the activation the checkpoint declares, cross-checked against mlx-vlm.
gelu_pytorch_tanhhidden_actgelu_pytorch_tanh(hard-coded)nn.GELU()(erf)hidden_act = gelu(erf)gelu_pytorch_tanhhidden_activationgelu_pytorch_tanhEvery changed family was smoke-tested through
mlxcel generate --imageand onemlxcel-serverrequest (DiffusionGemma CLI only; its server rejects image input).Open. About 20 remaining mismatched families (Gemma3n sparse, Gemma 4 bf16 dense MLP and drafter, Gemma 4 TP gate, RecurrentGemma, MiniCPM-o, Aya/Granite SigLIP, CLIP
quick_geludefaults, Phi, GPT-2, Moondream, LFM2-VL, DeepSeek-VL2 and others) are in #2548 with one acceptance item each. Separately,# prompt_tokensdisagree with mlx-vlm for InternVL (262 vs 268), MiniCPM-V 4.6 (78 vs 76) and Gemma 3 (273 vs 274), which invalidates those cross-runtime comparisons: #2547.docs: record the prefill cost of the f32 tanh GELU switch
update(core): compile the f32 tanh GELU for the switched vision towers
fix(cuda): keep PTX cache temporaries under NAME_MAX
docs: record the compiled tanh GELU prefill A/B per family
fix(vision): add GELU selector, use tanh where checkpoints declare it (…
inureyes commentedon Oct 22, 2026
PR #2552 merged as 58db639 with
Refs #2423; this issue stays open for the families it did not change.Landed: criteria 1 (inventory in
docs/gelu-call-sites.md), 2 (bf16 re-test: no NaN or Inf withx*x*xtanh in input dtype or f32 on gemma-3-4b SigLIP at 4096 patches, MiniCPM-V 4.6 bf16 and PaliGemma2; the original NaN came frompower(x, 3)), 5 (comments), and the selector and folds of criterion 4.Switched to tanh (computed in f32 through
compiled_gelu_tanh_f32), because the checkpoint and the mlx-vlm port both declare it: MiniCPM-V 4.6 encoder MLP and window merger, Gemma 3 VLM SigLIP and PaliGemma SigLIP (missinghidden_actresolves to the SigLIP default), Gemma 4 vision GeGLU. Each was smoke-tested throughmlxcel generate --imageandmlxcel-server(DiffusionGemma CLI only).The teacher-forced trace gate could not distinguish erf from tanh in any measured family: the reverted arm never fails (0 decided flips; logits move by at most the bf16 step), so these decisions rest on the declared activation, not on the gate.
Measured image-prefill cost against erf on GB10 (interleaved, null arm): Gemma 3 +2.8% pooled (same sign in both blocks), Gemma 4 +2.0%, MiniCPM-V 4.6 and PaliGemma within the null arm. The eager f32 form had cost Gemma 3 +17%; a single-kernel fusion for the remainder is tracked in #2548.
Open: every other family in #2548, and the mlxcel vs mlx-vlm prompt-length mismatch in #2547 that invalidates cross-runtime traces for InternVL, MiniCPM-V 4.6 and Gemma 3. Record:
docs/benchmark_results/gelu-erf-vs-tanh-2026-10-23.md.fix(models): settle the GELU formula for the families #2423 left open (…
inureyes commentedon Nov 3, 2026
Every remaining acceptance criterion here was moved to #2548, which closed on 2026-10-31 with PR #2785 and follow-up PR #2804 (
status:done). Criteria 1, 2, 5 and the selector half of 4 landed in PR #2552; the per-family measurement (3), the per-family decisions and private-copy folds (4), and the real CLI/server smoke tests (5) were completed in #2548.No implementation work remains, so this issue closes as done rather than shipping another PR.