The Quantization Accuracy Gap
When you compress a 16-bit float model to 4-bit, standard Post-Training Quantization (PTQ) chops off the numbers, introducing "quantization noise" that degrades accuracy.
Quantization-Aware Training (QAT) pre-conditions the model's weights during pre-training to model low-precision noise beforehand. When you compress a QAT model, it loses almost zero accuracy.
But training LoRA adapters on a QAT base model introduces a strict mathematical constraint: your adapters must not bypass the QAT paths. If you run LoRA using naive high-learning rates or incorrect export types (like standard q8_0 or q4_k_m), you lose the QAT benefit.
We updated the gemma4-tuning MLX pipeline to enforce strict QAT constraints and created a validation system to track the semantic drift.
The Fixes
1. Conditionally Overriding Attention Parameter Instantiation
Gemma 4 unquantized QAT checkpoints contain a strict structural loading mismatch in standard MLX, leading to key parsing failures on load. We patched the initialization path inside our CLI runner to conditionally capture and instantiate KV projections and normalization boundaries on the fly.
See the loader patch in src/gemmmma/cli.py.
2. Eliminating Repetition with Conservative Learning Rates
Our initial dry runs with QAT models showed validation loss climbing, causing the model to get stuck in infinite generation loops (e.g., repeating "plays it plays it" endlessly).
By lowering the learning rate to a conservative 2e-6 from a clean initialization, we stabilized the gradient paths, dropping validation loss from 2.26 to 0.70 with zero text repetition.
# Our stable, high-fidelity QAT training configuration
uv run mlxtune train \
--model ./model/google-gemma-4-E2B-it-qat-q4_0-unquantized \
--qat \
--iters 20 \
--batch-size 2 \
--rank 8 \
--lr 2e-63. Strict GGUF q4_0 Alignment
To ensure the model actually executes along the pre-conditioned QAT pathways, we configured our GGUF export command to override standard parameters and enforce the q4_0 quantization scheme whenever --qat is enabled:
# From src/gemmmma/cli.py
if qat:
if outtype != "q4_0":
print("==> [QAT Optimization] Overriding outtype to 'q4_0' for QAT alignment.")
outtype = "q4_0"Check out the GGUF compilation step in src/gemmmma/cli.py.
🔬 MLOps: Tracking Semantic Drift
To verify that our quantized outputs match the high-precision reference outputs, we implemented an automated drift analyzer. It calculates a word-level Jaccard Similarity and alerts the developer if the model is diverging:
# Our vocabulary similarity metric
intersection = len(set(ref_words) & set(gguf_words))
union = len(set(ref_words) | set(gguf_words))
jaccard_score = intersection / union if union > 0 else 0
semantic_drift = (1.0 - jaccard_score) * 100All parameters, GPU memory footprints (which peaked at 12.43 GB on our Mac), validation trends, and inference responses. Our local experiment journal dynamically logs... all parameters...
- Read our complete Experiment Journal.
- See the raw logged data in experiments.json.
- View the concepts guide in docs/PIPELINE.md.
- Check out active hardware workarounds in our Friction Log.
🤖 Teaching the Machines: agentskills.io Integration
Local model training involves dozens of steps and fine-grained parameters, making it easy for AI coding assistants to lose track of constraints.
We authored four repository-housed AI Agent Skills conforming strictly to the agentskills.io specification:
gemma-fine-tuning: Text-based LoRA training loops on custom ChatML data.gemma-qat-tuning: Alignment parameters, unquantized checkpoints, and strictq4_0GGUF compilation.gemma-multimodal-tuning: Multimodal (Vision & Audio) ChatML formatting,mlx-vlmtraining, and key-sanitized weights saving.gemma-model-export: Compiling fused models into multi-format mobile runtimes (GGUF, LiteRT-LM, MLX Swift) and implementing sandboxed persistent storage.
Our main README.md and GEMINI.md guidelines have been aligned so that any agent entering the workspace is immediately equipped with these capabilities.