From 665a8e4b1d6a818dd01a7b9cf78dafff4fa322c0 Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Sun, 13 Jul 2025 04:23:07 +0530
Subject: [PATCH] Fix falcon H1 dropout issue (#2938)

Because we don't have down and gate multipliers, the MLP output values are too huge, causing NaN and unstable training. To bypass that lets rely on HF's implementation for the time being
---
 unsloth/models/llama.py | 57 +++++++++++++++++++----------------------
 1 file changed, 26 insertions(+), 31 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index aae86080a8..0c6f6c1e86 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2768,40 +2768,35 @@ class FastLlamaModel:
         if lora_dropout == 0 and bias == "none":
             for idx, layer in enumerate(model.model.model.layers):
 
-                # Determine MLP module name (falcon_h1 has feed_forward, llama style has mlp)
-                if hasattr(layer, "mlp"):
-                    mlp_module_name = "mlp"
-                elif hasattr(layer, "feed_forward"):
-                    mlp_module_name = "feed_forward"
-                else:
-                    logger.warning_once(f"Unsloth: No MLP module found in layer {idx} so skipping peft mlp patching")
-                    continue
+                if model_type != "falcon_h1":
+                    # LoRAMLP.apply doesn't have functionality for gate and down mutlipliers yet.
+                    # Don't patch falcon h1 for the time being.
 
-                mlp_module = getattr(layer, mlp_module_name)
+                    # MLP patching
+                    mlp_module = layer.mlp
+                    gate_proj = mlp_module.gate_proj
+                    up_proj   = mlp_module.  up_proj
+                    down_proj = mlp_module.down_proj
 
-                # MLP patching
-                gate_proj = mlp_module.gate_proj
-                up_proj   = mlp_module.  up_proj
-                down_proj = mlp_module.down_proj
+                    if hasattr(gate_proj, "lora_A") and \
+                        hasattr(  up_proj, "lora_A") and \
+                        hasattr(down_proj, "lora_A") and \
+                        (getattr(gate_proj, "base_layer", gate_proj).bias is None) and \
+                        (getattr(  up_proj, "base_layer",   up_proj).bias is None) and \
+                        (getattr(down_proj, "base_layer", down_proj).bias is None) and \
+                        (len(getattr(gate_proj, "lora_magnitude_vector", []) or []) == 0) and \
+                        (len(getattr(  up_proj, "lora_magnitude_vector", []) or []) == 0) and \
+                        (len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0):
 
-                if  hasattr(gate_proj, "lora_A") and \
-                    hasattr(  up_proj, "lora_A") and \
-                    hasattr(down_proj, "lora_A") and \
-                    (getattr(gate_proj, "base_layer", gate_proj).bias is None) and \
-                    (getattr(  up_proj, "base_layer",   up_proj).bias is None) and \
-                    (getattr(down_proj, "base_layer", down_proj).bias is None) and \
-                    (len(getattr(gate_proj, "lora_magnitude_vector", []) or []) == 0) and \
-                    (len(getattr(  up_proj, "lora_magnitude_vector", []) or []) == 0) and \
-                    (len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0):
-
-                    # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
-                    mlp_module.forward = types.MethodType(_apply_lora_mlp, mlp_module)
-                    n_mlp += 1
-                else:
-                    logger.warning_once(
-                        "Not an error, but Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\n"\
-                        "are not enabled or a bias term (like in Qwen) is used."
-                    )
+                        # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
+                        mlp_module.forward = types.MethodType(_apply_lora_mlp, mlp_module)
+                        n_mlp += 1
+                    else:
+                        logger.warning_once(
+                            "Not an error, but Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\n"\
+                            "are not enabled or a bias term (like in Qwen) is used."
+                        )
+                    pass
                 pass
 
                 # QKV attention patching