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.
This commit is contained in:
Daniel Han
parent
f16d0df146
commit
6c76ae9f9d
1 changed files with 279 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue