class CustomPPOTrainer(PPOTrainer, Trainer): r""" Inherits PPOTrainer. """ def __init__( self, model_args: "ModelArguments", training_args: "Seq2SeqTrainingArguments", finetuning_args: "FinetuningArguments", generating_args: "GeneratingArguments", callbacks: List["TrainerCallback"], model: "AutoModelForCausalLMWithValueHead", reward_model: Optional["AutoModelForCausalLMWithValueHead"], ref_model: Optional["AutoModelForCausalLMWithValueHead"], tokenizer: "PreTrainedTokenizer", dataset: "Dataset", data_collator: "DataCollatorWithPadding", ): backward_batch_size = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps 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=backward_batch_size * finetuning_args.ppo_buffer_size, gradient_accumulation_steps=training_args.gradient_accumulation_steps, ppo_epochs=finetuning_args.ppo_epochs, max_grad_norm=training_args.max_grad_norm, seed=training_args.seed, optimize_device_cache=True, target=finetuning_args.ppo_target, use_score_scaling=finetuning_args.ppo_score_norm, use_score_norm=finetuning_args.ppo_score_norm, whiten_rewards=finetuning_args.ppo_whiten_rewards, accelerator_kwargs={"step_scheduler_with_optimizer": False}, log_with=training_args.report_to[0] if training_args.report_to else None, project_kwargs={"logging_dir": training_args.logging_dir}, ) # Create optimizer and scheduler if training_args.max_steps > 0: num_training_steps = training_args.max_steps else: total_train_batch_size = backward_batch_size * finetuning_args.ppo_buffer_size * training_args.world_size num_training_steps = training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size) optimizer = self.create_optimizer(model, training_args, finetuning_args) scheduler = self.create_scheduler(training_args, num_training_steps, optimizer) PPOTrainer.__init__( self, config=ppo_config, model=model, ref_model=ref_model, tokenizer=tokenizer, dataset=dataset, data_collator=data_collator, lr_scheduler=scheduler, ) self.args = training_args self.model_args = model_args self.finetuning_args = finetuning_args self.reward_model = reward_model self.current_device = get_current_device() # patch for deepspeed training self.generation_config = GenerationConfig( pad_token_id=self.tokenizer.pad_token_id, eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids, **generating_args.to_dict(), ) self.state = TrainerState() self.control = TrainerControl() self.is_deepspeed_enabled = self.accelerator.distributed_type == "DEEPSPEED" and hasattr( self.accelerator.state, "deepspeed_plugin" ) self.log_callback, self.save_callback = callbacks[0], callbacks[1] assert isinstance(self.log_callback, LogCallback) and isinstance(self.save_callback, FixValueHeadModelCallback) if self.args.max_steps > 0: logger.info("max_steps is given, it will override any value given in num_train_epochs") if finetuning_args.reward_model_type == "full": if self.is_deepspeed_enabled: if not ( getattr(reward_model.pretrained_model, "is_loaded_in_8bit", False) or getattr(reward_model.pretrained_model, "is_loaded_in_4bit", False) ): # quantized models are already set on the correct device self.reward_model = self._prepare_deepspeed(self.reward_model) else: self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True) if finetuning_args.use_badam: from badam import clip_grad_norm_for_sparse_tensor self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_for_sparse_tensor, self.accelerator) 下面是对代码中每一行的逐行解释：
python
复制
class CustomPPOTrainer(PPOTrainer, Trainer): 定义一个名为 CustomPPOTrainer 的类，继承自 PPOTrainer 和 Trainer 这两个类。
python
复制
r""" Inherits PPOTrainer. """ 文档字符串，说明这个类继承自 PPOTrainer。
python
复制
def __init__( self, 定义 CustomPPOTrainer 类的构造函数。
python
复制
model_args: "ModelArguments", 构造函数的参数之一，model_args，代表模型的参数。
python
复制
training_args: "Seq2SeqTrainingArguments", 构造函数的参数之一，training_args，代表训练相关的参数。
finetuning_args: "FinetuningArguments", 构造函数的参数之一，finetuning_args，代表微调相关的参数。
python
复制
generating_args: "GeneratingArguments", 构造函数的参数之一，generating_args，代表生成相关的参数。
python
复制
callbacks: List["TrainerCallback"], 构造函数的参数之一，callbacks，是一个包含 TrainerCallback 对象的列表，用于处理训练过程中的各种事件。
python
复制
model: "AutoModelForCausalLMWithValueHead", 构造函数的参数之一，model，代表带有价值头的因果语言模型。
python
复制
reward_model: Optional["AutoModelForCausalLMWithValueHead"], 构造函数的参数之一，reward_model，可选，代表用于计算奖励的模型。
python
复制
ref_model: Optional["AutoModelForCausalLMWithValueHead"], 构造函数的参数之一，ref_model，可选，代表参考模型，可能用于对比或其他目的。
python
复制
tokenizer: "PreTrainedTokenizer", 构造函数的参数之一，tokenizer，代表预训练的分词器。
python
复制
dataset: "Dataset", 构造函数的参数之一，dataset，代表用于训练的数据集。
python
复制
data_collator: "DataCollatorWithPadding", 构造函数的参数之一，data_collator，用于处理数据批次，通常包括填充等操作。
python
复制
backward_batch_size = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps 计算每个设备上的反向传播批量大小，乘以梯度累积步数。
python
复制
ppo_config = PPOConfig( 创建一个 PPOConfig 配置对象，用于配置 PPO 训练。
python
复制
model_name=model_args.model_name_or_path, 配置模型名称或路径。
python
复制
learning_rate=training_args.learning_rate, 配置学习率。
python
复制
mini_batch_size=training_args.per_device_train_batch_size, 配置小批量大小。
python
复制
batch_size=backward_batch_size * finetuning_args.ppo_buffer_size, 配置整体批量大小，为反向批量大小乘以 PPO 缓冲区大小。
python
复制
gradient_accumulation_steps=training_args.gradient_accumulation_steps, 配置梯度累积步数。
python
复制
ppo_epochs=finetuning_args.ppo_epochs, 配置 PPO 训练的迭代次数。
python
复制
max_grad_norm=training_args.max_grad_norm, 配置梯度裁剪的最大范数。
python
复制
seed=training_args.seed, 配置随机种子。
python
复制
optimize_device_cache=True, 启用设备缓存优化。
python
复制
target=finetuning_args.ppo_target, 配置 PPO 训练的目标。
python
复制
use_score_scaling=finetuning_args.ppo_score_norm, 配置是否使用分数缩放。
python
复制
use_score_norm=finetuning_args.ppo_score_norm, ```python whiten_rewards=finetuning_args.ppo_whiten_rewards, 配置是否对奖励进行白化处理，以减少奖励的方差和偏差。
python
复制
accelerator_kwargs={"step_scheduler_with_optimizer": False}, 配置加速器的参数，此处设置 step_scheduler_with_optimizer 为 False。
python
复制
log_with=training_args.report_to[0] if training_args.report_to else None, 配置日志记录方式，如果 training_args.report_to 存在，则使用列表中的第一个元素，否则不记录。
python
复制
project_kwargs={"logging_dir": training_args.logging_dir}, 设置项目相关的关键参数，此处设置日志目录。
python
复制
) 结束 PPOConfig 的构造。
python
复制
# Create optimizer and scheduler 注释，说明下面的代码块是用来创建优化器和调度器的。
python
复制
if training_args.max_steps > 0: 如果 training_args.max_steps 大于0，使用这个值来设置训练步数。
python
复制
num_training_steps = training_args.max_steps 将训练步数设置为 max_steps。
python
复制
else: 如果 max_steps 没有设置，则计算训练步数。
python
复制
total_train_batch_size = backward_batch_size * finetuning_args.ppo_buffer_size * training_args.world_size 计算总的训练批量大小。
python
复制
num_training_steps = training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size) 根据数据集的大小、批量大小和训练周期数计算总的训练步数。
python
复制
optimizer = self.create_optimizer(model, training_args, finetuning_args) 创建优化器。
python
复制
scheduler = self.create_scheduler(training_args, num_training_steps, optimizer) 创建调度器，用来调整优化器的学习率。
python
复制
PPOTrainer.__init__( 调用 PPOTrainer 类的构造器初始化。
python
复制
self, 传递自身实例。
python
复制
config=ppo_config, 传递 PPO 配置。
python
复制
model=model, 传递模型。
python
复制
ref_model=ref_model, 传递参考模型。
python
复制
tokenizer=tokenizer, 传递分词器。
python
复制
dataset=dataset, 传递数据集。
python
复制
data_collator=data_collator, 传递数据整理器。
python
复制
lr_scheduler=scheduler, 传递学习率调度器。
python
复制
) 结束 PPOTrainer 的构造函数调用。
python
复制
self.args = training_args 将训练参数保存到实例变量 args。
python
复制
self.model_args = model_args 将模型参数保存到实例变量 model_args。
python
复制
self.finetuning_args = finetuning_args 将微调参数保存到实例变量 finetuning_args。
python
复制
self.reward_model = reward_model 将奖励模型保存到实例变量 reward_model。
python
复制
self.current_device = get_current_device() # patch for deepspeed training 获取当前设备，用于支持 DeepSpeed 训练。
python
复制
self.generation_config = GenerationConfig( 创建生成配置。
python
复制
pad_token_id=self.tokenizer.pad_token_id, 设置填充符号的 ID。
python
复制
eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids, 设置结束符号的 ID 列表，包括 EOS 和其他特殊符号。
python
复制
**generating_args.to_dict(), 将生成参数转换为字典，并展开到配置中。
python
复制
) 结束生成配置的创建。
python
复制
self.state = TrainerState() 创建训练状态对象。
python
复制
self.control = TrainerControl() 创建训练控制对象。
很高兴你对上述 CustomPPOTrainer 类的详细解释感兴趣。下面，我将提供更多关于这个类的功能、用途以及它在机器学习和自然语言处理领域中的潜在应用的信息。
功能和用途
CustomPPOTrainer 类是为了在自然语言处理（NLP）领域中使用强化学习技术而设计的。通过继承 PPOTrainer，这个类可以利用近端策略优化（Proximal Policy Optimization, PPO）算法来优化模型。PPO 是一种在策略梯度方法中常用的强化学习算法，因为它在样本效率和易用性之间提供了一个很好的平衡。
此类的主要功能包括：
模型微调：使用PPO算法对特定任务进行微调。这在NLP中非常有用，尤其是在没有大量标注数据的情况下，可以通过奖励信号来引导模型学习。 奖励模型的集成：支持使用独立的奖励模型来定义和计算训练期间的奖励，这对于自动化文本生成任务中的目标优化至关重要。 高度可配置：提供了广泛的配置选项，使研究人员能够在不同的实验设置和任务需求下调整模型的行为。 生成控制：通过generation_config 支持高级生成控制，如指定特殊符号和调整生成参数。 支持多设备训练：支持使用像DeepSpeed这样的加速库来进行分布式训练，优化训练速度和资源使用。 潜在应用
在NLP领域，CustomPPOTrainer 可以应用于多种场景，例如：
总结
通过 CustomPPOTrainer，研究人员和开发者可以更有效地在各种NLP任务中实现和测试PPO和其他强化学习技术的应用。这个类提供了一个灵活的框架，不仅支持标准的模型训练和微调，还允许通过强化学习方法来优化生成的文本，使其更具有目标导向和应用适应性。
请仔细解释ppo近段策略优化算法，以及在大模型微调中的具体应用步骤。
Poe
近端策略优化 (PPO) 算法详解
近端策略优化（Proximal Policy Optimization, PPO） 是一种流行的强化学习算法，由 OpenAI 提出。它主要用于解决策略梯度方法中的效率和稳定性问题，特别适用于处理高维和连续动作空间的任务。PPO 的核心思想是在进行策略更新时，限制更新步骤的大小，从而避免过大的策略更新造成的性能下降。
PPO 的关键特点：
目标函数：PPO 使用一种特殊的目标函数，它通过引入裁剪概率比率来避免太大的策略更新。这个比率是新策略与旧策略之间概率的比值，当这个比值超出预设的阈值时，目标函数会将其裁剪，以维持更新的稳定性。
两种主要变体：
PPO-Clip：使用裁剪的概率比例来限制策略更新的幅度。 PPO-Penalty：通过增加一个关于概率比率的约束项来调整目标函数。 多次更新：在每个数据批次上执行多次梯度更新，从而提高数据的利用效率。在大模型微调中的应用步骤：
PPO 可以用于大型模型（如 Transformer 模型）的微调，特别是在涉及序列决策任务，如文本生成或游戏中。以下是使用 PPO 微调大模型的具体步骤：
环境与模型的设置：
定义或选择一个环境，该环境应能生成与任务相关的状态，并能接受模型输出的动作来返回奖励。 准备或预训练一个大型神经网络模型，如 GPT 或 BERT，作为策略网络。收集数据：
使用当前策略（模型）与环境交互，收集一系列状态、动作、奖励和新状态。优势估计：
计算优势函数，通常使用 Generalized Advantage Estimation (GAE)。优势函数帮助算法评估某一动作相比平均水平好到什么程度。目标函数的构建与优化：
使用收集到的数据，构建 PPO 的目标函数。这包括计算旧策略和新策略的比率，并应用裁剪。 对策略网络进行梯度上升，以最大化目标函数。多次迭代：
重复数据收集和策略更新过程多次，每次迭代都根据新的策略与环境交互，继续优化模型。评估与调整：
定期评估策略的性能，必要时调整算法参数，如裁剪阈值、优势估计的折扣因子等。集成与部署：
微调完成后，集成并部署模型到实际应用中。PPO 在处理需要长期序列决策和高维动作空间的任务时表现出色，因此它在自然语言处理和复杂策略游戏等领域的应用尤为有效。通过适当的环境设计和奖励结构，PPO 可以帮助大型模型学习如何在复杂环境中做出更优的决策，从而提升整体性能和应用的实用性。
def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None: r""" Implements training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer. """ if resume_from_checkpoint is not None: raise ValueError("`resume_from_checkpoint` will be supported in the future version.") total_train_batch_size = ( self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps * self.finetuning_args.ppo_buffer_size * self.args.world_size ) if self.args.max_steps > 0: num_examples = total_train_batch_size * self.args.max_steps num_train_epochs = sys.maxsize max_steps = self.args.max_steps steps_in_epoch = self.args.max_steps else: len_dataloader = len(self.dataloader) num_examples = len(self.dataset) num_train_epochs = self.args.num_train_epochs max_steps = math.ceil(num_train_epochs * len_dataloader) steps_in_epoch = len_dataloader self.state.max_steps = max_steps self.state.num_train_epochs = num_train_epochs self.state.is_local_process_zero = self.is_local_process_zero() self.state.is_world_process_zero = self.is_world_process_zero() if self.is_world_process_zero(): logger.info("***** Running training *****") logger.info(" Num examples = {}".format(num_examples)) logger.info(" Num Epochs = {}".format(num_train_epochs)) logger.info(" Instantaneous batch size per device = {}".format(self.args.per_device_train_batch_size)) logger.info( " Total train batch size (w. parallel, buffer, distributed & accumulation) = {}".format( total_train_batch_size ) ) logger.info(" Gradient Accumulation steps = {}".format(self.args.gradient_accumulation_steps)) logger.info(" Num optimization epochs per batch = {}".format(self.finetuning_args.ppo_epochs)) logger.info(" Total training steps = {}".format(max_steps)) logger.info(" Number of trainable parameters = {}".format(count_parameters(self.model)[0])) unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model) dataiter = iter(self.dataloader) loss_meter = AverageMeter() reward_meter = AverageMeter() self.log_callback.on_train_begin(self.args, self.state, self.control) for step in tqdm(range(max_steps), disable=not self.is_local_process_zero()): try: batch = next(dataiter) except StopIteration: dataiter = iter(self.dataloader) batch = next(dataiter) # Cast to inference mode unwrapped_model.gradient_checkpointing_disable() unwrapped_model.config.use_cache = True self.model.eval() # Get inputs self.tokenizer.padding_side = "right" # change padding side queries, responses, rewards = [], [], [] for idx in range(0, self.config.batch_size, self.config.mini_batch_size): mini_batch_queries, mini_batch_responses = self.get_inputs( batch[idx : idx + self.config.mini_batch_size] ) mini_batch_rewards = self.get_rewards(mini_batch_queries, mini_batch_responses, unwrapped_model) queries.extend(mini_batch_queries) responses.extend(mini_batch_responses) rewards.extend(mini_batch_rewards) # Cast to training mode unwrapped_model.gradient_checkpointing_enable() unwrapped_model.config.use_cache = False self.model.train() # Run PPO step stats = self.step(queries, responses, rewards) self.tokenizer.padding_side = "left" # restore padding side loss_meter.update(float(stats["ppo/loss/total"]), n=len(rewards)) reward_meter.update(torch.stack(rewards).mean().item(), n=len(rewards)) if self.config.log_with is not None: try: batch["query"] = self.tokenizer.batch_decode(queries, skip_special_tokens=True) batch["response"] = self.tokenizer.batch_decode(responses, skip_special_tokens=True) self.log_stats(stats, batch, rewards) except Exception: logger.warning("Failed to save stats due to unknown errors.") self.state.global_step += 1 self.log_callback.on_step_end(self.args, self.state, self.control) if self.is_local_process_zero() and (step + 1) % self.args.logging_steps == 0: logs = dict( loss=round(loss_meter.avg, 4), reward=round(reward_meter.avg, 4), learning_rate=stats["ppo/learning_rate"], epoch=round(step / steps_in_epoch, 2), ) tqdm.write(str(logs)) logs["step"] = step self.state.log_history.append(logs) self.log_callback.on_log(self.args, self.state, self.control) loss_meter.reset() reward_meter.reset() if (step + 1) % self.args.save_steps == 0: # save checkpoint self.save_model( os.path.join(self.args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, self.state.global_step)) ) self.save_callback.on_save( self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model) ) if self.control.should_epoch_stop or self.control.should_training_stop: break self.log_callback.on_train_end(self.args, self.state, self.control) self.save_callback.on_train_end( self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model) ) 这段代码是一个较为复杂的 PPO 训练循环实现，涵盖了从模型初始化到循环训练，以及日志记录和模型保存等功能。我们将逐行解释其核心内容：
python
复制
def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None: ppo_train 方法，该方法接受一个可选的参数 resume_from_checkpoint，用于未来版本支持从检查点恢复训练。 python
复制
if resume_from_checkpoint is not None: raise ValueError("`resume_from_checkpoint` will be supported in the future version.") resume_from_checkpoint 参数，当前版本会抛出异常，提示该功能将在未来版本中支持。 python
复制
total_train_batch_size = ( self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps * self.finetuning_args.ppo_buffer_size * self.args.world_size ) python
复制
if self.args.max_steps > 0: num_examples = total_train_batch_size * self.args.max_steps num_train_epochs = sys.maxsize max_steps = self.args.max_steps steps_in_epoch = self.args.max_steps else: len_dataloader = len(self.dataloader) num_examples = len(self.dataset) num_train_epochs = self.args.num_train_epochs max_steps = math.ceil(num_train_epochs * len_dataloader) steps_in_epoch = len_dataloader max_steps) 来决定训练的迭代方式和周期。如果指定了最大步数，则使用最大步数来计算总例数和训练周期，否则根据数据加载器和数据集的大小计算。 python
复制
self.state.max_steps = max_steps self.state.num_train_epochs = num_train_epochs self.state.is_local_process_zero = self.is_local_process_zero() self.state.is_world_process_zero = self.is_world_process_zero() python
复制
if self.is_world_process_zero(): ... python
复制
unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model) dataiter = iter(self.dataloader) loss_meter = AverageMeter() reward_meter = AverageMeter() self.log_callback.on_train_begin(self.args, self.state, self.control) python
复制
for step in tqdm(range(max_steps), disable=not self.is_local_process_zero()): ... tqdm 来展示进度条。循环将根据最大步数执行，并在本地主进程中显示进度。 接下来的代码块涉及处理数据批次、执行模型的推理和训练、记录统计数据、处理日志和保存模型的功能。由于篇幅限制，这里不对每一行代码进行详细解释，但大致流程包括：
从数据加载器中获取批次数据。 将模型置于推理模式，处理输入数据，获取奖励。 切换到训练模式，执行 PPO 训练步骤。 更新损失和奖励的度量，记录日志和统计信息。 根据设定的间隔保存模型和执行回调。这个过程循环，直到达到最大步数或满足特定的停止条件。整个训练循环详细地处理了模型的运行状态、性能监控以及与分布式训练相关的特定操作。
def create_optimizer( self, model: "AutoModelForCausalLMWithValueHead", training_args: "Seq2SeqTrainingArguments", finetuning_args: "FinetuningArguments", ) -> "torch.optim.Optimizer": optimizer = create_custom_optimzer(model, training_args, finetuning_args) if optimizer is None: decay_params, nodecay_params = [], [] decay_param_names = self.get_decay_parameter_names(model) for name, param in model.named_parameters(): if param.requires_grad: if name in decay_param_names: decay_params.append(param) else: nodecay_params.append(param) optim_class, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args) param_groups = [ dict(params=nodecay_params), dict(params=decay_params, weight_decay=training_args.weight_decay), ] optimizer = optim_class(param_groups, **optim_kwargs) return optimizer def create_scheduler( self, training_args: "Seq2SeqTrainingArguments", num_training_steps: int, optimizer: "torch.optim.Optimizer" ) -> "torch.optim.lr_scheduler.LRScheduler": create_custom_scheduler(training_args, num_training_steps, optimizer) lr_scheduler = get_scheduler( training_args.lr_scheduler_type, optimizer=optimizer, num_warmup_steps=training_args.get_warmup_steps(num_training_steps), num_training_steps=num_training_steps, ) return lr_scheduler @torch.no_grad() def get_inputs(self, batch: Dict[str, torch.Tensor]) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: r""" Generates model's responses given queries. """ if self.model_args.upcast_layernorm: layernorm_params = dump_layernorm(self.model) if batch["input_ids"].size(0) == 1: # handle llama2 ppo with gradient accumulation > 1 start_index = (batch["input_ids"][0] != self.tokenizer.pad_token_id).nonzero()[0].item() for k, v in batch.items(): batch[k] = v[:, start_index:] unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model) generate_output: torch.Tensor = unwrapped_model.generate( generation_config=self.generation_config, logits_processor=get_logits_processor(), **batch ) if self.model_args.upcast_layernorm: restore_layernorm(self.model, layernorm_params) query = batch["input_ids"].detach().cpu() response = generate_output[:, batch["input_ids"].size(-1) :].detach().cpu() queries, responses = [], [] for i in range(len(query)): query_start_index = (query[i] != self.tokenizer.pad_token_id).nonzero()[0].item() response_index = (response[i] != self.tokenizer.pad_token_id).nonzero() if len(response_index) == 0: response_length = 1 # allow empty response else: response_length = response_index[-1].item() + 1 queries.append(query[i, query_start_index:]) # remove padding from left responses.append(response[i, :response_length]) # remove padding from right return queries, responses @torch.no_grad() def get_rewards( self, queries: List[torch.Tensor], responses: List[torch.Tensor], unwrapped_model: "AutoModelForCausalLMWithValueHead", ) -> List[torch.Tensor]: r""" Computes scores using given reward model. Both inputs and outputs are put on CPU. """ if self.finetuning_args.reward_model_type == "api": token_ids = [torch.cat((q, r), dim=-1).tolist() for q, r in zip(queries, responses)] messages = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True) return get_rewards_from_server(self.reward_model, messages) if self.finetuning_args.reward_model_type == "lora": replace_model(unwrapped_model, target="reward") reward_model = self.model else: reward_model = self.reward_model batch = self.prepare_model_inputs(queries, responses) with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype): # support bf16 _, _, values = reward_model(**batch, output_hidden_states=True, return_dict=True, use_cache=False) if getattr(unwrapped_model.config, "model_type", None) == "chatglm": # assume same architecture values = torch.transpose(values, 0, 1) rewards = [] for i in range(values.size(0)): end_indexes = (batch["input_ids"][i] != self.tokenizer.pad_token_id).nonzero() end_index = end_indexes[-1].item() if len(end_indexes) else 0 rewards.append(values[i, end_index].float().detach().cpu()) # use fp32 type if self.finetuning_args.reward_model_type == "lora": replace_model(unwrapped_model, target="default") return rewards
