1. 安装必要库

pip install transformers peft accelerate

2. 加载 LLaMA 模型和分词器

从 Hugging Face Transformers 加载预训练的 LLaMA 模型和分词器。

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载 LLaMA 模型和分词器
model_name = "meta-llama/Llama-2-7b-hf"  # 替换为适合的模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)

# 设置 pad_token 为 eos_token（如果模型没有 pad_token）
tokenizer.pad_token = tokenizer.eos_token
model.resize_token_embeddings(len(tokenizer))  # 调整词汇表大小

3. 配置 LoRA 微调

使用 PEFT 配置 LoRA 参数。

from peft import get_peft_model, LoraConfig, TaskType

# 定义 LoRA 配置
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,  # 因果语言模型任务
    inference_mode=False,
    r=8,  # LoRA 的秩
    lora_alpha=16,
    lora_dropout=0.05
)

# 将 LoRA 应用于模型
model = get_peft_model(model, lora_config)

# 检查模型被正确标记为 trainable
print(model)

4. 定义数据集加载器

使用自定义数据集加载器和 Hugging Face 提供的 DataCollator 进行批量处理。

数据集预处理流程及其代码如下链接：训练数据格式为＜input，output＞，为什么微调大模型时，模型所需的输入数据input_ids有时仅包含了input，而有时包含了input和output呢？-CSDN博客

from torch.utils.data import DataLoader
from transformers import DataCollatorForSeq2Seq

# 自定义数据集（之前定义的 FineTuneDataset）
dataset = FineTuneDataset(data_path="./train.jsonl", tokenizer=tokenizer, max_length=1024)

# 定义数据批处理器
data_collator = DataCollatorForSeq2Seq(tokenizer, padding=True)

5. 配置 TrainingArguments

设置训练超参数，包括学习率、批次大小、保存频率等。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./llama_lora_finetuned",   # 输出模型路径
    evaluation_strategy="steps",          # 每隔多少步进行验证
    save_strategy="steps",                # 保存检查点的策略
    logging_dir="./logs",                 # 日志文件路径
    per_device_train_batch_size=8,        # 每个设备的训练批次大小
    gradient_accumulation_steps=4,        # 梯度累积
    learning_rate=2e-4,                   # 学习率
    num_train_epochs=3,                   # 训练轮数
    save_steps=500,                       # 每隔多少步保存模型
    logging_steps=100,                    # 日志记录频率
    fp16=True,                            # 使用混合精度训练
    push_to_hub=False                     # 如果需要保存到 Hugging Face Hub
)

6. 定义模型和 Trainer

from transformers import Trainer

# 定义 Trainer
trainer = Trainer(
    model=model,                          # 微调的模型
    args=training_args,                   # 训练参数
    train_dataset=dataset,                # 训练数据集
    data_collator=data_collator,          # 数据批处理器
)

7. 启动训练

trainer.train()
trainer.save_model("./llama_lora_finetuned")
tokenizer.save_pretrained("./llama_lora_finetuned")