Distilled and W4A16 quantized Gemma 4 E4B-it

This is a compressed version of google/gemma-4-E4B-it prepared for the Resilient AI Challenge 2026 (image-to-text track) by the ECE-ILAB team. The model can be served directly with vllm serve --config vllm_config.yaml on the evaluation stack.

The first submission we made (internally called lambda) failed to load on the evaluator stack vllm/vllm-openai:v0.23.0-cu129 with the error ValueError: no module 'audio_tower.layers.0.feed_forward1.ffw_layer_1.linear.weight_packed'. We narrowed the problem down to the fact that Gemma 4 uses custom Linear subclasses inside its audio_tower (the Conformer ffw_layer_1.linear) and inside its vision_tower (the input_proj of the patch embedder). When the compressed-tensors loader in vLLM 0.23 wraps these layers as ReplicatedLinear modules, the resulting objects do not expose a .weight attribute, but the Transformers code for Gemma 4 accesses self.input_proj.weight.dtype directly. The same loading path crashes on both towers, just at different lines.

We first tried a partial fix where only the audio_tower was kept in BF16. That version loaded past the audio stage but then crashed inside the vision profile run with AttributeError: 'ReplicatedLinear' object has no attribute 'weight' at gemma4/modeling_gemma4.py:613. The fix needed to cover both towers, which is what this version (eta) does. The language model and the 84 distilled Per-Layer Embedding linears stay in W4A16 pack-quantized form, while audio_tower, vision_tower, lm_head and the various embedding modules remain in BF16. Validation was done on an L4 GPU running the exact PEReN docker image: the model loads in 98 seconds, takes 9.73 GiB of GPU memory, and leaves 7.84 GiB of KV cache at max_model_len 2048, which corresponds to a maximum concurrency of about 70x.

The compressed checkpoint is 11.0 GB on disk, down from 15.88 GB for the BF16 student, a 1.44x compression ratio. The submitted config is tuned for low-energy inference on the L4 evaluation hardware. We set max-model-len to 2048 because the KV cache scales linearly with context length and the VQA prompts in the benchmark never get close to that limit anyway. We use kv-cache-dtype fp8 which halves KV memory again with no observable quality loss on our internal benchmark. CUDA graphs are disabled (enforce-eager true) to save VRAM, and gpu-memory-utilization is left at the challenge default of 0.85. The generation_config.json caps max_new_tokens at 40, a value we calibrated on 500 representative VQA samples drawn from the official mix (vqav2, chartqa, docvqa, textvqa, coco_captions, ok_vqa, scienceqa, infographicvqa, mathvista, ai2d). On that sample the median output length was 25 tokens, the mean 42, and the 95th percentile 48, so 40 covers the vast majority of correct answers while preventing Gemma 4 from drifting into the long verbose generations it tends to produce on open-ended prompts. The sampling parameters required by the challenge (temperature 1.0, top_p 0.95, top_k 64) are kept as-is.

The training pipeline is in two stages. The first is a Flash-Aware INT4 distillation of the Per-Layer Embeddings, where every PLE Linear has its weights snapped to an INT4 grid (per-row asymmetric min-max) while gradients flow through standard BF16 ops. Loss is a temperature-scaled KL on response-only positions computed in FP32 to avoid the 262144-vocab overflow. The student was trained for 1901 out of 2000 steps with early stopping on stable KL, warmup 100, grad accumulation 8, learning rate 5e-5. Final KL averaged 0.022 over the last 100 steps with zero NaN events recorded. Training took 73 minutes on a single A100-SXM4-80GB and cost 0.248 kWh measured with CodeCarbon. The second stage is the post-training W4A16 pack-quantized step, applied via llmcompressor on the 343 Linear modules of the language decoder. The recipe ignores lm_head, audio_tower, vision_tower and the embedding modules, which is what guarantees the model actually loads on the evaluator.

A companion dataset at Achta/gemma-4-E4B-it-distillation-artifacts contains the distillation script, the PLE checkpoint and the calibration cache, for anyone who wants to reproduce the pipeline end to end.

ECE-ILAB