From a86363eca972118e2c6c4bb91c42810851fc72d6 Mon Sep 17 00:00:00 2001 From: oKatanaaa Date: Thu, 11 Dec 2025 03:21:02 +0000 Subject: [PATCH 1/4] 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": From 1837de275165b5307b057036c420f2778c6d1343 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Dec 2025 03:31:41 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/llama.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d38018ee1b..e0d8cbcf25 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3011,6 +3011,7 @@ class FastLlamaModel: 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 @@ -3029,9 +3030,9 @@ class FastLlamaModel: 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" - ): + 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, From 7403104b0c05c0794bd8f74342624a22c930a535 Mon Sep 17 00:00:00 2001 From: oKatanaaa Date: Sat, 13 Dec 2025 00:02:48 +0000 Subject: [PATCH 3/4] fix: add a log instead of silent exception --- unsloth/models/llama.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index e0d8cbcf25..6e47907166 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3022,8 +3022,11 @@ class FastLlamaModel: if hasattr(target_module, "weight"): try: delattr(target_module, "weight") - except Exception: - pass + 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) From f7e0f4b152b67479f3b3b889f198daa3b9b28691 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 1 Jan 2026 12:54:21 +0000 Subject: [PATCH 4/4] Add TODO comment for ensure_weight_tying in vision models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- unsloth/models/vision.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 9f847f2837..b4ce718f46 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -930,7 +930,7 @@ class FastBaseModel: task_type = TaskType.CAUSAL_LM, temporary_location = "_unsloth_temporary_saved_buffers", qat_scheme = None, - ensure_weight_tying = False, + 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":