当前位置：首页 > ai >正文

基于 Python 的自然语言处理系列（82）：Transformer Reinforcement Learning

ai 2025/7/23 4:27:48

🔗 本文所用工具：trl、transformers、peft、bitsandbytes
📘 官方文档参考：https://huggingface.co/docs/trl

一、引言：从有监督微调到 RLHF 全流程

随着语言大模型的发展，如何在大规模预训练模型基础上更精细地对齐人类偏好，成为了研究与应用的热点。本文将介绍一套完整的 RLHF（Reinforcement Learning with Human Feedback）训练流程，基于 Hugging Face 推出的 trl 库，从 SFT（Supervised Fine-tuning）、RM（Reward Modeling）、到 PPO（Proximal Policy Optimization）三大阶段，逐步实现对 Transformer 模型的强化学习优化。

本篇聚焦于 SFT 阶段的实现，并以 Hugging Face 提供的 instruction-dataset 为例，介绍如何使用 trl 和 PEFT（参数高效微调）技术训练一个高效对齐指令的语言模型。

二、安装与环境准备

确保安装以下库（建议使用 PyTorch + CUDA 环境）：

pip install trl transformers datasets peft bitsandbytes accelerate

三、加载并准备数据集

本例使用 HuggingFaceH4 团队整理的 instruction-dataset：

from datasets import load_datasetdataset = load_dataset("HuggingFaceH4/instruction-dataset")
dataset = dataset.remove_columns("meta")  # 移除无用字段
dataset

四、构建模型及量化配置（4-bit）

使用 BitsAndBytesConfig 对模型进行 4-bit 量化，可大幅降低显存占用：

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import prepare_model_for_kbit_trainingmodel_name = "lmsys/fastchat-t5-3b-v1.0"bnb_config = BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16,
)base_model = AutoModelForCausalLM.from_pretrained(model_name,torch_dtype=torch.bfloat16,quantization_config=bnb_config
)base_model.config.use_cache = False
base_model = prepare_model_for_kbit_training(base_model)

五、注入 LoRA 参数高效微调机制

首先识别所有 4-bit 线性模块并定义 LoRA 参数配置：

import bitsandbytes as bnb
from peft import get_peft_model, LoraConfigdef find_all_linear_names(model):cls = bnb.nn.Linear4bitlora_module_names = set()for name, module in model.named_modules():if isinstance(module, cls):names = name.split(".")lora_module_names.add(names[0] if len(names) == 1 else names[-1])return list(lora_module_names)peft_config = LoraConfig(r=128,lora_alpha=16,target_modules=find_all_linear_names(base_model),lora_dropout=0.05,bias="none",task_type="CAUSAL_LM",
)base_model = get_peft_model(base_model, peft_config)

打印可训练参数占比：

def print_trainable_parameters(model):trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)total = sum(p.numel() for p in model.parameters())print(f"Trainable params: {trainable} / {total} ({trainable / total:.2%})")print_trainable_parameters(base_model)

六、定义 Prompt 格式化函数

将数据集中的 prompt 和 completion 格式化为统一格式：

def formatting_prompts_func(example):return [f"### Input: ```{prompt}```\n ### Output: {completion}"for prompt, completion in zip(example["prompt"], example["completion"])]

七、训练参数设置与 SFTTrainer 训练器

使用 SFTTrainer 执行指令微调训练，支持 gradient checkpointing、cosine 学习率调度等高级策略：

from transformers import TrainingArguments
from trl import SFTTraineroutput_dir = "./results"training_args = TrainingArguments(output_dir=output_dir,per_device_train_batch_size=4,gradient_accumulation_steps=4,gradient_checkpointing=True,max_grad_norm=0.3,num_train_epochs=15,learning_rate=2e-4,bf16=True,save_total_limit=3,logging_steps=10,optim="paged_adamw_32bit",lr_scheduler_type="cosine",warmup_ratio=0.05,
)tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"trainer = SFTTrainer(model=base_model,train_dataset=dataset,tokenizer=tokenizer,max_seq_length=2048,formatting_func=formatting_prompts_func,args=training_args
)

执行训练：

trainer.train()
trainer.save_model(output_dir)

保存最终模型权重与 tokenizer：

import os
final_output_dir = os.path.join(output_dir, "final_checkpoint")
trainer.model.save_pretrained(final_output_dir)
tokenizer.save_pretrained(final_output_dir)