diff --git a/tests/test_gemma4_moe_4bit_swap.py b/tests/test_gemma4_moe_4bit_swap.py new file mode 100644 index 0000000000..2093067624 --- /dev/null +++ b/tests/test_gemma4_moe_4bit_swap.py @@ -0,0 +1,142 @@ +"""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/unsloth/models/gemma4_moe_4bit.py b/unsloth/models/gemma4_moe_4bit.py new file mode 100644 index 0000000000..d8453beb98 --- /dev/null +++ b/unsloth/models/gemma4_moe_4bit.py @@ -0,0 +1,215 @@ +# 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", + "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 _per_expert_forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Replacement Gemma4TextExperts.forward using per-expert Linear4bit. + + Mirrors the reference forward in transformers.models.gemma4.modeling_gemma4 + but dispatches through swapped nn.ModuleList[Linear4bit] modules instead + of nn.functional.linear on the 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 + + +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) BF16 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)) + + # Drop the fused Parameters before attaching the ModuleLists so peak + # VRAM during the swap stays bounded by one expert at a time. + 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 forward bind so other Gemma4TextExperts instances + # (e.g. in a sibling model) keep the class-level method. + module.forward = MethodType(_per_expert_forward, 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/vision.py b/unsloth/models/vision.py index 648488b2b6..c459a153fe 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1051,6 +1051,39 @@ class FastBaseModel: # attn_implementation = attn_implementation, **kwargs, ) + # Opt-in per-expert Linear4bit swap for Gemma-4 MoE checkpoints + # whose fused 3D expert weights bnb cannot quantize (#5344). + # Off by default; users enable via UNSLOTH_GEMMA4_MOE_4BIT=1. + if load_in_4bit and not full_finetuning: + try: + from unsloth.models.gemma4_moe_4bit import ( + is_gemma4_moe_4bit_enabled, + swap_gemma4_experts_to_per_expert_linear4bit, + ) + if is_gemma4_moe_4bit_enabled(): + _swapped = swap_gemma4_experts_to_per_expert_linear4bit( + model, + compute_dtype = ( + bnb_config.bnb_4bit_compute_dtype + if bnb_config is not None + else torch.bfloat16 + ), + ) + 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: + 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, + ) + # Guardrail: see _warn_if_quantization_silently_dropped + #5344. _warn_if_quantization_silently_dropped( model,