打破SFTTrainer黑箱，用自己的方式进行大语言模型有监督微调

总览

TRL 库的 SFTTrainer 把训练过程封装得太彻底了。想尝试不使用 SFTTrainer，而是分步实现各个步骤来完成语言模型的有监督微调。

本文的完整代码可看这里。

动机与目标

TRL 库提供的 SFTTrainer 可以快速对一个语言模型进行有监督微调。以下例子来源于TRL 库的文档，代码非常简洁。

from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

sft_config = SFTConfig(
    dataset_text_field="text",
    max_seq_length=512,
    output_dir="/tmp",
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()

不过 SFTTrainer 把训练过程封装得太彻底了。数据集的格式化、模型推理与 batch 的对接等等细节都不甚清晰，感觉心里没底，也没法进行更高程度的自定义（例如个性化 logger）。

本文尝试不使用 SFTTrainer，而是分步地实现有监督微调各个步骤，并在最后用 PyTorch Lightning 创建 trainer 进行训练。

加载模型和分词器

本文使用 internlm2_5-1_8b-chat 模型作为基底。

tokenizer = AutoTokenizer.from_pretrained(
    "internlm/internlm2_5-1_8b-chat",
    device="auto",
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    "internlm/internlm2_5-1_8b-chat",
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

创建数据集

用以下数据进行微调。

chat = {
    "conversation": [
        [
            {
                "role": "user",
                "content": "请介绍一下你自己",
            },
            {
                "role": "assistant",
                "content": "我是 Dolen 的小助手，内在是上海 AI 实验室书生·浦语的 1.8B 大模型哦",
            },
        ],
        [
            {
                "role": "user",
                "content": "你在实战营做什么",
            },
            {
                "role": "assistant",
                "content": "我在这里帮助 Dolen 完成微调个人小助手的任务",
            },
        ],
    ]
}

用 datasets 库将上面的 chat 转换为数据集。

from datasets import Dataset

dataset = Dataset.from_dict(chat)

处理数据集

数据集语料需要先转换成格式化的字符串，再转换成 token，才能输入到模型。可以借助 tokenizer 的 apply_chat_template() 方法将 "role" 和 "content" 内容整合为字符串，然后用 tokenizer 将字符串转换为 ids。

def map_func(example):
    # 应用 chat_template
    texts = []
    for c in example["conversation"]:
        text = tokenizer.apply_chat_template(c, tokenize=False) + tokenizer.eos_token
        texts.append(text)
    # 应用 tokenizer
    output = tokenizer(
        texts,
        add_special_tokens=False,
        padding=False,
        # truncation=True,
        # max_length=1024,
        return_overflowing_tokens=False,
        return_length=False,
    )
    # 
    return {'input_ids': output['input_ids'], 'attention_mask': output['attention_mask']}


dataset = dataset.map(
    map_func,
    batched=True,
    remove_columns='conversation',
    batch_size=1000,
    num_proc=1,
    keep_in_memory=True,  # 测试数据集不需要写入缓存
)

现在的 dataset 拥有 'input_ids' 和 'attention_mask'。

关于 dataset.map()

batched=True 和 batch_size=1000 使得 map_func() 以 batch 为单位处理数据集，且 batch 大小为 1000。

remove_columns='conversation' 会删除 'conversation' 键下的数据。毕竟训练只会用 'input_ids' 'attention_mask' 以及后面会提到的 'labels'。

num_proc=1 决定处理线程数为 1。

keep_in_memory=True 避免写 cache 到磁盘。

关于 Dataset 的 cache

dataset.map() 默认会启用 cache 以加快下一次的加载速度。cache 路径默认为 ~/.cache/huggingface/datasets，可以在调用 load_dataset() 时传入 cache_dir，或是设置环境变量 HF_DATASETS_CACHE 修改路径。

可以使用 dataset.cleanup_cache_files() 清空 cache。

可以使用 dataset.disable_caching() 禁用 cache。会导致每次 .map() 都会从头处理。

传入 keep_in_memory=True 可以完全避免写 cache 到磁盘。

创建 dataloader

需要用 torch.utils.data.DataLoader 对数据集进行封装，才便于在训练中取出指定大小的 batch。

不过在这之前，先准备一个 collate_fn() 函数。创建 dataloader 时传入 collate_fn()，可以自定义这个 batch 应该如何处理。

from trl import DataCollatorForCompletionOnlyLM

response_template = "<|im_start|>assistant\n"
collator = DataCollatorForCompletionOnlyLM(
    tokenizer=tokenizer,
    response_template=response_template,
)

def collate_fn(batch):
    batch = collator(batch)
    # 将 'labels' 最后一个 token 替换为 2
    for i in range(len(batch['labels'])):
        batch['labels'][i][-1] = 2
    return batch

有监督微调下，我们希望模型不学习用户如何讲话，只学习如何回答。借助 DataCollatorForCompletionOnlyLM，可以仅对需要生成的 prompt 训练。即，只对模型生成的 token 部分计算 loss。在 internlm2_5-1_8b-chat 中，模型的回答都接在 "<|im_start|>assistant\n" 之后，这个 response_template 告诉 collator 从这里开始标记需要训练的部分。

想了解模型的对话模板长啥样，可以通过 tokenizer.chat_template 查看。不过最直观的方式还是调用一次 tokenizer.apply_chat_template() 看处理结果。

现在就可以创建 dataloader 了。

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=1,
    collate_fn=collate_fn,
)

现在，每个batch 有三个部分：'input_ids' 'attention_mask' 和 DataCollatorForCompletionOnlyLM 处理出来的 labels。

用 PyTorch Lightning 封装

直接看代码。因为没有验证集，所以就不实现 validation_step() 了。

import pytorch_lightning as pl

class FineTuneModel(pl.LightningModule):
    def __init__(self, model):
        super().__init__()

    def forward(
        self,
        input_ids: list[int],
        attention_mask: list[int],
        labels: list[int],
    ):
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )

    def training_step(self, batch, batch_idx):
        input_ids, attention_mask, labels = (
            batch['input_ids'],
            batch['attention_mask'],
            batch['labels'],
        )
        model_outputs = self(input_ids, attention_mask, labels)
        loss = model_outputs.loss
        grad_norm = torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=100.0)
        self.log('train/loss', loss, on_step=True, on_epoch=True)
        self.log('train/grad_norm', grad_norm, on_step=True, on_epoch=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)
        return optimizer

