feat(convert): quantize gemma4_unified (12B) — recognize + keep vision_embedder bf16 - #76
Merged
Merged
Conversation
Wire the Gemma 4 12B "unified" encoder-free checkpoint (model_type `gemma4_unified`, arch Gemma4UnifiedForConditionalGeneration) into the `mlx convert` pipeline so `--q-mode affine --q-bits 4 --q-group-size 32` produces a correct quantized model. The runtime loader already supports it (PR #74); only the converter did not recognize it and the default quant predicate would corrupt one vision weight. - Route `gemma4_unified` to the existing `Gemma4Recipe`: extend `model_types()`, `recipe_for`, and `CONVERTIBLE_MODEL_TYPES`; treat it identically to `gemma4` in the registry-consistency sym8 allowlist. - Extend the TS auto-detect so `model_type == "gemma4_unified"` maps to `gemma4`. - Exclude `vision_embedder.*` from `should_quantize`: the loader installs it as dense bf16 only (no `.scales` branch), so quantizing `patch_dense.weight` would corrupt the unified vision path. `embed_vision` / `embed_audio` keep quantizing (they have affine loader branches). - Keep `vision_embedder.*` at its bare key in `Gemma4Recipe::sanitize` (sibling of `embed_vision.*`) instead of mis-prefixing under `language_model.model.`. - Add focused unit tests for the quant-skip, recipe recognition, and bare key placement; document `gemma4_unified` in docs/cli.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The CLI auto-detect collapsed model_type 'gemma4_unified' to 'gemma4'
before calling native convertModel. That is wrong: native recipe_for
already resolves 'gemma4_unified' to the shared Gemma4Recipe, so passing
the raw string yields identical behavior for every recipe-keyed path
(sym8_supported, sanitizer, embed_quantizable, mtp_policy, etc.), while
collapsing it (a) made the native recipe_for("gemma4_unified") arm dead
code and (b) misrouted a unified checkpoint carrying gemma-QAT metadata
into the E2B-only prequantized importer, whose gate keys on the exact
string "gemma4" and then hard-errors in validate_e2b_qat_schedule.
Pass 'gemma4_unified' through unchanged; keep 'gemma4'/'gemma4_text'
collapsing to 'gemma4'. Add a regression test that drives run() against a
synthetic config.json and asserts the modelType handed to the mocked
native convertModel: 'gemma4_unified' stays 'gemma4_unified', while
'gemma4' and 'gemma4_text' still resolve to 'gemma4'.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ce4623db9b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
The runtime loader treats an architecture-only unified config (no `model_type`, only `architectures: ['Gemma4UnifiedForConditionalGeneration']`) as unified (model-loader.ts maps it to gemma4; persistence.rs parse_config sets is_unified on EITHER model_type == "gemma4_unified" OR that architecture). The CLI converter auto-detect inspected only `config.model_type`, so such a config left modelType undefined, skipped Gemma4Recipe::sanitize on the native side, and produced unloadable output. Extend the gemma4_unified auto-detect arm to also match the architecture-only shape, mapping to the same 'gemma4_unified' pass-through string (not 'gemma4', which would re-introduce the E2B-importer misroute the exact-"gemma4" gate guards against). Adds a regression test covering the architecture-only config. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The `e2e <family>` cargo legs run `cargo test … -- --ignored` with no thread cap, so libtest's default per-test thread pool (= nproc) launches several `#[ignore]` tests at once. Each loads 2-4 full models (lfm2's `lfm2_paged_vs_flat_parity` has 6 such tests, ~8 LFM2-1.2B live at peak), all driving the single shared Metal GPU. Under that oversubscription on the shared macos-latest runner a command buffer occasionally waits past the GPU watchdog -> kIOGPUCommandBufferCallbackErrorTimeout thrown at device.cpp:456 -> uncaught -> SIGABRT, flaking the leg (observed on `lfm2_paged_vs_flat_greedy_token_parity`; a no-change rerun went green, confirming a resource flake, not a regression). Serialize the cargo e2e tests with `--test-threads=1`. A single shared GPU means the tests never actually ran in parallel — parallelism only piled on GPU memory/queue pressure — so this costs ~no wall time. The parity/oracle comparisons and `temperature=0` greedy determinism are unchanged, so a real regression still fails. The gemma4 TS leg is single-instance Vitest and unaffected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
`runs-on: macos-latest` resolves to a HETEROGENEOUS pool (macOS 26 and macOS 15 both present). The Build job compiles the Metal shaders with `xcrun -sdk macosx metal` and no `-std=metal` pin, so it inherits the builder's SDK default Metal Shading Language version. When Build lands on macOS 26 it emits an MSL 4.0 `mlx.metallib` / `paged_attn.metallib`; that artifact is then downloaded by the consumer jobs (yarn test, and the `yarn mlx convert` step of the e2e legs, which use the prebuilt addon). A consumer that lands on macOS 15 cannot load it: Failed to load the default metallib. This library is using language version 4.0 which is not supported on this device. -> C++ exception across the Rust FFI boundary -> "fatal runtime error: Rust cannot catch foreign exceptions, aborting" -> convert never writes its output -> the lfm2_session tests panic "model path does not exist". Pure runner-pool roulette (which OS the builder vs each consumer draws), so it surfaces intermittently. Pin every macOS job to `macos-26` so the builder and all consumers share one Metal version (MSL 4.0 loads on every macos-26 runner). The single ubuntu publish job is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
1 participant
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.
What
Enable
mlx convertto quantizegemma4_unified(Gemma 4 12B — encoder-free text+vision+audio) checkpoints. The runtime loader already supportsgemma4_unified(PR #74); only the converter didn't recognize it, and the default quant predicate would have corrupted one vision weight.Motivating model:
google/gemma-4-12B-it-qat-q4_0-unquantized(dense 12B, bf16, QAT-trained for llama.cppq4_0). The best quality+speed target is 4-bit affine group_size 32 — near-bf16 quality at 4-bit memory/speed.Changes
gemma4_unified—Gemma4Recipe::model_types/recipe_for/CONVERTIBLE_MODEL_TYPESroute it to the existing sharedGemma4Recipe(it is dense — no MoE split, sym8-supported, no MTP — so it behaves identically togemma4on everyrecipe_for-keyed path).vision_embedder.*bf16 —should_quantize()excludes any key containingvision_embedder. The encoder-free unified vision patch projection is installed dense-bf16-only by the loader (apply_unified_vision_embedder_weights, no.scalesbranch), so quantizingpatch_dense.weightwould corrupt the vision path.embed_vision/embed_audioprojections still quantize (they have affine loader branches).vision_embedderis unified-only → zero collateral on other families.vision_embedder.*bare — sibling ofembed_vision.*, not mis-prefixed underlanguage_model.model..gemma4_unifiedthrough unchanged — not collapsed to'gemma4'. The driver'smodel_typeis the CLI value, and the gemma-QAT (wNa8o8) prequantized importer gate is exactSome("gemma4"); collapsing would (a) dead-code the nativerecipe_for("gemma4_unified")arm and (b) misroute a gemma-QAT unified checkpoint into the E2B-only importer.End-to-end validation
Converted the real 23GB checkpoint → 8.4GB
{bits:4, group_size:32, mode:affine}:gemma4_unified(passthrough confirmed)vision_embedder.*→ bf16, no.scales✓ · text /embed_vision/embed_audio→ quantized (U32 + scales) ✓Tests
should_quantizeexcludesvision_embedder(with positive controls thatembed_vision+ text layers still quantize);gemma4_unifiedresolves to a recipe + is inCONVERTIBLE_MODEL_TYPES; sanitize keepsvision_embedder.*bare; extendedrecipe_registry_reproduces_inline_flags(sym8 allowlist).convert-cmd.test.tsdrivesrun()and asserts themodelTypehanded to native isgemma4_unified(andgemma4/gemma4_textstill map togemma4).convert::106 pass, fmt + clippy clean, CLI test 4/4, typecheck clean.Review
Three
/codex:adversarial-reviewcycles → final verdict approve, no material findings. The CLI-collapse misroute (cycle #1) was fixed; a flagged sanitize-test SIGABRT (cycle #2) was confirmed to be a review-harness/metallib environment artifact (the pre-existinggemma4_recipe_sanitize_transformstest aborts identically; CI runs one process per binary).🤖 Generated with Claude Code
Note
Medium Risk
Changes conversion routing and quantization rules for a new model type (wrong passthrough could misroute QAT checkpoints); CI runner change is low risk but affects all macOS jobs.
Overview
Adds first-class
gemma4_unifiedsupport inmlx convertso encoder-free Gemma 4 unified checkpoints can be quantized without breaking vision or misrouting QAT imports.The native converter registers
gemma4_unifiedon the sharedGemma4Recipe, excludesvision_embedder.*from quantization (bf16-only loader path), and keeps those weights at bare keys during sanitize. The CLI auto-detectsgemma4_unified(including architecture-only configs) and passes the string through instead of collapsing togemma4, avoiding the E2B prequantized importer gate.CI moves macOS jobs to
macos-26and runs ignored cargo model e2e tests with--test-threads=1to avoid Metal GPU watchdog timeouts on shared runners. Docs and TS/Rust tests cover detection, quant exclusions, and sanitize routing.Reviewed by Cursor Bugbot for commit fb4e8e0. Bugbot is set up for automated code reviews on this repo. Configure here.