Skip to content

[MODEL] support gemma4_unified#2921

Merged
Qubitium merged 1 commit into
ModelCloud:mainfrom
HaozheZhang6:support-gemma4-unified
Jun 15, 2026
Merged

[MODEL] support gemma4_unified#2921
Qubitium merged 1 commit into
ModelCloud:mainfrom
HaozheZhang6:support-gemma4-unified

Conversation

@HaozheZhang6

Copy link
Copy Markdown
Contributor

Summary

gemma4_unified (multimodal Gemma 4, e.g. OpenYourMind/gemma-4-12B-it-abliterated-uncensored) fails with Unsupport model_type gemma4_unified, and failed to auto-detect module tree. It's a close sibling of the existing gemma4 composite model, but its text stack drops the per-layer input adapters, so it needs its own definition rather than reusing Gemma4ForConditionalGenerationGPTQ.

What Changed

  • New Gemma4UnifiedForConditionalGenerationGPTQ definition over model.language_model.layers: per-projection q/k/v norms + dual pre/post feed-forward norms (same as gemma4), but no per_layer_input_gate / post_per_layer_input_norm / per_layer_projection — this variant has none. Its language_model also has no project_per_layer_inputs, which is why reusing the gemma4 class AttributeErrors during the calibration generate hook.
  • prepare_layer_replay_kwargs refreshes the rope per layer: the model builds one (cos, sin) per attention layer_type and hands each decoder layer the tuple for its own type, so cached-layer replay has to regenerate it for the layer's layer_type (sliding vs full). Without it, replay reuses a stale tuple and hits a rope shape mismatch in apply_rotary_pos_emb.
  • Registered gemma4_unified in MODEL_MAP and re-exported from definitions/__init__.py.

Vision/audio towers stay as base (non-quantized) modules, same as gemma4.

Tests

  • I added a new simple/fast unit test for this change, or documented why that is not applicable.
  • I ran the new targeted test locally before opening this PR.
  • I ran any other directly relevant local tests.

tests/test_gemma4_unified_support.py (CPU-only): model_type → definition resolution, module tree excludes the per-layer-input paths, and the rope refresh returns the right tuple per layer_type.

$ pytest tests/test_gemma4_unified_support.py -v
tests/test_gemma4_unified_support.py::test_gemma4_unified_model_type_selects_definition PASSED
tests/test_gemma4_unified_support.py::test_gemma4_unified_module_tree_excludes_per_layer_input_paths PASSED
tests/test_gemma4_unified_support.py::test_gemma4_unified_replay_kwargs_refresh_position_embeddings_per_layer_type PASSED
3 passed in 3.93s

I also ran a full quantize → save → reload → generate on a tiny gemma4_unified (4-bit, group_size 128) on an A100: quantization completes across both the sliding and full attention layers, the saved model reloads with TritonV2Linear packed layers, and generation runs without error.

Review Requirements

  • I personally reviewed every file in this diff.
  • I checked that the code matches existing project structure, APIs, and conventions.
  • I avoided unnecessary monkeypatching and used the project's normal extension points where possible.

Notes

The definition follows gemma4.py's rope-replay approach but is kept as a separate class/file since the per-layer-input machinery there doesn't apply here. I didn't have the full 12B checkpoint on hand, so the end-to-end run used a tiny synthetic gemma4_unified — happy to validate against a specific checkpoint if you'd like.

Gemma 4 unified (multimodal) checkpoints fail to quantize with
"Unsupport model_type gemma4_unified". Add a dedicated definition that
reuses the composite Gemma 4 decoder layout (per-projection q/k/v norms,
dual pre/post feed-forward norms) but drops the per-layer input adapters
this variant does not have, and refreshes the rope (cos, sin) per layer
during cached replay so the sliding/full attention boundaries stay correct.

Fixes ModelCloud#2920

Signed-off-by: HaozheZhang6 <zhang.hz6666@gmail.com>

@Qubitium Qubitium left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@HaozheZhang6 Thank you! LGTM.

@Qubitium Qubitium merged commit 9de4b03 into ModelCloud:main Jun 15, 2026
1 check passed
Qubitium pushed a commit that referenced this pull request Jun 19, 2026
…2925)

The gemma4_unified family also has a standalone text stack (model_type
"gemma4_unified_text", Gemma4UnifiedTextModel) registered in transformers,
mirroring how gemma4_text parallels gemma4. #2921 added the composite
gemma4_unified but not the text-only variant, so a text-only checkpoint
still fails with "Unsupport model_type gemma4_unified_text".

This extracts the shared decoder layout into a constant and adds
Gemma4UnifiedTextQModel rooted at model.layers (model.norm /
model.rotary_emb), reusing the same per-layer rope replay helper.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants