model.save_pretrained_merged("merged", tokenizer) Not functioning properly

[SKillSeven]

My English is not good and I'm a novice in the AI field. Perhaps it's my fault, but I still hope to inform you. Even though it's my fault, it might still be able to help those who are in a similar situation to me.

After training the model, I used model.save_pretrained_merged("merged", tokenizer)
Attempting to save the fine-tuned model.
However, after reloading the model of this path, when the same question was asked, the answer obtained during the training process was not returned.
Ultimately, I discovered that
model.save_pretrained_merged("merged", tokenizer)
Requires coordination
model.load_adapter("xxxxx/xxx")
Be capable of effectively loading the trained model's performance

Attach the source code for my training and loading process
train：

from mlx_tune import FastLanguageModel, SFTTrainer


max_seq_length = 2048  # 可以自定义任意窗口长度，已根据RoPE编码自动伸缩模型窗口尺寸了。
dtype = None  # 设置为None自动获取。目前 Float16 支持GPU类型：Tesla T4, V100； Bfloat16 支持GPU类型： Ampere+
load_in_4bit = True  # 使用4bit量化以减少内存使用。可以设置为False。

# 调用 unsloth 预先量化好的4bit模型
fourbit_models = [
    "unsloth/mistral-7b-bnb-4bit",
    "unsloth/mistral-7b-instruct-v0.2-bnb-4bit",
    "unsloth/llama-2-7b-bnb-4bit",
    "unsloth/gemma-7b-bnb-4bit",
    "unsloth/gemma-7b-it-bnb-4bit",  # Instruct version of Gemma 7b
    "unsloth/gemma-2b-bnb-4bit",
    "unsloth/gemma-2b-it-bnb-4bit",  # Instruct version of Gemma 2b
    "unsloth/llama-3-8b-bnb-4bit",  # [NEW] 15 Trillion token Llama-3
]

# 调用 unsloth 预先量化好的4bit模型
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="mlx-community/llama-3.2-1b-Instruct-4bit",
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,  # 以4bit加载模型
)

# 以下是PEFT模型的默认参数，您可以根据需要进行调整。
model = FastLanguageModel.get_peft_model(
    model,
    r=16,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj", ],
    lora_alpha=16,
    lora_dropout=0,  # Supports any, but = 0 is optimized
    bias="none",  # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    # use_gradient_checkpointing="unsloth",  # True or "unsloth" for very long context
    use_gradient_checkpointing=True,
    random_state=3407,
    use_rslora=False,  # We support rank stabilized LoRA
    loftq_config=None,  # And LoftQ
)

alpaca_prompt = """
### Instruction:
{}

### Input:
{}

### Response:
{}"""

EOS_TOKEN = tokenizer.eos_token  # 必须添加 EOS_TOKEN 这个特殊符号，否则生成会无限循环。。


def formatting_prompts_func(examples):
    instructions = examples["instruction_zh"]
    inputs = examples["input_zh"]
    outputs = examples["output_zh"]
    texts = []
    for instruction, input, output in zip(instructions, inputs, outputs):
        text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN  # 必须添加 EOS_TOKEN 这个特殊符号，否则生成会无限循环。
        texts.append(text)
    return {"text": texts, }


# 从数据集中加载数据
from datasets import load_dataset

dataset = load_dataset(
    "json",
    data_files="/xxxxx/xxxxtempTrain.jsonl",
    split="train"
)
dataset = dataset.map(formatting_prompts_func, batched=True, )

# 以下是PEFT模型的默认参数，您可以根据需要进行调整。
# from trl import SFTTrainer
from transformers import TrainingArguments

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    packing=False,  # 可以让小上下文窗口训练速度增加5倍以上
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=100,  # 微调循环次数
        learning_rate=1e-5,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=1337,
        # output_dir="outputs",
    ),
)

# 训练模型
trainer_stats = trainer.train()

# 推理
FastLanguageModel.for_inference(model) # Enable native 2x faster inference


tempContent =  alpaca_prompt.format(
            "说明xxxx公司有哪些人员",  # instruction
            "",  # input
            "",  # output - leave this blank for generation!
        )

from mlx_lm import generate
response = generate(
    model.model,
    tokenizer,
    prompt=tempContent,
    max_tokens=300,
    verbose=False,
)
print(f"   Response: {response}")

# 实际测试有问题，无法保存进初始模型内部
model.save_pretrained_merged("merged", tokenizer)  # Full model

inference:

from mlx_tune import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="/xxx/merged",
)

alpaca_prompt = """
### Instruction:
{}

### Input:
{}

### Response:
{}"""

# 推理# Enable native 2x faster inference
# 推理
FastLanguageModel.for_inference(model)  # Enable native 2x faster inference


tempContent =  alpaca_prompt.format(
            "说明xxxxx公司有哪些人员",  # instruction
            "",  # input
            "",  # output - leave this blank for generation!
        )


from mlx_lm import generate
response = generate(
    model.model,
    tokenizer,
    prompt=tempContent,
    max_tokens=300,
    verbose=False,
)

print(f"   Response: {response}")

If it's my question and you have spare time, could you please explain to me why it's incorrect? It would be great if you could provide a brief example. Thank you very much.

[ARahim3]

Hey @SKillSeven, thanks for the super detailed writeup, and no, this isn't your fault. You found a real bug.

What's going on: you fine-tuned a 4-bit base (llama-3.2-1b-Instruct-4bit). When save_pretrained_merged fused your LoRA weights, it was re-quantizing the result back to 4-bit. Your training run was pretty light (100 steps at lr=1e-5), so the change LoRA learned was small, small enough that the 4-bit rounding basically erased it. That's why the merged model "forgot" everything. load_adapter worked for you because the adapter stays at full precision on top of the frozen 4-bit base, so nothing gets rounded off.

I just pushed a fix: save_pretrained_merged(..., save_method="merged_16bit") (the default) now dequantizes the base before merging and saves a proper 16-bit model, so the fine-tune is kept. After you update, your original code will just work — no load_adapter dance needed.

Two other things that'll help your results regardless:

• lr=1e-5 for 100 steps is quite low for LoRA. Try lr=2e-4, you'll see a much stronger effect.
• If you'd rather keep the small 4-bit file, save_pretrained_merged(..., save_method="merged_4bit") still works, just know a weak fine-tune may not survive it. For 4-bit, sticking with the adapter + load_adapter is the safest.

Thanks again for reporting it. This genuinely helps other people hitting the same thing. 🙏