diff --git a/swift/trainers/seq2seq_trainer.py b/swift/trainers/seq2seq_trainer.py index 3b35658fc1..78f51559b6 100644 --- a/swift/trainers/seq2seq_trainer.py +++ b/swift/trainers/seq2seq_trainer.py @@ -163,7 +163,15 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N if (self.args.enable_dft_loss or loss_scale is not None or self.args.enable_channel_loss or self.template.sequence_parallel_size > 1): if self.template.sequence_parallel_size > 1: - outputs.loss = per_token_loss_func_sp(outputs, labels, enable_dft_loss=self.args.enable_dft_loss) + sp_loss = per_token_loss_func_sp( + outputs, + labels, + enable_dft_loss=self.args.enable_dft_loss, + return_labels=self.args.enable_channel_loss) + if self.args.enable_channel_loss: + outputs.loss, channel_labels = sp_loss + else: + outputs.loss = sp_loss if loss_scale is not None: position_ids = sequence_parallel.real_position_ids if position_ids is not None: @@ -183,15 +191,19 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N if self.args.enable_channel_loss: metrics = self.custom_metrics[mode] - masks = torch.roll(labels, shifts=-1, dims=-1).view(-1) != -100 + if self.template.sequence_parallel_size == 1: + channel_labels = torch.roll(labels, shifts=-1, dims=-1) + channel_loss = outputs.loss.view(-1) + masks = channel_labels.view(-1) != -100 if self.template.padding_free: cu_seqlens = self.get_cu_seqlens(text_position_ids, inputs.get('logits_to_keep')) else: - cu_seqlens = torch.arange(0, labels.shape[0] + 1) * labels.shape[1] + seq_len = channel_labels.numel() // labels.shape[0] + cu_seqlens = torch.arange(0, labels.shape[0] + 1) * seq_len for i in range(cu_seqlens.shape[0] - 1): channel = None if channels is None else channels[i] slice_ = slice(cu_seqlens[i], cu_seqlens[i + 1]) - metrics[f'loss_{channel}'].update(outputs.loss[slice_][masks[slice_]]) + metrics[f'loss_{channel}'].update(channel_loss[slice_][masks[slice_]]) unwrapped_model = self.accelerator.unwrap_model(model) if is_peft_available() and isinstance(unwrapped_model, PeftModel): diff --git a/swift/trainers/utils.py b/swift/trainers/utils.py index 1c6442a7c5..8368979250 100644 --- a/swift/trainers/utils.py +++ b/swift/trainers/utils.py @@ -164,7 +164,7 @@ def is_instance_of_ms_model(model: Module) -> bool: return False -def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, **kwargs) -> torch.Tensor: +def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, return_labels=False, **kwargs): """Common loss function for sequence parallel training""" if hasattr(outputs, 'logits'): logits = outputs.logits @@ -192,7 +192,10 @@ def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, **kwargs) -> if position_ids is not None and position_ids.min() == -1: _pos_mask = position_ids >= 0 loss = loss[_pos_mask].contiguous() + labels = labels[_pos_mask].contiguous() + if return_labels: + return loss, labels return loss