[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

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.")

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.")

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

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,

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(