This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,233 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user