From 6c76ae9f9d3c3b308142fc16ab479862f52ae9da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 May 2026 14:09:04 +0000 Subject: [PATCH] gemma4 moe 4bit: active-only + cache + static_bf16 forward variants (#5344 follow-up) Follow-up to #5432. Closes ~half the speed gap between the 4-bit loop forward and the BF16 grouped_mm path by skipping dequantization of inactive experts. Three new forward variants, all gated by env vars on top of the existing UNSLOTH_GEMMA4_MOE_4BIT_GROUPED=1: 1. UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_ACTIVE_ONLY=1 Dequant only the experts touched by the current batch's top-k routing instead of all num_experts. Uses torch.unique(top_k_index) for the active set and remaps top_k_index from [0, num_experts) to [0, E_active). Bench on gemma-4-26B-A4B-it shows 2.71-3.01 tok/s vs the loop forward's 1.59 tok/s (~1.7-2.0x speedup), same VRAM (14.27 GB resident, +0.4 GB peak), logits cos_sim 0.998 vs loop. Bit-exact on the synthetic stub. 2. UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_CACHE=1 (requires _ACTIVE_ONLY=1) Per-module LRU of dequantized BF16 experts. Capped at UNSLOTH_GEMMA4_MOE_4BIT_CACHE_SIZE (default 8). Empirically break-even with active-only on the 26B-A4B bench because the torch.stack-into-compact-tensor cost remains; the cache only skips the dequant math. 3. UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_STATIC_BF16=1 Speed-ceiling experiment: dequant every expert ONCE on first forward, keep fused BF16 tensors live for the rest of the module's lifetime. Reaches 5.98 tok/s on 26B-A4B (within 6% of BF16 baseline 6.34) at the cost of peak forward VRAM 57 GB (the 4-bit storage is still resident, so net is BF16 + 4-bit cohabit). Proves per-forward dequant is the bottleneck; the remaining ~2x gap cannot be closed without persistent fused weights or a custom 4-bit grouped GEMM kernel. UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_COMPILE=1 also included as a negative result: torch.compile on the dequant helper regresses to 1.03 tok/s because bnb.functional.dequantize_4bit is a custom CUDA op (graph break) and Inductor can't fuse around it. Forward selection ladder in swap_gemma4_experts_to_per_expert_linear4bit: STATIC_BF16 > ACTIVE_ONLY+CACHE > ACTIVE_ONLY > GROUPED > loop. The env vars are mutually compatible; later flags layer on earlier ones. Unit equivalence tests in temp/sim_5344_{active_only,cached,static_bf16}_unit.py confirm bit-exactness (cos=1.0, max_abs_diff=0) on a synthetic Gemma4TextExperts stub. Real-model bench in temp/sim_5344_grouped_4bit_bench.py with selectable BENCH_MODES. Design notes + full Pareto frontier in outputs/moe_4bit_speed_summary.md. --- unsloth/models/gemma4_moe_4bit.py | 280 +++++++++++++++++++++++++++++- 1 file changed, 279 insertions(+), 1 deletion(-) diff --git a/unsloth/models/gemma4_moe_4bit.py b/unsloth/models/gemma4_moe_4bit.py index 44ed75e628..06d2bffc69 100644 --- a/unsloth/models/gemma4_moe_4bit.py +++ b/unsloth/models/gemma4_moe_4bit.py @@ -44,6 +44,8 @@ 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", ] @@ -53,6 +55,42 @@ def is_gemma4_moe_4bit_enabled() -> bool: 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_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, @@ -92,6 +130,223 @@ def _per_expert_forward( return final_hidden_states +_COMPILED_DEQUANT_STACK = None + + +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_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, @@ -195,7 +450,30 @@ def swap_gemma4_experts_to_per_expert_linear4bit( module.down_proj_4bit = down_list # Per-instance bind so sibling Gemma4TextExperts keep the class method. - module.forward = MethodType(_per_expert_forward, module) + # 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_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