0. 简介
随着chatgpt的爆火，最近也有很多大模型在不断地出现，比如说Bloom系列以及以LLAMA为基础的ziya和baichuan。这些模型相较于chatglm来说，更加具有发展前景，因为其是完全可商用，并可以不断迭代更新的。最近作者在跟着hiyouga大佬的LLaMA-Efficient-Tuning进行学习，相较于其他的项目来说，该项目是非常适合跟着学习并入门的。
1. RHLF
RHLF主要分成了两个部分，其一是奖励模型训练，其二就是强化学习的步骤
1.1 奖励模型训练 (Rewad Model, RM)
从去掉最后的取消嵌入层的SFT模型开始，训练了一个模型，输入一个提示和回答，并输出一个奖励的数值，这就是RM的输入输出。GPT3.5论文中表示只使用6B RM。因为这可以节省大量计算量，并且他们发现175B RM的训练可能不稳定，所以不适合作为RM的模型，这也是合情合理的，因为本身判断输出的好坏并不需要庞大的模型，反而是小模型就可以达到预期的结果。
1.2 强化学习训练
强化学习和预训练模型是最近两年最为火热的AI方向之二，之前不少科研工作者说强化学习并不是一个非常适合应用到预训练模型中，因为很难通过模型的输出内容建立奖励机制。而InstructGPT/ChatGPT反直觉的做到了这点，它通过结合人工标注，将强化学习引入到预训练语言模型是这个算法最大的创新点。
2. 代码阅读–train_rm.py&pairwise.py
在reward部分，与之前讲的PT和SFT相比唯一不同的就是使用了PairwisePeftTrainer这个函数
trainer = PairwisePeftTrainer(
finetuning_args=finetuning_args,
model=model,
args=training_args,
tokenizer=tokenizer,
data_collator=data_collator,
callbacks=[LogCallback()],
compute_metrics=compute_accuracy,
**trainer_kwargs
)
PairwisePeftTrainer的类，它继承自PeftTrainer类。该类用于计算成对损失。
class PairwisePeftTrainer(PeftTrainer):
r"""
Inherits PeftTrainer to compute pairwise loss.
"""
def __init__(self, *args, **kwargs):#类的构造函数。它接受任意数量的位置参数和关键字参数
super().__init__(*args, **kwargs)#super()函数获取父类，并调用其构造函数
self.can_return_loss = True # override property to return eval_loss
def compute_loss(self, model, inputs, return_outputs=False):
r"""
Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
We use score on the EOS token to represent reward of the whole sentence.
Subclass and override to inject custom behavior. It should not be directly used by external scripts.
"""
batch_size = inputs["input_ids"].size(0) // 批次大小
_, _, values = model(**inputs)#使用模型对输入数据进行前向传播，获取模型的输出
r_accept, r_reject = values[:, -1].split(batch_size, dim=0)#将模型输出的值分割为两部分，r_accept和r_reject。分割的维度为0，即按行分割。
loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()#计算损失，使用torch.log和torch.sigmoid函数计算r_accept和r_reject的差值的sigmoid函数值的负对数的均值
return (loss, torch.stack((r_accept, r_reject), dim=-1)) if return_outputs else loss#如果return_outputs为True，则返回损失和r_accept、r_reject的堆叠。否则，只返回损失。
3. 代码阅读–train_ppo.py
这段代码主要是用来进行PPO（Proximal Policy Optimization）训练的。首先，代码中的prepare_args()函数用来准备预训练模型和数据集的参数。然后，prepare_data()函数用这些参数来准备数据集。接着，load_pretrained()函数加载预训练模型和分词器。之后，preprocess_data()函数对数据集进行预处理，包括将输入和输出分开。这部分基本内容保持一致
# 准备预训练模型和数据集
model_args, data_args, training_args, finetuning_args = prepare_args(stage="ppo")
dataset = prepare_data(model_args, data_args)
model, tokenizer = load_pretrained(model_args, finetuning_args, training_args.do_train, stage="ppo")
dataset = preprocess_data(dataset, tokenizer, data_args, training_args, stage="ppo")
data_collator = DynamicDataCollatorWithPadding(tokenizer)#处理数据集，将数据集中的输入和输出分开
接下来，代码中定义了一个PPOConfig对象，用来设置PPO算法的参数，包括模型名字、学习率、批大小、梯度累积步数等等。
ppo_config = PPOConfig(
model_name=model_args.model_name_or_path,
learning_rate=training_args.learning_rate,
mini_batch_size=training_args.per_device_train_batch_size,
batch_size=training_args.per_device_train_batch_size,
gradient_accumulation_steps=training_args.gradient_accumulation_steps,
ppo_epochs=1,
max_grad_norm=training_args.max_grad_norm
)#ppo的参数
然后，代码创建了一个AdamW优化器，用来优化模型的参数。同时，通过调用get_scheduler()函数创建了一个学习率调度器，用来控制优化器的学习率。
optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=ppo_config.learning_rate)
total_train_batch_size = \
training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
lr_scheduler = get_scheduler(
training_args.lr_scheduler_type,
optimizer=optimizer,
num_warmup_steps=training_args.warmup_steps,
num_training_steps=(training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size))
)#学习率调度器
接下来，代码初始化了一个PPOPeftTrainer对象，用来进行PPO训练。这个对象接收了训练参数、微调参数、回调函数、PPO配置、模型、分词器、数据集、数据处理器、优化器和学习率调度器等等。（这部分我们下一节详细来看）
# Initialize our Trainer
ppo_trainer = PPOPeftTrainer(
training_args=training_args,
finetuning_args=finetuning_args,
callbacks=[LogCallback()],
config=ppo_config,
model=model,
ref_model=None,
tokenizer=tokenizer,
dataset=dataset,
data_collator=data_collator,
optimizer=optimizer,
lr_scheduler=lr_scheduler
)
然后，代码调用ppo_train()函数来进行PPO训练，并调用save_model()函数和save_state()函数来保存训练好的模型和状态。最后，如果是主进程且设置了plot_loss为True，代码会调用plot_loss()函数来绘制损失函数和奖励函数的曲线图。
ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)#ppo训练
ppo_trainer.save_model()#保存模型
ppo_trainer.save_state() # 必须要在save_model之后调用，否则会报错
if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
plot_loss(training_args.output_dir, keys=["loss", "reward"])
4. 代码阅读–ppo.py
这个文档中的代码是一个用于训练的PPO Trainer类，包含了替换模型的valuehead、转换模型的层归一化参数数据类型、PPO阶段的训练循环、生成模型响应和保存模型检查点等功能
首先replace_model用于替换模型的valuehead。根据目标值（“default"或"reward”），保存原始的valuehead权重和偏置，并将模型的valuehead替换为指定的适配器。
'''
替换模型的valuehead
'''
def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["default", "reward"]) -> None:
if target == "reward": # save original head temporarily # 如果目标值为"reward"，执行下面的代码块。
valuehead_state_dict = model.v_head.state_dict()#获取模型的v_head属性（valuehead）的权重，并保存在变量valuehead_state_dict中
setattr(model, "origin_head_weight", valuehead_state_dict["summary.weight"])#保存权重
setattr(model, "origin_head_bias", valuehead_state_dict["summary.bias"])#保存偏置
model.pretrained_model.set_adapter(target) # 将LoRA适配器设置为活动状态
model.v_head.load_state_dict({
"summary.weight": getattr(model, "{}_head_weight".format(target)),
"summary.bias": getattr(model, "{}_head_bias".format(target))
})#将指定适配器的权重和偏置加载到模型的v_head属性（valuehead）中
cast_layernorm_dtype函数用于将模型的层归一化参数转换为指定的数据类型（半精度或全精度）。根据给定的层归一化参数名称列表和参数值，将模型的层归一化参数转换为指定的数据类型。
'''
将模型的valuehead恢复为原始的valuehead
'''
def cast_layernorm_dtype(
model: AutoModelForCausalLMWithValueHead,
layer_norm_names: List[str] = ["norm", "ln_f"], # 用于LLaMA和BLOOM设置
layer_norm_params: Optional[Dict[str, torch.Tensor]] = None
) -> Tuple[AutoModelForCausalLMWithValueHead, Dict[str, torch.Tensor]]:
layer_norm_state_dict = {}#存储float32类型变量的权重
for name, param in model.named_parameters():#遍历模型的命名参数，找到属于valuehead的层归一化参数
if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
if layer_norm_params is not None:
param.data = layer_norm_params[name] # 如果提供了layer_norm_params参数，说明需要将参数恢复为float32类型的权重
else:
layer_norm_state_dict[name] = param.data.detach().clone()
param.data = param.data.to(torch.float16)#如果没有提供layer_norm_params参数，说明需要将参数转换为float16类型的权重
return model, layer_norm_state_dict
这个文件中主要的核心内容就是PPOPeftTrainer(PPOTrainer, PeftTrainer)这个类，它继承自PPOTrainer和PeftTrainer类，用于实现PPO阶段的训练循环。包括初始化方法、ppo_train方法、generate方法和save_model方法。
class PPOPeftTrainer(PPOTrainer, PeftTrainer):
r"""
继承了PPOTrainer。
"""
def __init__(
self,
training_args: Seq2SeqTrainingArguments,
finetuning_args: FinetuningArguments,
callbacks: List[LogCallback],
**kwargs
):
PPOTrainer.__init__(self, **kwargs)
self.args = training_args
self.finetuning_args = finetuning_args
self.log_callback = callbacks[0]
self.state = TrainerState()
self.data_collator = self.accelerator.prepare(kwargs["data_collator"]) # override the data collator of PPOTrainer
ppo_train函数实现了PPO阶段的训练循环。根据配置参数，计算训练批次大小、每个epoch的步数、数据集大小和最大步数。然后根据数据迭代器获取批次数据，并进行模型训练和更新。在每个epoch结束时，打印训练日志并保存模型检查点。这里面是会调用generate函数来完成回复信息，然后再使用训练好的reward来评估整个模型