# =========================================================================================
# Fine-tuning script based on https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_%281B_and_3B%29-Conversational.ipynb
# This script and HF checkpoint are only intended to showcase how to do finetuning in a way compatible with ExecuTorch
# Only 10 steps are done, and quality of the finetuned model is not evaluated
# =========================================================================================

from unsloth import FastLanguageModel
from unsloth.chat_templates import (
    get_chat_template,
    standardize_data_formats,
    standardize_sharegpt,
    train_on_responses_only,
)

from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
from transformers import DataCollatorForSeq2Seq
import torch
import torch.nn as nn

batch_size = 2
learning_rate = 2e-5
gradient_accumulation_steps = 4
max_steps = 10
full_finetuning = True
qat_scheme = "int8-int4"
output_dir = "/tmp/unsloth_example"


model_id = "unsloth/Qwen3-4B"
chat_template = "qwen3"
max_seq_length = 2048
dtype = torch.bfloat16
load_in_4bit = False

################################################################################
# Define model/tokenizer
################################################################################

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_id,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit =load_in_4bit,
    full_finetuning=full_finetuning,
    qat_scheme=qat_scheme,
)
tokenizer = get_chat_template(tokenizer, chat_template = chat_template)

print("MODEL AFTER LOADING")
print(model)

################################################################################
# Untie model weights
################################################################################

def untie_word_embeddings_(model):
    """Untie input and output embeddings in a Hugging Face causal LM."""
    # 1) Persist setting in config
    if hasattr(model.config, "tie_word_embeddings"):
        model.config.tie_word_embeddings = False

    # 2) Find input and output embeddings
    in_emb = model.get_input_embeddings()         # nn.Embedding
    out_proj = model.get_output_embeddings() or getattr(model, "lm_head", None)
    if out_proj is None:
        raise AttributeError("Couldn't locate output projection (lm_head).")

    # (Optional) sanity: shapes should match [vocab, hidden]
    assert out_proj.weight.shape == in_emb.weight.shape, (
        f"Shape mismatch: out_proj {out_proj.weight.shape} vs in_emb {in_emb.weight.shape}"
    )

    # 3) Only clone if they are actually tied (shared storage)
    if out_proj.weight.data_ptr() == in_emb.weight.data_ptr():
        with torch.no_grad():
            W = in_emb.weight.detach().clone()
        out_proj.weight = nn.Parameter(W)  # new storage, keeps dtype/device

    # 4) Prevent future automatic re-tying
    def _no_tie(self):
        return
    model.tie_weights = _no_tie.__get__(model, model.__class__)

    # 5) Verify no shared storage
    assert out_proj.weight.data_ptr() != in_emb.weight.data_ptr(), "Embeddings still tied!"

    return model

model = untie_word_embeddings_(model)

print("MODEL AFTER UNTYING")
print(model)
print(model.config)


################################################################################
# Process dataset
################################################################################

def formatting_prompts_func(examples):
    convos = examples["conversations"]
    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
    return { "text" : texts, }
dataset = load_dataset("mlabonne/FineTome-100k", split = "train")
dataset = standardize_sharegpt(dataset)
dataset = dataset.map(formatting_prompts_func, batched = True,)

print("DATASET ENTRY")
print(dataset[0])
print("\n\n")

################################################################################
# Define trainer
################################################################################

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    packing=False,
    args=SFTConfig(
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        warmup_steps=5,
        num_train_epochs=1,
        max_steps=max_steps,
        learning_rate=learning_rate,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
        report_to="none",
    ),
)

################################################################################
# Do fine tuning
################################################################################
print("DOING FINETUNING")
trainer_stats = trainer.train()
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(
    f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training."
)

################################################################################
# Save model
################################################################################
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

################################################################################
# Convert model
################################################################################
from torchao.quantization import Int8DynamicActivationIntxWeightConfig, IntxWeightOnlyConfig, ModuleFqnToConfig, quantize_
from torchao.quantization.qat import QATConfig
from torchao.quantization.granularity import PerGroup, PerAxis
from transformers import TorchAoConfig

base_config = Int8DynamicActivationIntxWeightConfig(weight_dtype=torch.int4, weight_granularity=PerGroup(32))
quantize_(model, QATConfig(base_config, step="convert"))
 
