LoRA微调7B模型：单卡24G显存下的数据准备、训练配置与推理效果对比

GPU: RTX 4090 24GB (单卡)
CUDA: 12.1
torch: 2.4.1+cu121
transformers: 4.45.2
peft: 0.13.2
trl: 0.12.1
accelerate: 1.0.1
bitsandbytes: 0.44.1  (备选QLoRA方案用)
datasets: 3.0.1

{
  "instruction": "你是工业设备诊断助手。请根据故障描述输出JSON，字段包括 fault_code(string)、cause(string)、solution(string)、severity(high/medium/low)。",
  "input": "伺服驱动器上电报E-042，电机抖动明显。",
  "output": "{\"fault_code\": \"E-042\", \"cause\": \"编码器信号干扰或接线松动\", \"solution\": \"检查编码器屏蔽线接地，重新插拔CN2接口并锁紧\", \"severity\": \"high\"}"
}

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)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

ds = load_dataset("json", data_files={"train": "data/train.jsonl",
                                      "valid": "data/valid.jsonl"})

PROMPT = ("system\n{instruction}\n"
          "user\n{input}\n"
          "assistant\n")

def preprocess(ex):
    prompt = PROMPT.format(instruction=ex["instruction"], input=ex["input"])
    answer = ex["output"] + tokenizer.eos_token
    p_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    a_ids = tokenizer(answer, add_special_tokens=False)["input_ids"]
    input_ids = p_ids + a_ids
    labels = [-100] * len(p_ids) + a_ids
    # 截断保护
    input_ids = input_ids[:2048]
    labels = labels[:2048]
    return {"input_ids": input_ids, "labels": labels,
            "attention_mask": [1] * len(input_ids)}

train_ds = ds["train"].map(preprocess, remove_columns=ds["train"].column_names)
valid_ds = ds["valid"].map(preprocess, remove_columns=ds["valid"].column_names)

import torch
from transformers import (AutoModelForCausalLM, TrainingArguments, Trainer)
from peft import LoraConfig, get_peft_model, TaskType
from dataclasses import dataclass

@dataclass
class Collator:
    pad_id: int
    def __call__(self, features):
        max_len = max(len(f["input_ids"]) for f in features)
        out = {"input_ids": [], "labels": [], "attention_mask": []}
        for f in features:
            pad = max_len - len(f["input_ids"])
            out["input_ids"].append(f["input_ids"] + [self.pad_id] * pad)
            out["labels"].append(f["labels"] + [-100] * pad)
            out["attention_mask"].append(f["attention_mask"] + [0] * pad)
        return {k: torch.tensor(v, dtype=torch.long) for k, v in out.items()}

model = AutoModelForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, device_map={"": 0},
    trust_remote_code=True, attn_implementation="flash_attention_2")

lora_cfg = LoraConfig(
    task_type=TaskType.CAUSAL_LM, r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj","k_proj","v_proj","o_proj",
                    "gate_proj","up_proj","down_proj"],
    bias="none")
model = get_peft_model(model, lora_cfg)
model.print_trainable_parameters()  # trainable: 40,370,176 || all: 7,655,986,688 || 0.53%

args = TrainingArguments(
    output_dir="out/lora-e042",
    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,
    eval_strategy="steps", eval_steps=200,
    save_strategy="steps", save_steps=200, save_total_limit=3,
    bf16=True, gradient_checkpointing=True,
    optim="adamw_torch", report_to="none",
    dataloader_num_workers=4, remove_unused_columns=False)

trainer = Trainer(model=model, args=args,
                  train_dataset=train_ds, eval_dataset=valid_ds,
                  data_collator=Collator(tokenizer.pad_token_id))
trainer.train()
trainer.save_model("out/lora-e042/final")