SoFarSoGoodya/DeepMath-R1-Distill-Qwen-7B

1#!/usr/bin/env python32"""3DeepMath — merge LoRA adapters into a single, directly-loadable full model.4 5DeepMath is produced by a TWO-STAGE LoRA pipeline (LLaMA-Factory, template `deepseekr1`):6 7    DeepSeek-R1-Distill-Qwen-7B  (base)8      + SFT LoRA adapter   --(trained on cleaned NuminaMath-CoT)-->9        DeepMath-SFT  (merged intermediate)10      + DPO LoRA adapter --(trained on math preference pairs)---->11        DeepMath (final)12 13This repository ships the two adapters (small). Run this script ONCE on a14machine with the base model available to reconstruct the final full model:15 16    python merge.py \17        --base  deepseek-ai/DeepSeek-R1-Distill-Qwen-7B \18        --sft   ./sft_adapter \19        --dpo   ./adapter \20        --out   ./DeepMath-merged21 22Notes23-----24* Merging is pure PEFT `merge_and_unload` (no LLaMA-Factory needed at merge time).25* CPU is fine but slow; a GPU with ~16GB+ is comfortable. Weights are bf16.26* The original training merged with `llamafactory-cli export --template deepseekr1`;27  both approaches yield the same weights for these adapters.28"""29 30import argparse31import torch32from transformers import AutoModelForCausalLM, AutoTokenizer33from peft import PeftModel34 35 36def merge(base_path: str, sft_adapter: str, dpo_adapter: str, out: str, device: str, dtype: str):37    torch_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[dtype]38    print(f"[1/4] Loading base model: {base_path}  ({dtype}, device={device})")39    model = AutoModelForCausalLM.from_pretrained(40        base_path, torch_dtype=torch_dtype, device_map=device41    )42    tokenizer = AutoTokenizer.from_pretrained(base_path)43 44    print(f"[2/4] Merging SFT adapter: {sft_adapter}")45    model = PeftModel.from_pretrained(model, sft_adapter)46    model = model.merge_and_unload()47 48    print(f"[3/4] Merging DPO adapter: {dpo_adapter}")49    model = PeftModel.from_pretrained(model, dpo_adapter)50    model = model.merge_and_unload()51 52    print(f"[4/4] Saving merged model -> {out}")53    model.save_pretrained(out, safe_serialization=True)54    tokenizer.save_pretrained(out)55    print("Done. Load with: AutoModelForCausalLM.from_pretrained('%s')" % out)56 57 58if __name__ == "__main__":59    ap = argparse.ArgumentParser(description="Merge DeepMath SFT+DPO LoRA adapters into the base model.")60    ap.add_argument("--base", default="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",61                    help="Base model path or HF id.")62    ap.add_argument("--sft", default="./sft_adapter", help="Path to the SFT LoRA adapter.")63    ap.add_argument("--dpo", default="./adapter", help="Path to the final DPO LoRA adapter.")64    ap.add_argument("--out", default="./DeepMath-merged", help="Output directory for the merged model.")65    ap.add_argument("--device", default="auto", help="device_map: 'auto', 'cpu', 'cuda:0', ...")66    ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"], help="Weight dtype.")67    args = ap.parse_args()68    merge(args.base, args.sft, args.dpo, args.out, args.device, args.dtype)69