Fine-tuning with LoRA (train.py)

{"prompt": "...", "response": "..."}

Merge LoRA into a base model (merge.py)

# Merge only
python .\merge.py --base Qwen/Qwen2.5-3B-Instruct --adapter C:\Source\research\qwen2.5-3b-lora --out C:\Source\research\qwen2.5-3b-merged --cpu-only --dtype fp32

# Merge and test generation
python .\merge.py --base Qwen/Qwen2.5-0.5B-Instruct --adapter C:\Source\research\qwen2.5-3b-lora --out C:\Source\research\qwen2.5-3b-merged --cpu-only --dtype fp32 --infer

Web UI for fine-tuning and merging

Fine-tuning

Merging

from __future__ import annotations

import argparse
import os

import torch
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Fine-tune a causal LM with LoRA/QLoRA.")
    parser.add_argument("--base-model", default=os.getenv("BASE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct"))
    parser.add_argument("--data", default=os.getenv("DATA_PATH", "data.jsonl"))
    parser.add_argument("--out", default=os.getenv("OUT_DIR", "qwen2.5-lora"))
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--max-steps", type=int, default=-1)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=2e-4)
    parser.add_argument("--lora-r", type=int, default=8)
    parser.add_argument("--lora-alpha", type=int, default=16)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument("--max-seq-length", type=int, default=1024)
    parser.add_argument("--use-4bit", action="store_true")
    return parser.parse_args()


def positive(value: float | int, name: str) -> None:
    if value <= 0:
        raise SystemExit(f"{name} must be greater than zero")


def main() -> None:
    args = parse_args()
    for value, name in (
        (args.epochs, "epochs"),
        (args.batch_size, "batch-size"),
        (args.gradient_accumulation, "gradient-accumulation"),
        (args.learning_rate, "learning-rate"),
        (args.lora_r, "lora-r"),
        (args.lora_alpha, "lora-alpha"),
        (args.max_seq_length, "max-seq-length"),
    ):
        positive(value, name)
    if not 0 <= args.lora_dropout < 1:
        raise SystemExit("lora-dropout must be in [0, 1)")

    use_cuda = torch.cuda.is_available()
    use_4bit = args.use_4bit and use_cuda
    print(f"Device: {'CUDA - ' + torch.cuda.get_device_name(0) if use_cuda else 'CPU'}")
    if args.use_4bit and not use_cuda:
        print("4-bit QLoRA requested but CUDA is unavailable; using regular LoRA on CPU.")

    quantization = (
        BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        if use_4bit
        else None
    )
    tokenizer = AutoTokenizer.from_pretrained(
        args.base_model, use_fast=True, trust_remote_code=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        quantization_config=quantization,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16 if use_cuda else torch.float32,
        device_map="auto" if use_4bit else None,
    )
    if use_4bit:
        model = prepare_model_for_kbit_training(model)
    model.config.use_cache = False
    model = get_peft_model(
        model,
        LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "v_proj"],
        ),
    )
    model.print_trainable_parameters()

    dataset = load_dataset("json", data_files=args.data, split="train")

    def tokenize(samples):
        texts = [
            f"### Instruction:\n{prompt}\n\n### Response:\n{response}{tokenizer.eos_token}"
            for prompt, response in zip(samples["prompt"], samples["response"])
        ]
        return tokenizer(
            texts,
            truncation=True,
            max_length=args.max_seq_length,
            padding=False,
        )

    dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
    training_args = TrainingArguments(
        output_dir=args.out,
        num_train_epochs=args.epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation,
        learning_rate=args.learning_rate,
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        logging_steps=1,
        save_strategy="epoch",
        bf16=use_cuda and torch.cuda.is_bf16_supported(),
        fp16=use_cuda and not torch.cuda.is_bf16_supported(),
        max_grad_norm=0.3,
        gradient_checkpointing=use_cuda,
        optim="paged_adamw_8bit" if use_4bit else "adamw_torch",
        dataloader_pin_memory=use_cuda,
        report_to="none",
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
    )
    trainer.train()
    trainer.model.save_pretrained(args.out)
    tokenizer.save_pretrained(args.out)
    print(f"Adapter saved to {args.out}")


if __name__ == "__main__":
    main()