From a86363eca972118e2c6c4bb91c42810851fc72d6 Mon Sep 17 00:00:00 2001
From: oKatanaaa <whitemarsstudios@gmail.com@>
Date: Thu, 11 Dec 2025 03:21:02 +0000
Subject: [PATCH] fix: weights tying

---
 unsloth/models/llama.py  | 48 ++++++++++++++++++++++++++++++++++++++++
 unsloth/models/vision.py |  1 +
 2 files changed, 49 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 4c9337ccf9..d38018ee1b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2601,6 +2601,7 @@ class FastLlamaModel:
         loftq_config = {},
         temporary_location = "_unsloth_temporary_saved_buffers",
         qat_scheme = None,
+        ensure_weight_tying = False,
         **kwargs,
     ):
         if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1":
@@ -2630,6 +2631,7 @@ class FastLlamaModel:
                 init_lora_weights = init_lora_weights,
                 loftq_config = loftq_config,
                 temporary_location = temporary_location,
+                ensure_weight_tying = ensure_weight_tying,
                 **kwargs,
             )
         if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":
@@ -2953,6 +2955,7 @@ class FastLlamaModel:
             loftq_config = loftq_config,
             use_rslora = use_rslora,
             modules_to_save = modules_to_save,
+            ensure_weight_tying = ensure_weight_tying,
             **kwargs,
         )
         if not SUPPORTS_LOFTQ:
@@ -3002,6 +3005,51 @@ class FastLlamaModel:
 
         model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing)
 
+        if ensure_weight_tying:
+            try:
+                input_embeddings = model.get_input_embeddings()
+                output_embeddings = model.get_output_embeddings()
+
+                if input_embeddings is not None and output_embeddings is not None:
+                    def _retie_parameter(target_module, source_module):
+                        if not hasattr(source_module, "weight"):
+                            return
+                        weight = source_module.weight
+                        # Remove existing registration to avoid "attribute already exists"
+                        if "weight" in getattr(target_module, "_parameters", {}):
+                            target_module._parameters.pop("weight")
+                        if hasattr(target_module, "weight"):
+                            try:
+                                delattr(target_module, "weight")
+                            except Exception:
+                                pass
+                        target_module.register_parameter("weight", weight)
+
+                    # Tie trainable copies created by ModulesToSaveWrapper first (these are used in forward)
+                    if hasattr(input_embeddings, "modules_to_save") and hasattr(
+                        output_embeddings, "modules_to_save"
+                    ):
+                        if hasattr(input_embeddings.modules_to_save, "default") and hasattr(
+                            output_embeddings.modules_to_save, "default"
+                        ):
+                            _retie_parameter(
+                                output_embeddings.modules_to_save.default,
+                                input_embeddings.modules_to_save.default,
+                            )
+
+                    # Tie original_module references as well if present
+                    if hasattr(input_embeddings, "original_module") and hasattr(
+                        output_embeddings, "original_module"
+                    ):
+                        _retie_parameter(
+                            output_embeddings.original_module,
+                            input_embeddings.original_module,
+                        )
+            except Exception as e:
+                logger.warning_once(
+                    f"Unsloth: Failed to ensure weight tying between embeddings and lm_head: {e}"
+                )
+
         if train_embed_tokens:
             print("Unsloth: Training embed_tokens in mixed precision to save VRAM")
             assert hasattr(model.get_input_embeddings(), "modules_to_save")
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index ed19f587cf..9f847f2837 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -930,6 +930,7 @@ class FastBaseModel:
         task_type = TaskType.CAUSAL_LM,
         temporary_location = "_unsloth_temporary_saved_buffers",
         qat_scheme = None,
+        ensure_weight_tying = False,
         **kwargs,
     ):
         if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":