LoRA微调7B大模型：从数据准备到QLoRA量化训练全记录

{
  "system": "你是一名工业设备运维专家，回答要准确、简洁、可执行。",
  "question": "XX型号轴承更换扭矩是多少？",
  "answer": "该型号轴承安装扭矩为45N·m，分三次拧紧：15N·m→30N·m→45N·m，每次间隔对角拧紧。"
}

from datasets import load_dataset
from transformers import AutoTokenizer

MODEL_PATH = "Qwen/Qwen2.5-7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

def format_sample(sample):
    messages = [
        {"role": "system", "content": sample["system"]},
        {"role": "user", "content": sample["question"]},
        {"role": "assistant", "content": sample["answer"]},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    # 只对assistant部分计算loss
    prompt_messages = messages[:2]
    prompt_text = tokenizer.apply_chat_template(prompt_messages, tokenize=False, add_generation_prompt=True)
    return {"text": text, "prompt_text": prompt_text}

dataset = load_dataset("json", data_files={"train": "train.json", "valid": "valid.json"})
dataset = dataset.map(format_sample, remove_columns=dataset["train"].column_names)

def tokenize_fn(sample):
    tokenized = tokenizer(sample["text"], truncation=True, max_length=1024, padding=False)
    prompt_ids = tokenizer(sample["prompt_text"], truncation=True, max_length=1024, padding=False)["input_ids"]
    labels = tokenized["input_ids"].copy()
    # mask掉prompt部分
    for i in range(min(len(prompt_ids), len(labels))):
        labels[i] = -100
    tokenized["labels"] = labels
    return tokenized

tokenized_dataset = dataset.map(tokenize_fn, remove_columns=["text", "prompt_text"])

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer

MODEL_PATH = "Qwen/Qwen2.5-7B-Instruct"

# 4-bit量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False

# LoRA配置
lora_config = LoraConfig(
    r=64,
    lora_alpha=128,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)

training_args = TrainingArguments(
    output_dir="./qwen2.5-7b-lora-ops",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    num_train_epochs=3,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    logging_steps=10,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    bf16=True,
    optim="paged_adamw_8bit",
    gradient_checkpointing=True,
    report_to="none",
    max_grad_norm=0.3,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["valid"],
    peft_config=lora_config,
    dataset_text_field="input_ids",  # 已tokenize，这里用data_collator
    max_seq_length=1024,
    packing=False,
)

trainer.train()
trainer.model.save_pretrained("./qwen2.5-7b-lora-ops/final")