diff --git a/tests/test_gemma4_moe_4bit_swap.py b/tests/test_gemma4_moe_4bit_swap.py new file mode 100644 index 0000000000..f12b592dc4 --- /dev/null +++ b/tests/test_gemma4_moe_4bit_swap.py @@ -0,0 +1,147 @@ +"""Unit tests for the Gemma-4 MoE per-expert Linear4bit swap (#5344). + +End-to-end correctness on the real 26B-A4B checkpoint requires a GPU + the +checkpoint on disk, so this file restricts itself to fast CPU-only tests +that exercise the swap helper's shape contract, idempotence, and gating +behaviour. The full repro (resident VRAM 46 GB -> 14.27 GB, cosine sim 0.994 +vs BF16) is documented in the PR description. +""" + +import importlib +import os + +import torch +import torch.nn as nn + + +def _stub_gemma4_module(): + """Construct a stub Gemma4TextExperts-like module without importing + transformers' Gemma4Config (which would force a fresh transformers + download in CPU-only CI).""" + try: + from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts + except Exception: + return None + + # The class init requires a config; build a tiny synthetic one and then + # overwrite the fused weights with shapes small enough for CPU tests. + class _StubConfig: + num_experts = 4 + hidden_size = 16 + moe_intermediate_size = 8 + hidden_activation = "gelu_pytorch_tanh" + + module = Gemma4TextExperts.__new__(Gemma4TextExperts) + nn.Module.__init__(module) + module.num_experts = _StubConfig.num_experts + module.hidden_dim = _StubConfig.hidden_size + module.intermediate_dim = _StubConfig.moe_intermediate_size + module.gate_up_proj = nn.Parameter( + torch.randn( + _StubConfig.num_experts, + 2 * _StubConfig.moe_intermediate_size, + _StubConfig.hidden_size, + dtype = torch.bfloat16, + ), + requires_grad = False, + ) + module.down_proj = nn.Parameter( + torch.randn( + _StubConfig.num_experts, + _StubConfig.hidden_size, + _StubConfig.moe_intermediate_size, + dtype = torch.bfloat16, + ), + requires_grad = False, + ) + from transformers.activations import ACT2FN + + module.act_fn = ACT2FN[_StubConfig.hidden_activation] + return module + + +def test_is_enabled_reads_env_var(): + from unsloth.models import gemma4_moe_4bit + + old = os.environ.pop("UNSLOTH_GEMMA4_MOE_4BIT", None) + try: + assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is False + os.environ["UNSLOTH_GEMMA4_MOE_4BIT"] = "1" + assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is True + os.environ["UNSLOTH_GEMMA4_MOE_4BIT"] = "0" + assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is False + finally: + if old is None: + os.environ.pop("UNSLOTH_GEMMA4_MOE_4BIT", None) + else: + os.environ["UNSLOTH_GEMMA4_MOE_4BIT"] = old + + +def test_swap_skips_models_without_gemma4_experts(): + from unsloth.models.gemma4_moe_4bit import ( + swap_gemma4_experts_to_per_expert_linear4bit, + ) + + model = nn.Sequential(nn.Linear(8, 8), nn.Linear(8, 8)) + assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0 + + +def test_swap_skips_when_transformers_lacks_gemma4(): + """If transformers does not expose Gemma4TextExperts, the helper must + return 0 without raising. We simulate the ImportError by patching.""" + import unsloth.models.gemma4_moe_4bit as g4m + + real_import = importlib.import_module + + def _broken_import(name, *args, **kwargs): + if name == "transformers.models.gemma4.modeling_gemma4": + raise ImportError("simulated absence") + return real_import(name, *args, **kwargs) + + try: + importlib.import_module = _broken_import + # Re-exercise via the public helper. It imports Gemma4TextExperts + # inside its try/except, so the simulated ImportError must yield 0. + model = nn.Sequential(nn.Linear(8, 8)) + assert g4m.swap_gemma4_experts_to_per_expert_linear4bit(model) == 0 + finally: + importlib.import_module = real_import + + +def test_swap_idempotent_on_stub_module_without_cuda(): + """On CPU we cannot exercise bnb (Linear4bit requires CUDA). Verify the + helper at least returns 0 for the no-bnb-experts case without raising, + and is idempotent across repeated calls.""" + from unsloth.models.gemma4_moe_4bit import ( + swap_gemma4_experts_to_per_expert_linear4bit, + ) + + if not torch.cuda.is_available(): + # CPU-only: bnb's Linear4bit init would fail. Validate the model-walk + # path on an empty Sequential to confirm the helper is side-effect-free. + model = nn.Sequential(nn.Linear(4, 4)) + assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0 + assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0 + return + + # GPU path: build the stub and run a real swap. + module = _stub_gemma4_module() + if module is None: + return # transformers without gemma4 module: nothing to test + model = nn.Sequential(module.to("cuda")) + n1 = swap_gemma4_experts_to_per_expert_linear4bit(model) + n2 = swap_gemma4_experts_to_per_expert_linear4bit(model) + assert n1 == 1 + assert n2 == 0 # idempotent: already-swapped modules are skipped + assert hasattr(module, "gate_up_proj_4bit") + assert hasattr(module, "down_proj_4bit") + assert len(module.gate_up_proj_4bit) == module.num_experts + assert len(module.down_proj_4bit) == module.num_experts + + +if __name__ == "__main__": + test_is_enabled_reads_env_var() + test_swap_skips_models_without_gemma4_experts() + test_swap_skips_when_transformers_lacks_gemma4() + test_swap_idempotent_on_stub_module_without_cuda() + print("All 4 swap tests passed.") diff --git a/tests/test_issue_5344_guardrail.py b/tests/test_issue_5344_guardrail.py new file mode 100644 index 0000000000..09599bb071 --- /dev/null +++ b/tests/test_issue_5344_guardrail.py @@ -0,0 +1,177 @@ +"""Unit tests for the unslothai/unsloth#5344 silent-quantization-bypass guardrail. + +Covers two failure modes the helper detects: + 1. total bypass: load_in_4bit was requested but no bnb modules exist. + 2. partial bypass: bnb quantized nn.Linear but a large fraction of weight + bytes live in non-nn.Linear Parameters (e.g. Gemma-4 MoE fused experts). +""" + +import warnings + +import torch +import torch.nn as nn + + +# unsloth must be imported before transformers per its loading order, but +# these tests do not exercise the real loader. Import the helper directly. +from unsloth.models.vision import _warn_if_quantization_silently_dropped + + +class _PretendLinear4bit(nn.Module): + """type(m).__name__ == 'Linear4bit' so the guardrail counts it as quantized.""" + + def __init__(self): + super().__init__() + self.weight = nn.Parameter( + torch.zeros(1, dtype = torch.uint8), + requires_grad = False, + ) + + +_PretendLinear4bit.__name__ = "Linear4bit" + + +def _unquantized_model(): + return nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4)) + + +def _quantized_model(): + return nn.Sequential(nn.Linear(4, 4), _PretendLinear4bit()) + + +def test_fires_when_4bit_requested_but_no_bnb_modules(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert any("load_in_4bit=True was requested" in m for m in msgs), msgs + assert any("issues/5344" in m for m in msgs), msgs + + +def test_silent_when_4bit_succeeded(): + model = _quantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert not any("load_in_4bit" in m for m in msgs), msgs + + +def test_silent_for_full_finetuning(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = False, + load_in_8bit = False, + full_finetuning = True, + ) + msgs = [str(w.message) for w in caught] + assert not any("load_in_4bit" in m or "load_in_8bit" in m for m in msgs), msgs + + +def test_silent_when_no_quantization_requested(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = False, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert not any("load_in_4bit" in m or "load_in_8bit" in m for m in msgs), msgs + + +def test_fires_for_8bit_silent_bypass(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = False, + load_in_8bit = True, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert any("load_in_8bit=True was requested" in m for m in msgs), msgs + + +class _MoEFusedExpertWrapper(nn.Module): + """Mimics Gemma4TextExperts: fused 3D weights stored as nn.Parameter, not + as separate nn.Linear instances. bnb's replace_with_bnb_linear skips this.""" + + def __init__(self, num_experts = 128, intermediate = 1408, hidden = 2816): + super().__init__() + self.gate_up_proj = nn.Parameter( + torch.zeros((num_experts, intermediate, hidden), dtype = torch.bfloat16), + requires_grad = False, + ) + + +def _partial_quant_model(): + return nn.Sequential(_PretendLinear4bit(), _MoEFusedExpertWrapper()) + + +def test_fires_on_partial_quant_moe_experts(): + model = _partial_quant_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert any("partially applied" in m for m in msgs), msgs + assert any("gate_up_proj" in m for m in msgs), msgs + + +class _NormParam(nn.Module): + """An RMSNorm-like module: large BF16 weight whose name is in the skip list.""" + + def __init__(self, dim = 8 * 1024 * 1024 + 10): + super().__init__() + self.norm_weight = nn.Parameter( + torch.zeros(dim, dtype = torch.bfloat16), + requires_grad = False, + ) + + +def test_silent_when_only_skip_list_tensors_unquantized(): + model = nn.Sequential(_PretendLinear4bit(), _NormParam()) + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert not any("partially applied" in m for m in msgs), msgs + + +if __name__ == "__main__": + test_fires_when_4bit_requested_but_no_bnb_modules() + test_silent_when_4bit_succeeded() + test_silent_for_full_finetuning() + test_silent_when_no_quantization_requested() + test_fires_for_8bit_silent_bypass() + test_fires_on_partial_quant_moe_experts() + test_silent_when_only_skip_list_tensors_unquantized() + print("All 7 guardrail tests passed.") diff --git a/unsloth/models/gemma4_moe_4bit.py b/unsloth/models/gemma4_moe_4bit.py new file mode 100644 index 0000000000..6d85412ffc --- /dev/null +++ b/unsloth/models/gemma4_moe_4bit.py @@ -0,0 +1,648 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-expert bitsandbytes Linear4bit swap for Gemma-4 MoE experts. + +Refs: https://github.com/unslothai/unsloth/issues/5344 + +Gemma4TextExperts stores all experts as two fused 3D Parameters +(gate_up_proj, down_proj) shaped (num_experts, out_dim, in_dim) so that +torch._grouped_mm can dispatch a single grouped matmul per layer. The +fused storage is great for forward throughput but breaks bnb 4-bit +quantization: bnb.nn.Linear4bit only swaps nn.Linear instances, so the +fused 3D Parameters stay in BF16, defeating QLoRA VRAM savings. + +This module swaps each Gemma4TextExperts module's fused weights for two +nn.ModuleList[Linear4bit] of length num_experts, and overrides forward to +dispatch per-expert. The trade-off is the loss of torch._grouped_mm +throughput in exchange for a ~4x reduction in expert weight VRAM +(measured on unsloth/gemma-4-26B-A4B-it: 46 GB -> 14.27 GB resident). + +Gated on UNSLOTH_GEMMA4_MOE_4BIT (default off) and on load_in_4bit=True. +Default off until the matching per-expert LoRA path lands; opt in via +the env var if you want the VRAM win without QLoRA training. +""" + +from __future__ import annotations + +import os +from types import MethodType + +import torch +import torch.nn as nn + + +__all__ = [ + "is_gemma4_moe_4bit_enabled", + "is_gemma4_moe_4bit_grouped_enabled", + "is_gemma4_moe_4bit_grouped_active_only_enabled", + "swap_gemma4_experts_to_per_expert_linear4bit", +] + + +def is_gemma4_moe_4bit_enabled() -> bool: + """Opt-in via UNSLOTH_GEMMA4_MOE_4BIT=1.""" + return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT", "0") == "1" + + +def is_gemma4_moe_4bit_grouped_enabled() -> bool: + """Opt-in via UNSLOTH_GEMMA4_MOE_4BIT_GROUPED=1. Requires the base swap to + also be enabled. Uses dequant-then-torch._grouped_mm per forward instead of + a per-expert Linear4bit loop; trades transient BF16 staging buffer for + grouped-GEMM throughput. See unslothai/unsloth#5344 follow-up.""" + return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_GROUPED", "0") == "1" + + +def is_gemma4_moe_4bit_grouped_active_only_enabled() -> bool: + """Opt-in via UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_ACTIVE_ONLY=1. Requires both + base swap + grouped. Dequantizes only the experts touched by the current + batch's top-k routing (~k*B active per layer) instead of all num_experts. + For Gemma-4 MoE (128 experts, top-k=4) this cuts dequant work an order of + magnitude in autoregressive decode.""" + return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_ACTIVE_ONLY", "0") == "1" + + +def is_gemma4_moe_4bit_grouped_cached_enabled() -> bool: + """Opt-in via UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_CACHE=1. Layered on active-only: + cache the BF16 dequantized expert weights in a per-module LRU. Decode + revisits to the same expert skip the dequant entirely. Cache cap via + UNSLOTH_GEMMA4_MOE_4BIT_CACHE_SIZE (default 8).""" + return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_CACHE", "0") == "1" + + +def is_gemma4_moe_4bit_grouped_pt_dequant_enabled() -> bool: + """Opt-in via UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_PT_DEQUANT=1. Layered on + grouped+active_only. Replaces bnb.functional.dequantize_4bit with a + pure-tensor NF4 dequant so torch.compile can fuse the unpack + codebook + lookup + stack + grouped_mm into one Inductor graph. Per-block absmax is + pre-dequantized at swap time and cached on each Linear4bit, so the + per-forward path is bnb-free.""" + return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_PT_DEQUANT", "0") == "1" + + +def is_gemma4_moe_4bit_grouped_static_bf16_enabled() -> bool: + """Opt-in via UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_STATIC_BF16=1. Dequant every + expert ONCE on the first forward and keep the fused (E, 2I, H) / (E, H, I) + BF16 tensors live for the lifetime of the module. Subsequent forwards skip + dequant entirely. Wins back grouped_mm throughput at the cost of holding + a permanent BF16 mirror -- i.e. peak VRAM rises back toward the BF16 + baseline. Useful as a speed-ceiling experiment and for inference workloads + that have spare VRAM.""" + return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_STATIC_BF16", "0") == "1" + + +def _per_expert_forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Replacement Gemma4TextExperts.forward dispatching through swapped + nn.ModuleList[Linear4bit] instead of fused 3D Parameters.""" + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot( + top_k_index, + num_classes = self.num_experts, + ) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim = (-1, -2)), 0).nonzero() + + for expert_idx in expert_hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate_up = self.gate_up_proj_4bit[expert_idx](current_state) + gate, up = gate_up.chunk(2, dim = -1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = self.down_proj_4bit[expert_idx](current_hidden_states) + current_hidden_states = ( + current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + ) + final_hidden_states.index_add_( + 0, + token_idx, + current_hidden_states.to(final_hidden_states.dtype), + ) + + return final_hidden_states + + +# NF4 codebook (16 entries, bnb convention). Used by the pure-PyTorch dequant +# path which torch.compile can fuse with the surrounding stack + grouped_mm. +_NF4_CODES = torch.tensor( + [ + -1.0, + -0.6961928009986877, + -0.5250730514526367, + -0.39491748809814453, + -0.28444138169288635, + -0.18477343022823334, + -0.09105003625154495, + 0.0, + 0.07958029955625534, + 0.16093020141124725, + 0.24611230194568634, + 0.33791524171829224, + 0.44070982933044434, + 0.5626170039176941, + 0.7229568362236023, + 1.0, + ], + dtype = torch.float32, +) + +_COMPILED_PT_DEQUANT_STACK = None +_COMPILED_DEQUANT_STACK = None + + +def _ensure_pt_dequant_state(layer): + """Cache the dequantized absmax + codebook on the layer so the per-forward + path needs no bnb calls. Idempotent. Called at swap time.""" + if getattr(layer, "_unsloth_pt_dequant_ready", False): + return + from bitsandbytes.functional import dequantize_blockwise + + qs = layer.weight.quant_state + if qs.nested: + absmax_fp32 = dequantize_blockwise(qs.absmax, qs.state2) + absmax_fp32 = (absmax_fp32 + qs.offset).to(torch.float32) + else: + absmax_fp32 = qs.absmax.to(torch.float32) + layer._unsloth_pt_absmax_fp32 = absmax_fp32.contiguous() + layer._unsloth_pt_blocksize = qs.blocksize + layer._unsloth_pt_shape = tuple(qs.shape) + layer._unsloth_pt_dtype = qs.dtype + layer._unsloth_pt_dequant_ready = True + + +def _pt_dequant_one(packed_uint8, absmax_fp32, blocksize, shape, dtype, codes): + """Pure-PyTorch NF4 dequant of one expert weight. Bit-exact vs bnb. + Pure tensor ops -> torch.compile-friendly.""" + packed = packed_uint8.reshape(-1) + high = (packed >> 4) & 0xF + low = packed & 0xF + indices = torch.stack([high, low], dim = -1).reshape(-1).to(torch.long) + values = codes[indices] # fp32 + n_elements = values.numel() + n_blocks = (n_elements + blocksize - 1) // blocksize + values = values.view(n_blocks, blocksize) * absmax_fp32.view(-1, 1) + target = shape[0] * shape[1] + return values.reshape(-1)[:target].view(shape).to(dtype) + + +def _pt_dequant_stack_subset(layers, indices_cpu, codes): + """Pure-PyTorch dequant of a subset of experts and stack into (E_active, out, in).""" + return torch.stack( + [ + _pt_dequant_one( + layers[i].weight.data, + layers[i]._unsloth_pt_absmax_fp32, + layers[i]._unsloth_pt_blocksize, + layers[i]._unsloth_pt_shape, + layers[i]._unsloth_pt_dtype, + codes, + ) + for i in indices_cpu + ], + dim = 0, + ) + + +def _get_compiled_pt_dequant_stack(): + """Lazy-compile the pure-PT dequant+stack helper. Re-used across forwards.""" + global _COMPILED_PT_DEQUANT_STACK + if _COMPILED_PT_DEQUANT_STACK is None: + _COMPILED_PT_DEQUANT_STACK = torch.compile( + _pt_dequant_stack_subset, + dynamic = True, + fullgraph = False, + ) + return _COMPILED_PT_DEQUANT_STACK + + +def _dequant_stack(layers): + """Dequantize each Linear4bit in a ModuleList and stack into (E, out, in).""" + from bitsandbytes.functional import dequantize_4bit + + return torch.stack( + [dequantize_4bit(L.weight.data, L.weight.quant_state) for L in layers], + dim = 0, + ) + + +def _dequant_stack_subset(layers, indices_cpu): + """Dequantize only experts whose CPU-int indices are in indices_cpu.""" + from bitsandbytes.functional import dequantize_4bit + + return torch.stack( + [ + dequantize_4bit(layers[i].weight.data, layers[i].weight.quant_state) + for i in indices_cpu + ], + dim = 0, + ) + + +def _get_compiled_dequant_stack(): + """Lazily compile the dequant+stack helper. Done once and cached so + successive forwards reuse the compiled graph (otherwise the per-call + compile cost dwarfs the runtime saving).""" + global _COMPILED_DEQUANT_STACK + if _COMPILED_DEQUANT_STACK is None: + _COMPILED_DEQUANT_STACK = torch.compile( + _dequant_stack, + dynamic = False, + fullgraph = False, + ) + return _COMPILED_DEQUANT_STACK + + +def _grouped_mm_forward_4bit( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Dequantize per-expert Linear4bit weights into a fused 3D BF16 tensor and + run unsloth_zoo.forward_native_grouped_mm. Buffer is freed at end of each + forward so resident VRAM stays at 4-bit. Transient peak per layer is + (E * 2I * H + E * H * I) * 2 bytes BF16. Opt-in compile of the dequant + helper via UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_COMPILE=1; bnb's CUDA dequant + is a graph break so the win is bounded by kernel-launch overhead saved + via CUDA-graph capture, not by Inductor fusion.""" + from unsloth_zoo.temporary_patches.moe_utils import ( + forward_native_grouped_mm, + ) + + if os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_COMPILE", "0") == "1": + dequant_stack = _get_compiled_dequant_stack() + else: + dequant_stack = _dequant_stack + + gate_up = dequant_stack(self.gate_up_proj_4bit) + down = dequant_stack(self.down_proj_4bit) + self.gate_up_proj = nn.Parameter(gate_up, requires_grad = False) + self.down_proj = nn.Parameter(down, requires_grad = False) + try: + return forward_native_grouped_mm( + self, + hidden_states, + top_k_index, + top_k_weights, + ) + finally: + del self.gate_up_proj + del self.down_proj + + +def _grouped_mm_forward_4bit_active_only( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Like _grouped_mm_forward_4bit but only dequants the experts that this + batch's top-k routing actually touches. Remaps top_k_index from the full + expert space [0, num_experts) to the compact active range [0, E_active) + so torch._grouped_mm sees only the populated groups. + + For decode (S*K << num_experts) the saving is large; for prefill where + most experts get hit the active set approaches num_experts and the + overhead of the unique+remap is dominated by the dequant savings.""" + from unsloth_zoo.temporary_patches.moe_utils import ( + forward_native_grouped_mm, + ) + + flat = top_k_index.reshape(-1) + active_experts, inverse = torch.unique(flat, return_inverse = True) + active_cpu = active_experts.tolist() + n_active = len(active_cpu) + + gate_up = _dequant_stack_subset(self.gate_up_proj_4bit, active_cpu) + down = _dequant_stack_subset(self.down_proj_4bit, active_cpu) + + compact_top_k = inverse.view_as(top_k_index) + + saved_n_experts = self.num_experts + self.num_experts = n_active + self.gate_up_proj = nn.Parameter(gate_up, requires_grad = False) + self.down_proj = nn.Parameter(down, requires_grad = False) + try: + return forward_native_grouped_mm( + self, + hidden_states, + compact_top_k, + top_k_weights, + ) + finally: + del self.gate_up_proj + del self.down_proj + self.num_experts = saved_n_experts + + +def _cached_dequant(module, attr_name, expert_idx, layer): + """LRU dequant cache for one expert's weight. Hit returns the cached BF16 + tensor; miss dequants, stores, and evicts oldest. Cache cap is per-attribute + per Gemma4TextExperts module so each layer maintains its own working set.""" + from bitsandbytes.functional import dequantize_4bit + + cache_attr = f"_unsloth_dequant_cache_{attr_name}" + cache = getattr(module, cache_attr, None) + if cache is None: + from collections import OrderedDict + + cache = OrderedDict() + setattr(module, cache_attr, cache) + cached = cache.get(expert_idx, None) + if cached is not None: + cache.move_to_end(expert_idx) + return cached + w = dequantize_4bit(layer.weight.data, layer.weight.quant_state) + cache[expert_idx] = w + cap = int(os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT_CACHE_SIZE", "8")) + while len(cache) > cap: + cache.popitem(last = False) + return w + + +def _grouped_mm_forward_4bit_pt_compiled( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Active-only grouped forward using pure-PyTorch NF4 dequant + torch.compile. + Pre-cached per-expert absmax means the per-forward path is bnb-free, so + Inductor can fuse unpack + codebook lookup + stack into one Triton + kernel.""" + from unsloth_zoo.temporary_patches.moe_utils import ( + forward_native_grouped_mm, + ) + + flat = top_k_index.reshape(-1) + active_experts, inverse = torch.unique(flat, return_inverse = True) + active_cpu = active_experts.tolist() + n_active = len(active_cpu) + + codes = _NF4_CODES.to(hidden_states.device) + dequant_fn = _get_compiled_pt_dequant_stack() + + gate_up = dequant_fn(self.gate_up_proj_4bit, active_cpu, codes) + down = dequant_fn(self.down_proj_4bit, active_cpu, codes) + + compact_top_k = inverse.view_as(top_k_index) + + saved_n_experts = self.num_experts + self.num_experts = n_active + self.gate_up_proj = nn.Parameter(gate_up, requires_grad = False) + self.down_proj = nn.Parameter(down, requires_grad = False) + try: + return forward_native_grouped_mm( + self, + hidden_states, + compact_top_k, + top_k_weights, + ) + finally: + del self.gate_up_proj + del self.down_proj + self.num_experts = saved_n_experts + + +def _grouped_mm_forward_4bit_static_bf16( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """First call: dequant every expert to a permanent (E, 2I, H) / (E, H, I) + BF16 fused tensor stored on the module. Subsequent calls reuse them and + skip dequant entirely. Speed-ceiling experiment: trades the 4-bit VRAM + win for grouped_mm throughput.""" + from unsloth_zoo.temporary_patches.moe_utils import ( + forward_native_grouped_mm, + ) + + if not hasattr(self, "_unsloth_static_bf16_gate_up"): + self._unsloth_static_bf16_gate_up = _dequant_stack(self.gate_up_proj_4bit) + self._unsloth_static_bf16_down = _dequant_stack(self.down_proj_4bit) + self.gate_up_proj = self._unsloth_static_bf16_gate_up + self.down_proj = self._unsloth_static_bf16_down + try: + return forward_native_grouped_mm( + self, + hidden_states, + top_k_index, + top_k_weights, + ) + finally: + del self.gate_up_proj + del self.down_proj + + +def _grouped_mm_forward_4bit_active_only_cached( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Active-only grouped forward with a per-module LRU cache of dequantized + experts. Decode patterns where the same experts get hit across consecutive + tokens skip the dequant on the second visit. Cache cap via + UNSLOTH_GEMMA4_MOE_4BIT_CACHE_SIZE (default 8 per module).""" + from unsloth_zoo.temporary_patches.moe_utils import ( + forward_native_grouped_mm, + ) + + flat = top_k_index.reshape(-1) + active_experts, inverse = torch.unique(flat, return_inverse = True) + active_cpu = active_experts.tolist() + n_active = len(active_cpu) + + gate_up = torch.stack( + [ + _cached_dequant(self, "gate_up", i, self.gate_up_proj_4bit[i]) + for i in active_cpu + ], + dim = 0, + ) + down = torch.stack( + [_cached_dequant(self, "down", i, self.down_proj_4bit[i]) for i in active_cpu], + dim = 0, + ) + + compact_top_k = inverse.view_as(top_k_index) + + saved_n_experts = self.num_experts + self.num_experts = n_active + self.gate_up_proj = nn.Parameter(gate_up, requires_grad = False) + self.down_proj = nn.Parameter(down, requires_grad = False) + try: + return forward_native_grouped_mm( + self, + hidden_states, + compact_top_k, + top_k_weights, + ) + finally: + del self.gate_up_proj + del self.down_proj + self.num_experts = saved_n_experts + + +def _quantize_one_expert_to_linear4bit( + weight_2d: torch.Tensor, + compute_dtype: torch.dtype, + quant_type: str = "nf4", +): + """Build a bnb.nn.Linear4bit from a single (out, in) weight slice. + Params4bit triggers on-the-fly quantization on .to(device).""" + import bitsandbytes as bnb + + out_features, in_features = weight_2d.shape + layer = bnb.nn.Linear4bit( + in_features, + out_features, + bias = False, + compute_dtype = compute_dtype, + quant_type = quant_type, + quant_storage = torch.uint8, + ) + layer.weight = bnb.nn.Params4bit( + data = weight_2d.detach().clone().contiguous(), + requires_grad = False, + quant_type = quant_type, + ) + return layer + + +def swap_gemma4_experts_to_per_expert_linear4bit( + model: nn.Module, + compute_dtype: torch.dtype = torch.bfloat16, + quant_type: str = "nf4", + verbose: bool = False, +) -> int: + """Find every Gemma4TextExperts module in `model`, replace its fused 3D + weights with two nn.ModuleList[Linear4bit] (per-expert), and patch + forward to dispatch per-expert. + + Returns the count of swapped modules. Zero if the model has no Gemma-4 + MoE experts or if transformers does not expose Gemma4TextExperts. + """ + try: + from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts + except Exception: + return 0 + + swapped = 0 + for module in model.modules(): + if not isinstance(module, Gemma4TextExperts): + continue + # Idempotent: once swapped, the fused 3D Parameters are gone. + if hasattr(module, "_unsloth_gemma4_moe_4bit_swapped"): + continue + if not hasattr(module, "gate_up_proj") or not hasattr(module, "down_proj"): + continue + + gate_up = module.gate_up_proj + down = module.down_proj + if not isinstance(gate_up, nn.Parameter) or gate_up.ndim != 3: + continue + if not isinstance(down, nn.Parameter) or down.ndim != 3: + continue + + num_experts, two_intermediate, hidden = gate_up.shape + num_experts_d, hidden_d, intermediate = down.shape + if ( + num_experts != num_experts_d + or hidden != hidden_d + or two_intermediate != 2 * intermediate + ): + # Unrecognised layout: skip rather than risk corrupting weights. + if verbose: + print( + f"Unsloth: skipping Gemma4TextExperts swap due to " + f"unexpected shapes gate_up={tuple(gate_up.shape)} " + f"down={tuple(down.shape)}" + ) + continue + + device = gate_up.device + + gate_up_list = nn.ModuleList() + down_list = nn.ModuleList() + for e in range(num_experts): + gu = _quantize_one_expert_to_linear4bit( + gate_up.data[e], + compute_dtype = compute_dtype, + quant_type = quant_type, + ) + dp = _quantize_one_expert_to_linear4bit( + down.data[e], + compute_dtype = compute_dtype, + quant_type = quant_type, + ) + gate_up_list.append(gu.to(device)) + down_list.append(dp.to(device)) + + # Per-module peak = fused BF16 + accumulated per-expert nf4; released here. + del module.gate_up_proj + del module.down_proj + + module.gate_up_proj_4bit = gate_up_list + module.down_proj_4bit = down_list + + # Per-instance bind so sibling Gemma4TextExperts keep the class method. + # Forward variant ladder (most specific wins). STATIC_BF16 is the + # speed-ceiling variant; it overrides ACTIVE_ONLY/CACHE if set since + # those become no-ops once weights are kept permanently dequantized. + if ( + is_gemma4_moe_4bit_grouped_enabled() + and is_gemma4_moe_4bit_grouped_static_bf16_enabled() + ): + _fwd = _grouped_mm_forward_4bit_static_bf16 + elif ( + is_gemma4_moe_4bit_grouped_enabled() + and is_gemma4_moe_4bit_grouped_pt_dequant_enabled() + ): + for L in list(module.gate_up_proj_4bit) + list(module.down_proj_4bit): + _ensure_pt_dequant_state(L) + _fwd = _grouped_mm_forward_4bit_pt_compiled + elif ( + is_gemma4_moe_4bit_grouped_enabled() + and is_gemma4_moe_4bit_grouped_active_only_enabled() + and is_gemma4_moe_4bit_grouped_cached_enabled() + ): + _fwd = _grouped_mm_forward_4bit_active_only_cached + elif ( + is_gemma4_moe_4bit_grouped_enabled() + and is_gemma4_moe_4bit_grouped_active_only_enabled() + ): + _fwd = _grouped_mm_forward_4bit_active_only + elif is_gemma4_moe_4bit_grouped_enabled(): + _fwd = _grouped_mm_forward_4bit + else: + _fwd = _per_expert_forward + module.forward = MethodType(_fwd, module) + module._unsloth_gemma4_moe_4bit_swapped = True + + swapped += 1 + + if swapped > 0 and torch.cuda.is_available(): + # Free the cached fused tensors so post-swap VRAM reflects 4-bit. + torch.cuda.empty_cache() + + return swapped diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d39f2588ef..dc1cf17482 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2478,6 +2478,16 @@ class FastLlamaModel: and not _head.weight.is_floating_point() ): _head.to(dtype) + # Guardrail: warn before dispatch hooks if quantization was silently dropped. + from unsloth.models.vision import _warn_if_quantization_silently_dropped + + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = kwargs.get("load_in_8bit", False), + full_finetuning = kwargs.get("full_finetuning", False), + quantization_config = kwargs.get("quantization_config"), + ) # Attach dispatch hooks for bnb multi-device loads. from unsloth.models.vision import _attach_bnb_multidevice_hooks @@ -2500,9 +2510,19 @@ class FastLlamaModel: attn_implementation = preferred_attn_impl, **kwargs, ) - # Attach dispatch hooks for bnb multi-device loads. - from unsloth.models.vision import _attach_bnb_multidevice_hooks + # Guardrail (#5344) + multi-device dispatch hooks share an import. + from unsloth.models.vision import ( + _warn_if_quantization_silently_dropped, + _attach_bnb_multidevice_hooks, + ) + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = kwargs.get("load_in_8bit", False), + full_finetuning = kwargs.get("full_finetuning", False), + quantization_config = kwargs.get("quantization_config"), + ) _attach_bnb_multidevice_hooks( model, load_in_4bit = load_in_4bit, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index df371e00c8..46858766d6 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -68,7 +68,6 @@ from unsloth_zoo.patching_utils import patch_model_and_tokenizer from unsloth_zoo.training_utils import prepare_model_for_training from unsloth_zoo.utils import Version -from transformers import __version__ as transformers_version import types import functools @@ -99,6 +98,140 @@ __all__ = [ ] +# bnb-quantized Linear class names (see unslothai/unsloth#5344 guardrail). +_BNB_QUANT_CLASS_NAMES = ("Linear4bit", "Linear8bitLt", "LinearNF4", "LinearFP4") + +# Substrings the guardrail treats as intentionally-not-quantized: embeddings, +# norms, biases, routers/gates that need fp16/fp32 precision, vision/audio +# towers, classification heads, rotary tables. +_GUARDRAIL_SKIP_PATTERNS = ( + "embed", + "embedding", + "norm", + "ln_", + "rms", + ".bias", + "lm_head", + "multi_modal_projector", + "merger", + "modality_projection", + "router", + "block_sparse_moe.gate", + "mamba", + "audio_tower", + "vision_tower", + "score", + "classifier", + "qa_outputs", + "rotary", +) + +# A floating Parameter larger than this, found outside the skip list, counts +# as bulk weight that should have been 4-bit. Tuned so head dims and small +# projections do not false-fire but MoE fused expert tensors do. +_GUARDRAIL_BULK_WEIGHT_NUMEL = 8 * 1024 * 1024 + + +def _warn_if_quantization_silently_dropped( + model, + load_in_4bit, + load_in_8bit, + full_finetuning, + quantization_config = None, +): + """Guardrail for unslothai/unsloth#5344. + + Two failure modes covered: + + 1. TOTAL bypass: load_in_4bit was requested but the model contains zero + bnb Linear4bit / Linear8bitLt modules. transformers / bnb / a backend + incompatibility dropped kwargs.quantization_config. + + 2. PARTIAL bypass: bnb quantized the nn.Linear modules but a large fraction + of weight bytes live in non-nn.Linear Parameters (e.g. Gemma-4 MoE fused + 3D expert tensors, custom Linear-like wrappers). bnb only swaps nn.Linear + instances; fused expert weights stay in BF16, defeating QLoRA savings + even though some Linear4bit modules exist. + + Warns rather than raises so non-bnb backends (CPU / MLX / AMD-without-bnb) + with legitimately no Linear4bit are not broken. + """ + if full_finetuning: + return + if quantization_config is not None: + if isinstance(quantization_config, dict): + load_in_4bit = load_in_4bit or bool(quantization_config.get("load_in_4bit")) + load_in_8bit = load_in_8bit or bool(quantization_config.get("load_in_8bit")) + else: + load_in_4bit = load_in_4bit or bool( + getattr(quantization_config, "load_in_4bit", False) + ) + load_in_8bit = load_in_8bit or bool( + getattr(quantization_config, "load_in_8bit", False) + ) + if not (load_in_4bit or load_in_8bit): + return + + has_bnb = any(type(m).__name__ in _BNB_QUANT_CLASS_NAMES for m in model.modules()) + + # Failure mode 1: total bypass. + if not has_bnb: + kind = "4bit" if load_in_4bit else "8bit" + bnb_class_name = "Linear4bit" if load_in_4bit else "Linear8bitLt" + warnings.warn( + f"Unsloth: load_in_{kind}=True was requested but no bitsandbytes " + f"{bnb_class_name} modules were produced. The runtime quantization " + f"config was silently dropped and the model is in full precision. " + f"See https://github.com/unslothai/unsloth/issues/5344 for known " + f"triggers (transformers/bnb version mismatch, MoE checkpoints " + f"without a -bnb-4bit sibling, multi-GPU dispatch). Workaround: " + f'pass device_map="cuda:0" and pin transformers/bitsandbytes to ' + f"a version known to work.", + stacklevel = 3, + ) + return + + # Failure mode 2: partial bypass. Walk named_parameters and find large + # floating tensors outside the skip list. If they aggregate to >= 2x the + # quantized payload, partial quant is essentially negating 4-bit savings. + quantized_bytes = 0 + suspect_bytes = 0 + suspect_samples = [] + for name, p in model.named_parameters(): + if p is None: + continue + nbytes = p.numel() * p.element_size() + if p.dtype in (torch.uint8, torch.int8): + # bnb stores 4-bit payloads as uint8 and 8-bit payloads (Int8Params) as int8. + quantized_bytes += nbytes + continue + if p.dtype not in (torch.bfloat16, torch.float16, torch.float32): + continue + if p.numel() < _GUARDRAIL_BULK_WEIGHT_NUMEL: + continue + lname = name.lower() + if any(pat in lname for pat in _GUARDRAIL_SKIP_PATTERNS): + continue + suspect_bytes += nbytes + if len(suspect_samples) < 3: + suspect_samples.append((name, str(p.dtype), tuple(p.shape))) + + if quantized_bytes > 0 and suspect_bytes >= 2 * quantized_bytes: + kind = "4bit" if load_in_4bit else "8bit" + suspect_human = ", ".join(f"{n} ({d}, {s})" for n, d, s in suspect_samples) + warnings.warn( + f"Unsloth: load_in_{kind}=True is partially applied. " + f"bitsandbytes quantized ~{quantized_bytes/1024**3:.2f} GB of " + f"nn.Linear weights, but ~{suspect_bytes/1024**3:.2f} GB of " + f"non-nn.Linear floating Parameters were left unquantized (e.g. " + f"fused MoE expert tensors, custom Linear-like wrappers). " + f"Examples: {suspect_human}. The model's effective VRAM " + f"footprint is close to its full-precision size. See " + f"https://github.com/unslothai/unsloth/issues/5344.", + stacklevel = 3, + ) + + def _infer_device_map_from_loaded_model(model): """Build a compact device_map by inspecting actual parameter placements.""" device_map = {} @@ -911,6 +1044,82 @@ class FastBaseModel: verify_fp8_support_if_applicable(model_config) + # Resolve 4-bit + Gemma4 swap parameters once (shared by both load paths). + _user_qcfg = kwargs.get("quantization_config", None) + if isinstance(_user_qcfg, dict): + _qcfg_4bit = bool(_user_qcfg.get("load_in_4bit", False)) + _qcfg_dtype = _user_qcfg.get("bnb_4bit_compute_dtype", None) + _qcfg_quant_type = _user_qcfg.get("bnb_4bit_quant_type", None) + elif _user_qcfg is not None: + _qcfg_4bit = bool(getattr(_user_qcfg, "load_in_4bit", False)) + _qcfg_dtype = getattr(_user_qcfg, "bnb_4bit_compute_dtype", None) + _qcfg_quant_type = getattr(_user_qcfg, "bnb_4bit_quant_type", None) + else: + _qcfg_4bit = False + _qcfg_dtype = None + _qcfg_quant_type = None + if isinstance(_qcfg_dtype, str): + _qcfg_dtype_str = _qcfg_dtype.removeprefix("torch.") + _maybe_dtype = getattr(torch, _qcfg_dtype_str, None) + _qcfg_dtype = ( + _maybe_dtype if isinstance(_maybe_dtype, torch.dtype) else None + ) + _effective_load_in_4bit = bool(load_in_4bit) or _qcfg_4bit + + def _maybe_swap_gemma4_moe_4bit(_target_model): + if not (_effective_load_in_4bit and not full_finetuning): + return + try: + from unsloth.models.gemma4_moe_4bit import ( + is_gemma4_moe_4bit_enabled, + swap_gemma4_experts_to_per_expert_linear4bit, + ) + + if not is_gemma4_moe_4bit_enabled(): + return + if bnb_config is not None: + _compute_dtype = bnb_config.bnb_4bit_compute_dtype + _quant_type = getattr(bnb_config, "bnb_4bit_quant_type", "nf4") + else: + _compute_dtype = ( + _qcfg_dtype if _qcfg_dtype is not None else torch.bfloat16 + ) + _quant_type = ( + _qcfg_quant_type if _qcfg_quant_type is not None else "nf4" + ) + _swapped = swap_gemma4_experts_to_per_expert_linear4bit( + _target_model, + compute_dtype = _compute_dtype, + quant_type = _quant_type, + ) + if _swapped > 0: + print( + f"Unsloth: swapped {_swapped} " + f"Gemma4TextExperts module(s) to per-expert " + f"Linear4bit (see " + f"https://github.com/unslothai/unsloth/issues/5344)." + ) + except Exception as _e: + _partial = sum( + 1 + for _m in _target_model.modules() + if getattr(_m, "_unsloth_gemma4_moe_4bit_swapped", False) + ) + if _partial: + raise RuntimeError( + f"Unsloth: Gemma-4 MoE 4-bit swap failed after " + f"converting {_partial} module(s); model is in a " + f"mixed 4-bit/BF16 state. Reload the model to " + f"recover. Original error: " + f"{type(_e).__name__}: {_e}" + ) from _e + warnings.warn( + f"Unsloth: Gemma-4 MoE 4-bit swap failed: " + f"{type(_e).__name__}: {_e}. Falling back to BF16 " + f"experts. Unset UNSLOTH_GEMMA4_MOE_4BIT to silence.", + stacklevel = 2, + ) + raise_handler = RaiseUninitialized() if not fast_inference: # Prevent load_in_fp8 from being forwarded into HF internal model loading @@ -935,6 +1144,16 @@ class FastBaseModel: # attn_implementation = attn_implementation, **kwargs, ) + _maybe_swap_gemma4_moe_4bit(model) + + # Guardrail: see _warn_if_quantization_silently_dropped + #5344. + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = load_in_8bit, + full_finetuning = full_finetuning, + quantization_config = kwargs.get("quantization_config"), + ) # Attach dispatch hooks for bnb multi-device loads. _attach_bnb_multidevice_hooks( model, @@ -1049,6 +1268,7 @@ class FastBaseModel: bnb_config, is_vision_model = is_vlm, ) + _maybe_swap_gemma4_moe_4bit(model) model.vllm_engine = llm model.fast_generate = model.vllm_engine.generate model.fast_generate_batches = functools.partial(