Add Gemma-4 (Block-AP) support - #1
Open
icnatspell wants to merge 1 commit into
Open
Conversation
Quantize google/gemma-4-E2B-it (Gemma-3n/MatFormer: per-layer embeddings, KV-sharing, alternating sliding/global attention, multimodal wrapper) via Block-AP. Model-specific logic is isolated in a Gemma4Adapter so block_ap.py stays model-agnostic. - quantize/model_adapters.py: BlockAPAdapter base + Gemma4Adapter (PLE with per-batch cache, KV producer/consumer wiring, static-kwarg capture). - quantize/block_ap.py: adapter-driven; follows the model's dtype (Gemma needs bf16 — fp16 overflows), device-selectable (cuda/cpu). - run_gemma4_w4.py: fp16 vs w4 PPL with BOS-prefixed deterministic windows; two eval backends (whole-model CPU, low-mem block-walk GPU); RTN ablation via --epochs 0. - datautils_block.py: stream C4 (new datasets lib can't load it by single-name id; full-shard download OOMs). - GEMMA4_QAT.md + README section; pyproject.toml/uv.lock; smoke test. Result (RTX 3080, w4g128, C4): fp16 109.09 -> w4 109.59 (+0.5% PPL), ~6GB VRAM. 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.
Summary
Adds Block-AP quantization-aware training for
google/gemma-4-E2B-it(Gemma-3n / MatFormer lineage: per-layer embeddings, KV-sharing, alternating sliding/global attention, multimodal wrapper). Model-specific logic is isolated in aGemma4Adaptersoquantize/block_ap.pystays model-agnostic — supporting another model is a new adapter, not a fork.Result
RTX 3080 (10GB), w4g128, C4 calib+eval: fp16 109.09 → w4 109.59 (+0.5% PPL) — near-lossless 4-bit, ~6GB peak VRAM, ~3s/block.
Changes
quantize/model_adapters.py—BlockAPAdapterbase +Gemma4Adapter(PLE with per-batch cache, KV producer/consumer wiring, static-kwarg capture).quantize/block_ap.py— adapter-driven; follows the model's loaded dtype (Gemma needs bf16 — fp16 overflows its activations), device-selectable (cuda/cpu).run_gemma4_w4.py— fp16 vs w4 PPL over BOS-prefixed deterministic windows; two eval backends (whole-model CPU, low-mem block-walk GPU validated to 0.066%); RTN ablation via--epochs 0.datautils_block.py— stream C4 (newdatasetslib can't load it by single-name id; full-shard download OOMs).GEMMA4_QAT.md+ README section + uv install instructions;pyproject.toml/uv.lock; smoke test.Notes / gotchas baked in
add_bos_token=Falseand the model is very BOS-sensitive (no-BOS PPL ~1000x inflated).-itmodel on raw corpora, not a bug; only the fp16→w4 delta measures quantization quality.Vision/audio encoders left in fp16; only the text decoder is quantized.
🤖 Generated with Claude Code