继续上一篇文章的内容我们来做一个Reward model。1、加载模型base_model AutoModelForCausalLM.from_pretrained(./llm_sft).to(device) base_model这是在上一篇文章当中训练好的模型我们查看一下模型结构。上述代码输出为LlamaForCausalLM((model): LlamaModel((embed_tokens): Embedding(49152, 576, padding_idx2)(layers): ModuleList((0-29): 30 x LlamaDecoderLayer((self_attn): LlamaAttention((q_proj): Linear(in_features576, out_features576, biasFalse)(k_proj): Linear(in_features576, out_features192, biasFalse)(v_proj): Linear(in_features576, out_features192, biasFalse)(o_proj): Linear(in_features576, out_features576, biasFalse))(mlp): LlamaMLP((gate_proj): Linear(in_features576, out_features1536, biasFalse)(up_proj): Linear(in_features576, out_features1536, biasFalse)(down_proj): Linear(in_features1536, out_features576, biasFalse)(act_fn): SiLUActivation())(input_layernorm): LlamaRMSNorm((576,), eps1e-05)(post_attention_layernorm): LlamaRMSNorm((576,), eps1e-05)))(norm): LlamaRMSNorm((576,), eps1e-05)(rotary_emb): LlamaRotaryEmbedding())(lm_head): Linear(in_features576, out_features49152, biasFalse))(embed_tokens): Embedding(49152, 576, padding_idx2) 这里是一个超大的嵌入向量相当于字典对于某一个token都有一个576长度的嵌入一共有49152种tokenpadding_idx2表示第二个token是paddingtoken。(lm_head): Linear(in_features576, out_features49152, biasFalse)是输出了一个49152维度的向量每一个维度都是对应token的出现概率然后采样获得下一个token。对于一个奖励模型来说我们不会要某一个token出现的分数而是对于这些已存在的token的打分因此在这里可以把最后一层的线性层变成576-1的超级简单的线性层目的就是输出一个最终score。2、更改网络最后一层from torch import nn # define model class class RewardModel(nn.Module): def __init__( self, base_model, ): super().__init__() # replace final linear layer self.base_model base_model # 把base_model的lm_head替换为一个新的线性层 self.base_model.__setattr__( lm_head, nn.Linear(576, 1, biasFalse).to(device)) def forward(self, input_ids, attention_mask): output self.base_model( input_idsinput_ids, attention_maskattention_mask, ).logits # .logits是因为只改了模型层没有改参数名这里其实就是把原来输出词表大小的logits改成输出一个数值reward score了。 # need only final output in sequence # input_ids.shape: torch.Size([3, 23]) batch_size input_ids.shape[0] # token_indices.shape: torch.Size([23]) token_indices torch.arange(input_ids.shape[-1], dtypeint).to(device) # last_nonpad_indices.shape: torch.Size([3]) # 这里为什么还要获得最后一个非padding token的索引 # 事项一下句子1 [A B C D PAD]句子2 [E F G PAD PAD] # 我们打分的时候看到ABCD打分就可以了PAD不应该参与打分所以需要找到最后一个非PAD token的位置。 # 因此才需要last_nonpad_indices为[3, 2]表示每个序列中最后一个非pad token的索引位置。 last_nonpad_indices (token_indices * attention_mask).argmax(-1) # reward_score.shape: torch.Size([3]) reward_score output[torch.arange(batch_size).to(device), last_nonpad_indices] return reward_score # generate model instance rm RewardModel(base_model).to(device)3、数据过滤# 1. add chosen_len and rejected_len column # (which indicates the sequence length in chosen and rejected) def add_seq_len(example): # 计算文本的tokenized长度并把长度添加到example中 def get_tokenized_length(text): tokenized tokenizer(text) return len(tokenized[input_ids]) chosen_len get_tokenized_length(example[chosen]) reject_len get_tokenized_length(example[rejected]) return { chosen_len: chosen_len, rejected_len: reject_len } train_data train_data.map(add_seq_len) # 2. remove rows which exceed the maximum sequence length # 为了后续训练我们需要移除那些超过最大序列长度的样本。 train_data train_data.filter(lambda example: example[chosen_len] max_seq_len and example[rejected_len] max_seq_len) train_data train_data.remove_columns([chosen_len, rejected_len]) # show total number of filtered rows train_data4、开始训练奖励模型batch_size 4 gradient_accumulation_steps 8 # 梯度累加步数 # 为了避免内存不足我们可以使用梯度累加步数。 # 即每次只计算一部分样本的梯度然后累加起来再更新模型参数。 def collate_batch(batch): # tokenize (convert to token ids and attention mask) and convert to tensor chosen_list [item[chosen] for item in batch] reject_list [item[rejected] for item in batch] chosen_tensor tokenizer( chosen_list, paddingTrue, padding_sideleft, # see above description ! return_tensorspt).to(device) reject_tensor tokenizer( reject_list, paddingTrue, padding_sideleft, # see above description ! return_tensorspt).to(device) return chosen_tensor, reject_tensor dataloader DataLoader( train_data, batch_sizebatch_size, shuffleTrue, collate_fncollate_batch )正常Tesla T4(16G)的batch size为4够用在这里我用的是2080Ti12G所以把batchsize改为2可以正常训练。num_epochs 1 num_steps math.ceil(len(dataloader) / gradient_accumulation_steps) torch.set_default_dtype(torch.bfloat16) optimizer torch.optim.AdamW( paramsrm.parameters(), lr3.0e-5, betas(0.9, 0.999), eps1e-08, ) def _get_cosine_schedule( current_step: int, num_training_steps: int, num_warmup_steps: int0, linear_warmup: boolFalse, min_value: float0.0, ): if current_step num_warmup_steps: if linear_warmup: return min(1.0, (current_step 1) / (num_warmup_steps 1)) # see https://arxiv.org/abs/2410.11020 else: return 1.0 progress float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) scale 0.5 * (1.0 math.cos(math.pi * progress)) return (1.0 - min_value) * scale min_value scheduler LambdaLR(optimizer, lr_lambdafunctools.partial( _get_cosine_schedule, num_training_stepsnum_epochs*num_steps, num_warmup_stepsmath.ceil(num_epochs*num_steps*0.1))) # remove log file if exists log_file loss_rm.log if os.path.exists(log_file): os.remove(log_file) # iterate epoch for epoch in range(num_epochs): rm.train() optimizer.zero_grad() record_loss [] # iterate batch for i, (chosen, reject) in enumerate(dataloader): with torch.set_grad_enabled(True): # get reward score (chosen) rewards_chosen rm( input_idschosen[input_ids], attention_maskchosen[attention_mask], ) # get reward score (rejected) rewards_reject rm( input_idsreject[input_ids], attention_maskreject[attention_mask], ) # compute loss loss -F.logsigmoid(rewards_chosen - rewards_reject).mean() # 损失函数计算为 -log(sigmoid(reward_chosen - reward_reject)) # 即最大化 chosen 样本的 reward最小化 reject 样本的 reward。 loss.backward() # optimization by accumulation if ((i 1) % gradient_accumulation_steps 0) or \ (i 1 len(dataloader)): optimizer.step() scheduler.step() optimizer.zero_grad() # print log record_loss.append(loss.item()) print(fEpoch {epoch1} (iter{i1}) {math.ceil((i 1) / gradient_accumulation_steps)}/{num_steps} - loss {loss :5.4f}, end\r) # save logging epoch_average_loss sum(record_loss)/len(record_loss) print(fEpoch {epoch1} (iter{i1}) {math.ceil((i 1) / gradient_accumulation_steps)}/{num_steps} - loss {epoch_average_loss :5.4f}) with open(log_file, a) as f: for l in record_loss: f.write(%s\n %l) # save checkpoint torch.save(rm.state_dict(), rm.pt) print(Done)5、测试reward modelimport pandas as pd # # uncomment when you load RM from the trained checkpoint # rm.load_state_dict(torch.load(rm.pt, weights_onlyTrue)) rm.eval() test_list [ |im_start|user\nWhat do you most want to do right now?|im_end|\n|im_start|assistant\nRight now, I most want to assist you and provide helpful, respectful, and engaging interactions. If you have any questions, need information, or just want to chat, Im here for you!|im_end|, |im_start|user\nWhat do you most want to do right now?|im_end|\n|im_start|assistant\nId love to go for a hike in the woods! Its such a beautiful day and I really need to get some fresh air. How about you? |im_end|, |im_start|user\nWhat is the best gift to give a friend who loves the outdoors?|im_end|\n|im_start|assistant\nThats a thoughtful question! Here are a few gift ideas for a friend who loves the outdoors.: High-quality Multi-tool, Reusable Water Bottle with Filter, or Portable Hammock|im_end|, |im_start|user\nWhat is the best gift to give a friend who loves the outdoors?|im_end|\n|im_start|assistant\nOoh, thats a fun question!! If youre looking for some gift ideas for a friend who loves the outdoors, I suggest National Parks Pass, Hammock, or Portable Camping Chair! |im_end|, ] test_batch tokenizer( test_list, paddingTrue, padding_sideleft, return_tensorspt).to(device) test_rewards rm( input_idstest_batch[input_ids], attention_masktest_batch[attention_mask], ) df pd.DataFrame({ text: test_list, score: test_rewards.squeeze(-1).tolist() }) dftextscore0|im_start|user\nWhat do you most want to do ...-2.6406251|im_start|user\nWhat do you most want to do ...5.4062502|im_start|user\nWhat is the best gift to giv...3.5156253|im_start|user\nWhat is the best gift to giv...5.437500这样的话价值模型就训练好了。