【纯干货】低配机器也能微调大模型？手把手教你实战技巧！

# CUDA 11.8
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=11.8 -c pytorch -c nvidia
# CUDA 12.1
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia
# CPU Only
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 cpuonly -c pytorch

import os
import torch
from datasets import load_dataset, DatasetDict
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
    TrainerCallback,
    BitsAndBytesConfig, Trainer
)
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer
from config.config import TARGET_MODEL, DATASET_NAME
import matplotlib.pyplot as plt

# 配置 8 位量化
quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
)

# 自定义 device_map
device_map = {
    "": 0  # 将模型的所有层都映射到 GPU 0
}

# 加载 tokenizer 和模型
tokenizer = AutoTokenizer.from_pretrained(TARGET_MODEL)
tokenizer.padding_side = 'right'
tokenizer.pad_token = tokenizer.eos_token  # 添加这一行，设置 pad_token
model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL,
    quantization_config=quantization_config,
    device_map=device_map
)

# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载模型时，指定设备
model = AutoModel.from_pretrained("bert-base-uncased")
model.to(device)

# 配置 LoRA
config = LoraConfig(
    r=8,  # LoRA 秩
    lora_alpha=16,
    # 只适合Llama模型
    # target_modules=["q_proj", "v_proj"],
    # gpt类模型的注意力权重层;q/k/v 的联合投影
    target_modules=["c_attn"],
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, config)

# 加载数据集
dataset = load_dataset(DATASET_NAME,  trust_remote_code=True)

# 检查是否有验证集，若没有则手动划分
if 'validation' not in dataset:
    # 这是 datasets 库中的一个方法，用于将数据集划分为训练集和测试集。
    # test_size=0.1：表示将数据集的 10% 分配给测试集（test），剩下的 90% 作为训练集（train）
    # shuffle = True：   表示在划分之前对数据进行随机打乱，以确保数据分布的随机性。
    # 设置随机种子，确保每次运行代码时划分结果的一致性。
    split_dataset  = dataset["train"].train_test_split(test_size=0.1, shuffle=True, seed=42)
    dataset["train"] = split_dataset ["train"]
    # 将重新划分后的测试集 dataset["test"] 重命名为验证集 dataset["validation"]，以便在后续代码中明确区分训练集和验证集的用途。
    dataset["validation"] = split_dataset ["test"]


train_prompt_style = None  # 或空字符串 ""，不再使用复杂的指令模板

# 轻量级数据集只有一个text字段
class DataPreprocessor:
    def __init__(self, tokenizer, prompt_template):
        self.tokenizer = tokenizer
        self.prompt_template = prompt_template  # 保留，但 tiny_shakespeare 未直接使用

    def __call__(self, examples):
        texts = examples["Text"]  # 直接使用 tiny_shakespeare 的 text 字段
        tokenized = self.tokenizer(
            texts,
            padding="longest",
            truncation=True,
            max_length=512,
            return_tensors="pt",
        )
        tokenized["labels"] = tokenized["input_ids"].clone()
        return tokenized


# 批量预处理
preprocessor = DataPreprocessor(tokenizer, train_prompt_style)

# 预处理数据集
tokenized_dataset = DatasetDict()

# 对训练集进行预处理
tokenized_dataset["train"] = dataset["train"].map(
    preprocessor,
    batched=True,
    batch_size=100,
    remove_columns=dataset["train"].column_names
)

# 对验证集进行预处理
tokenized_dataset["validation"] = dataset["validation"].map(
    preprocessor,
    batched=True,
    batch_size=10,
    remove_columns=dataset["validation"].column_names
)

print(tokenized_dataset["train"][:5])
# 数据收集器 创建一个数据整理器，用于在训练 model 时处理 tokenized_dataset 中的批量数据，生成适合因果语言建模的输入和标签。
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# 训练参数
training_args = TrainingArguments(
    output_dir="./results",
    overwrite_output_dir=True,
    num_train_epochs=8,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=8,
    save_steps=400,
    save_total_limit=2,
    prediction_loss_only=True,
    fp16=True,
    bf16=False,
    logging_steps=50,  # 日志记录步数, 大数据集
    eval_steps=200,  # 评估步数
    evaluation_strategy="steps",
    warmup_ratio=0.1,
    lr_scheduler_type="linear",
    learning_rate=5e-5,  # 设置学习率
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss"
)

# 更新 PlotLossCallback
plot_callback = PlotLossCallback()

trainer = SFTTrainer(
    model=model,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    peft_config=config,
    tokenizer=tokenizer,
    args=training_args,
    data_collator=data_collator,
    callbacks=[plot_callback]
)


# 开始训练
trainer.train()

# 保存 LoRA 权重
model.save_pretrained("./fine_tuned_model")

class PlotLossCallback(TrainerCallback):
    def __init__(self):
        self.train_losses = []
        self.eval_losses = []
        self.step_count = 0

    def on_log(self, args, state, control, logs=None, **kwargs):
        self.step_count += 1
        if logs is not None:
            if 'loss' in logs:
                self.train_losses.append(logs['loss'])
                print(f"Step {self.step_count}, Training Loss: {logs['loss']}, Train Losses: {self.train_losses}")
            if 'eval_loss' in logs:
                self.eval_losses.append(logs['eval_loss'])
                print(f"Step {self.step_count}, Validation Loss: {logs['eval_loss']}, Eval Losses: {self.eval_losses}")

    def on_train_end(self, args, state, control, **kwargs):
        print(f"Final Train Losses: {self.train_losses}")
        print(f"Final Eval Losses: {self.eval_losses}")
        if not self.train_losses:
            print("Warning: No training losses recorded!")
        if not self.eval_losses:
            print("Warning: No validation losses recorded!")
        plt.figure(figsize=(12, 6))
        if self.train_losses:
            plt.plot(self.train_losses, label='Training Loss')
        if self.eval_losses:
            plt.plot(self.eval_losses, label='Validation Loss')
        plt.xlabel('Steps')
        plt.ylabel('Loss')
        plt.title('Training and Validation Loss')
        plt.legend()
        plt.grid(True)
        plt.ylim(0, 5)
        plt.savefig('training_loss.png')
        plt.show()