Summary of Changes

Hello @oKatanaaa, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves an issue where weight tying between token embeddings and the language model head was being disrupted when using LoRA with Unsloth's offloading mechanism, particularly for models like Qwen3-4B. The changes introduce a new configuration option and explicit re-tying logic to preserve the shared storage of these weights, preventing their divergence during training and ensuring the integrity of models that rely on this architectural constraint.

Highlights

• New ensure_weight_tying Parameter: Introduced a new ensure_weight_tying parameter to the FastLanguageModel.get_peft_model function in both llama.py and vision.py, allowing explicit control over weight tying behavior during LoRA fine-tuning.
• Re-tying Logic for Embeddings and LM Head: Implemented a mechanism to re-tie embed_tokens (input embeddings) and lm_head (output embeddings) after PEFT wrapping. This ensures that Unsloth's offloading and duplication processes do not inadvertently break the shared storage required for tied weights, which was causing trainable copies to diverge.
• Vision Model Compatibility: Ensured that the vision.py module is compatible with the new ensure_weight_tying parameter by passing it through to the underlying LoRA configuration, maintaining consistent behavior across different model types.

Here is a simple script to reproduce the issue btw:

import torch
from unsloth import FastLanguageModel

# Assumes unsloth without the post-PEFT retie fix. ensure_weight_tying is passed via kwargs.
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Qwen3-4B-Base-unsloth-bnb-4bit",
    max_seq_length=256,
    dtype=None,
    load_in_4bit=True,
)

def ptr(mod):
    return None if mod is None else mod.weight.data_ptr()

def check(stage):
    in_emb = model.get_input_embeddings()
    out_emb = model.get_output_embeddings()
    in_def = getattr(getattr(in_emb, "modules_to_save", {}), "default", None)
    out_def = getattr(getattr(out_emb, "modules_to_save", {}), "default", None)
    print(f"{stage}: base tied? {ptr(in_emb)==ptr(out_emb)}, "
          f"default tied? {ptr(in_def)==ptr(out_def)}")

check("before PEFT")

model = FastLanguageModel.get_peft_model(
    model,
    r=8,
    target_modules=["q_proj","k_proj","v_proj","o_proj"],
    modules_to_save=["embed_tokens","lm_head"],
    ensure_weight_tying=True,              # flag is set, but no fix present
    use_gradient_checkpointing="unsloth",  # triggers offload
    max_seq_length=256,
)

check("after PEFT (expect untied without fix)")

hey @oKatanaaa can you please tell me the motivation to tie the embedding and lm_head weights for a model that didn't have it to begin with?

hey @oKatanaaa can you please tell me the motivation to tie the embedding and lm_head weights for a model that didn't have it to begin with?

I'm not sure. This PR is not enabling this functionality though. Even if ensure_weight_tying=True, PEFT won't tie weights if the model's embeddings are untied initially (huggingface/peft#2864).

Some advanced users could enforce weight tying to intentionally reduce model size (usually word embeddings take ~13% of total model size, tying would reduce it by ~6% which could be significant).

Thanks for the PR and Happy New Year!

I tested this on a B200 with the following configurations:

Test Results:

Weight tying verification:

• Before PEFT: base embeddings tied = True
• After PEFT with ensure_weight_tying=True: base tied = True, trainable copies tied = True
• After training: weights remain tied = True

The fix correctly re-ties the trainable copies after PEFT wrapping, which preserves weight sharing during Unsloth's offload/rebuild path.

I also added a commit with a TODO comment for vision.py since the parameter is added but not yet implemented for vision models.

Sidenote: This PR was reviewed automatically by the Unsloth Code Review Bot.

Thanks @oKatanaaa again :) Was trying out our new auto review system as well!