From 2bbb15994bb39c729cee00979a04ca5fbe9afcf8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 May 2026 14:52:24 +0000 Subject: [PATCH] gemma4 moe 4bit: pure-PyTorch NF4 dequant variant (negative result, kept for docs) Adds UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_PT_DEQUANT=1 as a 5th forward variant on top of #5432 + the follow-up active-only path. Implementation: - Pre-compute the dequantized per-block absmax (bnb's nested-blockwise scheme) once at swap time and cache on each Linear4bit as _unsloth_pt_absmax_fp32. Removes bnb from the per-forward path entirely. - Per-forward: nibble unpack + 16-entry NF4 codebook lookup + per-block absmax multiply + reshape to (out, in). All pure tensor ops. - The dequant+stack helper is wrapped in torch.compile so Inductor can fuse with the surrounding stack and grouped_mm. Numerical parity probe (temp/sim_5344_pt_nf4_probe.py) matches bnb bit-exactly on a synthetic stub (cos=1.0, max_abs_diff=0). Real-model swapped forward gives cos 0.996 vs BF16 baseline (vs active-only's 0.985) because the FP32 intermediate multiply is closer to ideal BF16 weights. Speed result on gemma-4-26B-A4B-it (B200): 2.17 tok/s vs active-only's 2.21 tok/s. Essentially break-even. Inductor's fusion across the per-expert iteration in the torch.stack list comprehension is bounded; the real bottleneck (per-expert dispatch + stack copy) survives the compile pass. Resident VRAM +1.33 GB for the cached FP32 absmax buffers. Kept in the codebase as a negative result + foundation for a future vectorised-across-experts dequant pass, which would stack packed uint8 and absmax into (E_active, ...) tensors BEFORE the dequant so Inductor sees a single batched op. Bit-exact loop-vs-pt_dequant equivalence test: temp/sim_5344_pt_dequant_unit.py (cos=1.0, max_abs_diff=0 on synthetic stub with 2 or 4 active experts). --- unsloth/models/gemma4_moe_4bit.py | 150 ++++++++++++++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/unsloth/models/gemma4_moe_4bit.py b/unsloth/models/gemma4_moe_4bit.py index d7ca5c5c0d..13727de986 100644 --- a/unsloth/models/gemma4_moe_4bit.py +++ b/unsloth/models/gemma4_moe_4bit.py @@ -80,6 +80,16 @@ def is_gemma4_moe_4bit_grouped_cached_enabled() -> bool: 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) @@ -130,9 +140,101 @@ def _per_expert_forward( 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 @@ -278,6 +380,47 @@ def _cached_dequant(module, attr_name, expert_idx, layer): 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, @@ -470,6 +613,13 @@ def swap_gemma4_experts_to_per_expert_linear4bit( 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()