################################################################################
# Quantize embeddings to 8-bit with PTQ since they are not supported by QAT yet
################################################################################

embedding_fqn = "model.embed_tokens"
embedding_config = IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0))
quantize_(model, embedding_config, lambda m, fqn: fqn == embedding_fqn)

################################################################################
# Attach quantization config to model
################################################################################

quant_config = ModuleFqnToConfig({"_default": base_config, embedding_fqn: embedding_config})
quantization_config = TorchAoConfig(quant_type=quant_config, include_input_output_embeddings=True, modules_to_not_convert=[])
model.config.quantization_config = TorchAoConfig(base_config)

print('MODEL AFTER CONVERT', model)

################################################################################
# Push converted model to hub
################################################################################
from huggingface_hub import get_token, whoami

def _get_username():
    token = get_token()
    username = whoami(token=token)["name"]
    return username

username = _get_username()
model_name = model_id.split("/")[-1]
save_to = f"{username}/{model_name}-{qat_scheme}-unsloth"
model.push_to_hub(save_to, safe_serialization=False)
tokenizer.push_to_hub(save_to)

################################################################################
# Load converted from hub and inspect
################################################################################
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(save_to)
print('model', model)
print("model.embed_tokens.weight", model.model.embed_tokens.weight)
print("model.layers[0].self_attn.q_proj.weight", model.model.layers[0].self_attn.q_proj.weight)
print("lm_head.weight", model.lm_head.weight)

Exporting to ExecuTorch

We can run the quantized model on a mobile phone using ExecuTorch. We first set up ExecuTorch:

git clone https://github.com/pytorch/executorch.git
pushd executorch
git checkout release/1.0
git submodule update --init --recursive
python install_executorch.py
popd

ExecuTorch's LLM export scripts require the checkpoint keys and parameters have certain names, which differ from those used in Hugging Face. So we first use a script that converts the Hugging Face checkpoint key names to ones that ExecuTorch expects: The following script does this for you.

Shell

pushd executorch
python -m executorch.examples.models.qwen3.convert_weights $(hf download metascroy/Qwen3-4B-int8-int4-unsloth) pytorch_model_converted.bin
popd

Once we have the checkpoint, we export it to ExecuTorch with a max_seq_length/max_context_length of 1024 to the XNNPACK backend as follows.

(Note: ExecuTorch LLM export script requires config.json have certain key names. The correct config to use for the LLM export script is located at examples/models/qwen3/config/4b_config.json within the ExecuTorch repo.)

Shell

pushd executorch
python -m executorch.examples.models.llama.export_llama \
  --model "qwen3_4b" \
  --checkpoint pytorch_model_converted.bin \
  --params examples/models/qwen3/config/4b_config.json \
  --output_name model.pte \
  -kv \
  --use_sdpa_with_kv_cache \
  -X \
  --xnnpack-extended-ops \
  --max_context_length 1024 \
  --max_seq_length 1024 \
  --dtype fp32 \
  --metadata '{"get_bos_id":199999, "get_eos_ids":[200020,199999]}'
popd

We can also uplaod the pte file to Hugging Face as follows:

hf upload metascroy/Qwen3-4B-int8-int4-unsloth executorch/model.pte

(We try to keep these instructions up-to-date, but if you find they do not work, check out our CI test in ExecuTorch for the latest source of truth, and let us know we need to update our model card.)

Running in a mobile app

After that you can run the model in an iOS mobile app using the executorch-examples repo.

First clone the repo and open the Xcode project:

git clone https://github.com/meta-pytorch/executorch-examples.git
open executorch-examples/llm/apple/etLLM.xcodeproj

Once open, connect your iPhone and select it as the device in Xcode. To build the app, press the play button. (This does require you sign the code.)

To run the model we just finetuned, you need to transfer the model.pte file and tokenizer.json to your phone. To do this, open finder, select your phone, and drag and drop the files to copy them over. Please rename the pte file to qwen3_model.pte before copying it over because the demo app requires the name begin with qwen3 to use the correct prompt template.

In the etLLM app, you can select the model and tokenizer to use by browsing your file system for the "qwen3_model.pte" and "tokenizer.json" files we just copied to the phone.