Merge pull request #3711 from oKatanaaa/ensure-weight-tying

FIX: weight tying for LoRA embeddings and lm_head

unsloth/models/llama.py

@@ -2600,6 +2600,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":
@@ -2629,6 +2630,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,55 @@ 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 as exc:
+                                logger.warning_once(
+                                    f"Unsloth: Could not delete existing weight attr during retie on "
+                                    f"{type(target_module).__name__}: {exc}"
+                                )
+                        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")

unsloth/models/vision.py

@@ -938,6 +938,7 @@ class FastBaseModel:
         task_type = TaskType.CAUSAL_LM,
         temporary_location = "_unsloth_temporary_saved_buffers",
         qat_scheme = None,
+        ensure_weight_tying = False,  # [TODO] Add `ensure_weight_tying` for `modules_to_save` for vision models
         **kwargs,
     ):
         if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":