LLaMA-Factory自定义权重loss_llamafactory自定义损失函数
动机：训练带思考过程的模型的时候想权重分配控制，比如思考部分loss权重0.8，答案部分权重1.2，这需要自定义重写LLaMA-Factory/src/llamafactory/train/sft/trainer.py里面的compute_loss方法，参考代码如下，需要注意如下几点：
1、gradient_accumulation_steps切记设置为1，我还不会考虑gradient_accumulation_steps大于1的情况。
2、注意我的训练数据使用</think>标签区分思考与答案的。
3、如下两个参数控制权重：weights[i][think_part] = 1.0 weights[i][answer_part] = 1.0
4、非deepseek模型不带think标签需要补充：--new_special_tokens \",\"
@override def compute_loss(self, model, inputs, return_outputs=False, *args, **kwargs): rets = super().compute_loss(model, inputs, *args, **kwargs) # # 获取模型输出 # model_output = model(**inputs) # # lm_logits shape: [2, 224, 151936] # lm_logits = model_output[\"logits\"] # loss = None # # labels shape:[2, 224] # labels = inputs[\"labels\"] # # 给labels加cuda，和lm_logits使用相同的cuda # labels = labels.to(lm_logits.device) # # shift_logits shape: [2, 223, 151936]，最后一个是结束字符 # shift_logits = lm_logits[..., :-1, :].contiguous() # # shift_labels shape: [2, 223]，第一个是开始字符 # shift_labels = labels[..., 1:].contiguous() # loss_fct = nn.CrossEntropyLoss(ignore_index=-100) # # shift_logits.size(-1) = 151936 # # shift_logits.view(-1, shift_logits.size(-1))= shift_logits.view(-1, 151936) 的作用 # # shape [2, 223, 151936] --> [446, 151936] # # shift_labels.view(-1)的作用：shape [2, 223] --> [446] # loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) # print(\'outputs[\"loss\"]=\', model_output[\"loss\"]) # print(\'loss=\', loss) # 获取模型输出 outputs = model(**inputs) logits = outputs.logits logits = logits[..., :-1, :].contiguous() labels = inputs[\"labels\"] labels = labels[..., 1:].contiguous() input_ids = inputs[\"input_ids\"] attention_mask = inputs[\"attention_mask\"] # 获取 的 token ID think_end_token = self.processing_class.convert_tokens_to_ids(\"\") # 初始化权重矩阵（默认权重为 0，忽略无效位置） batch_size, seq_len = labels.shape weights = torch.zeros_like(labels, dtype=torch.float32) for i in range(batch_size): # 获取当前样本的 input_ids 和 attention_mask sample_input_ids = input_ids[i] # sample_mask = attention_mask[i] sample_label = labels[i] # 找到有效 token 的位置（排除 padding） # valid_indices = (sample_mask == 1).nonzero(as_tuple=True)[0] valid_indices = (sample_label != -100).nonzero(as_tuple=True)[0] # print(\'valid_indices=\', valid_indices) # 在有效 token 中查找 的位置 try: # 查找第一个 的位置（假设每个样本只有一个） think_end_pos = (sample_input_ids[valid_indices] == think_end_token).nonzero()[0].item() # think_end_global_pos = valid_indices[think_end_pos] # 在完整序列中的位置 # print(\'think_end_pos=\', think_end_pos) # print(\'think_end_global_pos=\', think_end_global_pos) except IndexError: # 若未找到 ，全部视为答案部分（权重 1.0） weights[i][valid_indices] = 1.0 continue # 分割思考部分和答案部分 think_part = valid_indices[:think_end_pos + 1] # 包含 answer_part = valid_indices[think_end_pos + 1:] # 答案部分 # 设置权重（思考部分 0.8，答案部分 1.0） weights[i][think_part] = 1.0 weights[i][answer_part] = 1.0 # 计算加权损失（仅 labels != -100 的位置参与计算） loss = F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), reduction=\"none\", ignore_index=-100 # 自动忽略 labels=-100 的位置 ).view(batch_size, seq_len) # print(self.args.average_tokens_across_devices) # print((self.model_accepts_loss_kwargs or self.compute_loss_func)) # 应用权重矩阵并归一化 weighted_loss = (loss * weights).sum() / weights.sum() # print(\'rets= \', rets) # print(\'outputs[\"loss\"]=\', outputs[\"loss\"]) print(\'weighted_loss=\', weighted_loss) return (weighted_loss, logits) if return_outputs else weighted_loss
