wehub-resource-sync
234 lines
12 KiB
Python
234 lines
12 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
# Part of the implementation is borrowed from huggingface/transformers.
|
|
import inspect
|
|
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
from contextlib import contextmanager, nullcontext
|
|
from peft import PeftModel
|
|
from torch import nn
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
from transformers import Seq2SeqTrainer as HfSeq2SeqTrainer
|
|
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
|
|
from transformers.utils import is_peft_available
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
|
|
|
from swift.infer_engine import InferRequest, RequestConfig, TransformersEngine
|
|
from swift.sequence_parallel import sequence_parallel
|
|
from swift.utils import HfConfigFactory, JsonlWriter, Serializer, gc_collect, get_logger, unwrap_model_for_generation
|
|
from .arguments import Seq2SeqTrainingArguments
|
|
from .mixin import DataLoaderMixin, SwiftMixin
|
|
from .utils import per_token_loss_func, per_token_loss_func_sp
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class Seq2SeqTrainer(SwiftMixin, DataLoaderMixin, HfSeq2SeqTrainer):
|
|
args: Seq2SeqTrainingArguments
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.model_accepts_loss_kwargs = True # fix transformers>=4.46.2
|
|
if self.template.model_accepts_loss_kwargs is not None:
|
|
self.model_accepts_loss_kwargs = self.template.model_accepts_loss_kwargs
|
|
if self.args.predict_with_generate:
|
|
self.infer_engine = TransformersEngine(
|
|
self.model, template=self.template, max_batch_size=self.args.per_device_eval_batch_size)
|
|
self.jsonl_writer = JsonlWriter(os.path.join(self.args.output_dir, 'predict.jsonl'))
|
|
|
|
@staticmethod
|
|
def _predict_data_collator(batch):
|
|
return {'_data': batch}
|
|
|
|
@contextmanager
|
|
def _patch_predict_with_generate(self):
|
|
origin_data_collator = self.data_collator
|
|
self.data_collator = self._predict_data_collator
|
|
packing = self.template.packing
|
|
padding_free = self.template.padding_free
|
|
self.template.packing = False
|
|
self.template.padding_free = False
|
|
try:
|
|
yield
|
|
finally:
|
|
self.template.packing = packing
|
|
self.template.padding_free = padding_free
|
|
self.data_collator = origin_data_collator
|
|
|
|
def evaluate(self, *args, **kwargs):
|
|
context = self._patch_predict_with_generate() if self.args.predict_with_generate else nullcontext()
|
|
with context:
|
|
res = super().evaluate(*args, **kwargs)
|
|
gc_collect()
|
|
return res
|
|
|
|
def prediction_step(
|
|
self,
|
|
model: nn.Module,
|
|
inputs: Dict[str, Union[torch.Tensor, Any]],
|
|
prediction_loss_only: bool,
|
|
ignore_keys: Optional[List[str]] = None,
|
|
**gen_kwargs,
|
|
) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
|
|
if not self.args.predict_with_generate or prediction_loss_only:
|
|
with self.template.forward_context(self.model, inputs):
|
|
return super().prediction_step(
|
|
model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys)
|
|
data_list = inputs['_data']
|
|
labels_list = [InferRequest.remove_response(data['messages']) for data in data_list]
|
|
with unwrap_model_for_generation(
|
|
self.model_wrapped, self.accelerator,
|
|
gather_deepspeed3_params=self.args.ds3_gather_for_generation), self.template.generate_context():
|
|
resp_list = self.infer_engine.infer(
|
|
data_list,
|
|
RequestConfig(max_tokens=self.model.generation_config.max_new_tokens),
|
|
use_tqdm=False,
|
|
)
|
|
|
|
response_list = []
|
|
jsonl_cache = []
|
|
device = self.args.device
|
|
for data, resp, labels in zip(data_list, resp_list, labels_list):
|
|
response = resp.choices[0].message.content
|
|
jsonl_cache.append({'response': response, 'labels': labels, **data})
|
|
response_list.append(Serializer.to_tensor(resp.choices[0].message.content).to(device=device))
|
|
self.jsonl_writer.append(jsonl_cache, gather_obj=True)
|
|
labels_list = [Serializer.to_tensor(labels).to(device=device) for labels in labels_list]
|
|
response_list = pad_sequence(response_list, batch_first=True, padding_value=0)
|
|
labels_list = pad_sequence(labels_list, batch_first=True, padding_value=0)
|
|
return None, response_list, labels_list
|
|
|
|
def _prepare_inputs(self, inputs):
|
|
args = self.args
|
|
inputs = super()._prepare_inputs(inputs)
|
|
if self.template.sequence_parallel_size > 1:
|
|
sequence_parallel.prepare_inputs(inputs)
|
|
|
|
use_logits_to_keep = self.get_use_logits_to_keep(self.template.sequence_parallel_size == 1)
|
|
if use_logits_to_keep:
|
|
self.prepare_logits_to_keep(inputs)
|
|
if args.tuner_backend == 'unsloth' and isinstance(inputs['logits_to_keep'], torch.Tensor):
|
|
inputs['logits_to_keep'] = int(inputs['logits_to_keep'].sum())
|
|
|
|
base_model = self.template.get_base_model(self.model)
|
|
forward_params = inspect.signature(base_model.forward).parameters
|
|
if self.model.model_info.is_moe_model and any(key in forward_params
|
|
for key in ['output_router_logits', 'kwargs']):
|
|
HfConfigFactory.set_config_attr(base_model.config, 'router_aux_loss_coef', args.router_aux_loss_coef)
|
|
base_model.router_aux_loss_coef = args.router_aux_loss_coef
|
|
logger.info_once(f'router_aux_loss_coef: {args.router_aux_loss_coef}')
|
|
if args.router_aux_loss_coef > 0:
|
|
inputs['output_router_logits'] = True
|
|
inputs['compute_loss_func'] = self.compute_loss_func
|
|
return inputs
|
|
|
|
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
|
|
labels = None
|
|
compute_loss_func: Callable = inputs.pop('compute_loss_func', None)
|
|
loss_scale = inputs.pop('loss_scale', None)
|
|
text_position_ids = inputs.pop('text_position_ids', None)
|
|
if text_position_ids is None:
|
|
text_position_ids = inputs.get('position_ids')
|
|
channels = inputs.pop('channel', None)
|
|
|
|
if (self.label_smoother is not None or compute_loss_func is not None or loss_scale is not None
|
|
or self.args.enable_dft_loss or self.args.enable_channel_loss
|
|
or self.template.sequence_parallel_size > 1) and 'labels' in inputs:
|
|
if self.args.use_liger_kernel:
|
|
logger.warning_once('The cross_entropy loss function defined in Liger Kernel will not '
|
|
'take effect, potentially leading to increased GPU memory consumption.')
|
|
labels = inputs.pop('labels')
|
|
outputs = self.template.compute_sft_loss(model, inputs, num_items_in_batch=num_items_in_batch, trainer=self)
|
|
mode = 'train' if self.model.training else 'eval'
|
|
if getattr(outputs, 'aux_loss', None) is not None:
|
|
self.custom_metrics[mode]['aux_loss'].update(outputs.aux_loss)
|
|
# Save past state if it exists
|
|
# TODO: this needs to be fixed and made cleaner later.
|
|
if hasattr(self.args, 'past_index') and self.args.past_index >= 0:
|
|
self._past = outputs[self.args.past_index]
|
|
|
|
if labels is None:
|
|
labels = inputs['labels']
|
|
if isinstance(outputs, dict) and 'loss' not in outputs:
|
|
raise ValueError(
|
|
'The model did not return a loss from the inputs, only the following keys: '
|
|
f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}.")
|
|
# We don't use .loss here since the model may return tuples instead of ModelOutput.
|
|
loss = outputs['loss'] if isinstance(outputs, dict) else outputs[0]
|
|
else:
|
|
outputs.loss = None
|
|
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)
|
|
else:
|
|
outputs.loss = per_token_loss_func(outputs, labels, enable_dft_loss=self.args.enable_dft_loss)
|
|
|
|
if loss_scale is not None:
|
|
loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1).view(-1)
|
|
outputs.loss = outputs.loss * loss_scale
|
|
|
|
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.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]
|
|
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_]])
|
|
|
|
unwrapped_model = self.accelerator.unwrap_model(model)
|
|
if is_peft_available() and isinstance(unwrapped_model, PeftModel):
|
|
model_name = unwrapped_model.model._get_name()
|
|
else:
|
|
model_name = unwrapped_model._get_name()
|
|
# User-defined compute_loss function
|
|
if compute_loss_func is not None:
|
|
loss = compute_loss_func(
|
|
outputs, labels, num_items_in_batch=num_items_in_batch, loss_scale=loss_scale, trainer=self)
|
|
elif self.label_smoother is None:
|
|
# Handle the outputs.loss generated by loss_scale.
|
|
if num_items_in_batch is None:
|
|
# https://github.com/huggingface/transformers/blob/9dff7ca5c9693f4c02cdd2a9c2abc4772fcea5da/src/transformers/trainer.py#L2137
|
|
num_items_in_batch = (labels != -100).sum() # compat SP
|
|
if self.template.sequence_parallel_size > 1:
|
|
# labels are sharded by SP; outputs.loss was gathered
|
|
# to full length via GatherLoss. Reduce the denominator
|
|
# across the SP group so it matches the gathered loss.
|
|
dist.all_reduce(num_items_in_batch, op=dist.ReduceOp.SUM)
|
|
loss = outputs.loss.sum() / num_items_in_batch
|
|
else:
|
|
if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
|
|
loss = self.label_smoother(outputs, labels, shift_labels=True)
|
|
else:
|
|
loss = self.label_smoother(outputs, labels)
|
|
|
|
if self.model.model_info.is_moe_model and self.args.router_aux_loss_coef is not None:
|
|
aux_loss = outputs.get('aux_loss')
|
|
if aux_loss is not None:
|
|
if num_items_in_batch is not None:
|
|
aux_loss = aux_loss * ((labels[:, 1:] != -100).sum() / num_items_in_batch)
|
|
loss = loss + self.args.router_aux_loss_coef * aux_loss.to(loss.device)
|
|
|
|
if getattr(self.args, 'average_tokens_across_devices',
|
|
False) and self.model_accepts_loss_kwargs and num_items_in_batch is not None:
|
|
loss *= self.accelerator.num_processes
|
|
if mode == 'eval' and self.template.sequence_parallel_size > 1:
|
|
loss /= self.template.sequence_parallel_size
|
|
|
|
if (outputs.logits is not None and labels is not None and self.args.tuner_backend != 'unsloth'):
|
|
cu_seqlens = None
|
|
if self.template.padding_free and self.args.acc_strategy == 'seq':
|
|
cu_seqlens = self.get_cu_seqlens(text_position_ids, inputs.get('logits_to_keep'))
|
|
# Liger does not have logits
|
|
# Unsloth has a bug with output logits
|
|
self._compute_acc(outputs, labels, cu_seqlens=cu_seqlens)
|
|
return (loss, outputs) if return_outputs else loss
|
|
|
|
def training_step(self, model, inputs, *args, **kwargs):
|
|
with self.template.forward_context(self.model, inputs):
|
|
return super().training_step(model, inputs, *args, **kwargs)
|