From 1080d0c4dc15ec97e40eacf17e81ff04c8518c88 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 2 Jan 2026 07:19:08 +0000 Subject: [PATCH] Fix Gemma3 QAT training instability with int8-int4 scheme Gemma3 models have a large vocabulary (262144 tokens) which causes training loss to explode when using int8 embedding quantization. This fix auto-detects Gemma3 models and switches from int8-int4 (phone-deployment) to int4 weight-only QAT for stable training. --- unsloth/models/_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index ccb547f58e..3851e18f92 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2198,6 +2198,18 @@ def _prepare_model_for_qat( from torchao.quantization.granularity import PerGroup, PerAxis from torchao.quantization.qat import QATConfig + # Gemma3 models have issues with int8 embedding quantization due to their + # large vocabulary size (262144). Auto-switch to int4 weight-only instead. + if qat_scheme == "int8-int4": + model_types = get_transformers_model_type(model.config) + is_gemma3 = any("gemma3" in mt or "gemma_3" in mt for mt in model_types) + if is_gemma3: + print( + "Unsloth: Gemma3 has a large vocabulary causing int8 embedding issues. " + "Switching to int4 weight-only QAT for training stability." + ) + qat_scheme = "int4" + if not isinstance(qat_scheme, TorchAOConfig): torchao_config: Optional[TorchAOConfig] = None if qat_scheme == "fp8-int4":