从零开始实战：基于QLoRA微调Llama 3的完整指南

# 创建conda环境
conda create -n qlora python=3.10 -y
conda activate qlora

# 安装核心依赖（严格锁定版本避免兼容性问题）
pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121
pip install transformers==4.41.2 accelerate==0.30.1 peft==0.11.1 bitsandbytes==0.43.1
pip install datasets==2.19.1 trl==0.9.4 wandb==0.17.0

python -c "import bitsandbytes as bnb; print(bnb.__version__); print(bnb.cuda_available)"

from datasets import Dataset
import json

# 模拟100条医疗指令数据（实际项目请替换为真实数据）
raw_data = [
    {"instruction": "患者主诉胸闷气短，伴随心悸，该如何初步处理？", 
     "output": "建议立即休息，监测血压心率，若持续不缓解需急诊排查心梗或肺栓塞。可给予吸氧并做心电图。"},
    # ... 更多数据
]

# 转换为Alpaca模板
def format_example(example):
    return {
        "text": f"### 指令：{example['instruction']}\n### 回答：{example['output']}"
    }

dataset = Dataset.from_list(raw_data).map(format_example)
dataset = dataset.train_test_split(test_size=0.1)
print(dataset)

import torch
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    BitsAndBytesConfig
)

model_name = "meta-llama/Meta-Llama-3-8B"  # 需huggingface token

# 4-bit量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # NormalFloat 4-bit
    bnb_4bit_use_double_quant=True,      # 双重量化
    bnb_4bit_compute_dtype=torch.bfloat16 # 计算时使用bfloat16加速
)

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",                   # 自动分配到多GPU或CPU offload
    trust_remote_code=False,
    use_cache=False                      # 训练时关闭KV cache
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token  # 设置padding token
tokenizer.padding_side = "right"

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# 准备模型用于k-bit训练（启用梯度检查点，节省显存）
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

# LoRA配置
lora_config = LoraConfig(
    r=16,                     # 秩（rank），越大表达能力越强，但显存增加
    lora_alpha=32,            # 缩放系数，通常设为2*r
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Llama 3的注意力投影层
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 应显示约 0.1% 参数可训练（8B模型约8.4M参数）

from transformers import TrainingArguments, Trainer
from trl import SFTTrainer  # 或使用标准Trainer

training_args = TrainingArguments(
    output_dir="./llama3-qlora-medical",
    per_device_train_batch_size=2,      # 24GB显存建议2，若OOM可降至1
    gradient_accumulation_steps=8,       # 有效batch size = 2*8=16
    num_train_epochs=3,
    learning_rate=2e-4,
    fp16=True,                           # 混合精度
    logging_steps=10,
    save_steps=100,
    save_total_limit=2,
    optim="paged_adamw_8bit",            # 关键：分页优化器
    report_to="wandb",                   # 可选，用wandb监控
    remove_unused_columns=True,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
)

# 使用SFTTrainer（自动处理数据整理）
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    max_seq_length=512,
    dataset_text_field="text",
    packing=False,                       # 关闭packing，避免长度不一致影响
)

trainer.train()

import wandb
wandb.init(project="llama3-qlora", name="medical-finetune")
# 在TrainingArguments中已设置 report_to="wandb"

# 保存适配器
model.save_pretrained("./lora-adapter-medical")
tokenizer.save_pretrained("./lora-adapter-medical")

# 合并到原模型（生成完整权重，用于高效推理）
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
merged_model = PeftModel.from_pretrained(base_model, "./lora-adapter-medical")
merged_model = merged_model.merge_and_unload()  # 合并LoRA权重到基础模型

# 保存合并后的模型
merged_model.save_pretrained("./llama3-medical-merged")
tokenizer.save_pretrained("./llama3-medical-merged")

def generate_response(prompt, model, tokenizer, max_new_tokens=256):
    inputs = tokenizer(f"### 指令：{prompt}\n### 回答：", return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=0.7,
        do_sample=True,
        top_p=0.9,
        repetition_penalty=1.1
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 测试样例
test_prompt = "患者反复头痛一周，服用布洛芬无效，下一步建议？"
print("微调前（基础模型）：")
# 加载未微调模型对比...
print("微调后：")
print(generate_response(test_prompt, merged_model, tokenizer))

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch

app = FastAPI()
model = AutoModelForCausalLM.from_pretrained("./llama3-medical-merged", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("./llama3-medical-merged")

class Query(BaseModel):
    prompt: str
    max_tokens: int = 256

@app.post("/generate")
def generate(query: Query):
    inputs = tokenizer(query.prompt, return_tensors="pt").to("cuda")
    out = model.generate(**inputs, max_new_tokens=query.max_tokens)
    return {"response": tokenizer.decode(out[0], skip_special_tokens=True)}