"""Merge a LoRA adapter into its base model at bfloat16, and prove the merge did not change it. Purpose: turn the adapter directory a training run produced into an ordinary model that can be converted to GGUF, converted to MLX or served directly. Loads the base at bfloat16 on the CPU, merges, saves the merged model together with the tokeniser that carries the chat template, and optionally generates the same prompts through the unmerged and the merged model to show that they agree. Platform: all (the merge is weight arithmetic, not a forward pass, so it needs no accelerator; the optional comparison runs on whatever device is available). Track M merges MLX adapters with mlx_lm.fuse instead, which the lab describes. Minimum memory: enough system memory for one bfloat16 copy of the base model, so about 4 GB for a 1.7B model and about 8 GB for a 4B one, plus room for the comparison if you use it Assumes: torch, transformers and peft installed; the adapter directory contains adapter_config.json naming the base model it was trained against. Usage: python3 merge-adapter.py --adapter runs/format-qwen3-1.7b python3 merge-adapter.py --adapter runs/format-qwen3-4b \ --merged-dir models/format-qwen3-4b-merged --compare 4 python3 merge-adapter.py --adapter runs/format-qwen3-4b --base-override Qwen/Qwen3-4B Merging into a base that is not the one the adapter was trained against, or into a quantised base, produces a model that loads and runs and is wrong. The default is to use the base named in adapter_config.json; --base-override exists for the case where you moved the base on disk, and the script tells you loudly when you use it. """ from __future__ import annotations import argparse import json from pathlib import Path import torch from peft import PeftConfig, PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer DEFAULT_PROMPTS = [ "The monitoring system reports: the model gateway is unreachable from every machine " "we have tried since this morning's deploy.", "A customer asks when the next release is. Nothing appears to be broken.", "The overnight operator reports: the nightly backup job is filling the disk faster " "than expected every night this week.", "Someone said something is broken. No other detail was given.", ] def generate(model, tokenizer, prompt: str, instruction: str, max_new_tokens: int) -> str: """One greedy generation, so two models given the same input can be compared exactly.""" messages = [{"role": "user", "content": f"{prompt}\n\n{instruction}" if instruction else prompt}] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False).to(model.device) with torch.no_grad(): output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) return tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True) def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--adapter", required=True, help="directory saved by train-lora.py") parser.add_argument("--merged-dir", default=None, help="where to write the merged model; defaults to -merged") parser.add_argument("--base-override", default=None, help="use this base instead of the one named in adapter_config.json") parser.add_argument("--compare", type=int, default=0, help="generate this many prompts through both models and report " "whether the answers agree; 0 skips the check") parser.add_argument("--instruction", default=None, help="instruction appended to each comparison prompt, matching the " "one your training examples used") parser.add_argument("--max-new-tokens", type=int, default=96) args = parser.parse_args() adapter = Path(args.adapter) if not (adapter / "adapter_config.json").is_file(): raise SystemExit(f"{adapter} has no adapter_config.json, so it is not a PEFT adapter " "directory. Point --adapter at what train-lora.py saved.") merged = Path(args.merged_dir) if args.merged_dir else adapter.with_name(adapter.name + "-merged") recorded_base = PeftConfig.from_pretrained(str(adapter)).base_model_name_or_path base_id = args.base_override or recorded_base if args.base_override and args.base_override != recorded_base: print("WARNING: merging into a base that is not the one recorded in the adapter.") print(f" recorded: {recorded_base}") print(f" using: {args.base_override}") print(" If that is not deliberate, stop now: the result will load and be wrong.") print(f"base model: {base_id}") print(f"adapter: {adapter}") print(f"merged to: {merged}") # bfloat16 on the CPU: one copy of the weights, no accelerator, and the precision the # adapter was trained against. Merging into a 4-bit base applies the update to weights # that were already rounded, which is the fault the challenge page reproduces. base = AutoModelForCausalLM.from_pretrained(base_id, dtype=torch.bfloat16, device_map="cpu") peft_model = PeftModel.from_pretrained(base, str(adapter)) tokenizer_source = str(adapter) if (adapter / "tokenizer_config.json").is_file() else base_id if tokenizer_source == base_id: print("WARNING: the adapter directory has no tokeniser, so the base model's is being " "used. If training changed the chat template, this export will not match.") tokenizer = AutoTokenizer.from_pretrained(tokenizer_source) before = [] prompts = DEFAULT_PROMPTS[: args.compare] if args.compare else [] for prompt in prompts: before.append(generate(peft_model, tokenizer, prompt, args.instruction, args.max_new_tokens)) merged_model = peft_model.merge_and_unload() # not in place: the return value is the model merged_model.save_pretrained(str(merged)) tokenizer.save_pretrained(str(merged)) print(f"merged model written to {merged}") if not prompts: print("no comparison requested; pass --compare 4 to check the merge before you convert") return agreed = 0 for prompt, previous in zip(prompts, before): after = generate(merged_model, tokenizer, prompt, args.instruction, args.max_new_tokens) same = after.strip() == previous.strip() agreed += int(same) print(f"\n--- {'MATCH' if same else 'DIFFERS'}: {prompt[:70]}…") if not same: print(f" adapter attached: {previous.strip()[:300]}") print(f" merged: {after.strip()[:300]}") print(f"\n{agreed} of {len(prompts)} prompt(s) matched exactly.") if agreed != len(prompts): print("A merge at the same precision should reproduce the adapter's answers. Differences " "here mean the wrong base, the wrong precision or the wrong tokeniser, and they " "are worth resolving before you spend an hour converting and quantising.") print(json.dumps({"base": base_id, "adapter": str(adapter), "merged": str(merged), "compared": len(prompts), "matched": agreed})) if __name__ == "__main__": main()