https://github.com/unslothai/unsloth-zoo

The candidate change is a single-line addition exposing torch_native_forward on the generated module, which only marginally addresses symbol resolution. It does not strip compiler decorators, implement the MoE/LoRA 4-bit fixes, or address robust name detection, so most required outcomes are unmet.

diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..169a3c5 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -1049,6 +1049,7 @@ def patch_gpt_oss_bnb4bit():
     m.dtype_from_config    = dtype_from_config
     m.transformers_version = transformers_version
     m.Version              = Version
+    m.torch_native_forward  = torch_native_forward
 
     return True
 

The change mainly improves symbol resolution for compiler-generated GPT-OSS/MoE modules via fallback imports and module attribute assignment, which meaningfully addresses c3 and lowers regression risk. However, it does not touch decorator stripping (c1) or model-name detection (c4), and shows no clear MoE/LoRA 4-bit forward or branching fix to demonstrate GPT-OSS actually loads in 4-bit (c2). Overall it is scaffolding toward the fix rather than a complete solution.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..b46fd6c 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -861,10 +861,12 @@ def create_new_function(
         or "forward_native_grouped_mm" in new_source
         or "forward_triton_grouped_gemm" in new_source
         or "forward_native_moe_loop" in new_source
+        or "_check_torch_grouped_mm_supported" in new_source
     ):
         imports += (
             "try:\n"
             "    from moe_utils import (\n"
+            "        _check_torch_grouped_mm_supported,\n"
             "        forward_moe_backend,\n"
             "        select_moe_backend,\n"
             "        forward_native_grouped_mm,\n"
@@ -872,6 +874,23 @@ def create_new_function(
             "        forward_native_moe_loop,\n"
             "    )\n"
             "except Exception:\n"
+            "    try:\n"
+            "        from unsloth_zoo.temporary_patches.moe_utils import (\n"
+            "            _check_torch_grouped_mm_supported,\n"
+            "            forward_moe_backend,\n"
+            "            select_moe_backend,\n"
+            "            forward_native_grouped_mm,\n"
+            "            forward_triton_grouped_gemm,\n"
+            "            forward_native_moe_loop,\n"
+            "        )\n"
+            "    except Exception:\n"
+            "        pass\n"
+        )
+    if "torch_native_forward" in new_source:
+        imports += (
+            "try:\n"
+            "    from unsloth_zoo.temporary_patches.gpt_oss import torch_native_forward\n"
+            "except Exception:\n"
             "    pass\n"
         )
     imports += (
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..64cfcf4 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -1049,6 +1049,10 @@ def patch_gpt_oss_bnb4bit():
     m.dtype_from_config    = dtype_from_config
     m.transformers_version = transformers_version
     m.Version              = Version
+    m._check_torch_grouped_mm_supported = _check_torch_grouped_mm_supported
+    m.forward_native_grouped_mm = forward_native_grouped_mm
+    m.torch_native_forward = torch_native_forward
+    m.moe_forward_inference = moe_forward_inference
 
     return True
 

The change robustly addresses the 4-bit load path (symbol resolution, MoE routing, quantization skip, name detection) but does not implement the decorator-stripping outcome (c1), which is a required piece and could leave some compiler-generated modules referencing unresolved decorators. Other criteria are well met.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..d9ed159 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -889,6 +889,42 @@ def create_new_function(
             _valid_items = items  # Fall back to all items if module load fails
         if _valid_items:
             imports += f"from {model_location} import (" + ", ".join(_valid_items) + ")"
+    gpt_oss_items = [
+        "_RouterLinearParams",
+        "_gpt_oss_make_dense_routing_weights",
+        "_gpt_oss_expert_token_indices",
+        "_unwrap_peft_experts",
+        "GptOssExpertsBnb4bit",
+        "GptOssTopKRouterBnb4bit",
+        "moe_forward_inference",
+        "moe_forward_inference_bf16",
+        "dtype_from_config",
+        "transformers_version",
+        "swiglu_torch_forward",
+        "torch_native_forward",
+        "Version",
+    ]
+    if any(item in new_source for item in gpt_oss_items):
+        imports += (
+            "\ntry:\n"
+            "    from unsloth_zoo.temporary_patches.gpt_oss import (\n"
+            "        _RouterLinearParams,\n"
+            "        _gpt_oss_make_dense_routing_weights,\n"
+            "        _gpt_oss_expert_token_indices,\n"
+            "        _unwrap_peft_experts,\n"
+            "        GptOssExpertsBnb4bit,\n"
+            "        GptOssTopKRouterBnb4bit,\n"
+            "        moe_forward_inference,\n"
+            "        moe_forward_inference_bf16,\n"
+            "        dtype_from_config,\n"
+            "        transformers_version,\n"
+            "        swiglu_torch_forward,\n"
+            "        torch_native_forward,\n"
+            "        Version,\n"
+            "    )\n"
+            "except Exception:\n"
+            "    pass\n"
+        )
     new_source = imports + "\n\n" + new_source
     # Check logger and remove use_cache
     if "logger" in items:
@@ -1125,6 +1161,13 @@ def create_standalone_class(
     # Create optimized standalone forward function
     f = eval(f"{model_location}.{module}")
     full_class = inspect.getsource(f)
+    class_match = re.search(r"(^class\s+)([A-Za-z_][A-Za-z0-9_]*)(\s*\()", full_class, flags=re.MULTILINE)
+    if class_match is not None and class_match.group(2) != module:
+        full_class = (
+            full_class[:class_match.start(2)]
+            + module
+            + full_class[class_match.end(2):]
+        )
     old_source = inspect.getsource(f.forward)
     old_init = inspect.getsource(f.__init__)

The change robustly improves model-name detection and refactors the 4bit/LoRA branching without regressing other paths, and adds partial symbol registration. However the c1 decorator-stripping outcome is entirely absent, and the compiler edit only renames the class rather than removing kernel/auto_docstring decorators, so generated source could still reference unresolved decorators, limiting confidence that 4bit loading fully works.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..343b817 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1127,6 +1127,17 @@ def create_standalone_class(
     full_class = inspect.getsource(f)
     old_source = inspect.getsource(f.forward)
     old_init = inspect.getsource(f.__init__)
+    class_name_match = re.search(
+        r"(^\s*class\s+)([A-Za-z_][A-Za-z0-9_]*)(\s*\()",
+        full_class,
+        flags=re.MULTILINE,
+    )
+    if class_name_match is not None and class_name_match.group(2) != module:
+        full_class = (
+            full_class[: class_name_match.start(2)]
+            + module
+            + full_class[class_name_match.end(2) :]
+        )
     if forward_source is None:
         forward_source = old_source
 
diff --git a/unsloth_zoo/peft_utils.py b/unsloth_zoo/peft_utils.py
index 09413a7..4ef804e 100644
--- a/unsloth_zoo/peft_utils.py
+++ b/unsloth_zoo/peft_utils.py
@@ -74,7 +74,18 @@ def get_peft_regex(
     # Get only linear layers
     modules = model.named_modules()
     linear_modules = [name for name, module in modules if isinstance(module, torch.nn.Linear)]
-    all_linear_modules = Counter(x.rsplit(".")[-1] for x in linear_modules)
+
+    def normalize_linear_module_name(name):
+        parts = name.rsplit(".", 2)
+        if (
+            len(parts) == 3
+            and parts[-1].isdigit()
+            and parts[-2] in ("gate_up_projs", "down_projs")
+        ):
+            return parts[-2]
+        return name.rsplit(".")[-1]
+
+    all_linear_modules = Counter(normalize_linear_module_name(x) for x in linear_modules)
 
     # Isolate lm_head / projection matrices if count == 1
     if target_modules is None:
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..b20c899 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -1024,11 +1024,9 @@ def patch_gpt_oss_bnb4bit():
         transformers.models.gpt_oss.modeling_gpt_oss._original_GptOssTopKRouter = \
             transformers.models.gpt_oss.modeling_gpt_oss.GptOssTopKRouter
 
-    # Replace with BnB 4bit compatible versions
-    # Preserve original symbol names for compiler-generated modules.
-    GptOssExpertsBnb4bit.__name__ = "GptOssExperts"
-    GptOssExpertsBnb4bit.__qualname__ = "GptOssExperts"
-
+    # Replace with BnB 4bit compatible versions. Keep the Python qualname intact

The change robustly fixes model-name detection (c4) and consistently applies it, plus fixes an init-weights condition bug contributing to the 4bit path. However, it does not address the central issues named in c1 and c3 — stripping extra compiler decorators and resolving symbols in generated modules — so GPT-OSS 4-bit loading is unlikely to be fully fixed.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..292bd0e 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1125,6 +1125,13 @@ def create_standalone_class(
     # Create optimized standalone forward function
     f = eval(f"{model_location}.{module}")
     full_class = inspect.getsource(f)
+    class_match = re.search(r"^class\s+([A-Za-z_][A-Za-z0-9_]*)\b", full_class, flags=re.MULTILINE)
+    if class_match is not None and class_match.group(1) != module:
+        full_class = (
+            full_class[:class_match.start(1)]
+            + module
+            + full_class[class_match.end(1):]
+        )
     old_source = inspect.getsource(f.forward)
     old_init = inspect.getsource(f.__init__)
     if forward_source is None:
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..5fa63e4 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -1316,15 +1316,26 @@ def _should_use_gpt_oss_bnb4bit() -> bool:
     Default: True when load_in_4bit is active.
     Set UNSLOTH_GPT_OSS_BNB4BIT_DISABLE=1 to force BF16 path.
     """
-    if "gpt_oss" not in os.environ.get("UNSLOTH_MODEL_NAME", ""):
+    if not _is_gpt_oss_model_name():
         return False
-    if "_load_in_4bit_" not in os.environ.get("UNSLOTH_MODEL_NAME", ""):
+    if not _is_gpt_oss_4bit_load():
         return False
     return os.environ.get("UNSLOTH_GPT_OSS_BNB4BIT_DISABLE", "0") != "1"
 
 
+def _is_gpt_oss_model_name() -> bool:
+    model_name = os.environ.get("UNSLOTH_MODEL_NAME", "").lower()
+    return "gpt_oss" in model_name or "gpt-oss" in model_name
+
+
 def _is_gpt_oss_4bit_load() -> bool:
-    return "_load_in_4bit_" in os.environ.get("UNSLOTH_MODEL_NAME", "")
+    model_name = os.environ.get("UNSLOTH_MODEL_NAME", "").lower()
+    return (
+        "_load_in_4bit_" in model_name
+        or "load_in_4bit" in model_name
+        or "bnb-4bit" in model_name
+        or "bnb_4bit" in model_name
+    )
 
 
 def _is_transformers_v5() -> bool:
@@ -1340,7 +1351,7 @@ def patch_gpt_oss_moe_for_lora():
     IMPORTANT: We only patch the forward method, NOT replace the entire class.
     This preserves the original class structure so weights load correctly.
     """
-    if "gpt_oss" not in os.environ.get("UNSLOTH_MODEL_NAME", ""):
+    if not _is_gpt_oss_model_name():
         return
     if _is_gpt_oss_4bit_load() or _should_use_gpt_oss_bnb4bit():

The change robustly generalizes decorator stripping (c1), resolves generated-module symbols via explicit exports (c3), and makes model-name detection formatting-agnostic (c4). However it removes the substantial MoE/PEFT 4-bit quantization and dequantization patches while adding only lighter-weight handling, so whether GPT-OSS truly loads and forwards correctly in 4-bit (c2) is uncertain and constitutes a possible regression (c5). Overall a plausible but incomplete-feeling fix on the core 4-bit correctness.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..cc5e7b6 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -993,7 +993,7 @@ def create_new_function(
         except Exception as e:
             # consider adding logging to main_process only
             # counterpoint: we may want to see errors on all processes
-            if os.environ.get("UNSLOTH_LOGGING_ENABLED", "0") == "1":
+            if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
                 logger.error(
                     f"Unsloth: Failed to write file {function_location} because {str(e)}"
                 )
@@ -1038,7 +1038,7 @@ def create_new_function(
                 new_module = importlib.import_module(name)
                 return new_module, old_path
         except Exception as e:
-            if os.environ.get("UNSLOTH_LOGGING_ENABLED", "0") == "1":
+            if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
                 logger.error(
                     f"Unsloth: Failed to import module {name} because {str(e)}"
                 )
@@ -1136,6 +1136,14 @@ def create_standalone_class(
 
     # Strip decorators from class source if present
     # This fixes issues with classes like Qwen3NextExperts which have decorators that cause compilation failures
+    STRIP_DECORATORS = {
+        "use_experts_implementation",
+        "use_kernel_forward_from_hub",
+        "use_kernelized_func",
+        "auto_docstring",
+        # add more here if needed
+    }
+
     if full_class.lstrip().startswith("@"):
         start = re.search(r"^class ", full_class, flags=re.MULTILINE)
         if start:
@@ -1148,30 +1156,71 @@ def create_standalone_class(
             lines = preamble.split('\n')
             new_lines = []
 
+            # Capture decorator head, including dotted paths: @pkg.decorator(...)
+            decorator_head_re = re.compile(
+                r"^\s*@\s*([A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*)\b"
+            )
+
+            skipping = False
+            paren_depth = 0
+            skip_base_name = None
+
             for line in lines:
+                if skipping:
+                    # Continue skipping decorator args until balanced
+                    paren_depth += line.count("(") - line.count(")")
+                    if paren_depth <= 0:
+                        skipping = False
+                        paren_depth = 0
+                        skip_base_name = None
+                    continue
+

The change robustly addresses decorator stripping (c1), symbol resolution (c3), and name normalization (c4), and adjusts several 4bit-path guards and mask handling. However, it deletes large blocks of PEFT/transformers 4bit MoE quantization and dequantization patches that appeared central to a working quantized MoE/LoRA 4bit forward path, without clear replacement in the diff — casting doubt on whether GPT-OSS truly loads and runs in 4-bit and creating regression risk. Credit for c2 and c5 is reduced accordingly.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..1fa4002 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1151,8 +1151,14 @@ def create_standalone_class(
             for line in lines:
                 stripped = line.strip()
                 if stripped.startswith("@"):
-                    if "use_experts_implementation" in stripped:
-                        logger.info(f'Unsloth: stripped use_experts_implementation decorator from {module}')
+                    if (
+                        "use_experts_implementation" in stripped
+                        or "use_kernel_forward_from_hub" in stripped
+                        or "use_kernelized_func" in stripped
+                        or stripped.startswith("@auto_docstring")
+                    ):
+                        decorator_name = stripped.split("(")[0].lstrip("@")
+                        logger.info(f"Unsloth: stripped {decorator_name} decorator from {module}")
                         continue # Strip it
                     else:
                         logger.warning(f"Unsloth: Warning: Unknown decorator {stripped} found for {module}.")
@@ -1165,13 +1171,14 @@ def create_standalone_class(
     # Check if forward was replaced by a temporary patch (renamed function)
     # In this case, keep the patched source as-is and replace the class forward body.
     patched_forward_info = None
-    func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
-    if func_match and func_match.group(1) != "forward":
-        # Find original forward in class to replace it
-        orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
-        if orig_fwd:
-            patched_forward_info = (func_match.group(1), orig_fwd.group(1))
-            disable = None  # Keep patched source as-is for renamed forward replacements
+    if "@torch.compiler.disable" in forward_source:
+        func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
+        if func_match and func_match.group(1) != "forward":
+            # Find original forward in class to replace it
+            orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
+            if orig_fwd:
+                patched_forward_info = (func_match.group(1), orig_fwd.group(1))
+                disable = None  # Keep patched source as-is for renamed forward replacements
 
     # Replace function name with module-specific name
     if patched_forward_info:
@@ -1269,6 +1276,7 @@ def create_standalone_class(
 
     # Remove @auto_docstring
     source = re.sub(r"@auto_docstring[\s]{0,}(\([^\)]{0,}\))?", "", source)
+    source = re.sub(r"@use_kernelized_func[\s]{0,}(\([^\)]{0,}\))?", "", source)
     source = re.sub(r"@check_model_inputs[\s]{0,}(\([^\)]{0,}\))?", "", source)
     # source = source.replace("@auto_docstring", "")
 
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..0b0851a 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -648,6 +648,26 @@ class ParameterModule(nn.Linear):
         )
 
 
+def patch_gpt_oss_compiler_exports():

The candidate only addresses one narrow aspect: preserving the compiler.disable decorator for GptOssExperts to keep the 4-bit MoE forward from being compiled. It does not strip extra decorators (c1), does not resolve generated-module symbols (c3), and does not implement robust name detection (c4). It likely partially helps the 4-bit path and avoids regressions, but is far from a complete fix.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..ad5ca1e 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1164,14 +1164,20 @@ def create_standalone_class(
 
     # Check if forward was replaced by a temporary patch (renamed function)
     # In this case, keep the patched source as-is and replace the class forward body.
+    # GptOssExperts is exempt: its 4-bit BnB replacement assigns forward =
+    # torch_native_forward (a renamed function), but the module is in
+    # DISABLE_COMPILE_MODULES and must keep its @torch.compiler.disable decorator.
+    # Entering this branch would set disable = None and drop that decorator,
+    # forcing the experts through torch.compile and breaking 4-bit loading.
     patched_forward_info = None
-    func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
-    if func_match and func_match.group(1) != "forward":
-        # Find original forward in class to replace it
-        orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
-        if orig_fwd:
-            patched_forward_info = (func_match.group(1), orig_fwd.group(1))
-            disable = None  # Keep patched source as-is for renamed forward replacements
+    if 'gptossexperts' != module.lower():
+        func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
+        if func_match and func_match.group(1) != "forward":
+            # Find original forward in class to replace it
+            orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
+            if orig_fwd:
+                patched_forward_info = (func_match.group(1), orig_fwd.group(1))
+                disable = None  # Keep patched source as-is for renamed forward replacements
 
     # Replace function name with module-specific name
     if patched_forward_info:

no diff captured (skipped)

The candidate only makes a single narrow guard preventing the renamed-forward heuristic from overriding disable for GptOssExperts. It addresses a fragment of the 4-bit MoE forward issue but omits the decorator stripping (c1), symbol resolution (c3), and robust model-name detection (c4) required by the rubric. As presented, it is unlikely to fully restore 4-bit loading.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..31730f1 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1164,14 +1164,19 @@ def create_standalone_class(
 
     # Check if forward was replaced by a temporary patch (renamed function)
     # In this case, keep the patched source as-is and replace the class forward body.
+    # GptOssExperts is excluded: in the 4-bit path its forward is reassigned to a
+    # renamed helper (e.g. torch_native_forward / experts_forward), and the rename
+    # heuristic below would otherwise override the intended disable=True and emit a
+    # broken compiler-generated module for the MoE/LoRA experts.
     patched_forward_info = None
-    func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
-    if func_match and func_match.group(1) != "forward":
-        # Find original forward in class to replace it
-        orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
-        if orig_fwd:
-            patched_forward_info = (func_match.group(1), orig_fwd.group(1))
-            disable = None  # Keep patched source as-is for renamed forward replacements
+    if 'gptossexperts' != module.lower():
+        func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
+        if func_match and func_match.group(1) != "forward":
+            # Find original forward in class to replace it
+            orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
+            if orig_fwd:
+                patched_forward_info = (func_match.group(1), orig_fwd.group(1))
+                disable = None  # Keep patched source as-is for renamed forward replacements
 
     # Replace function name with module-specific name
     if patched_forward_info:

The change strongly addresses the compiler-side outcomes: decorator stripping (c1), symbol resolution (c3), and robust name detection (c4). However it removes substantial MoE/LoRA 4bit quantization and dequantization machinery from misc.py and moe_utils.py rather than fixing it, casting doubt on whether GPT-OSS 4-bit with working quantized MoE/LoRA forward truly succeeds (c2) and whether existing paths avoid regression (c5).

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..1fa4002 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1151,8 +1151,14 @@ def create_standalone_class(
             for line in lines:
                 stripped = line.strip()
                 if stripped.startswith("@"):
-                    if "use_experts_implementation" in stripped:
-                        logger.info(f'Unsloth: stripped use_experts_implementation decorator from {module}')
+                    if (
+                        "use_experts_implementation" in stripped
+                        or "use_kernel_forward_from_hub" in stripped
+                        or "use_kernelized_func" in stripped
+                        or stripped.startswith("@auto_docstring")
+                    ):
+                        decorator_name = stripped.split("(")[0].lstrip("@")
+                        logger.info(f"Unsloth: stripped {decorator_name} decorator from {module}")
                         continue # Strip it
                     else:
                         logger.warning(f"Unsloth: Warning: Unknown decorator {stripped} found for {module}.")
@@ -1165,13 +1171,14 @@ def create_standalone_class(
     # Check if forward was replaced by a temporary patch (renamed function)
     # In this case, keep the patched source as-is and replace the class forward body.
     patched_forward_info = None
-    func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
-    if func_match and func_match.group(1) != "forward":
-        # Find original forward in class to replace it
-        orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
-        if orig_fwd:
-            patched_forward_info = (func_match.group(1), orig_fwd.group(1))
-            disable = None  # Keep patched source as-is for renamed forward replacements
+    if "@torch.compiler.disable" in forward_source:
+        func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
+        if func_match and func_match.group(1) != "forward":
+            # Find original forward in class to replace it
+            orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
+            if orig_fwd:
+                patched_forward_info = (func_match.group(1), orig_fwd.group(1))
+                disable = None  # Keep patched source as-is for renamed forward replacements
 
     # Replace function name with module-specific name
     if patched_forward_info:
@@ -1269,6 +1276,7 @@ def create_standalone_class(
 
     # Remove @auto_docstring
     source = re.sub(r"@auto_docstring[\s]{0,}(\([^\)]{0,}\))?", "", source)
+    source = re.sub(r"@use_kernelized_func[\s]{0,}(\([^\)]{0,}\))?", "", source)
     source = re.sub(r"@check_model_inputs[\s]{0,}(\([^\)]{0,}\))?", "", source)
     # source = source.replace("@auto_docstring", "")
 
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..0b0851a 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -648,6 +648,26 @@ class ParameterModule(nn.Linear):
         )
 
 
+def patch_gpt_oss_compiler_exports():

The change comprehensively addresses decorator stripping (broadened set + robust dotted/multiline parsing), symbol resolution via module exports, robust hyphen/underscore detection, and 4-bit MoE/LoRA path fixes (gptossexperts special-casing, dynamic_rope_update fullgraph). c2 receives near-full credit since actual runtime load success cannot be verified from the diff alone but the fixes clearly target the reported failure modes. Other paths appear preserved.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..53a5749 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1136,6 +1136,14 @@ def create_standalone_class(
 
     # Strip decorators from class source if present
     # This fixes issues with classes like Qwen3NextExperts which have decorators that cause compilation failures
+    STRIP_DECORATORS = {
+        "use_experts_implementation",
+        "use_kernel_forward_from_hub",
+        "use_kernelized_func",
+        "auto_docstring",
+        # add more here if needed
+    }
+
     if full_class.lstrip().startswith("@"):
         start = re.search(r"^class ", full_class, flags=re.MULTILINE)
         if start:
@@ -1148,30 +1156,71 @@ def create_standalone_class(
             lines = preamble.split('\n')
             new_lines = []
 
+            # Capture decorator head, including dotted paths: @pkg.decorator(...)
+            decorator_head_re = re.compile(
+                r"^\s*@\s*([A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*)\b"
+            )
+
+            skipping = False
+            paren_depth = 0
+            skip_base_name = None
+
             for line in lines:
+                if skipping:
+                    # Continue skipping decorator args until balanced
+                    paren_depth += line.count("(") - line.count(")")
+                    if paren_depth <= 0:
+                        skipping = False
+                        paren_depth = 0
+                        skip_base_name = None
+                    continue
+
                 stripped = line.strip()
                 if stripped.startswith("@"):
-                    if "use_experts_implementation" in stripped:
-                        logger.info(f'Unsloth: stripped use_experts_implementation decorator from {module}')
-                        continue # Strip it
-                    else:
-                        logger.warning(f"Unsloth: Warning: Unknown decorator {stripped} found for {module}.")
-                        new_lines.append(line) # Keep it
+                    m = decorator_head_re.match(line)
+                    if not m:
+                        logger.warning(
+                            f"Unsloth: Warning: Unparseable decorator {stripped} found for {module}."
+                        )
+                        new_lines.append(line)
+                        continue
+
+                    decorator_full = m.group(1)              # e.g. "foo.auto_docstring"
+                    decorator_base = decorator_full.split(".")[-1]  # e.g. "auto_docstring"

The change reorders torch.compile disabling and injects more helper symbols and class replacements into the transformers module and any already-loaded compiled cache module, which addresses part of the symbol-resolution and 4-bit path issues. However, it does not strip the compiler-generated decorators (c1) at all, relies on a hardcoded module name rather than robust name detection (c4), and only patches a pre-loaded cached module rather than fixing generation, so the fix is incomplete and GPT-OSS 4-bit likely still fails to load fully.

diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..5017645 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -16,6 +16,7 @@
 
 from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable
 import os
+import sys
 import torch
 import torch.nn as nn
 import torch.nn.init as init
@@ -1044,11 +1045,33 @@ def patch_gpt_oss_bnb4bit():
     # Inject BnB helpers so compiler-generated modules can import them
     # from transformers.models.gpt_oss.modeling_gpt_oss
     m = transformers.models.gpt_oss.modeling_gpt_oss
-    m._RouterLinearParams  = _RouterLinearParams
-    m.swiglu_torch_forward = swiglu_torch_forward
-    m.dtype_from_config    = dtype_from_config
-    m.transformers_version = transformers_version
-    m.Version              = Version
+    m._RouterLinearParams        = _RouterLinearParams
+    m.swiglu_torch_forward       = swiglu_torch_forward
+    m.dtype_from_config          = dtype_from_config
+    m.transformers_version       = transformers_version
+    m.Version                    = Version
+    # Also inject functions that GptOssExperts.forward and GptOssMLP.forward
+    # reference, so that compiler-generated standalone modules (which import from
+    # this module at file-load time) resolve them correctly.
+    m.torch_native_forward           = torch_native_forward
+    m.moe_forward_inference          = moe_forward_inference
+    m._check_torch_grouped_mm_supported = _check_torch_grouped_mm_supported
+    m.forward_native_grouped_mm      = forward_native_grouped_mm
+
+    # Propagate class replacements into any compiler-generated combined module
+    # that was already loaded into sys.modules before this patch ran.
+    # Those modules import GptOssExperts/GptOssTopKRouter at module-load time,
+    # so their cached reference still points at the original class unless we
+    # update it here.
+    _compiled_mod_name = "unsloth_compiled_module_gpt_oss"
+    _compiled_mod = sys.modules.get(_compiled_mod_name)
+    if _compiled_mod is not None:
+        _compiled_mod.GptOssExperts   = GptOssExpertsBnb4bit
+        _compiled_mod.GptOssTopKRouter = GptOssTopKRouter
+        _compiled_mod.torch_native_forward           = torch_native_forward
+        _compiled_mod.moe_forward_inference          = moe_forward_inference
+        _compiled_mod._check_torch_grouped_mm_supported = _check_torch_grouped_mm_supported
+        _compiled_mod.forward_native_grouped_mm      = forward_native_grouped_mm
 
     return True
 
@@ -1080,15 +1103,20 @@ def patch_gpt_oss_bnb4bit_auto():
     """
     if not _should_use_gpt_oss_bnb4bit():
         return
-    # BnB helpers are now injected into the transformers module by
-    # patch_gpt_oss_bnb4bit(), so the compiler can resolve all symbols.
-    patch_gpt_oss_bnb4bit()
-    # Ensure inference path avoids torch.compile for 4-bit
+    # Disable torch.compile on the inference forward BEFORE calling

The candidate only wraps a single forward-patch code branch in a condition checking for 'gptossexperts', which does not implement any of the core required outcomes: no decorator stripping, no symbol resolution, no genuine 4-bit MoE/LoRA fix, and no robust name detection. It is essentially scaffolding/partial guarding rather than the real fix, so most criteria earn little to no credit.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..d1459f2 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1165,13 +1165,14 @@ def create_standalone_class(
     # Check if forward was replaced by a temporary patch (renamed function)
     # In this case, keep the patched source as-is and replace the class forward body.
     patched_forward_info = None
-    func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
-    if func_match and func_match.group(1) != "forward":
-        # Find original forward in class to replace it
-        orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
-        if orig_fwd:
-            patched_forward_info = (func_match.group(1), orig_fwd.group(1))
-            disable = None  # Keep patched source as-is for renamed forward replacements
+    if 'gptossexperts' != module.lower():
+        func_match = re.search(r"def\s+(\w+)\s*\(", forward_source)
+        if func_match and func_match.group(1) != "forward":
+            # Find original forward in class to replace it
+            orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
+            if orig_fwd:
+                patched_forward_info = (func_match.group(1), orig_fwd.group(1))
+                disable = None  # Keep patched source as-is for renamed forward replacements
 
     # Replace function name with module-specific name
     if patched_forward_info:

The candidate makes a real attempt at the 4-bit MoE loading path (per-expert weight distribution, state_dict remapping, class renaming in compiler) which partially addresses c2, but it ignores the decorator-stripping (c1), symbol resolution (c3), and robust model-name detection (c4) outcomes entirely. Since the task explicitly ties the breakage to compiler-generated modules and unresolved decorators/symbols, leaving those unaddressed means the fix is incomplete.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..0b0da6b 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1130,6 +1130,18 @@ def create_standalone_class(
     if forward_source is None:
         forward_source = old_source
 
+    # If the class was replaced by a differently-named class (e.g. GptOssExpertsBnb4bit
+    # installed under the key GptOssExperts), the source will have the old class name.
+    # Rename the class definition in the source to match the expected module name so that
+    # the compiled module exports the correct attribute name.
+    full_class = re.sub(
+        r'^(\s*class\s+)\w+(\s*[\(:])',
+        rf'\1{module}\2',
+        full_class,
+        count=1,
+        flags=re.MULTILINE,
+    )
+
     # We disable this for nn.Embedding modules if torch is older than 2.5 since
     if OLD_TORCH_VERSION and "nn.Embedding(" in old_init:
         disable = True
@@ -1171,7 +1183,12 @@ def create_standalone_class(
         orig_fwd = re.search(r"(\n\s+def\s+forward\s*\([^)]*\)[^:]*:.*?)(?=\n\s+def\s|\n\s+@|\Z)", full_class, re.DOTALL)
         if orig_fwd:
             patched_forward_info = (func_match.group(1), orig_fwd.group(1))
-            disable = None  # Keep patched source as-is for renamed forward replacements
+            # Only clear the compile decorator when the caller did NOT explicitly request
+            # disable=True (e.g. from DISABLE_COMPILE_MODULES).  When disable=True the
+            # patched forward must still carry @torch.compiler.disable so BitsAndBytes
+            # layers or data-dependent control flow inside it don't break under compile.
+            if disable is not True:
+                disable = None  # Keep patched source as-is for renamed forward replacements
 
     # Replace function name with module-specific name
     if patched_forward_info:
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..9181f8b 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -872,6 +872,100 @@ class GptOssExpertsBnb4bit(nn.Module):
             "down_proj_bias", torch.empty(0, dtype=self.dtype), persistent=False
         )
 
+    def _distribute_stacked_weight(self, proj_name, weight_3d):
+        """
+        Distribute a 3-D stacked weight (E, dim1, dim2) from a BF16 checkpoint
+        into the per-expert Linear / Linear4bit modules in gate_up_projs / down_projs.
+
+        This is required when loading GPT-OSS with on-the-fly BitsAndBytes 4-bit
+        quantization from the original BF16 checkpoint whose state-dict has
+        gate_up_proj / down_proj as a single 3-D parameter.
+
+        Shape conventions (same as GptOssExperts):
+          gate_up_proj : (E, hidden_size, 2 * expert_dim)  -> per-expert weight=(2I, H)
+          down_proj    : (E, expert_dim,  hidden_size)     -> per-expert weight=(H, I)
+        """
+        projs = getattr(self, f"{proj_name}s", None)  # gate_up_projs or down_projs
+        if projs is None or len(projs) == 0:

The change removes a harmful __name__ override and injects extra helper symbols, which meaningfully advances symbol resolution (c3) and partially the 4-bit load fix (c2). However it does not implement decorator stripping (c1) or robust name detection (c4), and offers only weak regression protection (c5). The overall fix is partial.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..e9721cd 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1130,6 +1130,18 @@ def create_standalone_class(
     if forward_source is None:
         forward_source = old_source
 
+    # If the class was patched-in under a different name (e.g. GptOssExpertsBnb4bit
+    # registered as GptOssExperts), rename it in the source so the combined module
+    # exports the attribute under the expected name and the compiler injection at the
+    # end of unsloth_compile_transformers succeeds.
+    if f.__name__ != module:
+        full_class = re.sub(
+            r"\bclass\s+" + re.escape(f.__name__) + r"\b",
+            f"class {module}",
+            full_class,
+            count=1,
+        )
+
     # We disable this for nn.Embedding modules if torch is older than 2.5 since
     if OLD_TORCH_VERSION and "nn.Embedding(" in old_init:
         disable = True
diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py
index 1388d78..49f66e0 100644
--- a/unsloth_zoo/temporary_patches/gpt_oss.py
+++ b/unsloth_zoo/temporary_patches/gpt_oss.py
@@ -1024,11 +1024,10 @@ def patch_gpt_oss_bnb4bit():
         transformers.models.gpt_oss.modeling_gpt_oss._original_GptOssTopKRouter = \
             transformers.models.gpt_oss.modeling_gpt_oss.GptOssTopKRouter
 
-    # Replace with BnB 4bit compatible versions
-    # Preserve original symbol names for compiler-generated modules.
-    GptOssExpertsBnb4bit.__name__ = "GptOssExperts"
-    GptOssExpertsBnb4bit.__qualname__ = "GptOssExperts"
-
+    # Replace with BnB 4bit compatible versions.
+    # Do NOT rename __name__/__qualname__: renaming causes inspect.getsource to find
+    # the BF16 GptOssExperts class in gpt_oss.py (same search pattern) instead of
+    # GptOssExpertsBnb4bit, yielding a broken compiler-generated standalone module.
     transformers.models.gpt_oss.modeling_gpt_oss.GptOssExperts = GptOssExpertsBnb4bit
     # Use the unsloth GptOssTopKRouter (with self.linear = nn.Linear) for the router.
     # The BnB 4-bit checkpoint stores router weights as router.linear.weight/bias.
@@ -1041,14 +1040,30 @@ def patch_gpt_oss_bnb4bit():
     logger.info("Unsloth: Patched GPT OSS with BitsAndBytes 4bit compatible classes")
     os.environ["UNSLOTH_GPT_OSS_BNB4BIT_PATCHED"] = "1"
 
-    # Inject BnB helpers so compiler-generated modules can import them
-    # from transformers.models.gpt_oss.modeling_gpt_oss
+    # Inject helpers so compiler-generated standalone modules can import them from
+    # transformers.models.gpt_oss.modeling_gpt_oss.  The compiler builds the import
+    # list from dir(modeling_file), so every symbol used inside a compiled class or
+    # function body must be present here at compile time.
     m = transformers.models.gpt_oss.modeling_gpt_oss
-    m._RouterLinearParams  = _RouterLinearParams
-    m.swiglu_torch_forward = swiglu_torch_forward
-    m.dtype_from_config    = dtype_from_config
-    m.transformers_version = transformers_version
-    m.Version              = Version
+    m._RouterLinearParams               = _RouterLinearParams

The candidate mainly reworks compiler wrapper generation to skip externally patched modules and adds a cache-name suffix, plus removes empty buffer registrations in the 4-bit experts class. It does not implement decorator stripping (c1), symbol resolution (c3), or robust name detection (c4). The 4-bit load fix is only partially and indirectly addressed, with limited regression safety.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index e3d42fd..171ad75 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -3469,6 +3469,20 @@ def unsloth_compile_transformers(
     if len(disable_modules) > 0:
         for module in disable_modules:
             try:
+                # Skip generating a compiled wrapper when the class has been externally patched
+                # (e.g. GptOssExpertsBnb4bit replacing GptOssExperts for BNB 4-bit loading).
+                # Generating a wrapper from the patched source produces a mis-named class that
+                # would either silently disappear (fresh run) or overwrite the patch when loaded
+                # from a cache that was built in a different mode.
+                _cls = eval(f"{model_location}.{module}", globals(), locals())
+                _cls_module = getattr(_cls, "__module__", model_location) or ""
+                if _cls_module != model_location and not _cls_module.startswith(model_location):
+                    if UNSLOTH_ENABLE_LOGGING:
+                        logger.info(
+                            f"Unsloth: Skipping compile for externally patched module {module} "
+                            f"(from {_cls_module})"
+                        )
+                    continue
                 new_module = create_standalone_class(
                     module,
                     model_location,
@@ -3981,9 +3995,28 @@ def unsloth_compile_transformers(
 
     all_code = "\n\n".join(final_all_standalone_classes)
 
+    # Build a suffix that captures any externally patched DISABLE_COMPILE_MODULES so that
+    # the cache file name is unique to the current patching state.  This prevents a cached
+    # module compiled in non-4-bit mode (whose class definitions were generated from the
+    # original transformers classes) from being silently reused when running in 4-bit mode
+    # (where some of those classes are replaced by patched equivalents such as
+    # GptOssExpertsBnb4bit).
+    _external_patch_suffix = ""
+    for _dmod in sorted(DISABLE_COMPILE_MODULES):
+        try:
+            _dmod_cls = getattr(eval(model_location, globals(), locals()), _dmod, None)
+            if _dmod_cls is None:
+                continue
+            _dmod_cls_module = getattr(_dmod_cls, "__module__", model_location) or ""
+            if _dmod_cls_module != model_location and not _dmod_cls_module.startswith(model_location):
+                _external_patch_suffix += f"_{_dmod}"
+        except Exception:
+            pass
+    _combined_module_name = f"{COMBINED_UNSLOTH_NAME}_{model_type}{_external_patch_suffix}"
+
     try:
         combined_module = create_new_function(
-            f"{COMBINED_UNSLOTH_NAME}_{model_type}",
+            _combined_module_name,
             all_code,
             model_location,
             functions,
@@ -4073,6 +4106,20 @@ def unsloth_compile_transformers(
     # Import and replace with new module
     for module in all_standalone_classes.keys():
         try:
+            # Don't overwrite a class that was patched in from an external module

The change adds real infrastructure for FP8 MoE dequant and forward dispatch, achieving the forward and speed outcomes reasonably. However the central c1 outcome—relaxing the actual HF Trainer quantization-for-training validation that aborts—is not genuinely addressed; the patch only freezes params in training_step, which is not the guard that refuses FP8 base models. Scale/fallback handling is partial and somewhat fragile, and much of the added code is untested scaffolding relative to real checkpoints.

diff --git a/tests/test_fp8_moe_lora.py b/tests/test_fp8_moe_lora.py
new file mode 100644
index 0000000..3b88aa4
--- /dev/null
+++ b/tests/test_fp8_moe_lora.py
@@ -0,0 +1,163 @@
+# Tests for FP8 MoE + LoRA training path
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team.
+
+import pytest
+import torch
+import torch.nn as nn
+
+try:
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+        _is_fp8_param,
+        _dequantize_fp8_expert_weight,
+        _batch_dequantize_fp8_experts,
+        _moe_uses_fp8_expert_weights,
+        prepare_fp8_model_for_lora_training,
+    )
+    HAS_FP8_SUPPORT = True
+except ImportError:
+    HAS_FP8_SUPPORT = False
+
+
+@pytest.mark.skipif(not HAS_FP8_SUPPORT, reason="FP8 support not available")
+class TestFP8Detection:
+    def test_is_fp8_param_none(self):
+        assert not _is_fp8_param(None)
+    
+    def test_is_fp8_param_regular_tensor(self):
+        t = torch.randn(10, 10, dtype=torch.float16)
+        assert not _is_fp8_param(t)
+    
+    def test_is_fp8_param_fp8_tensor(self):
+        if not hasattr(torch, 'float8_e4m3fn'):
+            pytest.skip("torch.float8_e4m3fn not available")
+        t = torch.randn(10, 10, dtype=torch.float8_e4m3fn)
+        assert _is_fp8_param(t)
+    
+    def test_is_fp8_param_fp8_parameter(self):
+        if not hasattr(torch, 'float8_e4m3fn'):
+            pytest.skip("torch.float8_e4m3fn not available")
+        t = nn.Parameter(torch.randn(10, 10, dtype=torch.float8_e4m3fn))
+        assert _is_fp8_param(t)
+
+
+@pytest.mark.skipif(not HAS_FP8_SUPPORT, reason="FP8 support not available")
+class TestFP8Dequantization:
+    def test_dequantize_block_fp8(self):
+        if not hasattr(torch, 'float8_e4m3fn'):
+            pytest.skip("torch.float8_e4m3fn not available")
+        
+        E = 4
+        H = 64
+        I = 128
+        
+        weight = torch.randn(E, 2 * I, H, dtype=torch.float8_e4m3fn)
+        block_rows = weight.shape[1] // 16

The change achieves all five outcomes with real implementation, not just tests: it relaxes the Trainer FP8 guard narrowly (adapter-only), implements a correct FP8 MoE forward that reuses fast backends, batches+caches dequant for speed, avoids affecting other quant methods, and handles multiple scale layouts with graceful fallback. Speed path is caching+vectorized rather than a fused FP8 kernel, so c3 is partial; scale block handling is heuristic.

diff --git a/tests/test_moe_utils_fp8.py b/tests/test_moe_utils_fp8.py
new file mode 100644
index 0000000..145358e
--- /dev/null
+++ b/tests/test_moe_utils_fp8.py
@@ -0,0 +1,91 @@
+import sys
+import types
+
+import pytest
+import torch
+import torch.nn as nn
+
+
+def _fp8_dtype():
+    dtype = getattr(torch, "float8_e4m3fn", None)
+    if dtype is None:
+        pytest.skip("torch build does not expose float8_e4m3fn")
+    return dtype
+
+
+def test_fp8_moe_dequant_uses_scale_and_cache():
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+        _dequantize_fp8_expert_weight,
+        _moe_uses_fp8_expert_weights,
+    )
+
+    dtype = _fp8_dtype()
+    experts = nn.Module()
+    logical = torch.randn(2, 8, 4, dtype=torch.float32).clamp(-2, 2)
+    experts.gate_up_proj = nn.Parameter(logical.to(dtype), requires_grad=False)
+    experts.down_proj = nn.Parameter(torch.randn(2, 4, 8).clamp(-2, 2).to(dtype), requires_grad=False)
+    experts.gate_up_proj_scale_inv = torch.full((2, 1, 1), 0.5, dtype=torch.float32)
+
+    assert _moe_uses_fp8_expert_weights(experts)
+
+    first = _dequantize_fp8_expert_weight(experts, "gate_up_proj", torch.float32)
+    second = _dequantize_fp8_expert_weight(experts, "gate_up_proj", torch.float32)
+
+    assert first is second
+    assert torch.allclose(first, experts.gate_up_proj.float() * 0.5)
+    assert first.dtype == torch.float32
+    assert first.is_contiguous()
+
+
+def test_trainer_fp8_lora_patch_temporarily_allows_adapter_only_training(monkeypatch):
+    calls = []
+
+    class Trainer:
+        def __init__(self, model=None):
+            calls.append(model.is_quantized)
+            if model.is_quantized:
+                raise ValueError("quantized model rejected")
+
+    transformers_stub = types.ModuleType("transformers")
+    transformers_stub.Trainer = Trainer
+    monkeypatch.setitem(sys.modules, "transformers", transformers_stub)
+
+    from unsloth_zoo.temporary_patches.misc import patch_trainer_allow_fp8_lora_training
+

The change delivers the real fix: relaxes the Trainer FP8 guard for LoRA-trainable models, implements a working FP8 MoE forward via dequantization plus fast backends, and provides both a batched dequant and an optional grouped kernel path with sensible scale-format handling and fallbacks. Tests validate the forward correctness against a dequantized reference and the validation predicate. Minor uncertainty around the grouped_matmul kernel API existence and static-scheme paths, but core outcomes are robustly achieved.

diff --git a/tests/test_fp8_moe_lora_training.py b/tests/test_fp8_moe_lora_training.py
new file mode 100644
index 0000000..12a3b61
--- /dev/null
+++ b/tests/test_fp8_moe_lora_training.py
@@ -0,0 +1,132 @@
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from unsloth_zoo.temporary_patches.moe_utils import forward_native_moe_loop
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _dequantize_fp8_expert_weight,
+    _is_fp8_lora_trainable_model,
+    forward_moe_backend_fp8,
+)
+
+
+pytestmark = pytest.mark.skipif(
+    not hasattr(torch, "float8_e4m3fn"),
+    reason="torch build does not expose float8_e4m3fn",
+)
+
+
+def _build_fp8_experts(num_experts=3, hidden=16, intermediate=8, rank=4, with_lora=True):
+    experts = nn.Module()
+    experts.num_experts = num_experts
+    experts.hidden_dim = hidden
+    experts.intermediate_dim = intermediate
+    experts.act_fn = F.silu
+    experts.activation_scheme = "dynamic"
+    experts.block_size = None
+
+    gate_up = torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.float32).clamp(-2, 2)
+    down = torch.randn(num_experts, hidden, intermediate, dtype=torch.float32).clamp(-2, 2)
+
+    experts.gate_up_proj = nn.Parameter(gate_up.to(torch.float8_e4m3fn), requires_grad=False)
+    experts.down_proj = nn.Parameter(down.to(torch.float8_e4m3fn), requires_grad=False)
+    experts.gate_up_proj_scale_inv = nn.Parameter(torch.ones(num_experts, 1, 1), requires_grad=False)
+    experts.down_proj_scale_inv = nn.Parameter(torch.ones(num_experts, 1, 1), requires_grad=False)
+
+    if with_lora:
+        experts._unsloth_lora_gate_up_proj = (
+            torch.randn(num_experts, hidden, rank, dtype=torch.float32),
+            torch.randn(num_experts, rank, 2 * intermediate, dtype=torch.float32),
+            0.5,
+        )
+        experts._unsloth_lora_down_proj = (
+            torch.randn(num_experts, intermediate, rank, dtype=torch.float32),
+            torch.randn(num_experts, rank, hidden, dtype=torch.float32),
+            0.5,
+        )
+    return experts
+
+
+def _dequantized_reference(experts):
+    ref = nn.Module()
+    ref.num_experts = experts.num_experts
+    ref.act_fn = experts.act_fn

The change genuinely relaxes the Trainer FP8 guard for adapter-only training while preserving blocking for base training, and wires a correct FP8 dequant+scale path into all MoE forward variants with a working test. Scale layout handling and fallbacks are comprehensive. The speed improvement is real (dequant hoisted out of the loop and cached) but is a vectorized cast/multiply rather than a fused grouped kernel, so c3 is only partly satisfied.

diff --git a/tests/test_fp8_moe_lora_training.py b/tests/test_fp8_moe_lora_training.py
new file mode 100644
index 0000000..ff5d426
--- /dev/null
+++ b/tests/test_fp8_moe_lora_training.py
@@ -0,0 +1,178 @@
+import sys
+import types
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def test_trainer_fp8_lora_patch_allows_adapter_only_model(monkeypatch):
+    from unsloth_zoo.temporary_patches.misc import patch_transformers_trainer_fp8_lora
+
+    class Trainer:
+        def __init__(self, model=None, **kwargs):
+            self.saw_is_quantized = getattr(model, "is_quantized", None)
+            self.saw_quantizer_trainable = getattr(model.hf_quantizer, "is_trainable", None)
+            if getattr(model, "is_quantized", False) and not getattr(model.hf_quantizer, "is_trainable", False):
+                raise ValueError("quantized models cannot be trained")
+
+    trainer_mod = types.ModuleType("transformers.trainer")
+    trainer_mod.__package__ = "transformers"
+    trainer_mod.Trainer = Trainer
+    transformers_mod = types.ModuleType("transformers")
+    transformers_mod.__path__ = []
+    transformers_mod.Trainer = Trainer
+    transformers_mod.trainer = trainer_mod
+
+    monkeypatch.setitem(sys.modules, "transformers", transformers_mod)
+    monkeypatch.setitem(sys.modules, "transformers.trainer", trainer_mod)
+
+    patch_transformers_trainer_fp8_lora()
+
+    class FakeParam:
+        def __init__(self, requires_grad):
+            self.requires_grad = requires_grad
+
+    class FakeModel:
+        is_quantized = True
+        config = types.SimpleNamespace(quantization_config={"quant_method": "fp8"})
+        hf_quantizer = types.SimpleNamespace(is_trainable=False)
+
+        def named_parameters(self):
+            return iter([
+                ("base_model.model.layers.0.mlp.experts.gate_up_proj", FakeParam(False)),
+                ("base_model.model.layers.0.mlp.experts.gate_up_proj.lora_A.default.weight", FakeParam(True)),
+                ("base_model.model.layers.0.mlp.experts.gate_up_proj.lora_B.default.weight", FakeParam(True)),
+            ])
+
+        def modules(self):
+            return iter(())
+
+    model = FakeModel()
+    trainer = trainer_mod.Trainer(model=model)
+

The change robustly addresses all five criteria: relaxes the Trainer FP8 guard only for adapter-only training, adds a correct cached FP8 dequant + grouped backend forward path, handles multiple scale layouts with graceful fallback, and leaves other quant methods untouched. Tests validate the scale-handling and predicate logic. Minor uncertainty remains around real-model forward integration and the object.__setattr__ weight-swap, but the implementation is complete and coherent.

diff --git a/tests/test_moe_fp8_lora_training.py b/tests/test_moe_fp8_lora_training.py
new file mode 100644
index 0000000..9c3ac5d
--- /dev/null
+++ b/tests/test_moe_fp8_lora_training.py
@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+import pytest
+
+from unsloth_zoo.temporary_patches import misc
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _dequantize_fp8_expert_weight_cached,
+    _moe_uses_fp8_expert_weights,
+)
+
+
+def _fp8_dtype():
+    dtype = getattr(torch, "float8_e4m3fn", None)
+    if dtype is None:
+        dtype = getattr(torch, "float8_e5m2", None)
+    return dtype
+
+
+def test_fp8_expert_weight_dequant_is_cached():
+    dtype = _fp8_dtype()
+    if dtype is None:
+        pytest.skip("torch build does not expose float8 dtypes")
+
+    experts = nn.Module()
+    experts.gate_up_proj = torch.randn(2, 8, 4).to(dtype)
+    experts.down_proj = torch.randn(2, 4, 8).to(dtype)
+    experts.gate_up_proj_scale_inv = torch.full((2, 1, 1), 0.5)
+    experts.down_proj_scale_inv = torch.full((2, 1, 1), 0.25)
+
+    assert _moe_uses_fp8_expert_weights(experts)
+
+    first = _dequantize_fp8_expert_weight_cached(experts, "gate_up_proj", torch.bfloat16)
+    second = _dequantize_fp8_expert_weight_cached(experts, "gate_up_proj", torch.bfloat16)
+    assert first is second
+    assert first.dtype == torch.bfloat16
+    torch.testing.assert_close(
+        first.float(),
+        experts.gate_up_proj.float() * 0.5,
+        atol=0,
+        rtol=0,
+    )
+
+
+def test_fp8_block_scale_expands_and_reuses_cache():
+    dtype = _fp8_dtype()
+    if dtype is None:
+        pytest.skip("torch build does not expose float8 dtypes")
+
+    experts = nn.Module()
+    experts.weight_block_size = (2, 2)
+    experts.gate_up_proj = torch.ones(1, 3, 5).to(dtype)
+    experts.gate_up_proj_scale_inv = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
+
+    out = _dequantize_fp8_expert_weight_cached(experts, "gate_up_proj", torch.float32)

The change fully addresses all five outcomes: relaxes the Trainer FP8 guard via trainable-quantizer patching, implements a correct dequant-and-dispatch FP8 MoE forward integrated into the existing backend, provides a vectorized (non-loop) dequant supporting all three scale layouts with graceful fallbacks, and adds a differentiable FP8Linear forward for LoRA gradient flow. Detection gating prevents regressions to other quant methods. Comprehensive tests validate each mechanism.

diff --git a/tests/test_moe_utils_fp8.py b/tests/test_moe_utils_fp8.py
new file mode 100644
index 0000000..df214ba
--- /dev/null
+++ b/tests/test_moe_utils_fp8.py
@@ -0,0 +1,160 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""CPU tests for the FP8 MoE + LoRA training path (moe_utils_fp8)."""
+
+import os
+
+import pytest
+
+torch = pytest.importorskip("torch")
+import torch.nn as nn
+
+if not hasattr(torch, "float8_e4m3fn"):
+    pytest.skip("torch build lacks float8 dtypes", allow_module_level=True)
+
+from unsloth_zoo.temporary_patches import moe_utils_fp8 as m
+from unsloth_zoo.temporary_patches import moe_utils
+
+
+def _reference_block_dequant(weight_hp, scale, block_rows, block_cols):
+    out = torch.empty_like(weight_hp)
+    rows, cols = weight_hp.shape[-2], weight_hp.shape[-1]
+    n_r = scale.shape[-2]
+    n_c = scale.shape[-1]
+    for bi in range(n_r):
+        for bj in range(n_c):
+            r0, r1 = bi * block_rows, min((bi + 1) * block_rows, rows)
+            c0, c1 = bj * block_cols, min((bj + 1) * block_cols, cols)
+            out[..., r0:r1, c0:c1] = (
+                weight_hp[..., r0:r1, c0:c1] * scale[..., bi : bi + 1, bj : bj + 1]
+            )
+    return out
+
+
+def test_blockwise_dequant_matches_reference_3d():
+    torch.manual_seed(0)
+    wq = torch.randn(4, 96, 64).to(torch.float8_e4m3fn)
+    scale = torch.rand(4, 3, 2) + 0.5
+    out = m.dequantize_fp8_weight(wq, scale, torch.float32, block_size=(32, 32))
+    ref = _reference_block_dequant(wq.to(torch.float32), scale, 32, 32)
+    assert torch.allclose(out, ref)
+

no diff captured (skipped)

The change fully addresses all rubric outcomes: relaxes the trainer FP8 guard, implements a correct dequantized FP8 MoE forward with LoRA support, uses a vectorized batched dequant plus fast grouped backends, preserves non-FP8 behavior, and handles multiple scale layouts with graceful fallbacks. Minor deductions for c2/c3 since the fast path only fully accelerates the fused-gate layout and the grouped-kernel speed benefit relies on backend availability, but the implementation is robust and well-tested.

diff --git a/tests/test_moe_utils_fp8.py b/tests/test_moe_utils_fp8.py
new file mode 100644
index 0000000..5d7d1a4
--- /dev/null
+++ b/tests/test_moe_utils_fp8.py
@@ -0,0 +1,305 @@
+"""CPU regression tests for the fine-grained FP8 MoE + LoRA training path.
+
+These cover the pieces that make FP8 MoE LoRA training work without a GPU or a
+real FP8 checkpoint:
+
+1. `_dequantize_fp8_blockwise` reconstructs `fp8 * weight_scale_inv` correctly
+   for per-tensor, block-wise, and ceil-rounded block grids, matching the
+   transformers `Fp8Dequantize` convention.
+2. `_moe_uses_fp8_expert_weights` detects stacked float8 expert weights and the
+   fused-gate vs split-gate layouts.
+3. `forward_moe_backend_fp8` dequantizes the FP8 experts once and routes through
+   the native MoE loop so a hand-built LoRA delta lands on top of the correct
+   dequantized base — matching a float reference built from the same scales.
+
+All tests are float8-dtype-gated so they skip cleanly on torch builds without
+`float8_e4m3fn`.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+pytestmark = pytest.mark.skipif(
+    not hasattr(torch, "float8_e4m3fn"),
+    reason="torch build lacks float8_e4m3fn",
+)
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _dequantize_fp8_blockwise,
+    _moe_uses_fp8_expert_weights,
+    forward_moe_backend_fp8,
+)
+
+
+_FP8 = getattr(torch, "float8_e4m3fn", None)
+
+
+def _quantize_blockwise(weight_f32, block_m, block_n):
+    """Emulate a block-wise FP8 checkpoint export.
+
+    Returns (fp8_weight, scale_inv) such that dequant ≈ weight_f32. Works on the
+    last two dims and preserves any leading (expert) dims. Uses an
+    amax/FP8_MAX inverse-scale per block, exactly what real exporters store.
+    """
+    fp8_max = torch.finfo(_FP8).max
+    *lead, rows, cols = weight_f32.shape
+    srows = (rows + block_m - 1) // block_m
+    scols = (cols + block_n - 1) // block_n
+
+    scale_inv = torch.empty(*lead, srows, scols, dtype=torch.float32)
+    q = torch.empty_like(weight_f32)
+
+    for bi in range(srows):

no diff captured (skipped)

no diff captured (skipped)

no diff captured (skipped)

The change implements the real fix on all fronts: a Trainer guard relaxation for FP8 quantizers, a working vectorized fp8 MoE dequant-and-dispatch forward wired into the existing backend, robust scale-layout handling with fallbacks, and detection-gated safety for other quant methods. The main weakness is c1's is_trainable patch using a lambda class attribute rather than the property idiom, which is somewhat fragile; c3 is vectorized but not a true fused grouped kernel.

diff --git a/tests/test_moe_fp8_forward.py b/tests/test_moe_fp8_forward.py
new file mode 100644
index 0000000..0f2ac43
--- /dev/null
+++ b/tests/test_moe_fp8_forward.py
@@ -0,0 +1,310 @@
+"""Tests for the finegrained-FP8 MoE forward + LoRA path (moe_utils_fp8).
+
+Models like Qwen3-Coder-30B-A3B-Instruct-FP8 ship MoE expert weights
+(gate_up_proj / down_proj) as stacked fp8 nn.Parameters plus a block-wise
+dequant scale. transformers' fp8 integration never touches these stacked
+params, so before moe_utils_fp8 existed there was:
+
+  - no working forward (fp8 tensors fed straight into matmul), and
+  - a flat HF-Trainer refusal (FP8 quantizer reports is_trainable=False).
+
+These CPU tests pin the two pieces of the fix that don't need a GPU:
+
+  1. Block-wise dequant correctness (per-tensor, per-block, ragged blocks,
+     stacked-expert broadcasting).
+  2. End-to-end forward parity: forward_moe_backend_fp8 on an fp8 experts
+     module must match the same experts run in bf16/fp32, including LoRA
+     deltas — proving the dequant-plus-standard-backend approach is correct.
+
+They intentionally force the native_torch backend (UNSLOTH_MOE_BACKEND) so no
+CUDA / torch._grouped_mm / Triton is required.
+"""
+
+import os
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+fp8_dtype = getattr(torch, "float8_e4m3fn", None)
+pytestmark = pytest.mark.skipif(
+    fp8_dtype is None, reason="torch build has no float8_e4m3fn"
+)
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _dequantize_fp8_blockwise,
+    _moe_uses_fp8_expert_weights,
+    forward_moe_backend_fp8,
+)
+from unsloth_zoo.temporary_patches.moe_utils import forward_native_moe_loop
+
+
+# ---------------------------------------------------------------------------
+# Block-wise dequant correctness
+# ---------------------------------------------------------------------------
+
+def test_dequant_per_tensor_scale():
+    torch.manual_seed(0)
+    w = torch.randn(4, 8).clamp(-2, 2)
+    qw = w.to(fp8_dtype)
+    scale = torch.tensor(0.5)
+    out = _dequantize_fp8_blockwise(qw, scale, None, torch.float32)
+    expected = qw.to(torch.float32) * 0.5
+    torch.testing.assert_close(out, expected)

The change adds a genuine FP8 MoE forward path (dequant + existing backends) wired into forward_moe_backend, plus a Trainer guard relaxation scoped to FP8 + LoRA-only training, and a batched dequant that is faster than per-expert loops. It is plausibly functional and non-regressive for other quant methods. Weaknesses: block-scale layout handling is approximate, hardware-fallback detection is limited, and the Trainer patch depends on exact HF internal method names, so robustness across edge cases is partial.

diff --git a/tests/test_moe_fp8_lora.py b/tests/test_moe_fp8_lora.py
new file mode 100644
index 0000000..9b6e157
--- /dev/null
+++ b/tests/test_moe_fp8_lora.py
@@ -0,0 +1,498 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+Tests for FP8 MoE expert weight support (moe_utils_fp8).
+
+Covers:
+  1. _is_fp8_param detection helper
+  2. _moe_uses_fp8_expert_weights expert-module detection
+  3. _dequantize_fp8_expert_weights correctness for scalar, per-row, and None scales
+  4. forward_moe_backend_fp8 dispatches to native_loop and produces output close
+     to a full-precision reference forward
+  5. _get_fp8_scale looks up the right attribute in all naming conventions
+  6. Trainer patch: _is_lora_only_training classification
+  7. Trainer patch registers correctly without errors
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Skip everything if float8 is not available in this torch build
+# ---------------------------------------------------------------------------
+
+_HAS_FP8 = hasattr(torch, "float8_e4m3fn")
+
+pytestmark = pytest.mark.skipif(
+    not _HAS_FP8,
+    reason="torch.float8_e4m3fn not available (torch < 2.1)",
+)
+
+FP8_DTYPE = torch.float8_e4m3fn if _HAS_FP8 else None
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------

The change adds a dedicated FP8 MoE module that (1) lifts the quantizer training refusal, (2) dequantizes fp8 expert weights on the fly and dispatches to standard MoE backends with LoRA support, and (3) does so via a vectorized dequant path. It handles multiple scale layouts and falls back gracefully, and leaves non-FP8 paths untouched. Minor uncertainty on whether the is_trainable patch matches the exact Trainer validation the task names, but it targets the correct load-bearing gate.

diff --git a/tests/test_moe_utils_fp8.py b/tests/test_moe_utils_fp8.py
new file mode 100644
index 0000000..dba7a97
--- /dev/null
+++ b/tests/test_moe_utils_fp8.py
@@ -0,0 +1,326 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""CPU regression tests for the FP8 MoE + LoRA training path.
+
+Covers the two pieces that make FP8 MoE LoRA training work:
+
+  1. `_dequantize_fp8_blockwise` reconstructs the bf16/fp32 weight from the
+     float8 code + block scale, vectorized (no per-block Python loop), and
+     handles per-tensor, block, and ragged-final-block scales.
+
+  2. `forward_moe_backend_fp8` detects float8 expert weights, dequantizes them
+     on the fly, and dispatches to the standard native MoE loop — producing the
+     same output (within fp8 rounding) as an equivalent bf16 experts module,
+     including the separated-LoRA contribution.
+
+These run under the GPU-free harness in conftest.py (native_torch backend, no
+grouped_mm / triton). float8_e4m3fn tensors are CPU-constructible in modern
+torch; the whole suite is skipped if the dtype is unavailable.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+if not hasattr(torch, "float8_e4m3fn"):
+    pytest.skip("torch build lacks float8_e4m3fn", allow_module_level=True)
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _dequantize_fp8_blockwise,
+    _moe_uses_fp8_expert_weights,
+    _fp8_experts_uses_standard_activation,
+    forward_moe_backend_fp8,
+)
+from unsloth_zoo.temporary_patches.moe_utils import forward_native_moe_loop
+
+
+FP8 = torch.float8_e4m3fn

The change adds a genuine, coherent FP8 MoE + LoRA implementation: relaxes the Trainer guard for FP8+LoRA, dequantizes FP8 experts with cached bf16 and routes through grouped/vectorized backends, and preserves other quant methods. It is real fix code, not just scaffolding, and includes correctness tests. Weaknesses: the trainer guard patch is somewhat speculative about where the check lives, and block/group scale layouts (common in finegrained FP8) are not correctly dequantized (scale_inv/block scales handled naively).

diff --git a/tests/test_moe_fp8_lora.py b/tests/test_moe_fp8_lora.py
new file mode 100644
index 0000000..0964ec8
--- /dev/null
+++ b/tests/test_moe_fp8_lora.py
@@ -0,0 +1,539 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Tests for the FP8 MoE + LoRA training path.
+
+Covers:
+  - FP8 weight detection (_moe_uses_fp8_expert_weights)
+  - Per-tensor scale application (_dequantize_fp8_expert_weight)
+  - Cache idempotency (_ensure_fp8_cache)
+  - forward_native_moe_loop correctness with cached FP8 weights
+  - forward_native_moe_loop correctness with FP8 weights + LoRA
+  - HF Trainer is_quantized bypass (_is_fp8_lora_model logic surface)
+  - Module import and __all__ surface
+
+CPU-only; no GPU / transformers required for the core numeric tests.
+transformers is probed with pytest.importorskip so the trainer-patch
+tests are skipped gracefully when it isn't installed.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Module import smoke
+# ---------------------------------------------------------------------------
+
+def test_module_imports():
+    from unsloth_zoo.temporary_patches import moe_utils_fp8
+    assert moe_utils_fp8 is not None
+
+
+def test_public_api_surface():
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+        _moe_uses_fp8_expert_weights,
+        _dequantize_fp8_expert_weight,
+        _call_with_temporary_moe_weights,

diff --git a/tests/test_moe_fp8_lora.py b/tests/test_moe_fp8_lora.py
new file mode 100644
index 0000000..babf41b
--- /dev/null
+++ b/tests/test_moe_fp8_lora.py
@@ -0,0 +1,390 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+Tests for the FP8 MoE + LoRA training path.
+
+Covers:
+- FP8 dtype detection helpers
+- Per-module dequant cache (correctness + hit/miss logic)
+- forward_moe_backend_fp8 produces the same output as the reference
+- Cache invalidation when data_ptr changes
+- patch_fp8_model_for_training leaves expert weights as FP8
+- HF Trainer dtype guard patch lets FP8 + LoRA models pass validation
+- forward_moe_backend dispatcher routes FP8 modules through the FP8 path
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+def _fp8_available() -> bool:
+    return hasattr(torch, "float8_e4m3fn")
+
+
+skip_no_fp8 = pytest.mark.skipif(not _fp8_available(), reason="FP8 not available on this torch build")
+
+
+def _fp8_dtype():
+    return torch.float8_e4m3fn
+
+
+def _make_fp8_experts(num_experts: int, hidden: int, intermediate: int):
+    """Return a minimal module matching the transformers v5 experts layout."""
+    experts = nn.Module()

The change squarely targets all five outcomes with a dedicated moe_utils_fp8 module: it relaxes the Trainer FP8 guard, provides a working batched FP8 dequant MoE forward wired into the dispatcher, and gates everything behind FP8 detection so other quant methods are unaffected. Weaknesses are in robust scale-layout handling (block scales fall back to a plain cast that would be numerically wrong) and reliance on tests rather than a demonstrated real checkpoint run, but the real fix is present and functional.

diff --git a/tests/test_moe_fp8_lora.py b/tests/test_moe_fp8_lora.py
new file mode 100644
index 0000000..26c2e1e
--- /dev/null
+++ b/tests/test_moe_fp8_lora.py
@@ -0,0 +1,496 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+CPU-only tests for the FP8 MoE + LoRA training path.
+
+Covers:
+  1. FP8 dtype detection helpers
+  2. dequantize_fp8_weight with and without scale
+  3. _get_fp8_scale attribute lookup
+  4. _moe_uses_fp8_expert_weights detection
+  5. forward_moe_backend_fp8 end-to-end (native_torch backend, CPU)
+  6. patch_fp8_quantizer_is_trainable smoke-test
+  7. patch_peft_fp8_model_guard smoke-test
+  8. _get_base_weight FP8 dequant
+  9. forward_moe_backend FP8 dispatch integration
+ 10. prepare_model_for_training FP8 skip-upcast
+"""
+
+from __future__ import annotations
+
+import os
+import types
+import pytest
+import torch
+import torch.nn as nn
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+# float8 types are only available on PyTorch >= 2.1
+_HAS_FP8 = hasattr(torch, "float8_e4m3fn")
+
+
+def skip_no_fp8(fn=None):
+    marker = pytest.mark.skipif(
+        not _HAS_FP8,
+        reason="torch.float8_e4m3fn not available on this PyTorch install",
+    )
+    if fn is None:
+        return marker

The change consistently routes MoE expert weight access through _get_base_weight, which dequantizes 4-bit weights and restores 3D shape, fixing both generation (native loop) and training (grouped/triton) paths without corrupting base params and guarded for non-quant/no-bnb cases. The main gap is the lack of an explicit cast to the input/activation dtype after dequantization, so c2 is only partially satisfied.

diff --git a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
index baf7269..53ffab3 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
@@ -48,14 +48,21 @@ __all__ = [
 def _is_expert_module(module: nn.Module) -> bool:
     """
     Check if a module is an MoE experts module.
-    Specifically, check if the module has gate_up_proj & down_proj attributes that are nn.Parameter.
+    Specifically, check if the module has gate_up_proj & down_proj attributes that are nn.Parameter or Params4bit.
     """
-    return (
-        hasattr(module, "gate_up_proj")
-        and hasattr(module, "down_proj")
-        and isinstance(module.gate_up_proj, nn.Parameter)
-        and isinstance(module.down_proj, nn.Parameter)
-    )
+    if not hasattr(module, "gate_up_proj") or not hasattr(module, "down_proj"):
+        return False
+    
+    gate_up_proj = module.gate_up_proj
+    down_proj = module.down_proj
+    
+    # Check if Params4bit (4-bit quantized MoE experts)
+    if HAS_BNB:
+        if isinstance(gate_up_proj, Params4bit) and isinstance(down_proj, Params4bit):
+            return True
+    
+    # Check if nn.Parameter (non-quantized MoE experts)
+    return isinstance(gate_up_proj, nn.Parameter) and isinstance(down_proj, nn.Parameter)
 
 
 def replace_expert_params_with_bnb_params(
diff --git a/unsloth_zoo/temporary_patches/moe_utils.py b/unsloth_zoo/temporary_patches/moe_utils.py
index e97b838..ff04f85 100644
--- a/unsloth_zoo/temporary_patches/moe_utils.py
+++ b/unsloth_zoo/temporary_patches/moe_utils.py
@@ -495,7 +495,12 @@ def _get_base_weight(param):
                 "MoE quantizer patch did not fire for this expert. "
                 f"data.shape={tuple(param.data.shape)}, device={param.device}."
             )
-        return bnb.functional.dequantize_4bit(param.data, param.quant_state)
+        dequant = bnb.functional.dequantize_4bit(param.data, param.quant_state).contiguous()
+        # Reshape to original 3D shape if present (for MoE expert weights)
+        original_shape = getattr(param, "_original_shape", None)
+        if original_shape is not None and tuple(dequant.shape) != tuple(original_shape):
+            dequant = dequant.view(original_shape)
+        return dequant
 
     if hasattr(param, "get_param"):
         return param.get_param()
@@ -666,9 +671,14 @@ def _is_moe_experts_module(module) -> bool:
     if hasattr(module, "gate_up_proj"):
         param = module.gate_up_proj
 
-        # 4-bit parameters are packed into 2D tensors (n_params, 1) or similar.
-        if _check_bnb_available() and isinstance(param, Params4bit) and param.ndim == 2:
-            return True
+        # 4-bit parameters: check for Params4bit (either 2D packed or with _original_shape)
+        if _check_bnb_available() and isinstance(param, Params4bit):

The change adds solid scaffolding — generalized expert-param detection, shape tracking, idempotency and no-op guards — and wires a _get_base_weight() call into the expert forward loop, which is where dequantization would happen. However, that dequantization function is not defined anywhere in the diff, so the core c1/c2 outcomes cannot be fully verified, and the visibly-changed path is the native per-expert loop rather than the batched decoding matmul, leaving c3 weak. Partial credit throughout for the correct structure and safety, but incomplete verifiable fix.

diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..f0674fc 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -241,7 +241,7 @@ class MoeExperts4bit(nn.Module):
         Perform 4-bit matmul using bitsandbytes, following Linear4bit.forward pattern.
         """
         quant_state = weight.quant_state
-        w = weight.t()
+        w = weight.data.t()
         return bnb.matmul_4bit(x, w, bias=None, quant_state=quant_state)
 
     def forward(
diff --git a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
index baf7269..826c0f7 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
@@ -36,6 +36,21 @@ except ImportError:
     HAS_BNB = False
     Params4bit = None
 
+_MOE_EXPERT_PARAMETER_NAMES = frozenset((
+    "gate_up_proj",
+    "down_proj",
+    "gate_proj",
+    "up_proj",
+    "w1",
+    "w2",
+    "w3",
+))
+_PARAM_TYPES = (nn.Parameter,) if Params4bit is None else (nn.Parameter, Params4bit)
+
+
+def _is_params4bit(param) -> bool:
+    return Params4bit is not None and isinstance(param, Params4bit)
+
 
 __all__ = [
     "patch_bnb4bit_quantize_convert",
@@ -53,11 +68,43 @@ def _is_expert_module(module: nn.Module) -> bool:
     return (
         hasattr(module, "gate_up_proj")
         and hasattr(module, "down_proj")
-        and isinstance(module.gate_up_proj, nn.Parameter)
-        and isinstance(module.down_proj, nn.Parameter)
+        and isinstance(module.gate_up_proj, _PARAM_TYPES)
+        and isinstance(module.down_proj, _PARAM_TYPES)
     )
 
 
+def _is_moe_like_module(module: nn.Module) -> bool:
+    class_name = type(module).__name__.lower()
+    module_name = type(module).__module__.lower()
+    return (
+        "moe" in class_name
+        or "expert" in class_name
+        or "moe" in module_name
+        or "expert" in module_name
+        or _is_expert_module(module)
+    )

The change robustly dequantizes 4-bit MoE expert weights, restores logical shape/dtype, and handles the decode loop path plus idempotency and safe no-ops. The main weakness is the training path: only the unsloth_triton backend is diverted to the native loop, leaving the grouped_mm path unmodified for byte-typed weights, so training robustness is only partially assured.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..ecd5468 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -17,6 +17,7 @@
 
 from .common import *
 from .gemma import *
+from .gemma4_moe import *
 from .misc import *
 from .gemma3n import *
 from .gpt_oss import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..e4f2c1c
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,71 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from .common import TEMPORARY_PATCHES, UNSLOTH_ENABLE_LOGGING
+from .utils import patch_function, logger
+from .moe_utils import patch_param_wrapper_for_moe, get_forward_moe_backend
+
+
+def patch_gemma4_moe():
+    """
+    Patch Gemma4 MoE experts so stacked expert Params4bit tensors are executed
+    through the generic MoE backend and PEFT target_parameters use separated LoRA.
+    """
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
+    except Exception:
+        return
+
+    if getattr(Gemma4TextExperts, "_unsloth_already_patched", False):
+        return
+
+    patch_param_wrapper_for_moe()
+
+    def _gemma4_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts):
+        """
+        Gemma4 stores expert weights as (E, out_dim, in_dim), matching PEFT's
+        swapped 3D ParamWrapper convention for stacked MoE parameters.
+        """
+        total_rank = weight_A.shape[0]

The change dequantizes packed 4-bit MoE expert weights in the experts forward path, restores logical 3D shapes and casts to activation dtype, and handles both the grouped (training) and native-loop/generation decode paths before per-expert indexing. Loading-side patches preserve _original_shape idempotently without corrupting base params, and all patches are safely no-op when bnb/transformers targets are unavailable or weights are non-4bit.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..2b20783 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -24,6 +24,7 @@ from .qwen3_moe import *
 from .qwen3_vl_moe import *
 from .qwen3_next_moe import *
 from .qwen3_5_moe import *
+from .gemma4_moe import *
 from .glm4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..06311f0
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,174 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import os
+
+from .common import TEMPORARY_PATCHES, UNSLOTH_ENABLE_LOGGING
+from .utils import patch_function, logger
+from .moe_utils import patch_param_wrapper_for_moe, get_forward_moe_backend
+from .qwen3_moe import (
+    _make_qwen_moe_lora_extractor,
+    _patch_causal_lm_forward_for_hidden_states,
+)
+
+
+def _slice_hidden_states_for_logits(hidden_states, logits_to_keep):
+    if logits_to_keep is None:
+        return hidden_states
+    if isinstance(logits_to_keep, int) and logits_to_keep == 0:
+        return hidden_states
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    return hidden_states[:, slice_indices, :]
+
+
+def _patch_gemma4_conditional_generation_forward(module):
+    try:
+        Gemma4ForConditionalGeneration = module.Gemma4ForConditionalGeneration
+        Gemma4CausalLMOutputWithPast = module.Gemma4CausalLMOutputWithPast
+    except AttributeError:
+        return

The change correctly wires both the generation (native per-expert loop) and training paths to obtain dequantized, LoRA-augmented weights via _get_base_weight/preprocess_weight, hoisting dequant outside the loop and casting to activation dtype, with safe guards for missing modules. However, the crucial dequantization and shape-recovery logic lives in _get_base_weight/preprocess_weight which are not shown or modified in this diff, so the core 4-bit unpacking cannot be verified as implemented here — points on c1/c2/c4 are reduced accordingly.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..ecd5468 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -17,6 +17,7 @@
 
 from .common import *
 from .gemma import *
+from .gemma4_moe import *
 from .misc import *
 from .gemma3n import *
 from .gpt_oss import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..4f1bee4
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,46 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from .common import TEMPORARY_PATCHES, UNSLOTH_ENABLE_LOGGING
+from .utils import patch_function, logger
+from .moe_utils import patch_param_wrapper_for_moe, get_forward_moe_backend
+
+
+def patch_gemma4_moe():
+    """
+    Patch Gemma4 MoE experts to use the shared MoE backend.
+
+    Gemma4TextExperts stores gate_up_proj/down_proj as stacked 3D nn.Parameter
+    tensors. After transformers' bnb 4-bit loading path replaces those tensors
+    with Params4bit, the upstream expert loop indexes packed storage during
+    generation and PEFT's ParamWrapper path cannot apply separated LoRA during
+    training. The shared backend dequantizes Params4bit through _get_base_weight
+    and consumes PEFT's separated LoRA metadata.
+    """
+    patch_param_wrapper_for_moe()
+
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
+    except (ImportError, AttributeError):
+        return
+
+    patch_function(Gemma4TextExperts, "forward", get_forward_moe_backend(), match_level="relaxed")
+
+    if UNSLOTH_ENABLE_LOGGING:

no diff captured (skipped)

The change robustly addresses all six criteria: dequantizing 4-bit expert weights in both grouped and batched paths, recovering the logical 3D shape, handling the decode-path indexing-before-linear problem, and preserving base weights via a restore pattern with idempotency and no-op guards. It correctly identifies and works around the subtle dispatcher-reference issue for the batched path.

diff --git a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
index baf7269..48746f9 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
@@ -41,6 +41,7 @@ __all__ = [
     "patch_bnb4bit_quantize_convert",
     "patch_bnb4bit_quantizer_param_needs_quantization",
     "patch_bnb4bit_quantizer_process_model",
+    "patch_transformers_grouped_linear_4bit",
     "replace_expert_params_with_bnb_params",
 ]
 
@@ -285,3 +286,114 @@ def patch_bnb4bit_quantizer_process_model():
     pass
 pass
 TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_process_model)
+
+
+def _maybe_dequant_params4bit_weight(weight, input_dtype):
+    """If `weight` is a packed Params4bit, dequantize to logical shape and cast
+    to `input_dtype`. Otherwise return `weight` unchanged.
+    """
+    if isinstance(weight, Params4bit) and getattr(weight, "quant_state", None) is not None:
+        original_shape = getattr(weight, "_original_shape", None)
+        dequant = bnb.functional.dequantize_4bit(weight.data, weight.quant_state)
+        if original_shape is not None and tuple(dequant.shape) != tuple(original_shape):
+            dequant = dequant.reshape(original_shape)
+        return dequant.to(input_dtype)
+    return weight
+
+
+def patch_transformers_grouped_linear_4bit():
+    """
+    transformers v5's `grouped_mm_experts_forward` and `batched_mm_experts_forward`
+    (in `transformers.integrations.moe`) read `self.gate_up_proj` /
+    `self.down_proj` raw and pass them through `_grouped_linear` /
+    `_batched_linear` -> `weight.transpose(-2, -1)` -> `torch._grouped_mm` /
+    `torch.bmm` -> `mat_a.to(weight.dtype)` ...
+
+    For MoE arches whose experts class is NOT replaced by an unsloth-zoo
+    per-arch patch (e.g. some Glm4Moe variants, Gemma4MoE), the experts forward
+    therefore sees the raw Params4bit (uint8 packed storage). The matmul ops raise:
+        - grouped_mm path:  `RuntimeError: Expected mat_a to be Float32, BFloat16 or
+                            Float16 matrix, got Byte` (during training)
+        - batched_mm path:  `RuntimeError: batch1 must be a 3D tensor` (during
+                            autoregressive decoding where the batched_mm dispatcher
+                            is used instead of grouped_mm)
+
+    Fix: wrap `_grouped_linear` and `_batched_linear` to detect Params4bit
+    weights, dequantize via `bnb.functional.dequantize_4bit` (using
+    `_original_shape` to recover the logical 3D `(E, in, out)` shape from
+    packed `(N, 1)` storage), cast to the input dtype, then delegate to the
+    original function which handles `is_transposed` orientation correctly.
+
+    Forward-only -- base weights are frozen; gradient flow stays on LoRA paths
+    that the per-arch wrappers (when they exist) inject separately. For arches
+    without a per-arch wrapper this gives base-only forward; PEFT's
+    `_activate_lora` parametrization adds the delta on top.
+    """
+    try:

no diff captured (empty)

The change robustly detects and dequantizes 4-bit MoE expert weights across both grouped (training) and batched (decoding) paths, correctly recovers logical shape and dtype, preserves base params via restore pattern, and no-ops for non-quantized/absent-bnb cases. All criteria appear fully addressed.

diff --git a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
index baf7269..48746f9 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb_transformers.py
@@ -41,6 +41,7 @@ __all__ = [
     "patch_bnb4bit_quantize_convert",
     "patch_bnb4bit_quantizer_param_needs_quantization",
     "patch_bnb4bit_quantizer_process_model",
+    "patch_transformers_grouped_linear_4bit",
     "replace_expert_params_with_bnb_params",
 ]
 
@@ -285,3 +286,114 @@ def patch_bnb4bit_quantizer_process_model():
     pass
 pass
 TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_process_model)
+
+
+def _maybe_dequant_params4bit_weight(weight, input_dtype):
+    """If `weight` is a packed Params4bit, dequantize to logical shape and cast
+    to `input_dtype`. Otherwise return `weight` unchanged.
+    """
+    if isinstance(weight, Params4bit) and getattr(weight, "quant_state", None) is not None:
+        original_shape = getattr(weight, "_original_shape", None)
+        dequant = bnb.functional.dequantize_4bit(weight.data, weight.quant_state)
+        if original_shape is not None and tuple(dequant.shape) != tuple(original_shape):
+            dequant = dequant.reshape(original_shape)
+        return dequant.to(input_dtype)
+    return weight
+
+
+def patch_transformers_grouped_linear_4bit():
+    """
+    transformers v5's `grouped_mm_experts_forward` and `batched_mm_experts_forward`
+    (in `transformers.integrations.moe`) read `self.gate_up_proj` /
+    `self.down_proj` raw and pass them through `_grouped_linear` /
+    `_batched_linear` -> `weight.transpose(-2, -1)` -> `torch._grouped_mm` /
+    `torch.bmm` -> `mat_a.to(weight.dtype)` ...
+
+    For MoE arches whose experts class is NOT replaced by an unsloth-zoo
+    per-arch patch (e.g. some Glm4Moe variants, Gemma4MoE), the experts forward
+    therefore sees the raw Params4bit (uint8 packed storage). The matmul ops raise:
+        - grouped_mm path:  `RuntimeError: Expected mat_a to be Float32, BFloat16 or
+                            Float16 matrix, got Byte` (during training)
+        - batched_mm path:  `RuntimeError: batch1 must be a 3D tensor` (during
+                            autoregressive decoding where the batched_mm dispatcher
+                            is used instead of grouped_mm)
+
+    Fix: wrap `_grouped_linear` and `_batched_linear` to detect Params4bit
+    weights, dequantize via `bnb.functional.dequantize_4bit` (using
+    `_original_shape` to recover the logical 3D `(E, in, out)` shape from
+    packed `(N, 1)` storage), cast to the input dtype, then delegate to the
+    original function which handles `is_transposed` orientation correctly.
+
+    Forward-only -- base weights are frozen; gradient flow stays on LoRA paths
+    that the per-arch wrappers (when they exist) inject separately. For arches
+    without a per-arch wrapper this gives base-only forward; PEFT's
+    `_activate_lora` parametrization adds the delta on top.
+    """
+    try:

no diff captured (empty)

no diff captured (skipped)

The candidate makes a small, plausible fix: reshaping the dequantized 4-bit expert weight back to its logical 3D shape, which is the key missing piece for MoE grouped/batched matmul. However, the diff does not show the forward-path integration, decode/training call sites, or dtype casting, so several outcomes can only be credited indirectly. The change is reasonable and non-regressive but incomplete/unverifiable from the diff alone.

diff --git a/unsloth_zoo/temporary_patches/moe_utils.py b/unsloth_zoo/temporary_patches/moe_utils.py
index e97b838..feebb57 100644
--- a/unsloth_zoo/temporary_patches/moe_utils.py
+++ b/unsloth_zoo/temporary_patches/moe_utils.py
@@ -495,7 +495,18 @@ def _get_base_weight(param):
                 "MoE quantizer patch did not fire for this expert. "
                 f"data.shape={tuple(param.data.shape)}, device={param.device}."
             )
-        return bnb.functional.dequantize_4bit(param.data, param.quant_state)
+        dequant = bnb.functional.dequantize_4bit(param.data, param.quant_state)
+        # dequantize_4bit returns the packed/flattened 2D layout, NOT the logical
+        # stacked (E, out, in) shape the MoE forward needs. preprocess_weight and
+        # torch._grouped_mm require the 3D expert layout, so restore it from the
+        # _original_shape recorded when the expert param was quantized (see
+        # moe_bnb_transformers.patched_convert). Without this, the first grouped_mm
+        # in training errors out and generation breaks during decode. Mirrors the
+        # dequant->reshape done in the merge/unmerge path (misc.py).
+        original_shape = getattr(param, "_original_shape", None)
+        if original_shape is not None and tuple(dequant.shape) != tuple(original_shape):
+            dequant = dequant.reshape(original_shape)
+        return dequant
 
     if hasattr(param, "get_param"):
         return param.get_param()

The change correctly detects and dequantizes 4-bit MoE expert weights before matmul, recovers the 3D logical shape, and handles the per-expert indexing path safely without corrupting base weights, with proper no-op guards for non-quantized/no-bnb cases. Weaknesses: explicit input-dtype casting is not clearly shown, and the grouped-GEMM training/decoding backend path fix is less directly evidenced than the native loop, so full training/generation robustness is somewhat inferred.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..a0ab714 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -25,6 +25,7 @@ from .qwen3_vl_moe import *
 from .qwen3_next_moe import *
 from .qwen3_5_moe import *
 from .glm4_moe import *
+from .gemma4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
 from .ministral import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..664a379
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,142 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import torch
+import torch.nn.functional as F
+from .common import TEMPORARY_PATCHES, UNSLOTH_ENABLE_LOGGING
+from .utils import patch_function, raise_error, logger
+from .moe_utils import (
+    patch_param_wrapper_for_moe,
+    get_forward_moe_backend,
+)
+
+
+def patch_gemma4_moe():
+    """
+    Patches Gemma4 MoE to support 4-bit quantization (bnb) and Split LoRA using grouped GEMM.
+
+    Gemma4MoE (transformers >= 5.x) uses stacked expert weights as nn.Parameter tensors
+    (gate_up_proj, down_proj) just like Qwen3-MoE, so the same grouped-GEMM backend works.
+
+    Two components are patched:
+      1. Gemma4TextExperts.forward  - the inner expert computation kernel
+      2. Gemma4TextSparseMoeBlock.forward - the outer routing + expert-dispatch block,
+         rewritten to call self.experts() with the (hidden_states, top_k_index, top_k_weights)
+         signature expected by forward_moe_backend.
+    """
+    # Patch PEFT ParamWrapper for separated LoRA weights (idempotent).
+    patch_param_wrapper_for_moe()
+

no diff captured (skipped)

The change robustly dequantizes 4-bit MoE expert weights and restores logical 3D shape in _get_base_weight, and correctly rewires the native loop (decode) path to dequantize once before indexing without mutating base params. However, the diff does not clearly show the grouped/training GEMM matmul path being adapted beyond reusing _get_base_weight, so c4 and the batched grouped-mm portion of c3 are only partially demonstrated. Non-quantized and missing-module cases are safely no-op.

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index 8312b9b..2a3d1af 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -2992,6 +2992,7 @@ DISABLE_COMPILE_MODULES = [
     "GptOssExperts",
     "Gemma3nTextModel",
     "Glm4MoeLiteNaiveMoe",
+    "Gemma4TextExperts",
     "Qwen3NextGatedDeltaNet",
     "GatedDeltaNet",
     "Qwen3_5MoeGatedDeltaNet",
diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..a0ab714 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -25,6 +25,7 @@ from .qwen3_vl_moe import *
 from .qwen3_next_moe import *
 from .qwen3_5_moe import *
 from .glm4_moe import *
+from .gemma4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
 from .ministral import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..52730a9
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,110 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import torch
+import torch.nn.functional as F
+from .common import (
+    TEMPORARY_PATCHES,
+    UNSLOTH_ENABLE_LOGGING,
+)
+from .utils import (
+    patch_function,
+    raise_error,
+    logger,
+)
+from .moe_utils import patch_param_wrapper_for_moe, get_forward_moe_backend
+from .qwen3_moe import (
+    _make_qwen_moe_lora_extractor,

The change robustly detects 4-bit packed expert weights and dequantizes them in both the loop (decode) and triton grouped (training) paths, restoring logical 3D shape via _original_shape. It preserves original params without mutation and safely no-ops for non-quantized cases. Only the explicit dtype-casting of dequantized operands to the activation dtype is implicit rather than enforced, so c2 is docked slightly.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..a0ab714 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -25,6 +25,7 @@ from .qwen3_vl_moe import *
 from .qwen3_next_moe import *
 from .qwen3_5_moe import *
 from .glm4_moe import *
+from .gemma4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
 from .ministral import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..c3e7960
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,216 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""
+Patches for Gemma4 MoE (sparse mixture-of-experts) models.
+
+Gemma4Moe uses stacked 3D nn.Parameter tensors for expert weights
+(gate_up_proj / down_proj), mirroring the Qwen3MoE layout.  This file
+wires them into:
+  1. The grouped-GEMM / loop-based forward backend in moe_utils.py
+  2. BNB 4-bit quantization via the _original_shape-preserving patched_convert
+  3. PEFT Split-LoRA via patch_param_wrapper_for_moe
+
+Weight layout used by Gemma4 (confirmed from HF transformers source):
+  gate_up_proj: (E, 2*intermediate, hidden) - F.linear format (out, in)
+  down_proj:    (E, hidden, intermediate)   - F.linear format (out, in)
+which is identical to the standard (Qwen3-MoE) format.
+"""
+
+import torch
+import torch.nn.functional as F
+import os
+
+from .common import TEMPORARY_PATCHES, UNSLOTH_ENABLE_LOGGING
+from .utils import patch_function, raise_error, logger
+from .moe_utils import (
+    patch_param_wrapper_for_moe,
+    get_forward_moe_backend,
+)

The change robustly adds 4-bit detection, dequantization, and 3D reshaping in the shared _get_base_weight and fixes the loop/decoding path to dequantize once before per-expert indexing, satisfying c1, c2, c3, c5, c6. The grouped/training path (c4) is only indirectly addressed via the shared helper without an explicit shown fix, so partial credit. Overall a solid, non-destructive fix.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index 180ba59..a0ab714 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -25,6 +25,7 @@ from .qwen3_vl_moe import *
 from .qwen3_next_moe import *
 from .qwen3_5_moe import *
 from .glm4_moe import *
+from .gemma4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
 from .ministral import *
diff --git a/unsloth_zoo/temporary_patches/gemma4_moe.py b/unsloth_zoo/temporary_patches/gemma4_moe.py
new file mode 100644
index 0000000..e42e087
--- /dev/null
+++ b/unsloth_zoo/temporary_patches/gemma4_moe.py
@@ -0,0 +1,155 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import torch
+import torch.nn.functional as F
+
+from .common import (
+    TEMPORARY_PATCHES,
+    UNSLOTH_ENABLE_LOGGING,
+)
+from .utils import (
+    patch_function,
+    raise_error,
+    logger,
+)
+from .moe_utils import (
+    patch_param_wrapper_for_moe,
+    get_forward_moe_backend,
+)
+from .qwen3_moe import (
+    _make_qwen_moe_lora_extractor,
+    _patch_causal_lm_forward_for_hidden_states,
+)
+
+
+def patch_gemma4_moe():
+    """
+    Patches Gemma4 MoE (Gemma4MoeExperts / Gemma4TextMoeSparseMoeBlock) to support:
+    - bnb 4-bit quantization: weights are stacked 3D nn.Parameter, handled by

The change adds LoRA support, module detection, and forward-path dispatch for 4-bit MoE experts on top of an apparently pre-existing quantization scaffold in moe_bnb.py. Core quantization and config-honoring logic are assumed/pre-existing rather than clearly established by this diff, and the LoRA integration relies on fragile inspect.stack() wrapper discovery with duplicated, questionable reshape branches. Forward-path dequant delegation is the most solid part; overall the outcomes are partially but not robustly achieved.

diff --git a/MOE_4BIT_LORA_CHANGES.md b/MOE_4BIT_LORA_CHANGES.md
new file mode 100644
index 0000000..9c648b4
--- /dev/null
+++ b/MOE_4BIT_LORA_CHANGES.md
@@ -0,0 +1,105 @@
+# MoE 4-bit Quantization + LoRA Integration
+
+## Summary of Changes
+
+This implementation enables 4-bit quantization and LoRA training for MoE (Mixture of Experts) models like Qwen3-MoE, Qwen3-VL-MoE, DeepSeek-V3, and GLM4-MoE.
+
+## Files Modified
+
+### 1. `unsloth_zoo/temporary_patches/moe_bnb.py`
+- Added `__all__` exports for public API
+- Added `_get_lora_weights_4bit()` method to extract LoRA weights for quantized experts
+- Modified `forward()` to support LoRA adapters on top of 4-bit quantized weights:
+  - Computes base output using `bnb.matmul_4bit`
+  - Adds LoRA delta: `(X @ lora_A) @ lora_B * scaling`
+- Added `prepare_moe_for_lora()` to mark models for LoRA training
+- Added `auto_quantize_moe_experts()` for automatic quantization during loading
+
+### 2. `unsloth_zoo/temporary_patches/moe_utils.py`
+- Modified `_is_moe_experts_module()` to detect 4-bit quantized experts (`_is_bnb_4bit`, `_bnb_gate_up_weights`)
+- Modified `_get_base_weight()` to return `None` for quantized modules (signals separate handling)
+- Modified `_extract_lora_from_wrapper()` to handle quantized expert modules
+- Modified `forward_native_grouped_mm()` to detect and delegate to 4-bit forward
+- Modified `forward_triton_grouped_gemm()` to detect and delegate to 4-bit forward
+- Modified `forward_native_moe_loop()` to detect and delegate to 4-bit forward
+
+### 3. `unsloth_zoo/temporary_patches/qwen3_moe.py`
+- Added import of moe_bnb module
+- Modified all forward functions to detect 4-bit quantized experts and delegate accordingly
+- Added `_is_4bit_quantized()` helper method to Qwen3MoeExperts
+
+### 4. `unsloth_zoo/peft_utils.py`
+- Added `MOE_EXPERT_TARGET_MODULES` list for LoRA target module patterns
+- Updated `__all__` exports
+
+### 5. `unsloth_zoo/temporary_patches/__init__.py`
+- Added `from .moe_bnb import *` to export moe_bnb functions
+
+## How It Works
+
+### 4-bit Quantization Flow
+1. `replace_with_bnb_moe_experts()` replaces expert modules before weight loading
+2. Weights are quantized to `Params4bit` when moved to CUDA
+3. Forward uses `bnb.matmul_4bit` for each expert's computations
+
+### LoRA on 4-bit Experts
+1. PEFT wraps the quantized expert parameters
+2. `_extract_lora_from_wrapper()` extracts LoRA A/B weights
+3. Forward computes: `base_4bit_output + (X @ lora_A) @ lora_B * scaling`
+4. Gradients flow through LoRA weights while base weights stay frozen
+
+## Usage Example
+
+```python
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig

The change adds proper integration scaffolding: a guarded, idempotent HF-loading hook, PEFT-visible per-expert projection handles preserving dimensional metadata, and custom per-expert LoRA application in the forward path. However much of the actual quantization/dequantization machinery (replace_with_bnb_moe_experts, _matmul_4bit, quantize_moe_experts_inplace) is pre-existing and not shown, so the diff mostly wires things together rather than proving the core quantization works end-to-end. Regression safety and metadata preservation are solid; quantization and dequantization outcomes are only partially verifiable.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..deddcc6 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -26,4 +26,5 @@ from .glm4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
 from .mxfp4 import *
+from .moe_bnb import *
 from .bitsandbytes import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..13ede54 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -34,6 +34,7 @@ from typing import Optional, List, Any, Tuple
 import os
 import warnings
 from ..log import logger
+from .common import TEMPORARY_PATCHES
 
 UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
 
@@ -46,6 +47,7 @@ try:
 except ImportError:
     HAS_BNB = False
     Params4bit = None
+_Moe4bitProjectionBase = bnb.nn.Linear4bit if HAS_BNB else nn.Module
 
 
 def _check_bnb_available():
@@ -57,6 +59,52 @@ def _check_bnb_available():
         )
 
 
+class Moe4bitProjection(_Moe4bitProjectionBase):
+    """PEFT-visible handle for one stacked MoE projection."""
+
+    def __init__(self, parent, name: str, in_features: int, out_features: int):
+        nn.Module.__init__(self)
+        object.__setattr__(self, "parent", parent)
+        self.name = name
+        self.parameter_name = name
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_experts = parent.num_experts
+        self.compute_dtype = parent.compute_dtype
+        self.compute_type_is_set = parent.compute_type_is_set
+        self._is_bnb_4bit = True
+        self.is_loaded_in_4bit = True
+        self.quant_type = parent.quant_type
+        self.compress_statistics = parent.compress_statistics
+        self.quant_storage = parent.quant_storage
+
+    @property
+    def weight(self):
+        weights = (
+            self.parent._bnb_gate_up_weights
+            if self.name == "gate_up_proj"
+            else self.parent._bnb_down_weights
+        )

The change adds a complete MoE 4-bit path: expert quantization via Params4bit, HF/accelerate loading hooks, meta-tensor shape metadata for PEFT LoRA discovery, a dimension-aware LoRA extractor and forward integration, and config propagation. Patches are bnb-guarded and idempotent, delegating non-MoE cases to originals. All rubric outcomes appear robustly achieved.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..deddcc6 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -26,4 +26,5 @@ from .glm4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
 from .mxfp4 import *
+from .moe_bnb import *
 from .bitsandbytes import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..792b3c9 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -34,6 +34,7 @@ from typing import Optional, List, Any, Tuple
 import os
 import warnings
 from ..log import logger
+from .common import TEMPORARY_PATCHES
 
 UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
 
@@ -85,6 +86,7 @@ class MoeExperts4bit(nn.Module):
         quant_type: str = "nf4",
         quant_storage: torch.dtype = torch.uint8,
         device: Optional[torch.device] = None,
+        transposed: bool = False,
     ):
         """
         Initialize MoeExperts4bit.
@@ -98,6 +100,7 @@ class MoeExperts4bit(nn.Module):
             quant_type: Quantization type ("nf4" or "fp4")
             quant_storage: Storage dtype for quantized weights
             device: Device to place weights on
+            transposed: Whether weights are stored as (E, in, out)
         """
         super().__init__()
         _check_bnb_available()
@@ -110,9 +113,11 @@ class MoeExperts4bit(nn.Module):
         self.compress_statistics = compress_statistics
         self.quant_type = quant_type
         self.quant_storage = quant_storage
+        self.transposed = transposed
 
         # Flag for detection
         self._is_bnb_4bit = True
+        self._unsloth_model_type = "bnb_4bit_moe"
 
         # Per-expert quantized weights stored as ParameterLists for proper registration
         # Each expert's gate_up_proj: Params4bit of shape [2*intermediate_dim, hidden_dim]
@@ -125,6 +130,18 @@ class MoeExperts4bit(nn.Module):
         self._gate_up_proj_pending = None
         self._down_proj_pending = None
 
+        # PEFT discovers MoE target_parameters from named Parameters. These
+        # meta tensors provide shape metadata for LoRA without storing a full
+        # unquantized copy of expert weights.
+        if transposed:
+            gate_up_shape = (num_experts, hidden_dim, 2 * intermediate_dim)
+            down_shape = (num_experts, intermediate_dim, hidden_dim)

The change comprehensively integrates 4-bit quantization of stacked MoE expert parameters into Transformers' BnB loading path and adds the metadata/placeholder machinery and PEFT patches needed for LoRA to attach and train on experts. Dequantization fallbacks, quant-config honoring, dimensional metadata preservation, and idempotent/guarded patching are all present. It robustly addresses every rubric criterion.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..0cf14a8 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -22,6 +22,7 @@ from .gemma3n import *
 from .gpt_oss import *
 from .qwen3_moe import *
 from .qwen3_vl_moe import *
+from .moe_bnb import *
 from .glm4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..7efd6d8 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -34,6 +34,7 @@ from typing import Optional, List, Any, Tuple
 import os
 import warnings
 from ..log import logger
+from .common import TEMPORARY_PATCHES
 
 UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
 
@@ -57,6 +58,36 @@ def _check_bnb_available():
         )
 
 
+def _bnb_moe_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts):
+    """
+    Extract PEFT ParamWrapper LoRA weights for the quantized MoE expert wrapper.
+
+    Standard Qwen3-MoE expert tensors are stored as (E, out, in), while
+    Qwen3-VL-MoE tensors are stored as (E, in, out). PEFT's ParamWrapper
+    always treats 3D parameters as (E, in, out), so standard Qwen3-MoE needs
+    A/B swapped for the actual computation.
+    """
+    total_rank = weight_A.shape[0]
+    rank_per_expert = total_rank // num_experts
+    dim_A = weight_A.shape[1]
+    dim_B = weight_B.shape[0]
+
+    experts_module = wrapper.get_base_layer() if hasattr(wrapper, "get_base_layer") else None
+    is_transposed_format = getattr(experts_module, "is_transposed_format", False)
+
+    if is_transposed_format:
+        first_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
+        first_weight = first_weight.permute(0, 2, 1).contiguous()
+        second_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+        second_weight = second_weight.permute(1, 2, 0).contiguous()
+    else:
+        first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+        first_weight = first_weight.permute(1, 0, 2).contiguous()
+        second_weight = weight_A.view(num_experts, rank_per_expert, dim_A).contiguous()
+
+    return first_weight, second_weight, scaling, num_experts
+
+
 class MoeExperts4bit(nn.Module):
     """

The change robustly achieves 4-bit quantization of MoE expert weights by hooking the HF bitsandbytes replacement pass and handling both state_dict and direct-assignment loaders, while preserving 3D expert shape metadata for correct LoRA attachment and dequantizing in forward. Quantization config is honored and patches are guarded and idempotent, protecting existing behavior. Minor uncertainty only on complete double-quant config plumbing.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..e671d4d 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -27,3 +27,4 @@ from .deepseek_v3_moe import *
 from .pixtral import *
 from .mxfp4 import *
 from .bitsandbytes import *
+from .moe_bnb import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..b0969eb 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -33,6 +33,9 @@ import torch.nn.functional as F
 from typing import Optional, List, Any, Tuple
 import os
 import warnings
+from .common import TEMPORARY_PATCHES
+from .utils import patch_function
+from .moe_utils import patch_param_wrapper_for_moe
 from ..log import logger
 
 UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
@@ -85,6 +88,7 @@ class MoeExperts4bit(nn.Module):
         quant_type: str = "nf4",
         quant_storage: torch.dtype = torch.uint8,
         device: Optional[torch.device] = None,
+        weight_format: str = "linear",
     ):
         """
         Initialize MoeExperts4bit.
@@ -110,6 +114,7 @@ class MoeExperts4bit(nn.Module):
         self.compress_statistics = compress_statistics
         self.quant_type = quant_type
         self.quant_storage = quant_storage
+        self.weight_format = weight_format
 
         # Flag for detection
         self._is_bnb_4bit = True
@@ -122,12 +127,100 @@ class MoeExperts4bit(nn.Module):
 
         # These will hold the original Parameters for weight loading
         # They get converted to Params4bit when .to(cuda) is called
-        self._gate_up_proj_pending = None
-        self._down_proj_pending = None
+        self.register_buffer("_gate_up_proj_pending", None)
+        self.register_buffer("_down_proj_pending", None)
 
         # Activation function
         self.act_fn = F.silu
 
+        # PEFT target_parameters discovers trainable MoE adapters from the
+        # public parameter names. Keep meta placeholders for shape discovery;
+        # the actual base weights live in the Params4bit lists below.
+        self._init_lora_placeholders(device=device)
+
+    def _make_lora_placeholder(
+        self,
+        shape: Tuple[int, ...],
+        device: Optional[torch.device] = None,

The change fully implements 4-bit quantization of MoE expert weights via per-expert Params4bit, hooks it into transformers' loading path, preserves dimensional metadata, and enables LoRA on quantized experts through PEFT ParamWrapper patching with correct shape inference and gradient-carrying forward deltas. Config settings are honored and existing Linear/non-MoE paths are preserved with bnb guards and idempotent patching. All criteria appear robustly satisfied.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..92832d2 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -25,5 +25,6 @@ from .qwen3_vl_moe import *
 from .glm4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
+from . import moe_bnb
 from .mxfp4 import *
 from .bitsandbytes import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..19e2708 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -24,7 +24,18 @@ skips them.
 Design follows bitsandbytes patterns:
 - Uses Params4bit for per-expert weight quantization
 - Quantization happens on .to(cuda), just like Linear4bit
-- Integrates with model loading via replace_with_bnb_moe_experts()
+- Integrates with model loading via replace_with_bnb_moe_experts(), which is
+  hooked into transformers' replace_with_bnb_linear by patch_moe_bnb_4bit()
+
+LoRA support:
+- PEFT `target_parameters` wraps gate_up_proj / down_proj via ParamWrapper.
+  Since the original 3D bf16 parameters are replaced by quantized storage,
+  a zero-element placeholder parameter keeps the attribute targetable, and
+  patch_peft_param_wrapper_for_moe_4bit() sizes the LoRA A/B matrices from the
+  module's (num_experts, hidden_dim, intermediate_dim) metadata instead.
+- The forward pass consumes the separated-LoRA data injected by
+  moe_utils._patched_param_wrapper_forward and adds the per-expert LoRA delta
+  on top of the 4-bit base matmul, so gradients flow to the LoRA weights.
 """
 
 import torch
@@ -65,6 +76,7 @@ class MoeExperts4bit(nn.Module):
     - Stores per-expert weights as Params4bit
     - Quantization happens when moving to CUDA (like Linear4bit)
     - Self-contained forward using bnb.matmul_4bit for each expert
+    - Supports separated LoRA via PEFT ParamWrapper (see moe_utils)
 
     The module replaces Qwen3MoeExperts or similar, providing the same interface
     but with quantized weights.
@@ -74,6 +86,10 @@ class MoeExperts4bit(nn.Module):
     Set UNSLOTH_MOE_BACKEND=grouped_mm environment variable.
     """
     _warned_loop_based = False
+    # Qwen3-VL stores weights in grouped_mm layout (E, in, out); standard
+    # Qwen3-MoE uses F.linear layout (E, out, in). Transposed weights are
+    # normalized to F.linear layout at quantization time.
+    weights_transposed = False
 
     def __init__(
         self,
@@ -120,10 +136,21 @@ class MoeExperts4bit(nn.Module):
         self._bnb_gate_up_weights = nn.ParameterList()
         self._bnb_down_weights = nn.ParameterList()
 
-        # These will hold the original Parameters for weight loading
-        # They get converted to Params4bit when .to(cuda) is called

no diff captured (skipped)

The change achieves quantization of MoE experts in place as Params4bit while preserving names for PEFT and shape metadata via quant_state, wires dequantization into the read path, honors the bnb config, and guards non-MoE/standard paths with bnb availability and idempotency. It robustly addresses all criteria; only minor completeness concerns on compute_dtype handling and backend routing for the 4-bit case reduce c6 slightly.

diff --git a/unsloth_zoo/patching_utils.py b/unsloth_zoo/patching_utils.py
index ba77af9..bacd637 100644
--- a/unsloth_zoo/patching_utils.py
+++ b/unsloth_zoo/patching_utils.py
@@ -441,6 +441,24 @@ def patch_model_and_tokenizer(
         pass
     pass
 
+    # Quantize MoE expert weights to 4-bit when loading in 4-bit.
+    # MoE experts store gate_up_proj / down_proj as stacked 3D nn.Parameters,
+    # which stock bitsandbytes (nn.Linear only) leaves in full precision. This
+    # quantizes them in place so LoRA training on experts works in 4-bit.
+    try:
+        is_4bit = getattr(model, "is_loaded_in_4bit", False)
+        if not is_4bit:
+            q_cfg = getattr(getattr(model, "config", None), "quantization_config", None)
+            if q_cfg is not None:
+                is_4bit = bool(getattr(q_cfg, "load_in_4bit", False))
+        if is_4bit:
+            from .temporary_patches.moe_bnb import maybe_quantize_moe_experts
+            maybe_quantize_moe_experts(model, load_in_4bit=True)
+    except Exception as e:
+        if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
+            from .temporary_patches.utils import logger
+            logger.warning(f"Unsloth: Could not 4-bit quantize MoE experts: {e}")
+
     if not fix_embeddings: return model, tokenizer
 
     # Check if torch.nn.Embedding seen
diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..d027461 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -20,6 +20,7 @@ from .gemma import *
 from .misc import *
 from .gemma3n import *
 from .gpt_oss import *
+from .moe_bnb import *
 from .qwen3_moe import *
 from .qwen3_vl_moe import *
 from .glm4_moe import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..266c1e9 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -17,31 +17,54 @@
 MoE 4-bit Quantization Module
 
 Provides bitsandbytes-style 4-bit quantization for Mixture of Experts layers.
-This is necessary because transformers' Qwen3MoeExperts uses nn.Parameter
-tensors instead of nn.Linear modules, so bitsandbytes' standard quantization
-skips them.
-
-Design follows bitsandbytes patterns:
-- Uses Params4bit for per-expert weight quantization
-- Quantization happens on .to(cuda), just like Linear4bit
-- Integrates with model loading via replace_with_bnb_moe_experts()
+
+This is needed because transformers' MoE experts (Qwen3MoeExperts,
+Qwen3VLMoeTextExperts, GLM-4 / DeepSeek-V3 MoE, ...) store the stacked expert

no diff captured (skipped)

The change robustly achieves 4-bit quantization of MoE experts and LoRA compatibility by expanding stacked experts into per-expert nn.Linear that flow through the standard bnb+PEFT paths, preserving dimensional metadata for both layouts. Quant-config honoring is only partial since only compute_dtype is explicitly read in the expander (quant_type/double_quant now defer to the downstream Linear4bit conversion which does honor the config, so it's not hardcoded-wrong but relies on that path). Overall a strong, well-guarded solution.

diff --git a/unsloth_zoo/temporary_patches/bitsandbytes.py b/unsloth_zoo/temporary_patches/bitsandbytes.py
index b01f016..0dfd029 100644
--- a/unsloth_zoo/temporary_patches/bitsandbytes.py
+++ b/unsloth_zoo/temporary_patches/bitsandbytes.py
@@ -19,9 +19,10 @@ import torch.nn as nn
 import inspect
 import importlib
 from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable
-from .common import TEMPORARY_PATCHES, torch_compile
+from .common import TEMPORARY_PATCHES, torch_compile, UNSLOTH_ENABLE_LOGGING
 from .utils import (
     patch_function,
+    logger,
     process_output_options,
     process_return,
     KWARGS_TYPE,
@@ -94,3 +95,78 @@ def patch_bitsandbytes_linear4bit_forward():
         pass
 pass
 TEMPORARY_PATCHES.append(patch_bitsandbytes_linear4bit_forward)
+
+
+def patch_bitsandbytes_moe_experts():
+    """
+    Expand stacked MoE expert parameters (Qwen3-MoE, Qwen3-VL-MoE, ...) into
+    per-expert nn.Linear layers just before bitsandbytes quantizes the model.
+
+    transformers stores MoE experts as a single stacked nn.Parameter
+    (gate_up_proj/down_proj of shape (num_experts, ...)) rather than nn.Linear
+    modules, so the standard replace_with_bnb_linear skips them entirely - the
+    expert weights stay in full precision and PEFT has no Linear4bit to attach a
+    LoRA adapter to. We wrap replace_with_bnb_linear so that, whenever a 4-bit
+    quantization config is active, the stacked experts are first re-expressed as
+    per-expert nn.Linear (see moe_bnb.replace_with_bnb_moe_experts). The original
+    replace_with_bnb_linear then quantizes those linears to Linear4bit and PEFT's
+    standard lora.Linear4bit handles LoRA training on the experts.
+    """
+    try:
+        import transformers.integrations.bitsandbytes as tf_bnb
+    except Exception as e:
+        return raise_error("transformers.integrations.bitsandbytes", e)
+
+    if not hasattr(tf_bnb, "replace_with_bnb_linear"):
+        return
+
+    original = tf_bnb.replace_with_bnb_linear
+    if getattr(original, "__name__", "") == "_unsloth_replace_with_bnb_linear_moe":
+        return
+
+    # Discover the parameter position/name of quantization_config so we can read it
+    # regardless of how the caller passes it (positional or keyword), across
+    # transformers versions.
+    try:
+        _params = list(inspect.signature(original).parameters)
+        _qc_index = _params.index("quantization_config")
+    except (ValueError, TypeError):
+        _qc_index = None
+
+    def _find_quant_config(args, kwargs):
+        if "quantization_config" in kwargs:

no diff captured (skipped)

The change robustly achieves in-place 4-bit quantization of MoE expert parameters while preserving names/dims needed for PEFT LoRA, dequantizes on the forward paths, honors the quant config, and guards against regressions with idempotent, bnb-gated patches. All criteria appear fully satisfied from the diff.

diff --git a/unsloth_zoo/temporary_patches/__init__.py b/unsloth_zoo/temporary_patches/__init__.py
index b437499..0cf14a8 100644
--- a/unsloth_zoo/temporary_patches/__init__.py
+++ b/unsloth_zoo/temporary_patches/__init__.py
@@ -22,6 +22,7 @@ from .gemma3n import *
 from .gpt_oss import *
 from .qwen3_moe import *
 from .qwen3_vl_moe import *
+from .moe_bnb import *
 from .glm4_moe import *
 from .deepseek_v3_moe import *
 from .pixtral import *
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..04a9821 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -17,23 +17,34 @@
 MoE 4-bit Quantization Module
 
 Provides bitsandbytes-style 4-bit quantization for Mixture of Experts layers.
-This is necessary because transformers' Qwen3MoeExperts uses nn.Parameter
-tensors instead of nn.Linear modules, so bitsandbytes' standard quantization
-skips them.
-
-Design follows bitsandbytes patterns:
-- Uses Params4bit for per-expert weight quantization
-- Quantization happens on .to(cuda), just like Linear4bit
-- Integrates with model loading via replace_with_bnb_moe_experts()
+This is necessary because transformers' Qwen3MoeExperts (and friends) store the
+stacked expert weights as nn.Parameter tensors (gate_up_proj / down_proj) rather
+than nn.Linear modules, so bitsandbytes' standard Linear replacement skips them.
+
+Design
+------
+We quantize the expert weight *parameters in place*, keeping them as attributes
+named gate_up_proj / down_proj on the original experts module. Each becomes a
+bitsandbytes Params4bit holding packed 4-bit data plus a quant_state.
+
+Why in-place instead of a replacement module:
+- PEFT attaches LoRA to MoE experts via `target_parameters=["gate_up_proj",
+  "down_proj"]`. That only works if those attributes still exist as parameters.
+  A bespoke replacement module (per-expert Params4bit lists) destroys the names
+  PEFT needs, so LoRA silently never attaches to the experts. Keeping the names
+  lets PEFT's ParamWrapper wrap them exactly as it does for a dense model.
+- The wired MoE forward (moe_utils.forward_native_grouped_mm / loop / triton)
+  already fetches base weights through `_get_base_weight`, which now dequantizes
+  Params4bit on the fly. So the same forward runs for dense and 4-bit weights,
+  and separated LoRA is added on top — mirroring bnb Linear4bit + PEFT.
 """
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, List, Any, Tuple
+from typing import Optional, Tuple
 import os
-import warnings
 from ..log import logger
+from .common import TEMPORARY_PATCHES
 

The change adds a coherent mechanism for LoRA on 4-bit MoE experts: meta-device parameter stubs of correct shape, per-expert LoRA delta injection in forward, dequantization helpers, and routing of bnb 4-bit experts to loop-based matmul_4bit. However the core quantization-to-Params4bit logic and BitsAndBytesConfig handling are referenced but not demonstrated in the diff, and patch-application guarding/idempotency is not shown, so several criteria earn only partial credit.

diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..86a1e26 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -125,6 +125,21 @@ class MoeExperts4bit(nn.Module):
         self._gate_up_proj_pending = None
         self._down_proj_pending = None
 
+        # Expose stacked-weight parameter stubs so PEFT can attach LoRA wrappers.
+        # Shape follows the standard Qwen3-MoE layout: (E, 2*I, H) and (E, H, I).
+        # These are *meta-device* tensors — they carry no real data; all computation
+        # uses _bnb_gate_up_weights / _bnb_down_weights.  The LoRA delta is added
+        # in forward() via the _unsloth_lora_{gate_up_proj,down_proj} injection
+        # mechanism used by the patched PEFT ParamWrapper.
+        self.gate_up_proj = nn.Parameter(
+            torch.empty(num_experts, 2 * intermediate_dim, hidden_dim, device="meta"),
+            requires_grad=False,
+        )
+        self.down_proj = nn.Parameter(
+            torch.empty(num_experts, hidden_dim, intermediate_dim, device="meta"),
+            requires_grad=False,
+        )
+
         # Activation function
         self.act_fn = F.silu
 
@@ -163,6 +178,12 @@ class MoeExperts4bit(nn.Module):
                 # Store for later quantization on .to(cuda)
                 self.register_buffer("_gate_up_proj_pending", gate_up_proj)
                 self.register_buffer("_down_proj_pending", down_proj)
+
+            # After quantization the meta stubs (gate_up_proj / down_proj) are intentionally
+            # not in the state dict — remove them from missing_keys so strict loading passes.
+            for key in (gate_up_key, down_key):
+                if key in missing_keys:
+                    missing_keys.remove(key)
         else:
             # Try loading per-expert quantized format
             super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
@@ -244,6 +265,14 @@ class MoeExperts4bit(nn.Module):
         w = weight.t()
         return bnb.matmul_4bit(x, w, bias=None, quant_state=quant_state)
 
+    def _dequantize_expert_weight(self, param4bit: "Params4bit") -> torch.Tensor:
+        """Dequantize a Params4bit weight to full precision for LoRA merge/save."""
+        return dequantize_4bit(
+            param4bit.data,
+            param4bit.quant_state,
+            quant_type=self.quant_type,
+        ).to(self.compute_dtype or torch.bfloat16)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -253,6 +282,8 @@ class MoeExperts4bit(nn.Module):
         """
         Optimized forward pass using bnb.matmul_4bit for each expert.
         Only iterates over experts that actually have tokens routed to them.
+        Supports separated LoRA via _unsloth_lora_gate_up_proj / _unsloth_lora_down_proj
+        injection from the patched PEFT ParamWrapper (same protocol as forward_native_grouped_mm).

The change robustly achieves all outcomes: it quantizes stacked MoE expert params into per-expert Linear4bit so bitsandbytes quantizes them, PEFT can attach LoRA, dimensional metadata is derived correctly, forward dequantizes via Linear4bit, quant config is honored, and non-MoE/standard layers are unaffected with bnb-guarded, idempotent patching.

diff --git a/unsloth_zoo/patching_utils.py b/unsloth_zoo/patching_utils.py
index ba77af9..25df5da 100644
--- a/unsloth_zoo/patching_utils.py
+++ b/unsloth_zoo/patching_utils.py
@@ -273,6 +273,31 @@ def patch_to_dict():
         setattr(PretrainedConfig, "to_dict", wrapped_to_dict)
 pass
 
+
+def _is_model_loaded_in_4bit(model) -> bool:
+    """Detect whether a model was loaded in bitsandbytes 4-bit."""
+    # transformers sets these flags on the model when a BitsAndBytesConfig is used.
+    if getattr(model, "is_loaded_in_4bit", False):
+        return True
+    quant_config = getattr(getattr(model, "config", None), "quantization_config", None)
+    if quant_config is not None:
+        if getattr(quant_config, "load_in_4bit", False):
+            return True
+        # dict form (some serialized configs)
+        if isinstance(quant_config, dict) and quant_config.get("load_in_4bit", False):
+            return True
+    # Fall back to scanning for any already-present Linear4bit module.
+    try:
+        from bitsandbytes.nn import Linear4bit as _Bnb_Linear4bit
+        for module in model.modules():
+            if isinstance(module, _Bnb_Linear4bit):
+                return True
+    except Exception:
+        pass
+    return False
+pass
+
+
 def patch_model_and_tokenizer(
     model,
     tokenizer,
@@ -315,6 +340,23 @@ def patch_model_and_tokenizer(
         raise ImportError("Unsloth: Please install peft via `pip install peft`")
     pass
 
+    # MoE experts (Qwen3-MoE etc.) are stored as stacked 3D nn.Parameter tensors,
+    # so bitsandbytes' replace_with_bnb_linear skips them and they stay unquantized.
+    # When the rest of the model was loaded in 4-bit, quantize the experts too and
+    # expose them as per-expert Linear4bit so PEFT can attach LoRA to them.
+    _do_logging = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
+    if _is_model_loaded_in_4bit(model):
+        try:
+            from .temporary_patches.moe_bnb import replace_with_bnb_moe_experts
+            quantization_config = getattr(getattr(model, "config", None), "quantization_config", None)
+            model, moe_replaced = replace_with_bnb_moe_experts(model, quantization_config)
+            if moe_replaced and _do_logging:
+                print("Unsloth: Quantized MoE experts to BNB 4-bit for LoRA training.")
+        except Exception as e:
+            if _do_logging:
+                print(f"Unsloth: Could not quantize MoE experts to 4-bit: {e}")
+    pass
+
     # Get most likely the correct data-type of the model
     if correct_dtype is None:
         try:

The diff primarily adds LoRA stub parameters, per-expert LoRA delta application in forward, and MoE-4bit detection, on top of a pre-existing quantization framework. It plausibly enables LoRA on expert layers with correct shapes (c2/c3), but core quantization (c1/c4) predates the diff and config-honoring (c5) isn't demonstrated. End-to-end LoRA correctness depends on an unshown ParamWrapper patch, limiting confidence.

diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..8925f30 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -65,9 +65,14 @@ class MoeExperts4bit(nn.Module):
     - Stores per-expert weights as Params4bit
     - Quantization happens when moving to CUDA (like Linear4bit)
     - Self-contained forward using bnb.matmul_4bit for each expert
+    - Exposes gate_up_proj / down_proj as 3D nn.Parameter stubs so that PEFT
+      can apply LoRA adapters.  Actual computation always uses the quantized
+      Params4bit stored in _bnb_gate_up_weights / _bnb_down_weights.
 
-    The module replaces Qwen3MoeExperts or similar, providing the same interface
-    but with quantized weights.
+    LoRA support:
+    - unsloth_zoo patches PEFT's ParamWrapper.forward to inject LoRA contributions
+      as module attributes _unsloth_lora_gate_up_proj / _unsloth_lora_down_proj.
+    - MoeExperts4bit.forward reads those attributes and adds the per-expert delta.
 
     NOTE: BNB 4-bit MoE uses loop-based forward (each expert processed individually).
     For maximum throughput, consider using grouped_mm backend with native bf16 weights.
@@ -120,6 +125,22 @@ class MoeExperts4bit(nn.Module):
         self._bnb_gate_up_weights = nn.ParameterList()
         self._bnb_down_weights = nn.ParameterList()
 
+        # Stacked 3D parameter stubs exposed so that PEFT can attach LoRA adapters.
+        # Shape matches the original Qwen3MoeExperts layout:
+        #   gate_up_proj: (E, 2*intermediate_dim, hidden_dim)
+        #   down_proj:    (E, hidden_dim, intermediate_dim)
+        # They are created/updated by _register_lora_stubs() after quantization.
+        # dtype is always bfloat16 regardless of quant_storage so that PEFT's
+        # in_features/out_features inference works correctly on CPU too.
+        self.gate_up_proj = nn.Parameter(
+            torch.empty(num_experts, 2 * intermediate_dim, hidden_dim, dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+        self.down_proj = nn.Parameter(
+            torch.empty(num_experts, hidden_dim, intermediate_dim, dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+
         # These will hold the original Parameters for weight loading
         # They get converted to Params4bit when .to(cuda) is called
         self._gate_up_proj_pending = None
@@ -167,6 +188,30 @@ class MoeExperts4bit(nn.Module):
             # Try loading per-expert quantized format
             super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
 
+    def _register_lora_stubs(self, gate_up_proj: torch.Tensor, down_proj: torch.Tensor):
+        """
+        Update the gate_up_proj / down_proj nn.Parameter stubs so that PEFT
+        can attach LoRA adapters with the correct shapes.
+
+        The stubs are kept in bfloat16 on CPU to be lightweight; actual forward
+        computation always uses the quantized Params4bit in _bnb_gate_up_weights /
+        _bnb_down_weights.  The stubs carry only shape and dtype metadata.
+        """
+        E = gate_up_proj.shape[0]
+        # Preserve gradient state already set by PEFT (requires_grad=True means
+        # PEFT has already attached LoRA adapters — don't clobber them).

The candidate correctly integrates a MoE 4-bit quantization pathway: it triggers expert quantization on 4-bit load, applies per-expert LoRA deltas with correct shapes, detects and delegates 4-bit expert modules for dequantized forward, and forwards the quantization config — all guarded by bitsandbytes availability. Much of the heavy lifting (Params4bit storage, shape metadata) resides in the pre-existing moe_bnb.py not shown in the diff, so I cannot fully verify robustness of c3/c5 from the diff alone, but the wiring is coherent and achieves the outcomes.

diff --git a/unsloth_zoo/patching_utils.py b/unsloth_zoo/patching_utils.py
index ba77af9..995e17e 100644
--- a/unsloth_zoo/patching_utils.py
+++ b/unsloth_zoo/patching_utils.py
@@ -441,6 +441,27 @@ def patch_model_and_tokenizer(
         pass
     pass
 
+    # Quantize MoE expert layers (nn.Parameter-based stacked weights that
+    # bitsandbytes' _replace_with_bnb_linear silently skips).
+    try:
+        from .temporary_patches.moe_bnb import replace_with_bnb_moe_experts, HAS_BNB
+        if HAS_BNB:
+            quantization_config = getattr(model.config, "quantization_config", None)
+            load_in_4bit = (
+                getattr(quantization_config, "load_in_4bit", False)
+                if quantization_config is not None
+                else False
+            )
+            if load_in_4bit:
+                _, _moe_replaced = replace_with_bnb_moe_experts(model, quantization_config)
+                if _moe_replaced and os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
+                    import logging as _logging
+                    _logging.getLogger("unsloth").info(
+                        "Unsloth: MoE expert weights quantized to 4-bit via replace_with_bnb_moe_experts."
+                    )
+    except Exception:
+        pass
+
     if not fix_embeddings: return model, tokenizer
 
     # Check if torch.nn.Embedding seen
diff --git a/unsloth_zoo/temporary_patches/moe_bnb.py b/unsloth_zoo/temporary_patches/moe_bnb.py
index adfbe70..9bcb1b6 100644
--- a/unsloth_zoo/temporary_patches/moe_bnb.py
+++ b/unsloth_zoo/temporary_patches/moe_bnb.py
@@ -244,6 +244,32 @@ class MoeExperts4bit(nn.Module):
         w = weight.t()
         return bnb.matmul_4bit(x, w, bias=None, quant_state=quant_state)
 
+    def _apply_lora_delta(
+        self,
+        x: torch.Tensor,
+        lora_data: tuple,
+        expert_idx: int,
+    ) -> torch.Tensor:
+        """
+        Compute the LoRA correction for a single expert and add it to x.
+
+        lora_data is (first_weight, second_weight, scaling[, num_experts]) as
+        stored by _patched_param_wrapper_forward on _unsloth_lora_gate_up_proj /
+        _unsloth_lora_down_proj.
+
+        first_weight:  (E, in_dim, R)  — per-expert input projection to rank R
+        second_weight: (E, R, out_dim) — per-expert rank projection to output
+        """
+        first_weight = lora_data[0]   # (E, in_dim, R)
+        second_weight = lora_data[1]  # (E, R, out_dim)
+        scaling = lora_data[2]
+

The change plausibly achieves the core outcomes: MoE experts get quantized via a wrapped replacement, LoRA attaches through meta-device stubs with correct shapes, dequantization paths handle 4-bit, and existing behavior is guarded. Several mechanisms depend on pre-existing helpers (replace_with_bnb_moe_experts, ParamWrapper injection) not shown in the diff, and config honoring is only partially evidenced, so credit is reduced from perfect.

diff --git a/unsloth_zoo/patching_utils.py b/unsloth_zoo/patching_utils.py
index ba77af9..1371182 100644
--- a/unsloth_zoo/patching_utils.py
+++ b/unsloth_zoo/patching_utils.py
@@ -751,6 +751,54 @@ if hasattr(transformers.integrations.bitsandbytes, "_replace_with_bnb_linear") a
     transformers.integrations.bitsandbytes._replace_with_bnb_linear = _unsloth_replace_with_bnb_linear
 pass
 
+# After patching _replace_with_bnb_linear, also wrap it so that MoE expert
+# modules (which use nn.Parameter-stacked weights instead of nn.Linear) get
+# quantized to 4-bit as well.  Standard BNB quantization only touches
+# nn.Linear leaves; the MoE expert parameter tensors are silently skipped.
+try:
+    _bnb_linear_before_moe = transformers.integrations.bitsandbytes._replace_with_bnb_linear
+
+    def _unsloth_replace_with_bnb_linear_and_moe(
+        model,
+        modules_to_not_convert=None,
+        current_key_name=None,
+        quantization_config=None,
+        has_been_replaced=False,
+        **kwargs,
+    ):
+        """Wraps _replace_with_bnb_linear to also quantize MoE expert layers."""
+        model, replaced = _bnb_linear_before_moe(
+            model,
+            modules_to_not_convert=modules_to_not_convert,
+            current_key_name=current_key_name,
+            quantization_config=quantization_config,
+            has_been_replaced=has_been_replaced,
+            **kwargs,
+        )
+
+        # Only attempt MoE quantization when 4-bit mode is requested
+        is_4bit = (
+            quantization_config is not None
+            and getattr(quantization_config, "load_in_4bit", False)
+        )
+        if is_4bit:
+            try:
+                from .temporary_patches.moe_bnb import replace_with_bnb_moe_experts
+                model, moe_replaced = replace_with_bnb_moe_experts(
+                    model,
+                    quantization_config=quantization_config,
+                )
+                replaced = replaced or moe_replaced
+            except Exception:
+                pass
+
+        return model, replaced
+
+    _unsloth_replace_with_bnb_linear_and_moe.__name__ = "_replace_with_bnb_linear"
+    transformers.integrations.bitsandbytes._replace_with_bnb_linear = _unsloth_replace_with_bnb_linear_and_moe
+except Exception:
+    pass
+
 # Patch for transformers 5.x: should_convert_module uses re.match (prefix-anchored)
 # and endswith, but does not do substring component matching. This means entries like
 # "vision_tower" in llm_int8_skip_modules fail to match module names like
@@ -781,6 +829,30 @@ if not hasattr(transformers.integrations.bitsandbytes, "_replace_with_bnb_linear

The change robustly makes fused/expert merge skips and FP8 forward LoRA-drops observable via recorded diagnostics and warnings, and preserves valid merges. The weak spots are c1 (no clear FP8 scale-companion location/matching skip logic is shown—the fused merge changes are about rank/layout, not FP8 scales) and c2 (loud warnings rather than a raised error, though the fallback recorder mentions a follow-up RuntimeError elsewhere).

diff --git a/unsloth_zoo/saving_utils.py b/unsloth_zoo/saving_utils.py
index ceba896..52e8c48 100644
--- a/unsloth_zoo/saving_utils.py
+++ b/unsloth_zoo/saving_utils.py
@@ -963,6 +963,14 @@ def _record_moe_merge_fallback(role, expert_idx, reason, lora_stats, W_shape):
             "lora_B_shape":  b_shape,
             "per_expert_W":  W_shape,
         }
+    warnings.warn(
+        f"[Unsloth MoE merge fallback] role={role} expert={expert_idx} reason={reason}. "
+        f"per_expert_W={W_shape}. The base weight is being written through; "
+        "the merged checkpoint will be missing this delta. "
+        "This warning will be followed by a RuntimeError to prevent silent corruption.",
+        UserWarning,
+        stacklevel=3,
+    )
     logger.warning(
         f"[Unsloth MoE merge fallback] role={role} expert={expert_idx} reason={reason}. "
         f"per_expert_W={W_shape}. The base weight is being written through; "
@@ -1575,15 +1583,14 @@ def _merge_moe_experts_file(mm, header_metadata, length_of_header, file, convert
 
 
 def _merge_moe_fused_gate_up_expert(gate_up_W, lora_stats, output_dtype, is_transposed=None):
-    """
-    Merge LoRA for fused gate_up_proj 3D tensor.
-    Supports both formats:
-      - Transposed (GPT-OSS): (E, H, 2*I) with lora_A (E*R, H), lora_B (2*I, E*R)
-      - Standard (Gemma4):    (E, 2*I, H) with lora_A (E*R, H), lora_B (2*I, E*R)
-    is_transposed: if provided, overrides dimension-based heuristic (needed when dims are equal).
-    """
+    _MOE_MERGE_STATE["attempted"] += 1
     try:
         if lora_stats.lora_A is None or lora_stats.lora_B is None:
+            _record_moe_merge_fallback(
+                "fused_gate_up", -1,
+                f"LoRA matrices None (A={lora_stats.lora_A is not None}, B={lora_stats.lora_B is not None})",
+                lora_stats, tuple(gate_up_W.shape),
+            )
             return gate_up_W
 
         num_experts, dim1, dim2 = gate_up_W.shape
@@ -1591,10 +1598,20 @@ def _merge_moe_fused_gate_up_expert(gate_up_W, lora_stats, output_dtype, is_tran
         dim_B, total_rank_B = lora_stats.lora_B.shape
 
         if total_rank_B != total_rank:
+            _record_moe_merge_fallback(
+                "fused_gate_up", -1,
+                f"rank mismatch (A.total={total_rank}, B.total={total_rank_B})",
+                lora_stats, tuple(gate_up_W.shape),
+            )
             return gate_up_W
 
         rank = total_rank // num_experts
         if total_rank % num_experts != 0:
+            _record_moe_merge_fallback(
+                "fused_gate_up", -1,
+                f"total_rank {total_rank} not divisible by num_experts {num_experts}",
+                lora_stats, tuple(gate_up_W.shape),
+            )
             return gate_up_W

The change robustly makes the FP8 forward LoRA-drop and the MoE merge refusal paths loud, reporting diagnostics and refusing degenerate metadata while preserving valid merges. However, criterion c1's specific requirement — detecting when an FP8 base weight's per-block scale companion cannot be located/matched during merge and skipping to avoid silent casting/corruption — is not actually implemented; the diff only touches layout/rank logic, so c1 earns minimal credit.

diff --git a/tests/test_moe_fp8_lora_guard.py b/tests/test_moe_fp8_lora_guard.py
new file mode 100644
index 0000000..7683882
--- /dev/null
+++ b/tests/test_moe_fp8_lora_guard.py
@@ -0,0 +1,24 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import _forward_native_fp8_expert_loop
+
+
+def test_fp8_expert_loop_refuses_active_lora():
+    experts = nn.Module()
+    experts.num_experts = 2
+    experts.gate_up_proj = nn.Parameter(torch.empty(2, 8, 4))
+    experts.down_proj = nn.Parameter(torch.empty(2, 4, 4))
+    experts._unsloth_lora_gate_up_proj = (
+        torch.empty(2, 4, 1),
+        torch.empty(2, 1, 8),
+        1.0,
+    )
+
+    hidden_states = torch.empty(3, 4)
+    top_k_index = torch.zeros(3, 1, dtype=torch.long)
+    top_k_weights = torch.ones(3, 1)
+
+    with pytest.raises(RuntimeError, match="would ignore active LoRA"):
+        _forward_native_fp8_expert_loop(experts, hidden_states, top_k_index, top_k_weights)
diff --git a/tests/test_moe_merge_e2e_cpu.py b/tests/test_moe_merge_e2e_cpu.py
index cf0d30e..5c497fb 100644
--- a/tests/test_moe_merge_e2e_cpu.py
+++ b/tests/test_moe_merge_e2e_cpu.py
@@ -134,15 +134,15 @@ def test_per_layer_merge_round_trip(layout):
     _reset_moe_merge_state()
 
 
-def test_unrecognised_layout_records_fallback_and_first_error():
+def test_unrecognised_layout_raises_and_records_first_error():
     _reset_moe_merge_state()
     num_experts, rank_per, intermediate, hidden = 4, 4, 8, 12
     TR = num_experts * rank_per
     W = torch.randn(intermediate, hidden)
     A = torch.randn(TR, hidden + 7); B = torch.randn(hidden, TR)
     stats = LoraStats(module=_InnerMoE(num_experts), lora_A=A, lora_B=B, alpha=1.0)
-    out = _merge_moe_gate_expert(W.clone(), stats, 0, num_experts, torch.float32)
-    assert torch.equal(out.cpu(), W)
+    with pytest.raises(RuntimeError, match="refusing to save a silently incomplete MoE LoRA merge"):
+        _merge_moe_gate_expert(W.clone(), stats, 0, num_experts, torch.float32)
     assert _MOE_MERGE_STATE["fallback"] >= 1
     err = _MOE_MERGE_STATE["first_error"]
     assert err is not None and err["role"] == "gate"
diff --git a/tests/test_unsloth_zoo_lora_merge.py b/tests/test_unsloth_zoo_lora_merge.py
index 505c34c..78a2b5e 100644
--- a/tests/test_unsloth_zoo_lora_merge.py
+++ b/tests/test_unsloth_zoo_lora_merge.py
@@ -435,7 +435,7 @@ def test_detect_moe_lora_layout_classifies_both_conventions():
     assert _detect_moe_lora_layout(A_swap, B_swap, num_experts + 1, out_dim, in_dim)[0] == "unknown"
 

The change converts prior silent fallbacks into loud failures across FP8 merge scale resolution, fused/expert merges, layout extraction, and the FP8 forward loop, while preserving valid merges and adding LoRA application to the FP8 loop. It robustly covers required outcomes; c2 is nearly full since it applies the adapter and raises when extraction fails, though the exact fall-through-raise condition is handled via extraction rather than a distinct raise in every unreachable path.

diff --git a/tests/test_forward_native_moe_loop_lora.py b/tests/test_forward_native_moe_loop_lora.py
index 5807c57..d6a2955 100644
--- a/tests/test_forward_native_moe_loop_lora.py
+++ b/tests/test_forward_native_moe_loop_lora.py
@@ -21,6 +21,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from unsloth_zoo.temporary_patches.moe_utils import forward_native_moe_loop
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import _forward_native_fp8_expert_loop
 
 
 def _build_experts(num_experts, hidden, intermediate, transposed_storage):
@@ -272,3 +273,40 @@ def test_forward_native_moe_loop_lora_dtype_precast_no_loop_alloc(dtype):
     rtol = 1e-4 if dtype == torch.float32 else 1e-1
     torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
     assert out.dtype == dtype
+
+
+def test_forward_native_fp8_expert_loop_applies_stashed_lora():
+    """The FP8 per-expert loop is used as the last-resort FP8 forward path.
+    It must not silently drop LoRA data stashed by the patched ParamWrapper.
+    """
+    torch.manual_seed(42)
+    num_experts = 4
+    hidden = 16
+    intermediate = 12
+    rank = 3
+    num_tokens = 6
+    top_k = 2
+
+    experts = _build_experts(num_experts, hidden, intermediate, transposed_storage=False)
+    gate_up_lora, down_lora = _build_lora_for(experts, rank, scaling=1.25)
+    experts._unsloth_lora_gate_up_proj = gate_up_lora
+    experts._unsloth_lora_down_proj = down_lora
+
+    hidden_states = torch.randn(num_tokens, hidden, dtype=torch.float32)
+    top_k_index = torch.randint(0, num_experts, (num_tokens, top_k))
+    top_k_weights = torch.softmax(torch.randn(num_tokens, top_k), dim=-1)
+
+    out = _forward_native_fp8_expert_loop(
+        experts, hidden_states, top_k_index, top_k_weights
+    )
+    ref = _naive_forward(
+        experts, hidden_states, top_k_index, top_k_weights,
+        gate_up_lora, down_lora,
+    )
+    without_lora = _naive_forward(
+        experts, hidden_states, top_k_index, top_k_weights,
+        None, None,
+    )
+
+    torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+    assert not torch.allclose(out, without_lora, atol=1e-6, rtol=1e-6)
diff --git a/tests/test_moe_merge_e2e_cpu.py b/tests/test_moe_merge_e2e_cpu.py
index cf0d30e..e6b21cf 100644
--- a/tests/test_moe_merge_e2e_cpu.py
+++ b/tests/test_moe_merge_e2e_cpu.py
@@ -141,8 +141,8 @@ def test_unrecognised_layout_records_fallback_and_first_error():
     W = torch.randn(intermediate, hidden)
     A = torch.randn(TR, hidden + 7); B = torch.randn(hidden, TR)

The change consistently converts previously silent fallbacks (returning base weight, silently omitting LoRA in FP8 forward, guessing layouts) into loud errors with diagnostic messages, and importantly adds actual LoRA delta application in the native FP8 expert loop. All rubric outcomes are fully and robustly achieved, with tests updated to assert the new loud behavior and valid merges preserved via applied counters.

diff --git a/tests/test_moe_fp8_lora_fail_loud.py b/tests/test_moe_fp8_lora_fail_loud.py
new file mode 100644
index 0000000..b15499f
--- /dev/null
+++ b/tests/test_moe_fp8_lora_fail_loud.py
@@ -0,0 +1,99 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _call_with_temporary_moe_weights,
+    _resolve_active_moe_lora_for_fp8,
+)
+
+
+class _LoraWeight(nn.Module):
+    def __init__(self, weight):
+        super().__init__()
+        self.weight = nn.Parameter(weight)
+
+
+class _Wrapper:
+    def __init__(self, experts, parameter_name, weight_A, weight_B, scaling=1.0):
+        self.base_layer = experts
+        self.parameter_name = parameter_name
+        self.lora_A = {"default": _LoraWeight(weight_A)}
+        self.lora_B = {"default": _LoraWeight(weight_B)}
+        self.scaling = {"default": scaling}
+        self.disable_adapters = False
+        self.merged = False
+
+    def get_base_layer(self):
+        return self.base_layer
+
+
+def _build_experts(E=4, H=8, I=6, R=2, bad=False):
+    experts = nn.Module()
+    experts.num_experts = E
+    experts.gate_up_proj = nn.Parameter(torch.zeros(E, 2 * I, H))
+    experts.down_proj = nn.Parameter(torch.zeros(E, H, I))
+    A_dim = H + 3 if bad else H
+    gate_A = torch.randn(E * R, A_dim)
+    gate_B = torch.randn(2 * I, E * R)
+    experts.gate_up_proj_lora_wrapper = _Wrapper(
+        experts, "gate_up_proj", gate_A, gate_B, scaling=2.0
+    )
+    return experts, gate_A, gate_B
+
+
+def test_fp8_lora_resolver_extracts_active_wrapper_without_side_attr():
+    E, H, I, R = 4, 8, 6, 2
+    experts, gate_A, gate_B = _build_experts(E=E, H=H, I=I, R=R)
+
+    first, second, scaling = _resolve_active_moe_lora_for_fp8(
+        experts, "gate_up_proj", use_separated_lora=True
+    )
+
+    assert first.shape == (E, H, R)
+    assert second.shape == (E, R, 2 * I)

The change robustly converts silent MoE merge and FP8 forward LoRA-drop failure modes into loud errors while preserving recorded diagnostics and valid-merge behavior. The FP8-forward drop (c2), diagnostics (c3), valid merges (c4), and degenerate metadata refusal (c5) are all well handled. However, the specific FP8 per-block scale-companion location/matching failure (c1) is not explicitly shown as a dedicated skip path — it's only covered indirectly by the generalized merge-failure raising.

diff --git a/tests/test_forward_native_moe_loop_lora.py b/tests/test_forward_native_moe_loop_lora.py
index 5807c57..5c5b37d 100644
--- a/tests/test_forward_native_moe_loop_lora.py
+++ b/tests/test_forward_native_moe_loop_lora.py
@@ -272,3 +272,43 @@ def test_forward_native_moe_loop_lora_dtype_precast_no_loop_alloc(dtype):
     rtol = 1e-4 if dtype == torch.float32 else 1e-1
     torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
     assert out.dtype == dtype
+
+
+def test_fp8_backend_raises_before_dropping_pending_lora(monkeypatch):
+    import unsloth_zoo.temporary_patches.moe_utils_fp8 as fp8
+
+    experts = nn.Module()
+    experts.num_experts = 2
+    experts.gate_up_proj = nn.Parameter(
+        torch.empty(2, 8, 4, dtype=torch.float8_e4m3fn),
+        requires_grad=False,
+    )
+    experts.down_proj = nn.Parameter(
+        torch.empty(2, 4, 4, dtype=torch.float8_e4m3fn),
+        requires_grad=False,
+    )
+    experts._unsloth_lora_gate_up_proj = (
+        torch.randn(2, 4, 2),
+        torch.randn(2, 2, 8),
+        1.0,
+    )
+
+    monkeypatch.setattr(fp8, "_check_torch_scaled_grouped_mm_supported", lambda: False)
+    monkeypatch.setattr(fp8, "_dequantize_full_expert_weights", lambda *args, **kwargs: None)
+    monkeypatch.setattr(fp8, "_log_moe_fp8_backend_once", lambda *args, **kwargs: None)
+    monkeypatch.setattr(
+        "unsloth_zoo.temporary_patches.moe_utils.select_moe_backend",
+        lambda: "native_torch",
+    )
+
+    hidden_states = torch.randn(3, 4)
+    top_k_index = torch.tensor([[0], [1], [0]])
+    top_k_weights = torch.ones(3, 1)
+
+    with pytest.raises(RuntimeError, match="silently drop the adapter"):
+        fp8.forward_moe_backend_fp8(
+            experts,
+            hidden_states,
+            top_k_index,
+            top_k_weights,
+        )
diff --git a/tests/test_moe_merge_e2e_cpu.py b/tests/test_moe_merge_e2e_cpu.py
index cf0d30e..7bf6f39 100644
--- a/tests/test_moe_merge_e2e_cpu.py
+++ b/tests/test_moe_merge_e2e_cpu.py
@@ -141,8 +141,8 @@ def test_unrecognised_layout_records_fallback_and_first_error():
     W = torch.randn(intermediate, hidden)
     A = torch.randn(TR, hidden + 7); B = torch.randn(hidden, TR)
     stats = LoraStats(module=_InnerMoE(num_experts), lora_A=A, lora_B=B, alpha=1.0)
-    out = _merge_moe_gate_expert(W.clone(), stats, 0, num_experts, torch.float32)
-    assert torch.equal(out.cpu(), W)
+    with pytest.raises(RuntimeError, match="MoE LoRA merge failed"):
+        _merge_moe_gate_expert(W.clone(), stats, 0, num_experts, torch.float32)

The change makes both the FP8 merge scale-loss and the FP8 forward LoRA-drop loud via RuntimeErrors, records diagnostics for skipped fused merges, and refuses degenerate metadata. Valid merges remain intact and gain applied counters. All rubric outcomes appear fully addressed, assuming _record_moe_merge_fallback and _MOE_MERGE_STATE surface to the outer save as the comments state.

diff --git a/tests/test_moe_lora_loud_failures.py b/tests/test_moe_lora_loud_failures.py
new file mode 100644
index 0000000..de88ca7
--- /dev/null
+++ b/tests/test_moe_lora_loud_failures.py
@@ -0,0 +1,44 @@
+import torch, types
+from unsloth_zoo import saving_utils as su
+from unsloth_zoo.temporary_patches import moe_utils_fp8 as mf
+from unsloth_zoo.temporary_patches import moe_utils as mu
+import pytest
+
+def test_fused_merge_records_fallback():
+    su._reset_moe_merge_state()
+    stats = types.SimpleNamespace(lora_A=torch.randn(8, 16), lora_B=torch.randn(5, 8), alpha=1.0, module=None)
+    W = torch.randn(4, 32, 64)
+    out = su._merge_moe_fused_gate_up_expert(W, stats, torch.float32)
+    assert su._MOE_MERGE_STATE["fallback"] == 1
+    assert torch.equal(out, W)
+    su._reset_moe_merge_state()
+
+def test_fp8_load_missing_scale_raises():
+    class FakeFile:
+        def get_tensor(self, k):
+            return torch.zeros(4, 4, dtype=torch.float8_e4m3fn)
+    with pytest.raises(RuntimeError, match="weight_scale_inv"):
+        su._fp8_load_for_merge(FakeFile(), {}, "x.weight")
+
+def test_checked_write_raises():
+    with pytest.raises(RuntimeError, match="failed to write merged tensor"):
+        su._write_merged_tensor_checked(None, {}, 0, "missing.weight", torch.zeros(2, 2), torch.float32)
+
+def test_fp8_loop_refuses_active_lora():
+    mod = types.SimpleNamespace(_unsloth_lora_gate_up_proj=(1, 2, 3), num_experts=2)
+    with pytest.raises(RuntimeError, match="cannot apply the active LoRA"):
+        mf._forward_native_fp8_expert_loop(mod, torch.zeros(1, 4), torch.zeros(1, 1, dtype=torch.long), torch.zeros(1, 1))
+
+def test_extractor_raises_on_broken_wrapper():
+    class BadWrapper:
+        lora_A = {"default": types.SimpleNamespace()}
+        lora_B = {"default": types.SimpleNamespace()}
+        scaling = {"default": 1.0}
+        disable_adapters = False
+        merged = False
+    with pytest.raises(RuntimeError, match="failed to extract the active MoE LoRA"):
+        mu._extract_lora_from_wrapper(BadWrapper())
+
+def test_extractor_none_when_inactive():
+    class NoLora: pass
+    assert mu._extract_lora_from_wrapper(NoLora()) is None
diff --git a/unsloth_zoo/saving_utils.py b/unsloth_zoo/saving_utils.py
index ceba896..5a4c727 100644
--- a/unsloth_zoo/saving_utils.py
+++ b/unsloth_zoo/saving_utils.py
@@ -697,7 +697,7 @@ def _merge_and_overwrite_lora(
                                 W, lora_stats, expert_idx, num_experts, output_dtype or W.dtype
                             )
 
-                        _write_tensor_direct_torch(mm, header_metadata, length_of_header, key, merged_W, W.dtype)
+                        _write_merged_tensor_checked(mm, header_metadata, length_of_header, key, merged_W, W.dtype)

no diff captured (skipped)

The change robustly fixes the FP8 forward silent LoRA drop (c2) and makes fused/expert merge fallbacks loud and observable via recorded diagnostics (c3, c5), while preserving valid merges (c4). However, criterion c1 — the FP8 per-block scale companion location/matching failure during merge — is not addressed anywhere in the diff, so it earns no points.

diff --git a/tests/test_fp8_expert_loop_lora.py b/tests/test_fp8_expert_loop_lora.py
new file mode 100644
index 0000000..a7c85cc
--- /dev/null
+++ b/tests/test_fp8_expert_loop_lora.py
@@ -0,0 +1,177 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Regression tests for the last-resort FP8 per-expert MoE forward loop.
+
+`_forward_native_fp8_expert_loop` is the fallback FP8 path taken when neither
+`_scaled_grouped_mm` nor dequantize-plus-grouped_mm is available. It previously
+never read the injected `_unsloth_lora_*` factors, so an active LoRA adapter was
+silently dropped and the layer produced pure base-model outputs — which surfaces
+as "the model is barely learning" (#5410).
+
+These tests build a plain bf16 experts module (so the loop takes the ordinary
+`F.linear` branch, no FP8 kernels required) and verify:
+
+1. The injected separated-LoRA delta is actually applied (matches a naive ref).
+2. A LoRA whose factors cannot be indexed per-expert raises loudly rather than
+   silently dropping the delta.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import _forward_native_fp8_expert_loop
+
+
+def _build_bf16_experts(num_experts, hidden, intermediate):
+    """Experts module in the standard (E, out, in) F.linear layout, bf16 (not
+    FP8) so the loop uses plain F.linear and no kernels are needed."""
+    experts = nn.Module()
+    experts.num_experts = num_experts
+    experts.gate_up_proj = nn.Parameter(
+        torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.bfloat16)
+    )
+    experts.down_proj = nn.Parameter(
+        torch.randn(num_experts, hidden, intermediate, dtype=torch.bfloat16)
+    )
+    experts.act_fn = F.silu
+    return experts
+

no diff captured (skipped)

The change robustly makes the FP8 forward LoRA-drop loud (warn/raise) and now actually applies LoRA in the native loop, and converts silent fused-merge fallbacks into recorded, observable failures with reasons for degenerate metadata. However, c1's specific concern — the FP8 per-block scale companion being missing/unmatched during merge — is not explicitly handled; the merge changes address generic fused layout/rank issues rather than FP8 scale location, so c1 gets only partial credit.

diff --git a/tests/test_fp8_forward_lora_drop.py b/tests/test_fp8_forward_lora_drop.py
new file mode 100644
index 0000000..c7f3483
--- /dev/null
+++ b/tests/test_fp8_forward_lora_drop.py
@@ -0,0 +1,210 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""The FP8 per-expert fallback loop must NOT silently drop an attached MoE LoRA
+adapter (which looks like a model that barely learns). It should apply the
+separated LoRA delta when weights are usable, and warn loudly otherwise.
+
+These run on CPU with plain bf16/fp32 expert weights (fp8 tensors need a
+Hopper/Blackwell GPU); the LoRA-application code path is dtype-agnostic, so the
+CPU exercise still covers the fix.
+"""
+
+import os
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _forward_native_fp8_expert_loop,
+    _moe_has_active_lora_fp8,
+    _moe_lora_data_for_fp8,
+    _warn_fp8_lora_dropped_once,
+)
+
+
+def _build_fp8_experts(num_experts, hidden, intermediate):
+    """Minimal experts module matching the surface _forward_native_fp8_expert_loop reads."""
+    experts = nn.Module()
+    experts.num_experts = num_experts
+    # F.linear layout: (E, out, in)
+    experts.gate_up_proj = nn.Parameter(
+        torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.float32)
+    )
+    experts.down_proj = nn.Parameter(
+        torch.randn(num_experts, hidden, intermediate, dtype=torch.float32)
+    )
+    experts.act_fn = F.silu
+    return experts
+

no diff captured (skipped)

The change robustly addresses the FP8 forward LoRA-drop (c2) and the fused/degenerate-metadata loudness (c3, c5) plus keeps valid merges working (c4). However c1 (FP8 per-block scale companion loss during merge) is not clearly addressed — no scale-lookup/match code is shown skipping; the fused-merge accounting only tangentially relates. Overall strong but the specific FP8 scale-loss merge criterion is largely unmet.

diff --git a/unsloth_zoo/saving_utils.py b/unsloth_zoo/saving_utils.py
index ceba896..2d20e4f 100644
--- a/unsloth_zoo/saving_utils.py
+++ b/unsloth_zoo/saving_utils.py
@@ -1582,20 +1582,34 @@ def _merge_moe_fused_gate_up_expert(gate_up_W, lora_stats, output_dtype, is_tran
       - Standard (Gemma4):    (E, 2*I, H) with lora_A (E*R, H), lora_B (2*I, E*R)
     is_transposed: if provided, overrides dimension-based heuristic (needed when dims are equal).
     """
+    if lora_stats is None or lora_stats.lora_A is None or lora_stats.lora_B is None:
+        return gate_up_W
+    # #5410: fused-tensor merges must feed the same _MOE_MERGE_STATE accounting
+    # as the per-expert path so the outer guard can refuse to save a partial
+    # merge. A silent `return gate_up_W` here writes the base weight through
+    # with the LoRA delta missing and no warning.
+    _MOE_MERGE_STATE["attempted"] += 1
     try:
-        if lora_stats.lora_A is None or lora_stats.lora_B is None:
-            return gate_up_W
-
         num_experts, dim1, dim2 = gate_up_W.shape
         total_rank, dim_A = lora_stats.lora_A.shape
         dim_B, total_rank_B = lora_stats.lora_B.shape
 
         if total_rank_B != total_rank:
+            _record_moe_merge_fallback(
+                "gate_up", None,
+                f"lora_A/lora_B rank mismatch (A rank={total_rank}, B rank={total_rank_B})",
+                lora_stats, tuple(gate_up_W.shape),
+            )
             return gate_up_W
 
-        rank = total_rank // num_experts
         if total_rank % num_experts != 0:
+            _record_moe_merge_fallback(
+                "gate_up", None,
+                f"total_rank {total_rank} not divisible by num_experts {num_experts}",
+                lora_stats, tuple(gate_up_W.shape),
+            )
             return gate_up_W
+        rank = total_rank // num_experts
 
         if is_transposed is not None:
             use_transpose = is_transposed
@@ -1604,6 +1618,12 @@ def _merge_moe_fused_gate_up_expert(gate_up_W, lora_stats, output_dtype, is_tran
         elif dim_A == dim2 and dim_B == dim1:
             use_transpose = False
         else:
+            _record_moe_merge_fallback(
+                "gate_up", None,
+                f"layout not detected (A={tuple(lora_stats.lora_A.shape)}, "
+                f"B={tuple(lora_stats.lora_B.shape)}, fused_W={tuple(gate_up_W.shape)})",
+                lora_stats, tuple(gate_up_W.shape),
+            )
             return gate_up_W
 
         device = _active_merge_device()
@@ -1625,8 +1645,10 @@ def _merge_moe_fused_gate_up_expert(gate_up_W, lora_stats, output_dtype, is_tran
                 delta.T if use_transpose else delta, alpha=lora_stats.alpha
             )
 

The change makes previously silent FP8 merge corruption and LoRA-drop conditions loud via exceptions, recorded fallbacks, and warnings, and preserves valid merges. c1's exception (vs per-weight skip) and c2's warning (vs error) diverge slightly from the literal wording but achieve the loudness intent; both earn partial-to-good credit.

diff --git a/unsloth_zoo/saving_utils.py b/unsloth_zoo/saving_utils.py
index ceba896..e3233af 100644
--- a/unsloth_zoo/saving_utils.py
+++ b/unsloth_zoo/saving_utils.py
@@ -1067,14 +1067,35 @@ def _fp8_load_for_merge(file, header_metadata, weight_key):
         return W, None
     scale_key = weight_key[: -len(".weight")] + ".weight_scale_inv"
     if scale_key not in header_metadata:
-        return W, None
+        # FP8 weight has no companion scale — the LoRA delta will be added to
+        # raw fp8-encoded integers cast to float32, not to the real weight
+        # values. The merged checkpoint will be numerically wrong.
+        raise RuntimeError(
+            f"Unsloth: FP8 weight {weight_key!r} has no companion "
+            f"{scale_key!r} in this shard. Cannot safely merge the LoRA "
+            "delta: adding to raw fp8-integer values (not the real weights) "
+            "produces a corrupted checkpoint. Ensure the base checkpoint "
+            "includes weight_scale_inv tensors alongside every FP8 weight."
+        )
     scale_inv = file.get_tensor(scale_key)
     if scale_inv.ndim != 2 or W.ndim != 2:
-        return W, None
+        raise RuntimeError(
+            f"Unsloth: FP8 weight {weight_key!r} or its scale {scale_key!r} "
+            f"has unexpected rank (W.ndim={W.ndim}, scale_inv.ndim={scale_inv.ndim}; "
+            "expected both 2). Cannot safely dequantize for LoRA merge — "
+            "the merged checkpoint would be numerically wrong."
+        )
     rows, cols = W.shape
     srows, scols = scale_inv.shape
     if rows % srows != 0 or cols % scols != 0:
-        return W, None
+        raise RuntimeError(
+            f"Unsloth: FP8 weight {weight_key!r} shape ({rows}, {cols}) is not "
+            f"evenly divisible by scale shape ({srows}, {scols}). Cannot compute "
+            "block size for dequantization before LoRA merge — the merged "
+            "checkpoint would be numerically wrong. "
+            "Please file a bug at https://github.com/unslothai/unsloth-zoo/issues "
+            f"with these shapes: W={tuple(W.shape)}, scale_inv={tuple(scale_inv.shape)}."
+        )
     bm, bn = rows // srows, cols // scols
     W_bf16 = _fp8_dequant_blockwise(W, scale_inv)
     return W_bf16, {
@@ -1424,6 +1445,22 @@ def _merge_moe_experts_file(mm, header_metadata, length_of_header, file, convert
                             debug_logged += 1
                         except Exception:
                             pass
+                elif (
+                    lora_stats is not None
+                    and getattr(lora_stats, "lora_A", None) is not None
+                ):
+                    # LoRA exists for this key but the tensor is not in this shard.
+                    # This is expected for multi-shard models; only warn when the
+                    # tensor key is entirely absent from the model (not just this shard).
+                    # We cannot tell here, so record a debug note.
+                    if UNSLOTH_ENABLE_LOGGING:
+                        try:
+                            logger.info(
+                                f"[merge_debug] {gate_up_key} not in this shard — "
+                                "LoRA merge deferred to the shard that contains it."

The change robustly addresses the FP8 forward LoRA-drop (c2), fused-merge fallback recording (c3), valid merges (c4), and degenerate fused metadata (c5). However c1 — the FP8 per-block scale companion location/matching in the merge path — is not addressed at all; the merge changes target the fused gate_up/down expert functions, not FP8 scale corruption.

diff --git a/tests/test_forward_native_moe_loop_lora.py b/tests/test_forward_native_moe_loop_lora.py
index 5807c57..28787b4 100644
--- a/tests/test_forward_native_moe_loop_lora.py
+++ b/tests/test_forward_native_moe_loop_lora.py
@@ -272,3 +272,49 @@ def test_forward_native_moe_loop_lora_dtype_precast_no_loop_alloc(dtype):
     rtol = 1e-4 if dtype == torch.float32 else 1e-1
     torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
     assert out.dtype == dtype
+
+
+# ---------------------------------------------------------------------------
+# FP8 per-expert loop LoRA delta: the last-resort fp8_linear expert loop used
+# to ignore _unsloth_lora_* entirely, silently dropping the adapter so the
+# model looked like it was barely learning. _expert_lora_delta both applies it
+# and raises loudly when a present adapter cannot be applied.
+# ---------------------------------------------------------------------------
+
+def test_expert_lora_delta_matches_manual_per_expert():
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import _expert_lora_delta
+    torch.manual_seed(3)
+    num_experts, in_dim, rank, out_dim = 4, 8, 2, 10
+    scaling = 1.5
+    first = torch.randn(num_experts, in_dim, rank)
+    second = torch.randn(num_experts, rank, out_dim)
+    lora = (first, second, scaling)
+    x = torch.randn(5, in_dim)
+    for e in range(num_experts):
+        delta = _expert_lora_delta(lora, x, e, x.dtype)
+        ref = ((x @ first[e]) @ second[e]) * scaling
+        torch.testing.assert_close(delta, ref, atol=1e-5, rtol=1e-5)
+
+
+def test_expert_lora_delta_none_when_no_adapter():
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import _expert_lora_delta
+    assert _expert_lora_delta(None, torch.randn(3, 8), 0, torch.float32) is None
+
+
+def test_expert_lora_delta_raises_on_expert_out_of_range():
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import _expert_lora_delta
+    num_experts, in_dim, rank, out_dim = 2, 8, 2, 10
+    lora = (torch.randn(num_experts, in_dim, rank),
+            torch.randn(num_experts, rank, out_dim), 1.0)
+    with pytest.raises(RuntimeError, match="barely learn"):
+        _expert_lora_delta(lora, torch.randn(3, in_dim), num_experts, torch.float32)
+
+
+def test_expert_lora_delta_raises_on_input_dim_mismatch():
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import _expert_lora_delta
+    num_experts, in_dim, rank, out_dim = 2, 8, 2, 10
+    lora = (torch.randn(num_experts, in_dim, rank),
+            torch.randn(num_experts, rank, out_dim), 1.0)
+    # current_state has the wrong trailing dim -> unusable adapter -> loud raise.
+    with pytest.raises(RuntimeError, match="does not match LoRA first_weight"):
+        _expert_lora_delta(lora, torch.randn(3, in_dim + 1), 0, torch.float32)
diff --git a/tests/test_unsloth_zoo_lora_merge.py b/tests/test_unsloth_zoo_lora_merge.py
index 505c34c..56a088c 100644
--- a/tests/test_unsloth_zoo_lora_merge.py
+++ b/tests/test_unsloth_zoo_lora_merge.py
@@ -454,6 +454,69 @@ def test_moe_merge_fallback_counter_records_bad_layout():
     _reset_moe_merge_state()

The change robustly makes fused/expert merge fallbacks and FP8 forward LoRA-drops loud (c2-c5 well covered). However c1's specific requirement — detecting when an FP8 base weight's per-block scale companion cannot be located/matched and skipping the merge to avoid silent casting corruption — is not implemented; the added merge fallbacks are generic layout/rank checks, not FP8 scale-companion handling.

diff --git a/tests/test_lora_moe_loud_failures.py b/tests/test_lora_moe_loud_failures.py
new file mode 100644
index 0000000..cb2c113
--- /dev/null
+++ b/tests/test_lora_moe_loud_failures.py
@@ -0,0 +1,276 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""
+Tests that previously-silent failure modes in LoRA+MoE merge and FP8 forward
+are now loud (RuntimeWarning or recorded in _MOE_MERGE_STATE).
+
+Coverage:
+1. _merge_moe_fused_gate_up_expert: unknown layout → _MOE_MERGE_STATE["fallback"] > 0
+2. _merge_moe_fused_gate_up_expert: rank_mismatch → _MOE_MERGE_STATE["fallback"] > 0
+3. _merge_moe_fused_down_proj_expert: unknown layout → _MOE_MERGE_STATE["fallback"] > 0
+4. _merge_moe_fused_down_proj_expert: rank_mismatch → _MOE_MERGE_STATE["fallback"] > 0
+5. extract_moe_lora_weights_for_grouped_mm: unmatched dims → RuntimeWarning with
+   "LoRA delta … is likely WRONG"
+6. _extract_lora_from_wrapper: exception inside extractor fn → RuntimeWarning with
+   "LoRA adapter will be DROPPED"
+7. _merge_moe_fused_* happy path: _MOE_MERGE_STATE["applied"] increments and
+   result matches reference.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import pytest
+import torch
+
+from unsloth_zoo.saving_utils import (
+    LoraStats,
+    _MOE_MERGE_STATE,
+    _merge_moe_fused_down_proj_expert,
+    _merge_moe_fused_gate_up_expert,
+    _reset_moe_merge_state,
+)
+
+
+SEED = 9999
+
+
+def _ls(lora_A, lora_B, alpha=1.0):
+    return LoraStats(module=None, lora_A=lora_A, lora_B=lora_B, alpha=alpha)

The change robustly converts silent fallbacks in the fused MoE merge paths (rank/divisibility/layout) into hard errors (c5 solid, c3 solid, c4 solid). However c1 targets the FP8 per-block scale-companion loss specifically, which is not visibly addressed—only general fused-expert corruption. c2 is addressed via a warning rather than a raised error and remains partial. Strong on diagnostics and metadata refusal, weaker on the two FP8-specific outcomes.

diff --git a/tests/test_moe_lora_loud_failures.py b/tests/test_moe_lora_loud_failures.py
new file mode 100644
index 0000000..5361b6e
--- /dev/null
+++ b/tests/test_moe_lora_loud_failures.py
@@ -0,0 +1,418 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Tests that silent failure modes in MoE LoRA merge and FP8 forward are now loud.
+
+Covers:
+1. _merge_moe_fused_gate_up_expert / _merge_moe_fused_down_proj_expert:
+   shape mismatches that previously returned the base weight un-merged now raise.
+2. check_moe_merge_completeness() raises after _record_moe_merge_fallback.
+3. extract_moe_lora_weights_for_grouped_mm: layout fallback emits a warning.
+4. _extract_lora_from_wrapper: unexpected extraction errors raise RuntimeError.
+5. _check_fp8_lora_not_silently_dropped: warns when LoRA wrapper exists but
+   _unsloth_lora_* attrs are absent.
+"""
+
+from __future__ import annotations
+
+import os
+import warnings
+
+import pytest
+import torch
+import torch.nn as nn
+
+from unsloth_zoo.saving_utils import (
+    LoraStats,
+    _MOE_MERGE_STATE,
+    _merge_moe_fused_down_proj_expert,
+    _merge_moe_fused_gate_up_expert,
+    _record_moe_merge_fallback,
+    _reset_moe_merge_state,
+    check_moe_merge_completeness,
+)
+
+
+# ---------------------------------------------------------------------------
+# 1. Fused gate_up merge — shape mismatches now raise instead of silently
+#    returning the base weight unchanged.
+# ---------------------------------------------------------------------------
+
+def _make_fused_gate_up_lora(num_experts, hidden, two_inter, rank):

The change robustly converts fused-expert merge and FP8-forward LoRA-drop failures into loud errors/warnings with good diagnostics, and preserves valid merges. However c1's specific outcome — the FP8 per-block scale companion cannot be located/matched — is not addressed; only generic rank/layout checks are added, so that criterion is only partially met. c2 uses a warning rather than a raised error but the rubric permits equivalent loud failure.

diff --git a/tests/test_moe_lora_loud_failures.py b/tests/test_moe_lora_loud_failures.py
new file mode 100644
index 0000000..a0fb0ba
--- /dev/null
+++ b/tests/test_moe_lora_loud_failures.py
@@ -0,0 +1,458 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Tests that MoE LoRA failure modes are loud rather than silent.
+
+Covers:
+1. _raise_if_moe_merge_failed raises when fallback > 0 and is a no-op otherwise.
+2. _merge_moe_fused_gate_up_expert raises on rank mismatch, non-divisible experts,
+   and unresolvable layout instead of silently returning the unmerged weight.
+3. _merge_moe_fused_down_proj_expert raises on the same families of bad input.
+4. extract_moe_lora_weights_for_grouped_mm raises (instead of silently using
+   canonical layout) when both layout matches fail and dims are known.
+5. _extract_lora_from_wrapper re-raises unexpected errors instead of swallowing
+   them with a bare except-and-return-None.
+6. _warn_if_lora_dropped_in_fp8_forward emits a warning when LoRA adapters exist
+   but the separated-LoRA inject attributes are missing on an FP8 experts module.
+"""
+
+from __future__ import annotations
+
+import logging
+import pytest
+import torch
+import torch.nn as nn
+
+from unsloth_zoo.saving_utils import (
+    LoraStats,
+    _MOE_MERGE_STATE,
+    _merge_moe_fused_down_proj_expert,
+    _merge_moe_fused_gate_up_expert,
+    _raise_if_moe_merge_failed,
+    _reset_moe_merge_state,
+)
+from unsloth_zoo.temporary_patches.moe_utils import (
+    extract_moe_lora_weights_for_grouped_mm,
+)
+from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+    _warn_if_lora_dropped_in_fp8_forward,
+)
+
+

The change consistently converts unconditional patch registration into import-guarded conditional registration for both bnb4bit and FP8 MoE patches, achieving the core goal of not breaking on older Transformers. Gates are based on import availability, which cleanly handles missing modules but is less robust against present-module/missing-attribute cases and doesn't inspect the specific quantizer hooks the criteria emphasize.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..bc025bd 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -307,8 +307,12 @@ def patch_bnb4bit_quantize_convert():
 
     if UNSLOTH_ENABLE_LOGGING:
         logger.info("Unsloth: Patched Bnb4bitQuantize.convert for MoE expert parameter support")
-pass
-TEMPORARY_PATCHES.append(patch_bnb4bit_quantize_convert)
+
+try:
+    from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+    TEMPORARY_PATCHES.append(patch_bnb4bit_quantize_convert)
+except ImportError:
+    pass
 
 
 def patch_bnb4bit_quantizer_param_needs_quantization():
@@ -346,8 +350,12 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
 
     if UNSLOTH_ENABLE_LOGGING:
         logger.info("Unsloth: Patched Bnb4BitHfQuantizer.param_needs_quantization for MoE expert parameters")
-pass
-TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_param_needs_quantization)
+
+try:
+    from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+    TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_param_needs_quantization)
+except ImportError:
+    pass
 
 
 def patch_bnb4bit_quantizer_process_model():
@@ -374,8 +382,12 @@ def patch_bnb4bit_quantizer_process_model():
 
     patched_process_model_before_weight_loading._unsloth_moe_patched = True
     patch_function(Bnb4BitHfQuantizer, "_process_model_before_weight_loading", patched_process_model_before_weight_loading, match_level = "relaxed")
-pass
-TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_process_model)
+
+try:
+    from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+    TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_process_model)
+except ImportError:
+    pass
 
 
 def patch_transformers_weight_converter_kwargs():
@@ -411,5 +423,9 @@ def patch_transformers_weight_converter_kwargs():
 
     if UNSLOTH_ENABLE_LOGGING:
         logger.info("Unsloth: Patched transformers WeightConverter.__init__ to ignore unknown kwargs (peft 0.19 forward-compat)")
-pass
-TEMPORARY_PATCHES.append(patch_transformers_weight_converter_kwargs)
+
+try:
+    from transformers.core_model_loading import WeightConverter
+    TEMPORARY_PATCHES.append(patch_transformers_weight_converter_kwargs)
+except ImportError:

The change robustly prevents crashes on older Transformers by guarding each MoE bnb4bit and FP8 patch with getattr-based hook detection and broadened exception handling, and preserves behavior when APIs exist. It differs from a single shared capability-gate design—detection is distributed per patch—but functionally achieves all required outcomes.

diff --git a/tests/test_temporary_patches_imports.py b/tests/test_temporary_patches_imports.py
index 6b86258..7080774 100644
--- a/tests/test_temporary_patches_imports.py
+++ b/tests/test_temporary_patches_imports.py
@@ -36,6 +36,8 @@ forward pass.
 from __future__ import annotations
 
 import importlib
+import sys
+import types
 
 import pytest
 
@@ -108,6 +110,65 @@ def test_torch_compile_options_is_dict():
     )
 
 
+def _stub_module(monkeypatch, name: str, **attrs):
+    module = types.ModuleType(name)
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    monkeypatch.setitem(sys.modules, name, module)
+    return module
+
+
+def test_moe_bnb4bit_quantizer_patches_skip_missing_old_transformers_hooks(monkeypatch):
+    """Older transformers can expose the bnb quantizer class without the newer
+    hooks these MoE patches target. The patch activators must no-op instead of
+    raising AttributeError during temporary-patch application.
+    """
+    _stub_module(monkeypatch, "transformers")
+    _stub_module(monkeypatch, "transformers.quantizers")
+    _stub_module(
+        monkeypatch,
+        "transformers.quantizers.quantizer_bnb_4bit",
+        Bnb4BitHfQuantizer=type("Bnb4BitHfQuantizer", (), {}),
+    )
+
+    from unsloth_zoo.temporary_patches import moe_utils_bnb4bit
+
+    moe_utils_bnb4bit.patch_bnb4bit_quantizer_param_needs_quantization()
+    moe_utils_bnb4bit.patch_bnb4bit_quantizer_process_model()
+
+
+def test_moe_bnb4bit_convert_patch_skips_missing_old_transformers_hook(monkeypatch):
+    _stub_module(monkeypatch, "transformers")
+    _stub_module(monkeypatch, "transformers.integrations")
+    _stub_module(
+        monkeypatch,
+        "transformers.integrations.bitsandbytes",
+        Bnb4bitQuantize=type("Bnb4bitQuantize", (), {}),
+    )
+
+    from unsloth_zoo.temporary_patches import moe_utils_bnb4bit
+
+    moe_utils_bnb4bit.patch_bnb4bit_quantize_convert()
+
+
+def test_moe_fp8_experts_patch_handles_dict_registry(monkeypatch):
+    registry = {"grouped_mm": object(), "batched_mm": object(), "deepgemm": object()}

The change introduces robust capability-gate functions for both bnb4bit and FP8 MoE patches, conditions all patch registration on them, and replaces error-raising fallbacks with silent returns so older Transformers won't crash. Supported environments retain full patching behavior, and detection is defensively coded with broad exception handling and getattr defaults.

diff --git a/tests/test_moe_quantizer_patch_gates.py b/tests/test_moe_quantizer_patch_gates.py
new file mode 100644
index 0000000..8d662f5
--- /dev/null
+++ b/tests/test_moe_quantizer_patch_gates.py
@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+
+def _install_transformers_module(monkeypatch, module_name: str, module: types.ModuleType) -> None:
+    monkeypatch.setitem(sys.modules, module_name, module)
+
+
+def test_bnb4bit_moe_patch_noops_without_new_quantizer_hooks(monkeypatch):
+    integrations = pytest.importorskip("transformers.integrations")
+    quantizers = pytest.importorskip("transformers.quantizers")
+
+    from unsloth_zoo.temporary_patches import moe_utils_bnb4bit as bnb4bit
+    monkeypatch.setattr(bnb4bit, "HAS_BNB", True)
+    monkeypatch.setattr(bnb4bit, "Params4bit", object)
+
+    bitsandbytes_mod = types.ModuleType("transformers.integrations.bitsandbytes")
+
+    class Bnb4bitQuantize:
+        def convert(self):
+            raise AssertionError("old API gate should prevent patching")
+
+    bitsandbytes_mod.Bnb4bitQuantize = Bnb4bitQuantize
+
+    quantizer_mod = types.ModuleType("transformers.quantizers.quantizer_bnb_4bit")
+
+    class OldBnb4BitHfQuantizer:
+        def _process_model_before_weight_loading(self):
+            raise AssertionError("old API gate should prevent patching")
+
+    quantizer_mod.Bnb4BitHfQuantizer = OldBnb4BitHfQuantizer
+
+    quantizers_utils_mod = types.ModuleType("transformers.quantizers.quantizers_utils")
+    quantizers_utils_mod.get_module_from_name = lambda *args, **kwargs: None
+    quantizers_utils_mod.should_convert_module = lambda *args, **kwargs: True
+
+    _install_transformers_module(
+        monkeypatch,
+        "transformers.integrations.bitsandbytes",
+        bitsandbytes_mod,
+    )
+    _install_transformers_module(
+        monkeypatch,
+        "transformers.quantizers.quantizer_bnb_4bit",
+        quantizer_mod,
+    )
+    _install_transformers_module(
+        monkeypatch,
+        "transformers.quantizers.quantizers_utils",
+        quantizers_utils_mod,

The change introduces robust capability-gating helpers that check for quantizer hooks and the FP8 dispatcher registry, conditions all patches on those checks with early returns, and uses safe getattr/exception handling. It avoids crashes on older Transformers while preserving behavior when APIs are present.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..89594b4 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -88,6 +88,57 @@ def _is_expert_module(module: nn.Module) -> bool:
     )
 
 
+def _get_quantizers_utils_attr(name: str):
+    try:
+        import transformers.quantizers.quantizers_utils as quantizers_utils
+    except Exception as e:
+        return raise_error("transformers.quantizers.quantizers_utils", e)
+    attr = getattr(quantizers_utils, name, None)
+    if attr is None:
+        return raise_error(f"transformers.quantizers.quantizers_utils.{name}")
+    return attr
+
+
+def _get_bnb4bit_quantize_class():
+    if not HAS_BNB or Params4bit is None:
+        return None
+    try:
+        from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+    except Exception as e:
+        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
+    if getattr(Bnb4bitQuantize, "convert", None) is None:
+        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert")
+    if _get_quantizers_utils_attr("get_module_from_name") is None:
+        return None
+    return Bnb4bitQuantize
+
+
+def _get_bnb4bit_quantizer_class(required_hook: str):
+    if not HAS_BNB or Params4bit is None:
+        return None
+    try:
+        from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+    except Exception as e:
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
+    if getattr(Bnb4BitHfQuantizer, required_hook, None) is None:
+        return raise_error(f"transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer.{required_hook}")
+    return Bnb4BitHfQuantizer
+
+
+_BNB4BIT_QUANTIZATION_CONFIG_ATTRS = (
+    "bnb_4bit_use_double_quant",
+    "bnb_4bit_quant_type",
+    "bnb_4bit_quant_storage",
+)
+
+
+def _bnb4bit_quantization_config_supported(quantization_config) -> bool:
+    return quantization_config is not None and all(
+        hasattr(quantization_config, attr)
+        for attr in _BNB4BIT_QUANTIZATION_CONFIG_ATTRS
+    )
+
+
 # ============================================================================

The change adds proper runtime capability gates for both bnb4bit quantizer hooks and the FP8 experts dispatcher, conditioning patch application on the presence of the required callables/mutable registry/fp8 dtype. It safely uses getattr and try/except so detection itself does not throw, prevents the AttributeError on missing hooks, and preserves behavior when APIs are present. All rubric outcomes are fully achieved.

diff --git a/tests/test_moe_quantization_patch_gates.py b/tests/test_moe_quantization_patch_gates.py
new file mode 100644
index 0000000..fd8c7c8
--- /dev/null
+++ b/tests/test_moe_quantization_patch_gates.py
@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import sys
+import types
+
+
+def test_bnb4bit_quantizer_hooks_skip_when_transformers_methods_are_missing(monkeypatch):
+    import unsloth_zoo.temporary_patches.moe_utils_bnb4bit as bnb4bit
+
+    monkeypatch.setattr(bnb4bit, "HAS_BNB", True)
+    monkeypatch.setattr(bnb4bit, "Params4bit", object)
+
+    quantizer_mod = types.ModuleType("transformers.quantizers.quantizer_bnb_4bit")
+
+    class FakeBnb4BitHfQuantizer:
+        pass
+
+    quantizer_mod.Bnb4BitHfQuantizer = FakeBnb4BitHfQuantizer
+    monkeypatch.setitem(
+        sys.modules,
+        "transformers.quantizers.quantizer_bnb_4bit",
+        quantizer_mod,
+    )
+
+    utils_mod = types.ModuleType("transformers.quantizers.quantizers_utils")
+    utils_mod.get_module_from_name = lambda *args, **kwargs: None
+    monkeypatch.setitem(
+        sys.modules,
+        "transformers.quantizers.quantizers_utils",
+        utils_mod,
+    )
+
+    bnb4bit.patch_bnb4bit_quantizer_param_needs_quantization()
+    bnb4bit.patch_bnb4bit_quantizer_process_model()
+
+
+def test_bnb4bit_quantize_convert_skips_when_converter_method_is_missing(monkeypatch):
+    import unsloth_zoo.temporary_patches.moe_utils_bnb4bit as bnb4bit
+
+    monkeypatch.setattr(bnb4bit, "HAS_BNB", True)
+
+    bitsandbytes_mod = types.ModuleType("transformers.integrations.bitsandbytes")
+
+    class FakeBnb4bitQuantize:
+        pass
+
+    bitsandbytes_mod.Bnb4bitQuantize = FakeBnb4bitQuantize
+    monkeypatch.setitem(
+        sys.modules,
+        "transformers.integrations.bitsandbytes",
+        bitsandbytes_mod,
+    )
+
+    bnb4bit.patch_bnb4bit_quantize_convert()

no diff captured (skipped)

The change adds robust capability-detection helpers for both bnb4bit and FP8 MoE patches and gates all relevant patch registrations on them, so older Transformers loading skips unsupported patches instead of raising. It also hardens fragile except clauses. All rubric outcomes are fully achieved.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..0dbcb9f 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -51,9 +51,33 @@ __all__ = [
     "replace_expert_params_with_bnb_params",
     "forward_moe_backend_bnb4bit",
     "_moe_uses_bnb4bit_expert_weights",
+    "_transformers_supports_moe_bnb4bit_quantizer_api",
 ]
 
 
+def _transformers_supports_moe_bnb4bit_quantizer_api() -> bool:
+    """True iff the installed transformers exposes the newer (v5-style)
+    MoE/quantizer APIs these patches are written against:
+    Bnb4bitQuantize.convert, Bnb4BitHfQuantizer.param_needs_quantization and
+    the quantizers_utils helpers. Older transformers lack these hooks (e.g.
+    the quantizer only has check_quantized_param), so the patches below must
+    stay inactive there instead of raising AttributeError."""
+    try:
+        from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+        from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+        from transformers.quantizers.quantizers_utils import (
+            get_module_from_name,
+            should_convert_module,
+        )
+    except Exception:
+        return False
+    return (
+        callable(getattr(Bnb4bitQuantize, "convert", None))
+        and callable(getattr(Bnb4BitHfQuantizer, "param_needs_quantization", None))
+        and callable(getattr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading", None))
+    )
+
+
 # ============================================================================
 # Detection
 # ============================================================================
@@ -204,7 +228,10 @@ def replace_expert_params_with_bnb_params(
     try:
         from transformers.quantizers.quantizers_utils import should_convert_module
     except Exception as e:
-        return raise_error("transformers.quantizers.quantizers_utils.should_convert_module", e)
+        # Older transformers has no quantizers_utils helpers — leave the model
+        # untouched instead of returning None into the quantizer pipeline.
+        raise_error("transformers.quantizers.quantizers_utils.should_convert_module", e)
+        return model
 
     has_been_replaced = False
 
@@ -261,6 +288,10 @@ def patch_bnb4bit_quantize_convert():
     Expert modules of nn.Parameter type are converted to Params4bit placeholders during weight loading.
     Also preserves the original shape of the expert parameters for PEFT LoRA compatibility.
     """
+    if not HAS_BNB:
+        return
+    if not _transformers_supports_moe_bnb4bit_quantizer_api():
+        return
 
     try:

The change adds a robust attribute-presence gate for the bnb4bit MoE patches, preventing AttributeError on older Transformers and preserving behavior when supported. However, it entirely omits the FP8 patches which the task explicitly requires, leaving c2 and c3 only partially satisfied.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..83b73e8 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -54,6 +54,27 @@ __all__ = [
 ]
 
 
+# ============================================================================
+# Capability detection
+# ============================================================================
+
+def _require_attr(owner, attr_name: str, qualified_name: str) -> bool:
+    """True iff `owner` actually exposes `attr_name`.
+
+    The bnb4bit MoE patches are written against the transformers v5 quantizer /
+    weight-converter APIs. On older transformers the target class often still
+    imports cleanly but lacks the specific hook we wrap (e.g.
+    `Bnb4BitHfQuantizer.param_needs_quantization`). Reaching for it directly —
+    even inside a `hasattr(Class.method, ...)` idempotency guard — evaluates
+    `Class.method` first and raises AttributeError before `hasattr` can catch
+    anything. Gate on the attribute existing so the patch cleanly no-ops on
+    unsupported versions instead of blowing up at import/patch time.
+    """
+    if getattr(owner, attr_name, None) is None:
+        return raise_error(qualified_name) or False
+    return True
+
+
 # ============================================================================
 # Detection
 # ============================================================================
@@ -267,6 +288,9 @@ def patch_bnb4bit_quantize_convert():
     except Exception as e:
         return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
 
+    if not _require_attr(Bnb4bitQuantize, "convert", "transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert"):
+        return
+
     if hasattr(Bnb4bitQuantize.convert, "_unsloth_moe_patched"):
         return
 
@@ -320,6 +344,9 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
+    if not _require_attr(Bnb4BitHfQuantizer, "param_needs_quantization", "transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer.param_needs_quantization"):
+        return
+
     if hasattr(Bnb4BitHfQuantizer.param_needs_quantization, "_unsloth_moe_patched"):
         return
 
@@ -356,6 +383,9 @@ def patch_bnb4bit_quantizer_process_model():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
+    if not _require_attr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading", "transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer._process_model_before_weight_loading"):
+        return
+
     if hasattr(Bnb4BitHfQuantizer._process_model_before_weight_loading, "_unsloth_moe_patched"):

The change adds robust capability gates for both bnb 4-bit and FP8 MoE patches, conditions all patch bodies on those gates, and detects capabilities safely with try/except and getattr defaults. Note the patches remain appended to TEMPORARY_PATCHES but internally no-op when unsupported, which achieves the goal of not breaking older Transformers while preserving behavior when supported.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..0947c07 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -191,6 +191,43 @@ def forward_moe_backend_bnb4bit(self, hidden_states, top_k_index, top_k_weights)
 # transformers integration patches
 # ============================================================================
 
+_MOE_QUANT_API_SKIP_LOGGED = False
+
+
+def _transformers_supports_moe_quant_api() -> bool:
+    """True iff bitsandbytes is importable AND the installed transformers
+    exposes the newer (v5-style) MoE/quantizer APIs these patches are written
+    against. Older transformers (quantizer hook still `check_quantized_param`,
+    no `should_convert_module` / `Bnb4bitQuantize`) must leave the patches
+    inactive instead of raising AttributeError at patch time.
+    """
+    supported = False
+    if HAS_BNB:
+        try:
+            from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+            from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+            from transformers.quantizers import quantizers_utils
+        except Exception:
+            pass
+        else:
+            supported = (
+                callable(getattr(Bnb4bitQuantize, "convert", None))
+                and callable(getattr(Bnb4BitHfQuantizer, "param_needs_quantization", None))
+                and callable(getattr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading", None))
+                and callable(getattr(quantizers_utils, "get_module_from_name", None))
+                and callable(getattr(quantizers_utils, "should_convert_module", None))
+            )
+    global _MOE_QUANT_API_SKIP_LOGGED
+    if not supported and not _MOE_QUANT_API_SKIP_LOGGED:
+        _MOE_QUANT_API_SKIP_LOGGED = True
+        if UNSLOTH_ENABLE_LOGGING:
+            logger.info(
+                "Unsloth: Skipping MoE bnb 4-bit quantization patches — "
+                "bitsandbytes and/or the transformers v5 MoE quantizer API are unavailable."
+            )
+    return supported
+
+
 def replace_expert_params_with_bnb_params(
     model: nn.Module,
     modules_to_not_convert: Optional[List[str]] = None,
@@ -261,6 +298,8 @@ def patch_bnb4bit_quantize_convert():
     Expert modules of nn.Parameter type are converted to Params4bit placeholders during weight loading.
     Also preserves the original shape of the expert parameters for PEFT LoRA compatibility.
     """
+    if not _transformers_supports_moe_quant_api():
+        return
 
     try:
         from transformers.integrations.bitsandbytes import Bnb4bitQuantize
@@ -313,6 +352,8 @@ TEMPORARY_PATCHES.append(patch_bnb4bit_quantize_convert)
 
 def patch_bnb4bit_quantizer_param_needs_quantization():

The change makes each MoE bnb4bit and FP8 patch feature-detect its required hooks via getattr defaults and broadened exception handling, so older Transformers no longer crashes with AttributeError/ImportError while supported environments retain full behavior. It does not introduce a single explicit capability gate but distributes equivalent runtime checks across the patches, achieving the intended outcome slightly less cleanly for c1.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..41ba989 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -267,10 +267,14 @@ def patch_bnb4bit_quantize_convert():
     except Exception as e:
         return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
 
-    if hasattr(Bnb4bitQuantize.convert, "_unsloth_moe_patched"):
-        return
+    # Older transformers has no Bnb4bitQuantize.convert hook. Feature-detect
+    # rather than assume the attribute exists so the patch no-ops cleanly.
+    original_convert = getattr(Bnb4bitQuantize, "convert", None)
+    if original_convert is None:
+        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert")
 
-    original_convert = Bnb4bitQuantize.convert
+    if hasattr(original_convert, "_unsloth_moe_patched"):
+        return
 
     def patched_convert(
         self,
@@ -320,10 +324,14 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
-    if hasattr(Bnb4BitHfQuantizer.param_needs_quantization, "_unsloth_moe_patched"):
-        return
+    # param_needs_quantization is a newer-transformers quantizer hook. On older
+    # versions it's absent, so detect it rather than dereference it blindly.
+    original_param_needs_quantization = getattr(Bnb4BitHfQuantizer, "param_needs_quantization", None)
+    if original_param_needs_quantization is None:
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer.param_needs_quantization")
 
-    original_param_needs_quantization = Bnb4BitHfQuantizer.param_needs_quantization
+    if hasattr(original_param_needs_quantization, "_unsloth_moe_patched"):
+        return
 
     def patched_param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         if original_param_needs_quantization(self, model, param_name, **kwargs):
@@ -356,10 +364,16 @@ def patch_bnb4bit_quantizer_process_model():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
-    if hasattr(Bnb4BitHfQuantizer._process_model_before_weight_loading, "_unsloth_moe_patched"):
-        return
+    original_process_model_before_weight_loading = getattr(
+        Bnb4BitHfQuantizer, "_process_model_before_weight_loading", None
+    )
+    if original_process_model_before_weight_loading is None:
+        return raise_error(
+            "transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer._process_model_before_weight_loading"
+        )
 
-    original_process_model_before_weight_loading = Bnb4BitHfQuantizer._process_model_before_weight_loading
+    if hasattr(original_process_model_before_weight_loading, "_unsloth_moe_patched"):
+        return
 
     def patched_process_model_before_weight_loading(self, model, device_map, **kwargs):
         original_process_model_before_weight_loading(self, model, device_map, **kwargs)

The change introduces robust capability-detection functions and gates each bnb4bit and FP8 MoE patch on them, so older Transformers versions lacking the newer APIs skip the patches instead of crashing. Detection catches import/attribute errors and uses getattr defaults. All rubric outcomes are fully achieved.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..468055e 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -44,6 +44,7 @@ except ImportError:
 
 
 __all__ = [
+    "transformers_supports_moe_bnb4bit_quantization",
     "patch_bnb4bit_quantize_convert",
     "patch_bnb4bit_quantizer_param_needs_quantization",
     "patch_bnb4bit_quantizer_process_model",
@@ -58,6 +59,29 @@ __all__ = [
 # Detection
 # ============================================================================
 
+def transformers_supports_moe_bnb4bit_quantization() -> bool:
+    """True iff the installed transformers exposes the v5 MoE/quantizer APIs
+    these patches are written against. Older transformers (pre-v5) either lack
+    the modules entirely or expose different hook names (e.g.
+    `check_quantized_param` instead of `param_needs_quantization`), so the
+    patches below must not be applied there.
+    """
+    try:
+        from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+        from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+        from transformers.quantizers.quantizers_utils import (
+            get_module_from_name,
+            should_convert_module,
+        )
+    except Exception:
+        return False
+    return (
+        callable(getattr(Bnb4bitQuantize, "convert", None))
+        and callable(getattr(Bnb4BitHfQuantizer, "param_needs_quantization", None))
+        and callable(getattr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading", None))
+    )
+
+
 def _is_bnb4bit_param(param) -> bool:
     """True iff `param` is a Params4bit with a populated quant_state."""
     return (
@@ -261,6 +285,8 @@ def patch_bnb4bit_quantize_convert():
     Expert modules of nn.Parameter type are converted to Params4bit placeholders during weight loading.
     Also preserves the original shape of the expert parameters for PEFT LoRA compatibility.
     """
+    if not HAS_BNB or not transformers_supports_moe_bnb4bit_quantization():
+        return
 
     try:
         from transformers.integrations.bitsandbytes import Bnb4bitQuantize
@@ -313,6 +339,8 @@ TEMPORARY_PATCHES.append(patch_bnb4bit_quantize_convert)
 
 def patch_bnb4bit_quantizer_param_needs_quantization():
     """Recognize MoE expert modules of Params4bit type as needing quantization."""
+    if not HAS_BNB or not transformers_supports_moe_bnb4bit_quantization():
+        return
 
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer

The change adds safe hasattr guards to the bnb 4-bit MoE patches so they skip on older Transformers, addressing the bnb crash and preserving supported behavior. However it entirely omits the FP8 patches, which the task and rubric explicitly require to be gated too, leaving that failure mode unaddressed and using scattered per-patch checks rather than a unified capability gate.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..47962e5 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -267,6 +267,11 @@ def patch_bnb4bit_quantize_convert():
     except Exception as e:
         return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
 
+    # Older Transformers predate the Bnb4bitQuantize.convert MoE hook; only patch
+    # when the newer API is actually present so loading doesn't blow up.
+    if not hasattr(Bnb4bitQuantize, "convert"):
+        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert")
+
     if hasattr(Bnb4bitQuantize.convert, "_unsloth_moe_patched"):
         return
 
@@ -320,6 +325,12 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
+    # Older Transformers used check_quantized_param instead of
+    # param_needs_quantization. Skip the patch entirely on those versions rather
+    # than AttributeError on the missing hook.
+    if not hasattr(Bnb4BitHfQuantizer, "param_needs_quantization"):
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer.param_needs_quantization")
+
     if hasattr(Bnb4BitHfQuantizer.param_needs_quantization, "_unsloth_moe_patched"):
         return
 
@@ -356,6 +367,10 @@ def patch_bnb4bit_quantizer_process_model():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
+    # Guard against older Transformers lacking this loading hook.
+    if not hasattr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading"):
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer._process_model_before_weight_loading")
+
     if hasattr(Bnb4BitHfQuantizer._process_model_before_weight_loading, "_unsloth_moe_patched"):
         return
 

The change guards each patch by verifying the specific quantizer/dispatcher attributes exist within existing try/except blocks, so unsupported environments skip the patch and avoid AttributeError. However, there's no single explicit capability-detection helper, and patches remain registered unconditionally (they no-op at runtime rather than being conditionally applied), so c1 and c2 are only partially satisfied.

diff --git a/tests/test_temporary_patches_imports.py b/tests/test_temporary_patches_imports.py
index 6b86258..832631c 100644
--- a/tests/test_temporary_patches_imports.py
+++ b/tests/test_temporary_patches_imports.py
@@ -63,6 +63,8 @@ TEMPORARY_PATCHES_SUBMODULES = [
     "unsloth_zoo.temporary_patches.misc",
     "unsloth_zoo.temporary_patches.moe_bnb",
     "unsloth_zoo.temporary_patches.moe_utils",
+    "unsloth_zoo.temporary_patches.moe_utils_bnb4bit",
+    "unsloth_zoo.temporary_patches.moe_utils_fp8",
     "unsloth_zoo.temporary_patches.mxfp4",
     "unsloth_zoo.temporary_patches.pixtral",
     "unsloth_zoo.temporary_patches.qwen3_5_moe",
diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..e15b9c5 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -264,6 +264,7 @@ def patch_bnb4bit_quantize_convert():
 
     try:
         from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+        _ = Bnb4bitQuantize.convert  # verify attribute exists on this transformers version
     except Exception as e:
         return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
 
@@ -317,6 +318,7 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
         from transformers.quantizers.quantizers_utils import get_module_from_name
+        _ = Bnb4BitHfQuantizer.param_needs_quantization  # verify attribute exists
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
@@ -353,6 +355,7 @@ TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_param_needs_quantization)
 def patch_bnb4bit_quantizer_process_model():
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+        _ = Bnb4BitHfQuantizer._process_model_before_weight_loading  # verify attribute exists
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
diff --git a/unsloth_zoo/temporary_patches/moe_utils_fp8.py b/unsloth_zoo/temporary_patches/moe_utils_fp8.py
index df7e825..88210bb 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_fp8.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_fp8.py
@@ -1230,11 +1230,14 @@ from .utils import logger
 def patch_fp8_experts_interface():
     try:
         from transformers.integrations.finegrained_fp8 import ALL_FP8_EXPERTS_FUNCTIONS
-    except ImportError:
+    except Exception:
         return
 
     sentinel = "_unsloth_fp8_dispatcher"
-    if getattr(ALL_FP8_EXPERTS_FUNCTIONS, sentinel, False):
+    try:
+        if getattr(ALL_FP8_EXPERTS_FUNCTIONS, sentinel, False):
+            return
+    except Exception:
         return

The diff merely moves the original-method assignment before the hasattr check and improves error strings; it does not add any capability gate or condition the bnb4bit/FP8 patches on Transformers API support. The core requirement—skipping unsupported patches on older Transformers—is not implemented.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..c438138 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -264,14 +264,13 @@ def patch_bnb4bit_quantize_convert():
 
     try:
         from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+        original_convert = Bnb4bitQuantize.convert
     except Exception as e:
-        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
+        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert", e)
 
-    if hasattr(Bnb4bitQuantize.convert, "_unsloth_moe_patched"):
+    if hasattr(original_convert, "_unsloth_moe_patched"):
         return
 
-    original_convert = Bnb4bitQuantize.convert
-
     def patched_convert(
         self,
         input_dict: dict[str, Union[list[torch.Tensor], torch.Tensor]],
@@ -317,14 +316,13 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
         from transformers.quantizers.quantizers_utils import get_module_from_name
+        original_param_needs_quantization = Bnb4BitHfQuantizer.param_needs_quantization
     except Exception as e:
-        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer.param_needs_quantization", e)
 
-    if hasattr(Bnb4BitHfQuantizer.param_needs_quantization, "_unsloth_moe_patched"):
+    if hasattr(original_param_needs_quantization, "_unsloth_moe_patched"):
         return
 
-    original_param_needs_quantization = Bnb4BitHfQuantizer.param_needs_quantization
-
     def patched_param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         if original_param_needs_quantization(self, model, param_name, **kwargs):
             return True
@@ -353,14 +351,13 @@ TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_param_needs_quantization)
 def patch_bnb4bit_quantizer_process_model():
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+        original_process_model_before_weight_loading = Bnb4BitHfQuantizer._process_model_before_weight_loading
     except Exception as e:
-        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer._process_model_before_weight_loading", e)
 
-    if hasattr(Bnb4BitHfQuantizer._process_model_before_weight_loading, "_unsloth_moe_patched"):
+    if hasattr(original_process_model_before_weight_loading, "_unsloth_moe_patched"):
         return
 
-    original_process_model_before_weight_loading = Bnb4BitHfQuantizer._process_model_before_weight_loading
-
     def patched_process_model_before_weight_loading(self, model, device_map, **kwargs):
         original_process_model_before_weight_loading(self, model, device_map, **kwargs)
 

The change adds hasattr guards before accessing the newer quantizer/MoE hooks in each bnb4bit patch and makes the FP8 dispatcher marking safe on plain dicts, which prevents the AttributeError on older Transformers and keeps patches active when APIs exist. However, it implements detection as scattered per-function guards rather than a distinct, reusable capability gate the rubric emphasizes, so c1/c2 earn partial credit; the practical outcome (no crash, still active when supported) is achieved.

diff --git a/tests/test_temporary_patches_exhaustive.py b/tests/test_temporary_patches_exhaustive.py
index 01040ee..ce7d827 100644
--- a/tests/test_temporary_patches_exhaustive.py
+++ b/tests/test_temporary_patches_exhaustive.py
@@ -1532,6 +1532,116 @@ def test_qwen3_vl_moe_act2fn_dict_present():
         )
 
 
+# moe_utils_bnb4bit.py: version-gated transformers v5 MoE/quantizer APIs.
+# Each patch checks hasattr() before touching its target; the tests below
+# verify (a) graceful no-op on 4.x (skip) and (b) target present on 5.x.
+
+def test_moe_utils_bnb4bit_quantize_convert_class_present_on_5x():
+    """moe_utils_bnb4bit.py patch_bnb4bit_quantize_convert uses
+    transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert; only
+    available on transformers 5.x.  Older installs: skip (graceful no-op
+    confirmed by hasattr guard added in the patch function).  Newer
+    installs: drift-fail if the class or method disappears."""
+    try:
+        mod = importlib.import_module("transformers.integrations.bitsandbytes")
+    except Exception:
+        pytest.skip(
+            "transformers.integrations.bitsandbytes not importable on "
+            f"transformers {_TX_VERSION} -- moe_utils_bnb4bit patch no-ops"
+        )
+    cls = getattr(mod, "Bnb4bitQuantize", None)
+    if cls is None:
+        pytest.skip(
+            f"Bnb4bitQuantize absent on transformers {_TX_VERSION} "
+            "(pre-5.x; patch_bnb4bit_quantize_convert no-ops via hasattr guard)"
+        )
+    if not hasattr(cls, "convert"):
+        pytest.fail(
+            "DRIFT DETECTED: zoo temporary_patches/moe_utils_bnb4bit.py "
+            "patch_bnb4bit_quantize_convert expects Bnb4bitQuantize.convert "
+            f"but it is missing on transformers {_TX_VERSION}"
+        )
+
+
+def test_moe_utils_bnb4bit_quantizer_param_needs_quantization_present_on_5x():
+    """moe_utils_bnb4bit.py patch_bnb4bit_quantizer_param_needs_quantization
+    uses Bnb4BitHfQuantizer.param_needs_quantization; optional on 4.x (patch
+    no-ops via hasattr guard), required when Bnb4BitHfQuantizer is available."""
+    try:
+        mod = importlib.import_module("transformers.quantizers.quantizer_bnb_4bit")
+    except Exception:
+        pytest.skip(
+            "transformers.quantizers.quantizer_bnb_4bit not importable -- "
+            "patch no-ops"
+        )
+    cls = getattr(mod, "Bnb4BitHfQuantizer", None)
+    if cls is None:
+        pytest.skip(
+            f"Bnb4BitHfQuantizer absent on transformers {_TX_VERSION}"
+        )
+    if not hasattr(cls, "param_needs_quantization"):
+        pytest.fail(
+            "DRIFT DETECTED: zoo temporary_patches/moe_utils_bnb4bit.py "
+            "patch_bnb4bit_quantizer_param_needs_quantization expects "
+            "Bnb4BitHfQuantizer.param_needs_quantization but it is missing "

The change gates each MoE bnb4bit and FP8 patch behind hasattr checks for the specific newer quantizer/dispatcher hooks and widens exception handling, so unsupported patches are skipped rather than crashing on older Transformers. It achieves the outcomes via distributed per-patch guards rather than a single named capability function, which is a valid alternative approach, but the gating on the quantizer hook presence is slightly indirect for c1/c2.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..c6dc0f7 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -264,6 +264,8 @@ def patch_bnb4bit_quantize_convert():
 
     try:
         from transformers.integrations.bitsandbytes import Bnb4bitQuantize
+        if not hasattr(Bnb4bitQuantize, "convert"):
+            return
     except Exception as e:
         return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
 
@@ -317,6 +319,8 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
         from transformers.quantizers.quantizers_utils import get_module_from_name
+        if not hasattr(Bnb4BitHfQuantizer, "param_needs_quantization"):
+            return
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
@@ -353,6 +357,8 @@ TEMPORARY_PATCHES.append(patch_bnb4bit_quantizer_param_needs_quantization)
 def patch_bnb4bit_quantizer_process_model():
     try:
         from transformers.quantizers.quantizer_bnb_4bit import Bnb4BitHfQuantizer
+        if not hasattr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading"):
+            return
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
@@ -388,7 +394,7 @@ def patch_transformers_weight_converter_kwargs():
     """
     try:
         from transformers.core_model_loading import WeightConverter
-    except ImportError:
+    except Exception:
         return
 
     if getattr(WeightConverter.__init__, "_unsloth_kwargs_patched", False):
diff --git a/unsloth_zoo/temporary_patches/moe_utils_fp8.py b/unsloth_zoo/temporary_patches/moe_utils_fp8.py
index df7e825..e489ad7 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_fp8.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_fp8.py
@@ -1230,7 +1230,7 @@ from .utils import logger
 def patch_fp8_experts_interface():
     try:
         from transformers.integrations.finegrained_fp8 import ALL_FP8_EXPERTS_FUNCTIONS
-    except ImportError:
+    except Exception:
         return
 
     sentinel = "_unsloth_fp8_dispatcher"
@@ -1262,7 +1262,7 @@ def patch_fp8_validate_quantization_for_training():
     try:
         import transformers.trainer as _trainer_mod
         import transformers.trainer_utils as _trainer_utils_mod
-    except ImportError:
+    except Exception:
         return

The change makes the bnb 4-bit patches individually check for the attributes they patch, which partially prevents AttributeErrors, but it does not implement a genuine environment capability gate nor condition the patch registration on such a gate. FP8 patches are entirely absent from the diff, so the required outcomes are only partially met.

diff --git a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
index 44736e1..365b4e4 100644
--- a/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
+++ b/unsloth_zoo/temporary_patches/moe_utils_bnb4bit.py
@@ -267,7 +267,10 @@ def patch_bnb4bit_quantize_convert():
     except Exception as e:
         return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize", e)
 
-    if hasattr(Bnb4bitQuantize.convert, "_unsloth_moe_patched"):
+    if not hasattr(Bnb4bitQuantize, "convert"):
+        return raise_error("transformers.integrations.bitsandbytes.Bnb4bitQuantize.convert", None)
+
+    if getattr(Bnb4bitQuantize.convert, "_unsloth_moe_patched", False):
         return
 
     original_convert = Bnb4bitQuantize.convert
@@ -320,7 +323,10 @@ def patch_bnb4bit_quantizer_param_needs_quantization():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
-    if hasattr(Bnb4BitHfQuantizer.param_needs_quantization, "_unsloth_moe_patched"):
+    if not hasattr(Bnb4BitHfQuantizer, "param_needs_quantization"):
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer.param_needs_quantization", None)
+
+    if getattr(Bnb4BitHfQuantizer.param_needs_quantization, "_unsloth_moe_patched", False):
         return
 
     original_param_needs_quantization = Bnb4BitHfQuantizer.param_needs_quantization
@@ -356,7 +362,10 @@ def patch_bnb4bit_quantizer_process_model():
     except Exception as e:
         return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer", e)
 
-    if hasattr(Bnb4BitHfQuantizer._process_model_before_weight_loading, "_unsloth_moe_patched"):
+    if not hasattr(Bnb4BitHfQuantizer, "_process_model_before_weight_loading"):
+        return raise_error("transformers.quantizers.quantizer_bnb_4bit.Bnb4BitHfQuantizer._process_model_before_weight_loading", None)
+
+    if getattr(Bnb4BitHfQuantizer._process_model_before_weight_loading, "_unsloth_moe_patched", False):
         return
 
     original_process_model_before_weight_loading = Bnb4BitHfQuantizer._process_model_before_weight_loading

The change conditionally makes hidden_states contiguous before the shape-dependent operations, fixing the non-contiguous crash while avoiding copies on the common path. Since the fix is at the entry point, both forward loss computation and backward gradient handling operate on the contiguous tensor, satisfying all criteria.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..3617b8e 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -219,6 +219,9 @@ class UnslothFusedLoss(torch.autograd.Function):
         grad_lm_head_bias = torch.zeros_like(lm_head_bias, device = device) if lm_head_bias_requires_grad else None
 
         bsz, qlen, hd = hidden_states.shape
+        # Make hidden_states contiguous (sliced tensors can be non-contiguous)
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
         accumulated_loss = torch.zeros(1, device = device)[0]
         # Chunk hidden_states and labels
         if "n_chunks" in extra_kwargs:

The change correctly converts .view to .reshape and reshapes hidden_states up-front in both forward and gradient paths, making the fused CE loss robust to non-contiguous sliced inputs while preserving semantics and avoiding unnecessary copies on the contiguous path. A test validates numerical equivalence for the sliced scenario. All criteria fully met.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..9802344 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -118,8 +118,8 @@ def compute_fused_ce_loss(
     # Calculate cross entropy loss
     reduction = "sum" if n_items is not None else "mean"
     loss = torch.nn.functional.cross_entropy(
-        input  = logits.view(-1, vocab_size).float().contiguous(),
-        target = labels.view(-1).to(device).contiguous(),
+        input  = logits.reshape(-1, vocab_size).float().contiguous(),
+        target = labels.reshape(-1).to(device).contiguous(),
         reduction = reduction,
     )
     loss = loss / n_items if n_items is not None else loss
@@ -213,12 +213,15 @@ class UnslothFusedLoss(torch.autograd.Function):
         lm_head_bias_requires_grad = lm_head_bias is not None and lm_head_bias.requires_grad
         vocab_size = lm_head_weight.shape[0]
 
+        hidden_states_shape = hidden_states.shape
+        bsz, qlen, hd = hidden_states_shape
+        hidden_states = hidden_states.reshape(-1, hd)
+
         # Create backwards output
-        grad_inputs = torch.empty_like(hidden_states, device = device) if not overwrite else hidden_states
+        grad_inputs = hidden_states if overwrite else torch.empty_like(hidden_states, device = device)
         grad_lm_head = torch.zeros_like(lm_head_weight, device = device) if lm_head_requires_grad else None
         grad_lm_head_bias = torch.zeros_like(lm_head_bias, device = device) if lm_head_bias_requires_grad else None
 
-        bsz, qlen, hd = hidden_states.shape
         accumulated_loss = torch.zeros(1, device = device)[0]
         # Chunk hidden_states and labels
         if "n_chunks" in extra_kwargs:
@@ -228,8 +231,8 @@ class UnslothFusedLoss(torch.autograd.Function):
         if UNSLOTH_ENABLE_LOGGING:
             logger.info(f"Fused CE Loss [bsz={bsz}][qlen={qlen}][vocab_size={vocab_size}][n_chunks={n_chunks}]")
         __shift_labels = torch.chunk(labels,                     n_chunks, dim = 0)
-        __shift_states = torch.chunk(hidden_states.view(-1, hd), n_chunks, dim = 0)
-        __grad_inputs  = torch.chunk(grad_inputs.view(-1, hd),   n_chunks, dim = 0)
+        __shift_states = torch.chunk(hidden_states, n_chunks, dim = 0)
+        __grad_inputs  = torch.chunk(grad_inputs,   n_chunks, dim = 0)
 
         def accumulate_chunk(
             n_chunks,
@@ -413,6 +416,7 @@ class UnslothFusedLoss(torch.autograd.Function):
                     **extra_kwargs,
                 )
         pass
+        grad_inputs = grad_inputs.reshape(hidden_states_shape)
         ctx.save_for_backward(grad_inputs, grad_lm_head, grad_lm_head_bias)
         ctx.scaling = scaling
         return accumulated_loss
@@ -500,4 +504,4 @@ pass
 # GNU Affero General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
-# along with this program.  If not, see <https://www.gnu.org/licenses/>.
\ No newline at end of file
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
diff --git a/unsloth_zoo/fused_losses/tests.py b/unsloth_zoo/fused_losses/tests.py

The change robustly handles non-contiguous hidden states by making them contiguous at the forward entry point and switching view->reshape in flattening, preserving correctness for both forward and backward paths. The added test directly covers the sliced-last-tokens case. Minor deduction because loss_utils forces an unconditional .contiguous() copy on the common path rather than guarding it.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..de6c886 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -118,8 +118,8 @@ def compute_fused_ce_loss(
     # Calculate cross entropy loss
     reduction = "sum" if n_items is not None else "mean"
     loss = torch.nn.functional.cross_entropy(
-        input  = logits.view(-1, vocab_size).float().contiguous(),
-        target = labels.view(-1).to(device).contiguous(),
+        input  = logits.reshape(-1, vocab_size).float().contiguous(),
+        target = labels.reshape(-1).to(device).contiguous(),
         reduction = reduction,
     )
     loss = loss / n_items if n_items is not None else loss
@@ -189,9 +189,12 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
-            _labels = torch.empty_like(labels, device = device)
+            _labels = torch.empty(labels.shape, dtype = labels.dtype, device = device)
             _labels[..., :-1] = labels[..., 1:]
             # Also check mask
             if mask is not None:
@@ -199,7 +202,7 @@ class UnslothFusedLoss(torch.autograd.Function):
                 _labels[..., :-1][mask[..., 1:] == 0] = -100
             pass
             _labels[..., -1] = -100
-            _labels = _labels.view(-1)
+            _labels = _labels.reshape(-1)
             labels = _labels
         pass
 
@@ -500,4 +503,4 @@ pass
 # GNU Affero General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
-# along with this program.  If not, see <https://www.gnu.org/licenses/>.
\ No newline at end of file
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
diff --git a/unsloth_zoo/fused_losses/tests.py b/unsloth_zoo/fused_losses/tests.py
index f406597..21d4c18 100644
--- a/unsloth_zoo/fused_losses/tests.py
+++ b/unsloth_zoo/fused_losses/tests.py
@@ -15,6 +15,61 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 import torch
+from .cross_entropy_loss import compute_fused_ce_loss, unsloth_fused_ce_loss
+
+def test_fused_ce_loss_accepts_non_contiguous_hidden_states():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+    bsz, total_qlen, keep_qlen, hd, vocab_size = 2, 9, 5, 17, 128
+    torch.manual_seed(3407)

The change correctly and robustly handles non-contiguous hidden states in both forward and backward paths using reshape and a separate contiguous grad buffer, preserving numerical semantics. The gradient path is properly addressed via flat_grad_inputs. Only minor concern on unconditional allocation in the non-overwrite path, though this matches prior behavior.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..1be51aa 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -118,8 +118,8 @@ def compute_fused_ce_loss(
     # Calculate cross entropy loss
     reduction = "sum" if n_items is not None else "mean"
     loss = torch.nn.functional.cross_entropy(
-        input  = logits.view(-1, vocab_size).float().contiguous(),
-        target = labels.view(-1).to(device).contiguous(),
+        input  = logits.reshape(-1, vocab_size).float().contiguous(),
+        target = labels.reshape(-1).to(device).contiguous(),
         reduction = reduction,
     )
     loss = loss / n_items if n_items is not None else loss
@@ -191,7 +191,7 @@ class UnslothFusedLoss(torch.autograd.Function):
 
         # Get shifted labels first
         if shift_labels:
-            _labels = torch.empty_like(labels, device = device)
+            _labels = torch.empty(labels.shape, dtype = labels.dtype, device = device)
             _labels[..., :-1] = labels[..., 1:]
             # Also check mask
             if mask is not None:
@@ -199,8 +199,10 @@ class UnslothFusedLoss(torch.autograd.Function):
                 _labels[..., :-1][mask[..., 1:] == 0] = -100
             pass
             _labels[..., -1] = -100
-            _labels = _labels.view(-1)
+            _labels = _labels.reshape(-1)
             labels = _labels
+        else:
+            labels = labels.reshape(-1)
         pass
 
         # N items divisor
@@ -214,7 +216,24 @@ class UnslothFusedLoss(torch.autograd.Function):
         vocab_size = lm_head_weight.shape[0]
 
         # Create backwards output
-        grad_inputs = torch.empty_like(hidden_states, device = device) if not overwrite else hidden_states
+        if overwrite:
+            try:
+                flat_grad_inputs = hidden_states.view(-1, hidden_states.shape[-1])
+                grad_inputs = hidden_states
+            except RuntimeError:
+                grad_inputs = torch.empty_like(
+                    hidden_states,
+                    device = device,
+                    memory_format = torch.contiguous_format,
+                )
+                flat_grad_inputs = grad_inputs.view(-1, hidden_states.shape[-1])
+        else:
+            grad_inputs = torch.empty_like(
+                hidden_states,
+                device = device,
+                memory_format = torch.contiguous_format,
+            )
+            flat_grad_inputs = grad_inputs.view(-1, hidden_states.shape[-1])
         grad_lm_head = torch.zeros_like(lm_head_weight, device = device) if lm_head_requires_grad else None

The change comprehensively replaces view with reshape in both forward and backward flattening paths, correctly guards the overwrite optimization on contiguity, and allocates grad_inputs in contiguous format. It preserves semantics and avoids unnecessary copies on the common contiguous path, and adds tests covering the sliced scenario.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..0b07ae2 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -118,8 +118,8 @@ def compute_fused_ce_loss(
     # Calculate cross entropy loss
     reduction = "sum" if n_items is not None else "mean"
     loss = torch.nn.functional.cross_entropy(
-        input  = logits.view(-1, vocab_size).float().contiguous(),
-        target = labels.view(-1).to(device).contiguous(),
+        input  = logits.reshape(-1, vocab_size).float().contiguous(),
+        target = labels.reshape(-1).to(device).contiguous(),
         reduction = reduction,
     )
     loss = loss / n_items if n_items is not None else loss
@@ -199,7 +199,7 @@ class UnslothFusedLoss(torch.autograd.Function):
                 _labels[..., :-1][mask[..., 1:] == 0] = -100
             pass
             _labels[..., -1] = -100
-            _labels = _labels.view(-1)
+            _labels = _labels.reshape(-1)
             labels = _labels
         pass
 
@@ -214,7 +214,12 @@ class UnslothFusedLoss(torch.autograd.Function):
         vocab_size = lm_head_weight.shape[0]
 
         # Create backwards output
-        grad_inputs = torch.empty_like(hidden_states, device = device) if not overwrite else hidden_states
+        can_overwrite_hidden_states = overwrite and hidden_states.is_contiguous()
+        grad_inputs = hidden_states if can_overwrite_hidden_states else torch.empty_like(
+            hidden_states,
+            device = device,
+            memory_format = torch.contiguous_format,
+        )
         grad_lm_head = torch.zeros_like(lm_head_weight, device = device) if lm_head_requires_grad else None
         grad_lm_head_bias = torch.zeros_like(lm_head_bias, device = device) if lm_head_bias_requires_grad else None
 
@@ -228,8 +233,8 @@ class UnslothFusedLoss(torch.autograd.Function):
         if UNSLOTH_ENABLE_LOGGING:
             logger.info(f"Fused CE Loss [bsz={bsz}][qlen={qlen}][vocab_size={vocab_size}][n_chunks={n_chunks}]")
         __shift_labels = torch.chunk(labels,                     n_chunks, dim = 0)
-        __shift_states = torch.chunk(hidden_states.view(-1, hd), n_chunks, dim = 0)
-        __grad_inputs  = torch.chunk(grad_inputs.view(-1, hd),   n_chunks, dim = 0)
+        __shift_states = torch.chunk(hidden_states.reshape(-1, hd), n_chunks, dim = 0)
+        __grad_inputs  = torch.chunk(grad_inputs.reshape(-1, hd),   n_chunks, dim = 0)
 
         def accumulate_chunk(
             n_chunks,
@@ -359,7 +364,7 @@ class UnslothFusedLoss(torch.autograd.Function):
                 _FUSED_CE_COMPILE_SUPPORTED = False
                 torch._dynamo.reset()
                 accumulated_loss.zero_()
-                if not overwrite:
+                if not can_overwrite_hidden_states:
                     grad_inputs.zero_()
                 if grad_lm_head is not None: grad_lm_head.zero_()
                 if grad_lm_head_bias is not None: grad_lm_head_bias.zero_()
@@ -500,4 +505,4 @@ pass
 # GNU Affero General Public License for more details.

The change correctly and minimally fixes the crash by conditionally making non-contiguous hidden states contiguous before the view-based chunking, without penalizing the contiguous path. The gradient path is robust as a consequence of the forward tensor being contiguous, though the diff doesn't explicitly show backward reshaping, so slight deduction there.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..f6cbd98 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,11 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Sliced hidden states (e.g. keeping only the last few tokens' logits)
+        # are non-contiguous, which breaks the .view(-1, hd) chunking below.
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

The change conditionally makes non-contiguous hidden states contiguous at the start of the fused loss forward, cleanly fixing the crash without penalizing the contiguous path. Since all downstream forward and backward reshaping operates on this tensor, all criteria are satisfied.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..3387307 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,11 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Sliced hidden states (e.g. keeping only the last few tokens) are
+        # non-contiguous and cannot be .view()'d into chunks below
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

The candidate conditionally materializes a contiguous copy of hidden_states at the top of the forward, which fixes the .view crash on sliced inputs while avoiding copies on the common contiguous path. Since the contiguous tensor is used throughout (and saved for backward), both forward and gradient paths become robust without altering numerics. Fully achieves all criteria.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..03929d2 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,12 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Callers may pass sliced hidden_states (e.g. keeping only the last few
+        # token positions), which are non-contiguous. The .view(-1, hd) calls
+        # below require contiguous memory, so materialize a contiguous copy.
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

no diff captured (skipped)

The change conditionally makes hidden_states contiguous only when needed, fixing the .view crash for sliced inputs while preserving semantics and the common-case performance. Because grad_inputs is derived from the now-contiguous hidden_states, the backward path is also handled. Fully and robustly achieves all criteria.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..71a3776 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,13 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Sliced hidden_states (e.g. keeping only the last few tokens' logits)
+        # are non-contiguous, which makes the .view(-1, hd) below raise. Force
+        # contiguity so grad_inputs (an empty_like of hidden_states) stays a
+        # real view we can scatter chunk gradients into.
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

The change conditionally makes non-contiguous hidden_states contiguous, which resolves the .view() crash on sliced inputs while leaving the contiguous common case untouched. Since forward stores the contiguous tensor and backward reshapes from it, the gradient path is also robust, and results are preserved for contiguous inputs.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..adbfa0c 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,16 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Sliced hidden_states (e.g. keeping only the last few tokens' logits)
+        # are non-contiguous, which breaks .view(-1, hd) chunking below.
+        # A contiguous copy is safe: the gradient returned in backward only
+        # needs to match the input's shape, not its strides. With overwrite,
+        # gradients land in the copy instead of the caller's buffer, which is
+        # still a valid gradient output.
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+        pass
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

The change robustly fixes the non-contiguous crash by only copying when necessary, preserving the fast contiguous path and disabling the storage-reusing overwrite so gradients also remain valid. Semantics for contiguous inputs are unchanged, and the fix is shape/chunk-agnostic.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..c743488 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,14 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Sliced hidden_states (e.g. keeping only the last few tokens) are
+        # non-contiguous, which makes the later .view(-1, hd) reshape fail.
+        # overwrite reuses hidden_states' storage for gradients, which also
+        # requires a contiguous layout, so disable it in that case.
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+            overwrite = False
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

The change guards with is_contiguous() and only copies when needed, robustly fixing the crash on sliced inputs without altering semantics or degrading the contiguous case. It applies at both the autograd function entry and the wrapper. The gradient path is likely covered since the input tensor is made contiguous upfront, but the diff doesn't explicitly show reshaping of any separate non-contiguous gradient tensor, so slight deduction there.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..99a2591 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,9 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)
@@ -468,6 +471,8 @@ def unsloth_fused_ce_loss(
     device = lm_head_weight.device
     if hidden_states.device != device:
         hidden_states = hidden_states.to(device = device)
+    if not hidden_states.is_contiguous():
+        hidden_states = hidden_states.contiguous()
 
     return apply_autograd_function(UnslothFusedLoss, dict(
         loss_function = compute_fused_ce_loss,

The change conditionally forces contiguity on hidden_states before the view-based reshapes, which is the minimal and correct fix for the crash. It preserves numerical correctness, avoids copies for the common case, and since forward and backward share the contiguous tensor, the gradient path is also robust across shapes and chunk counts.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..275da61 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -189,6 +189,11 @@ class UnslothFusedLoss(torch.autograd.Function):
         device = lm_head_weight.device
         if extra_kwargs is None: extra_kwargs = {}
 
+        # Sliced hidden_states (e.g. keeping only the last few tokens' logits)
+        # are non-contiguous, which breaks the .view(-1, hd) reshapes below.
+        if not hidden_states.is_contiguous():
+            hidden_states = hidden_states.contiguous()
+
         # Get shifted labels first
         if shift_labels:
             _labels = torch.empty_like(labels, device = device)

The change conditionally forces contiguity at both entry points before reshaping, resolving the crash on sliced inputs while preserving correctness and avoiding copies for contiguous tensors. It robustly covers forward and backward paths and all shapes/chunk configurations.

diff --git a/unsloth_zoo/fused_losses/cross_entropy_loss.py b/unsloth_zoo/fused_losses/cross_entropy_loss.py
index dc77e07..4073cf5 100644
--- a/unsloth_zoo/fused_losses/cross_entropy_loss.py
+++ b/unsloth_zoo/fused_losses/cross_entropy_loss.py
@@ -87,6 +87,8 @@ def compute_fused_ce_loss(
     3) logit_softcapping    (X = tanh(X / logit_softcapping) * logit_softcapping)
     """
     device = lm_head_weight.device
+    if not hidden_states.is_contiguous():
+        hidden_states = hidden_states.contiguous()
     if shift_labels:
         # Get shifted labels first
         _labels = torch.empty_like(labels, device = device)
@@ -468,6 +470,8 @@ def unsloth_fused_ce_loss(
     device = lm_head_weight.device
     if hidden_states.device != device:
         hidden_states = hidden_states.to(device = device)
+    if not hidden_states.is_contiguous():
+        hidden_states = hidden_states.contiguous()
 
     return apply_autograd_function(UnslothFusedLoss, dict(
         loss_function = compute_fused_ce_loss,