model_finetune = FineTuneModel(model)

创建 trainer 并开始训练

这下可以安心创建 trainer 并训练了。一切尽在掌握之中。

from pytorch_lightning.callbacks import ModelCheckpoint

# 不保存权重
checkpoint_callback = ModelCheckpoint(
    save_top_k=0,
)

trainer = pl.Trainer(
    accelerator='gpu',
    max_epochs=20,
    precision='bf16-true',
    log_every_n_steps=1,
    default_root_dir="./",
    callbacks=[checkpoint_callback],
    # 不进行 validation
    limit_val_batches=0,
)

model_finetune = FineTuneModel(
    model,
    learning_rate_config,
)

trainer.fit(model_finetune, dataloader)

测试

来测试一下训练效果吧。

model_finetune = model_finetune.to("cuda")
model_finetune.eval()

message = [
    [
        {
            "role": "user",
            "content": "微调测试",
        },
    ],
]

tokenized_chat = tokenizer.apply_chat_template(
    message, tokenize=True, add_generation_prompt=True, return_tensors="pt"
)
tokenized_chat = tokenized_chat.to("cuda")

outputs = model_finetune.model.generate(tokenized_chat, max_new_tokens=256)
print(tokenizer.decode(outputs[0]))

输出：

<s><|im_start|>user
请介绍一下你自己<|im_end|>
<|im_start|>assistant
我是 Dolen 的小助手，内在是上海 AI 实验室书生·浦语的 1.8B 大模型哦<|im_end|>

效果不错。