Llama2-chat微调：少量数据注入特定知识的局限性与有效策略

import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
from peft import LoraConfig, PeftModel, get_peft_model
from trl import SFTTrainer
import bitsandbytes as bnb

# 1. 数据集加载
dataset_name = "celsowm/guanaco-llama2-1k1" # 包含新增问题的自定义数据集
dataset = load_dataset(dataset_name, split="train")

# 2. 模型与分词器加载
model_id = "NousResearch/Llama-2-7b-chat-hf"
compute_dtype = getattr(torch, "float16")

# 配置4位量化以节省显存
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token # 设置pad token为eos token
tokenizer.padding_side = "right" # 设置padding方向

# 动态分配显存
n_gpus = torch.cuda.device_count()
max_memory = torch.cuda.get_device_properties(0).total_memory
max_memory = f'{max_memory // (1024**2)}MB' # 转换为MB字符串
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map='auto',
    max_memory={i: max_memory for i in range(n_gpus)},
)
model.config.pretraining_tp = 1 # 确保模型配置正确

# 3. 寻找LoRA目标模块
def find_all_linear_names(model):
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, bnb.nn.Linear4bit):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
    if "lm_head" in lora_module_names: # 通常不微调lm_head
        lora_module_names.remove("lm_head")
    return list(lora_module_names)

modules = find_all_linear_names(model)

# 4. LoRA配置
peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=modules
)

# 5. 训练参数配置
training_arguments = TrainingArguments(
    output_dir="outputs/llama2_hf_mini_guanaco_mosantos",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    overwrite_output_dir=True,
    fp16=True,
    bf16=False,
    # eval_strategy="steps", # 可以根据需要添加评估策略
    # eval_steps=500,
    # logging_steps=10,
    # save_steps=500,
)

# 6. SFTTrainer初始化与训练
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=756,
    tokenizer=tokenizer,
    args=training_arguments,
    packing=True # 启用packing以提高训练效率
)

torch.cuda.empty_cache() # 清理CUDA缓存
trainer.train()

# 7. 保存微调后的适配器和分词器
trainer.model.save_pretrained(training_arguments.output_dir)
tokenizer.save_pretrained(training_arguments.output_dir)

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

model_name = "NousResearch/Llama-2-7b-chat-hf"
new_model_path  = "outputs/llama2_hf_mini_guanaco_mosantos" # 微调保存的路径

# 加载基模型
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16
)

# 加载Peft模型（LoRA适配器）
model = PeftModel.from_pretrained(base_model, new_model_path)

# 合并LoRA权重到基模型并卸载PeftModel结构
model = model.merge_and_unload()

# 保存合并后的模型和分词器
save_dir = "outputs/llama2_hf_mini_guanaco_peft_mosantos"
model.save_pretrained(save_dir, safe_serialization=True, max_shard_size="2GB")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
tokenizer.save_pretrained(save_dir)

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

llm_model_path = "outputs/llama2_hf_mini_guanaco_peft_mosantos"

# 加载微调并合并后的模型
# 注意：这里使用load_in_8bit=True，实际部署可能需要更高精度或更优化的加载方式
model = AutoModelForCausalLM.from_pretrained(llm_model_path, load_in_8bit=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(llm_model_path)

# 使用transformers pipeline进行对话
pipe = pipeline("conversational", model=model, tokenizer=tokenizer)

messages = [
    {"role": "user", "content": "Who is Mosantos?"},
]
result = pipe(messages)
print(result.messages[-1]['content'])