wehub-resource-sync
1347 lines
63 KiB
Python
1347 lines
63 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
# Part of the implementation is borrowed from huggingface/transformers.
|
|
import collections
|
|
import datasets
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import numpy as np
|
|
import os
|
|
import random
|
|
import re
|
|
import safetensors
|
|
import shutil
|
|
import time
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.utils.checkpoint
|
|
import transformers
|
|
import warnings
|
|
from contextlib import contextmanager
|
|
from copy import copy
|
|
from datasets import Dataset as HfDataset
|
|
from functools import partial, wraps
|
|
from modelscope import check_local_model_is_latest
|
|
from packaging import version
|
|
from peft import PeftModel
|
|
from torch.utils.data import DataLoader
|
|
from transformers import PreTrainedModel
|
|
from transformers.integrations import is_deepspeed_zero3_enabled
|
|
from transformers.modeling_utils import unwrap_model
|
|
from transformers.trainer import OPTIMIZER_NAME, PREFIX_CHECKPOINT_DIR, SCHEDULER_NAME, TRAINER_STATE_NAME, ParallelMode
|
|
from transformers.trainer import Trainer as HfTrainer
|
|
from transformers.trainer import reissue_pt_warnings
|
|
from transformers.trainer_utils import IntervalStrategy
|
|
|
|
try:
|
|
from transformers.trainer_utils import sort_checkpoints
|
|
except ImportError:
|
|
sort_checkpoints = None
|
|
from types import MethodType
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
from swift.callbacks import callbacks_map
|
|
from swift.dataloader import BatchSamplerShard, DataLoaderDispatcher, DataLoaderShard
|
|
from swift.hub import get_hub
|
|
from swift.loss import loss_map
|
|
from swift.metrics import MeanMetric, compute_acc, eval_metrics_map
|
|
from swift.model import get_llm_model, get_lm_head_model, save_checkpoint
|
|
from swift.model.patcher import gather_sequence_parallel_outputs, revert_padding_free, transformers_seq_cls_forward
|
|
from swift.optimizers import OptimizerCallback, optimizers_map
|
|
from swift.sequence_parallel import SequenceParallelDispatcher, SequenceParallelSampler, sequence_parallel
|
|
from swift.template import Template, update_generation_config_eos_token
|
|
from swift.tuner_plugin import tuners_map
|
|
from swift.tuners import SwiftModel
|
|
from swift.utils import (HfConfigFactory, copy_files_by_pattern, deep_getattr, get_current_device, get_logger,
|
|
get_packed_seq_params, is_dist, is_mp, is_mp_ddp, ms_logger_context, seed_worker)
|
|
from .arguments import TrainingArguments
|
|
from .utils import (can_return_loss, dynamic_gradient_checkpointing, find_labels, get_function, get_resume_dir,
|
|
is_instance_of_ms_model, patch_modelscope_hub_timeout, replace_index_file)
|
|
|
|
logger = get_logger()
|
|
|
|
transformers_5 = version.parse(transformers.__version__) >= version.parse('5.0.0')
|
|
|
|
|
|
class SwiftMixin:
|
|
FLASH_CKPT_WAIT_TIMEOUT = 1800
|
|
|
|
def __init__(self,
|
|
model: PreTrainedModel,
|
|
args: TrainingArguments,
|
|
template: Template,
|
|
train_dataset: HfDataset,
|
|
eval_dataset: Optional[HfDataset] = None,
|
|
**kwargs) -> None:
|
|
if not hasattr(train_dataset, '__len__') and args.dataloader_num_workers > 1:
|
|
args.dataloader_num_workers = 1
|
|
logger.warning('Using IterableDataset, setting args.dataloader_num_workers to 1.')
|
|
self.compute_loss_func = None # Compatible with the older version of transformers
|
|
self.template = template
|
|
|
|
self.is_encoder_decoder = self.template.is_encoder_decoder
|
|
self.padding_free = self.template.padding_free
|
|
self.task_type = self.template.task_type
|
|
self.problem_type = getattr(model.config, 'problem_type', None)
|
|
self.optimizer_callback = optimizers_map[args.optimizer or 'default'](args, self)
|
|
if args.check_model and hasattr(model, 'model_dir'):
|
|
with ms_logger_context(logging.CRITICAL), patch_modelscope_hub_timeout():
|
|
config_info = self._collect_config_info()
|
|
config_info.update({
|
|
'invoked_by': 'local_trainer',
|
|
'third_party': 'swift',
|
|
'trainer_class': self.__class__.__name__,
|
|
})
|
|
check_local_model_is_latest(model.model_dir, user_agent=config_info)
|
|
if eval_dataset is None and args:
|
|
if getattr(args, 'eval_dataset', None):
|
|
# Avoid trainer throwing errors.
|
|
eval_dataset = []
|
|
else:
|
|
args.evaluation_strategy = IntervalStrategy.NO
|
|
args.eval_strategy = IntervalStrategy.NO
|
|
|
|
def _get_mean_metric():
|
|
return MeanMetric(nan_value=None, device=args.device)
|
|
|
|
self.custom_metrics = {
|
|
'train': collections.defaultdict(_get_mean_metric),
|
|
'eval': collections.defaultdict(_get_mean_metric)
|
|
}
|
|
self.hub = get_hub()
|
|
|
|
self.model_meta = model.model_meta
|
|
self.model_info = model.model_info
|
|
|
|
data_collator = self._get_data_collator(args, template)
|
|
kwargs.update(self.create_loss_and_eval_metric(args))
|
|
trainer_parameters = inspect.signature(HfTrainer.__init__).parameters
|
|
tokenizer_key = 'processing_class' if 'processing_class' in trainer_parameters else 'tokenizer'
|
|
kwargs[tokenizer_key] = template.tokenizer
|
|
# Pass callbacks in __init__ to correctly invoke on_init_end
|
|
callbacks = self._get_callbacks(args)
|
|
with self.hub.patch_hub():
|
|
super().__init__(
|
|
model=model,
|
|
args=args,
|
|
data_collator=data_collator,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
callbacks=callbacks,
|
|
**kwargs)
|
|
# fix https://github.com/huggingface/transformers/pull/43919
|
|
if transformers_5:
|
|
self.accelerator.gradient_state.plugin_kwargs['num_steps'] = 1
|
|
if get_function(model.__class__.forward) is not get_function(model.forward):
|
|
self.label_names = find_labels(model)
|
|
self.can_return_loss = can_return_loss(model)
|
|
self.label_names = self.label_names or ['labels']
|
|
self.start_time = time.time()
|
|
self._fix_gradient_checkpointing()
|
|
self._patch_tasks()
|
|
update_generation_config_eos_token(self.model.generation_config, self.template)
|
|
if getattr(self.model, 'origin_generation_config', None):
|
|
self.model.origin_generation_config.eos_token_id = self.model.generation_config.eos_token_id
|
|
if self.args.resume_only_model and self.args.ignore_data_skip:
|
|
# The weights have already been loaded outside the trainer,
|
|
# so reading train_state is skipped here.
|
|
self.args.resume_from_checkpoint = None
|
|
|
|
def _get_data_collator(self, args, template):
|
|
padding_to = template.max_length if args.tuner_type == 'longlora' else None
|
|
return partial(template.data_collator, padding_to=padding_to)
|
|
|
|
def _get_callbacks(self, args):
|
|
callbacks = []
|
|
for callback in args.callbacks:
|
|
callbacks.append(callbacks_map[callback](args, self))
|
|
return callbacks
|
|
|
|
def _collect_config_info(self) -> Dict[str, str]:
|
|
"""
|
|
Collects trainer-specific configuration details.
|
|
|
|
Subclasses can override this method to provide additional configuration
|
|
information for model compatibility verification.
|
|
|
|
Returns:
|
|
Dict[str, str]: Configuration parameters as key-value pairs.
|
|
"""
|
|
if self.__class__.__name__ == 'Seq2SeqTrainer':
|
|
if not self.template.use_chat_template:
|
|
return {
|
|
'seq2seq_mode': 'pt',
|
|
}
|
|
else:
|
|
return {
|
|
'seq2seq_mode': 'sft',
|
|
}
|
|
return {}
|
|
|
|
@property
|
|
def tokenizer(self):
|
|
# compat transformers5.0
|
|
return self.processing_class
|
|
|
|
@contextmanager
|
|
def _patch_deepspeed_load_checkpoint(self):
|
|
from transformers import trainer
|
|
if not self.args.resume_from_checkpoint or not self.args.resume_only_model or not hasattr(
|
|
trainer, 'deepspeed_load_checkpoint'):
|
|
yield
|
|
return
|
|
origin_deepspeed_load_checkpoint = trainer.deepspeed_load_checkpoint
|
|
|
|
def deepspeed_load_checkpoint(*args, **kwargs):
|
|
try:
|
|
return origin_deepspeed_load_checkpoint(*args, **kwargs)
|
|
except Exception as e:
|
|
logger.warning('Failed to call deepspeed_load_checkpoint function. '
|
|
f'If `--resume_only_model true` is set, this warning can be ignored. {e}.')
|
|
|
|
trainer.deepspeed_load_checkpoint = deepspeed_load_checkpoint
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
trainer.deepspeed_load_checkpoint = origin_deepspeed_load_checkpoint
|
|
|
|
def get_use_logits_to_keep(self, default_value: bool = True):
|
|
use_logits_to_keep = self.args.use_logits_to_keep
|
|
if use_logits_to_keep is None:
|
|
base_model = self.template.get_base_model(self.model)
|
|
if self.model.model_meta.is_multimodal and not transformers_5:
|
|
use_logits_to_keep = False
|
|
elif 'logits_to_keep' not in inspect.signature(base_model.forward).parameters:
|
|
use_logits_to_keep = False
|
|
else:
|
|
use_logits_to_keep = default_value
|
|
self.args.use_logits_to_keep = use_logits_to_keep
|
|
logger.info_once(f'use_logits_to_keep: {use_logits_to_keep}')
|
|
return use_logits_to_keep
|
|
|
|
def _save_initial_model(self, output_dir):
|
|
# pissa/olora/lora-ga
|
|
model = unwrap_model(self.model)
|
|
if isinstance(model, PeftModel):
|
|
config = model.peft_config.get('default')
|
|
init_lora_weights = getattr(config, 'init_lora_weights', None)
|
|
if (isinstance(init_lora_weights, str)
|
|
and any(s in init_lora_weights for s in ('pissa', 'olora', 'lora-ga'))):
|
|
config.init_lora_weights = True
|
|
model.save_pretrained(os.path.join(output_dir, 'initial_model'))
|
|
config.init_lora_weights = init_lora_weights
|
|
|
|
def _save_converted_model(self, output_dir):
|
|
# pissa/olora/lora-ga
|
|
model = unwrap_model(self.model)
|
|
if isinstance(model, PeftModel):
|
|
config = model.peft_config.get('default')
|
|
init_lora_weights = getattr(config, 'init_lora_weights', None)
|
|
if isinstance(init_lora_weights, str):
|
|
config = copy(config)
|
|
# Save requires_grad state to protect against peft inject_adapter side effects
|
|
# (peft >= 0.18.1 incorrectly freezes active adapter when loading a temporary adapter)
|
|
requires_grad_state = {n: p.requires_grad for n, p in model.named_parameters()}
|
|
os.makedirs(os.path.join(output_dir, 'converted'), exist_ok=True)
|
|
if 'lora-ga' in init_lora_weights:
|
|
try:
|
|
from lora_ga.entrypoint import LoraGAContext
|
|
with LoraGAContext(model):
|
|
model.save_pretrained(
|
|
os.path.join(output_dir, 'converted', 'default'),
|
|
path_initial_model_for_weight_conversion=os.path.join(
|
|
os.path.dirname(output_dir), 'initial_model'),
|
|
)
|
|
model.peft_config['default'] = config
|
|
except ImportError as e:
|
|
error_message = """
|
|
Since 'LoRA-GA' is not implemented by PEFT, you will need to install it directly from GitHub.
|
|
Command: 'pip install git+https://github.com/lxline/LoRA-GA.git'.
|
|
"""
|
|
logger.info(error_message)
|
|
raise RuntimeError(error_message) from e
|
|
elif 'pissa' in init_lora_weights or 'olora' in init_lora_weights:
|
|
model.save_pretrained(
|
|
os.path.join(output_dir, 'converted', 'default'),
|
|
path_initial_model_for_weight_conversion=os.path.join(
|
|
os.path.dirname(output_dir), 'initial_model'),
|
|
)
|
|
model.peft_config['default'] = config
|
|
# Restore requires_grad state after conversion to prevent peft side effects
|
|
for n, p in model.named_parameters():
|
|
if n in requires_grad_state:
|
|
p.requires_grad = requires_grad_state[n]
|
|
|
|
def _load_rng_state(self, *args, **kwargs):
|
|
if self.args.resume_only_model:
|
|
return
|
|
return super()._load_rng_state(*args, **kwargs)
|
|
|
|
def _load_optimizer_and_scheduler(self, *args, **kwargs):
|
|
if self.args.resume_only_model:
|
|
return
|
|
super()._load_optimizer_and_scheduler(*args, **kwargs)
|
|
callbacks = set(getattr(self.args, 'callbacks', []))
|
|
ds_config = getattr(self.args, 'deepspeed', None) or {}
|
|
checkpoint_config = ds_config.get('checkpoint') if isinstance(ds_config, dict) else None
|
|
load_universal = isinstance(checkpoint_config, dict) and checkpoint_config.get('load_universal', False)
|
|
if 'deepspeed_elastic' in callbacks and load_universal:
|
|
self._fix_optimizer_step_device(self.optimizer)
|
|
if is_mp_ddp():
|
|
# fix mp+ddp adamw
|
|
for v in self.optimizer.state.values():
|
|
if 'step' in v:
|
|
# not on the same device
|
|
device_set = set([t.device for t in v.values()]) - {v['step'].device, torch.device('cpu')}
|
|
if len(device_set) >= 1:
|
|
v['step'] = v['step'].to('cpu')
|
|
|
|
@staticmethod
|
|
def _fix_optimizer_step_device(optimizer):
|
|
state = getattr(optimizer, 'state', None)
|
|
if not isinstance(state, dict):
|
|
return
|
|
for value in state.values():
|
|
if not isinstance(value, dict):
|
|
continue
|
|
step = value.get('step')
|
|
if not isinstance(step, torch.Tensor):
|
|
continue
|
|
target_device = None
|
|
for state_key, state_value in value.items():
|
|
if state_key == 'step':
|
|
continue
|
|
if isinstance(state_value, torch.Tensor) and state_value.device.type != 'cpu':
|
|
target_device = state_value.device
|
|
break
|
|
if target_device is not None and step.device != target_device:
|
|
value['step'] = step.to(target_device)
|
|
|
|
def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
|
|
# If template defines a save_callback, delegate to it
|
|
if hasattr(self, 'template') and hasattr(self.template, 'save_callback'):
|
|
self.template.save_callback(self.model, output_dir)
|
|
return
|
|
# model
|
|
supported_classes = (SwiftModel, PreTrainedModel, PeftModel)
|
|
supported_names = ('SentenceTransformer', )
|
|
safe_serialization = self.args.safe_serialization
|
|
use_flash_ckpt = self.args.use_flash_ckpt
|
|
|
|
if not isinstance(self.model, supported_classes) and self.model.__class__.__name__ not in supported_names:
|
|
if state_dict is None:
|
|
state_dict = self.model.state_dict()
|
|
|
|
_unwrap_model = unwrap_model(self.model)
|
|
if isinstance(_unwrap_model, supported_classes):
|
|
save_kwargs = {'state_dict': state_dict, 'max_shard_size': self.args.max_shard_size}
|
|
if isinstance(_unwrap_model, PeftModel):
|
|
save_kwargs['selected_adapters'] = ['default']
|
|
if use_flash_ckpt:
|
|
_unwrap_model.save_pretrained(
|
|
output_dir,
|
|
safe_serialization=False,
|
|
save_function=self.flash_checkpointer.ckpt_agent.save,
|
|
**save_kwargs)
|
|
else:
|
|
_unwrap_model.save_pretrained(output_dir, safe_serialization=safe_serialization, **save_kwargs)
|
|
else:
|
|
logger.info('Trainer.model is not a `PreTrainedModel`, only saving its state dict.')
|
|
if use_flash_ckpt:
|
|
self.flash_checkpointer.ckpt_agent.save(state_dict, os.path.join(output_dir, 'pytorch_model.bin'))
|
|
else:
|
|
if safe_serialization:
|
|
safetensors.torch.save_file(state_dict, os.path.join(output_dir, 'model.safetensors'))
|
|
else:
|
|
torch.save(state_dict, os.path.join(output_dir, 'pytorch_model.bin'))
|
|
elif is_instance_of_ms_model(self.model):
|
|
if use_flash_ckpt:
|
|
PreTrainedModel.save_pretrained(
|
|
self.model,
|
|
output_dir,
|
|
state_dict=state_dict,
|
|
safe_serialization=False,
|
|
save_function=self.flash_checkpointer.ckpt_agent.save)
|
|
else:
|
|
# modelscope save_pretrained does not support safe_serialization
|
|
PreTrainedModel.save_pretrained(
|
|
self.model, output_dir, state_dict=state_dict, safe_serialization=safe_serialization)
|
|
elif self.args.tuner_type in tuners_map:
|
|
tuners_map[self.args.tuner_type].save_pretrained(
|
|
self.model, output_dir, state_dict=state_dict, safe_serialization=safe_serialization)
|
|
else:
|
|
if self.model.__class__.__name__ != 'SentenceTransformer':
|
|
save_kwargs = {'state_dict': state_dict, 'max_shard_size': self.args.max_shard_size}
|
|
if isinstance(self.model, PeftModel):
|
|
save_kwargs['selected_adapters'] = ['default']
|
|
if use_flash_ckpt:
|
|
self.model.save_pretrained(
|
|
output_dir,
|
|
safe_serialization=False,
|
|
save_function=self.flash_checkpointer.ckpt_agent.save,
|
|
**save_kwargs)
|
|
else:
|
|
self.model.save_pretrained(output_dir, safe_serialization=safe_serialization, **save_kwargs)
|
|
else:
|
|
|
|
@contextmanager
|
|
def save_context():
|
|
save_pretrained = self.model[0].auto_model.save_pretrained
|
|
_state_dict = {
|
|
key[len('0.auto_model.'):] if 'auto_model' in key else key: value
|
|
for key, value in state_dict.items()
|
|
}
|
|
self.model[0].auto_model.save_pretrained = partial(
|
|
self.model[0].auto_model.save_pretrained, state_dict=_state_dict)
|
|
yield
|
|
self.model[0].auto_model.save_pretrained = save_pretrained
|
|
|
|
with save_context():
|
|
if use_flash_ckpt:
|
|
self.model.save_pretrained(
|
|
output_dir,
|
|
state_dict=state_dict,
|
|
safe_serialization=False,
|
|
save_function=self.flash_checkpointer.ckpt_agent.save)
|
|
else:
|
|
self.model.save_pretrained(output_dir, safe_serialization=safe_serialization)
|
|
# copy sentencetransformers files
|
|
copy_files_by_pattern(
|
|
self.model.model_dir, output_dir, '*.py', exclude_patterns=['model.safetensors.index.json'])
|
|
copy_files_by_pattern(
|
|
self.model.model_dir, output_dir, '*.json', exclude_patterns=['model.safetensors.index.json'])
|
|
|
|
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
|
"""Compatible with swift and peft"""
|
|
# If we are executing this function, we are the process zero, so we don't check for that.
|
|
output_dir = output_dir if output_dir is not None else self.args.output_dir
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
self._save_model(output_dir, state_dict)
|
|
# training_args.bin
|
|
torch.save(self.args, os.path.join(output_dir, 'training_args.bin'))
|
|
self._save_converted_model(output_dir)
|
|
# args.json
|
|
args_path = os.path.join(os.path.dirname(output_dir), 'args.json')
|
|
if os.path.exists(args_path):
|
|
shutil.copy(args_path, os.path.join(output_dir, 'args.json'))
|
|
# predict.jsonl
|
|
predict_jsonl = os.path.join(os.path.dirname(output_dir), 'predict.jsonl')
|
|
if os.path.exists(predict_jsonl):
|
|
shutil.move(predict_jsonl, os.path.join(output_dir, 'predict.jsonl'))
|
|
|
|
is_adapter = isinstance(self.model, (SwiftModel, PeftModel))
|
|
# tokenizer
|
|
if not is_adapter:
|
|
additional_saved_files = self.model_meta.additional_saved_files
|
|
save_checkpoint(
|
|
None,
|
|
self.template.processor,
|
|
output_dir,
|
|
model_dirs=[self.model.model_dir],
|
|
additional_saved_files=additional_saved_files)
|
|
if getattr(self.model, 'origin_generation_config', None):
|
|
self.model.origin_generation_config.save_pretrained(output_dir)
|
|
|
|
def _rotate_flash_checkpoints(self, use_mtime=False, output_dir=None) -> None:
|
|
if (self.args.save_total_limit is None or self.args.save_total_limit <= 0):
|
|
return
|
|
|
|
last_step = self._get_last_checkpoint_step()
|
|
|
|
# Check if we should delete older checkpoint(s)
|
|
if hasattr(self, '_sorted_checkpoints'):
|
|
checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
|
|
else:
|
|
output_dir = output_dir if output_dir is not None else self.args.output_dir
|
|
if sort_checkpoints is not None:
|
|
checkpoints_sorted = sort_checkpoints(
|
|
output_dir=output_dir,
|
|
checkpoint_prefix=PREFIX_CHECKPOINT_DIR,
|
|
use_mtime=use_mtime,
|
|
best_model_checkpoint=self.state.best_model_checkpoint,
|
|
)
|
|
else:
|
|
checkpoints = []
|
|
for path in os.listdir(output_dir) if os.path.isdir(output_dir) else []:
|
|
if re.match(f'^{PREFIX_CHECKPOINT_DIR}-([0-9]+)$', path):
|
|
checkpoints.append(os.path.join(output_dir, path))
|
|
ordering = os.path.getmtime if use_mtime else lambda path: int(path.rsplit('-', 1)[-1])
|
|
checkpoints_sorted = sorted(checkpoints, key=ordering)
|
|
|
|
valid_checkpoints = []
|
|
for path in checkpoints_sorted:
|
|
regex_match = re.match(f'.*{PREFIX_CHECKPOINT_DIR}-([0-9]+)', path)
|
|
if regex_match is not None and regex_match.groups() is not None:
|
|
step = int(regex_match.groups()[0])
|
|
if step <= last_step:
|
|
valid_checkpoints.append(path)
|
|
|
|
if len(valid_checkpoints) <= self.args.save_total_limit:
|
|
return
|
|
|
|
# If save_total_limit=1 with load_best_model_at_end=True,
|
|
# we could end up deleting the last checkpoint, which
|
|
# should be avoided and allow resuming
|
|
save_total_limit = self.args.save_total_limit
|
|
if (self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1
|
|
and valid_checkpoints[-1] != self.state.best_model_checkpoint):
|
|
save_total_limit = 2
|
|
|
|
number_of_checkpoints_to_delete = max(0, len(valid_checkpoints) - save_total_limit)
|
|
checkpoints_to_be_deleted = valid_checkpoints[:number_of_checkpoints_to_delete]
|
|
for checkpoint in checkpoints_to_be_deleted:
|
|
logger.info(f'Deleting older checkpoint [{checkpoint}] '
|
|
f'due to save_total_limit = {self.args.save_total_limit}.')
|
|
shutil.rmtree(checkpoint, ignore_errors=True)
|
|
|
|
def get_last_checkpoint(self):
|
|
"""
|
|
Get the path of the last complete checkpoint. Some latter directories
|
|
may not have the complete checkpoint because the asynchronous
|
|
persistence may not finish. The step in the `dlrover_latest.txt` is
|
|
the last step of complete checkpoint. We can get the path by the step.
|
|
"""
|
|
step = self._get_last_checkpoint_step()
|
|
if step == 0:
|
|
return False
|
|
checkpoint_folder = f'{PREFIX_CHECKPOINT_DIR}-{step}'
|
|
ckpt_dir = os.path.join(self.args.output_dir, checkpoint_folder)
|
|
return ckpt_dir
|
|
|
|
def _get_last_checkpoint_step(self):
|
|
tracer_file = os.path.join(self.args.output_dir, 'dlrover_latest.txt')
|
|
if not os.path.exists(tracer_file):
|
|
return 0
|
|
with open(tracer_file, 'r') as f:
|
|
step = int(f.read())
|
|
return step
|
|
|
|
def get_resume_checkpoint(self):
|
|
"""
|
|
Get the path of the last complete checkpoint. Some latter directories
|
|
may not have the complete checkpoint because the asynchronous
|
|
persistence may not finish. The step in the `dlrover_latest.txt` is
|
|
the last step of complete checkpoint. We can get the path by the step.
|
|
"""
|
|
resume_dir = get_resume_dir(self.args.output_dir)
|
|
if resume_dir is None:
|
|
return None
|
|
tracer_file = os.path.join(resume_dir, 'dlrover_latest.txt')
|
|
if not os.path.exists(tracer_file):
|
|
return None
|
|
with open(tracer_file, 'r') as f:
|
|
step = int(f.read())
|
|
checkpoint_folder = f'{PREFIX_CHECKPOINT_DIR}-{step}'
|
|
|
|
ckpt_dir = os.path.join(resume_dir, checkpoint_folder)
|
|
with open(os.path.join(ckpt_dir, TRAINER_STATE_NAME), 'r', encoding='utf-8') as f:
|
|
train_state = json.load(f)
|
|
if train_state is not None and train_state.get('max_steps') == step:
|
|
return None
|
|
return ckpt_dir
|
|
|
|
def get_resume_checkpoint_until_find_ucp(self):
|
|
resume_dir = get_resume_dir(self.args.output_dir)
|
|
if resume_dir is None:
|
|
return None
|
|
tracer_file = os.path.join(resume_dir, 'ucp.txt')
|
|
if not os.path.exists(tracer_file):
|
|
step = 0
|
|
if step == 0:
|
|
return None
|
|
with open(tracer_file, 'r') as f:
|
|
step = int(f.read())
|
|
checkpoint_folder = f'{PREFIX_CHECKPOINT_DIR}-{step}'
|
|
ckpt_dir = os.path.join(resume_dir, checkpoint_folder)
|
|
return ckpt_dir
|
|
|
|
def wait_latest_checkpoint(self, timeout=None, max_steps=None):
|
|
"""
|
|
Wait for the latest checkpoint.
|
|
Args:
|
|
timeout (second): The timeout to wait.
|
|
"""
|
|
self.flash_checkpointer.async_save_engine.wait_latest_checkpoint(timeout, max_steps)
|
|
|
|
def _fix_zero3_gather_all_parameters(self) -> None:
|
|
if is_deepspeed_zero3_enabled() and not hasattr(self.deepspeed, '_zero3_consolidated_16bit_state_dict_origin'):
|
|
parameters = inspect.signature(self.deepspeed._zero3_consolidated_16bit_state_dict).parameters
|
|
if 'exclude_frozen_parameters' in parameters:
|
|
|
|
def _zero3_consolidated_16bit_state_dict(model, exclude_frozen_parameters=False):
|
|
unwrapped = unwrap_model(model)
|
|
exclude_frozen_parameters = False
|
|
if isinstance(unwrapped, SwiftModel) and unwrapped.has_additional_modules:
|
|
exclude_frozen_parameters = True
|
|
if isinstance(unwrapped, PeftModel):
|
|
exclude_frozen_parameters = True
|
|
return model._zero3_consolidated_16bit_state_dict_origin(exclude_frozen_parameters)
|
|
|
|
self.deepspeed._zero3_consolidated_16bit_state_dict_origin = (
|
|
self.deepspeed._zero3_consolidated_16bit_state_dict)
|
|
self.deepspeed._zero3_consolidated_16bit_state_dict = MethodType(_zero3_consolidated_16bit_state_dict,
|
|
self.deepspeed)
|
|
|
|
def _save_checkpoint(self, *args, **kwargs):
|
|
self.state.last_model_checkpoint = os.path.join(self.args.output_dir, f'checkpoint-{self.state.global_step}')
|
|
self._fix_zero3_gather_all_parameters()
|
|
|
|
if self.args.use_flash_ckpt:
|
|
result = self._save_flash_checkpoint(*args, **kwargs)
|
|
else:
|
|
result = super()._save_checkpoint(*args, **kwargs)
|
|
logger.info(f'Saving model checkpoint to {self.state.last_model_checkpoint}')
|
|
return result
|
|
|
|
def _save_flash_checkpoint(self, model, trial, metrics=None):
|
|
from dlrover.trainer.torch.flash_checkpoint.hf_trainer import HfDdpCheckpointer, HfDeepSpeedCheckpointer
|
|
from transformers.trainer import DeepSpeedSchedulerWrapper
|
|
from transformers.trainer_utils import SaveStrategy
|
|
run_dir = self._get_output_dir(trial=trial)
|
|
|
|
torch_native_save = torch.save
|
|
|
|
# Save model checkpoint
|
|
checkpoint_folder = f'{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}'
|
|
output_dir = os.path.join(run_dir, checkpoint_folder)
|
|
|
|
if not hasattr(self, 'flash_checkpointer'):
|
|
if self.is_deepspeed_enabled:
|
|
self.flash_checkpointer = HfDeepSpeedCheckpointer(self.model_wrapped, run_dir)
|
|
elif not self.is_deepspeed_enabled and not self.is_fsdp_enabled:
|
|
self.flash_checkpointer = HfDdpCheckpointer(run_dir)
|
|
else:
|
|
raise ValueError('Flash Checkpoint only supports DeepSpeed or DDP.')
|
|
|
|
if self.hp_search_backend is None and trial is None:
|
|
self.store_flos()
|
|
|
|
torch.save = self.flash_checkpointer.ckpt_agent.save
|
|
self.save_model(output_dir, _internal_call=True)
|
|
if self.is_deepspeed_enabled:
|
|
self.model_wrapped.save_checkpoint(output_dir)
|
|
|
|
elif (self.args.should_save and not self.is_deepspeed_enabled and not self.is_fsdp_enabled):
|
|
# deepspeed.save_checkpoint above saves model/optim/sched
|
|
torch.save(
|
|
self.optimizer.state_dict(),
|
|
os.path.join(output_dir, OPTIMIZER_NAME),
|
|
)
|
|
|
|
# Save SCHEDULER & SCALER
|
|
is_deepspeed_custom_scheduler = (
|
|
self.is_deepspeed_enabled and not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper))
|
|
if self.args.should_save and (not self.is_deepspeed_enabled or is_deepspeed_custom_scheduler):
|
|
with warnings.catch_warnings(record=True) as caught_warnings:
|
|
torch.save(
|
|
self.lr_scheduler.state_dict(),
|
|
os.path.join(output_dir, SCHEDULER_NAME),
|
|
)
|
|
reissue_pt_warnings(caught_warnings)
|
|
if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH] and self.state.best_global_step:
|
|
best_checkpoint_folder = f'{PREFIX_CHECKPOINT_DIR}-{self.state.best_global_step}'
|
|
best_checkpoint_dir = os.path.join(run_dir, best_checkpoint_folder)
|
|
|
|
if os.path.exists(best_checkpoint_dir):
|
|
self.state.best_model_checkpoint = best_checkpoint_dir
|
|
|
|
# Save the Trainer state
|
|
if self.args.should_save:
|
|
# Update `ExportableState` callbacks and `TrainerControl` state to where we are currently
|
|
from transformers.trainer_callback import ExportableState
|
|
for cb in [
|
|
cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
|
|
]:
|
|
cb_name = cb.__class__.__name__
|
|
cb_state = cb.state()
|
|
if isinstance(self.state.stateful_callbacks[cb_name], list):
|
|
self.state.stateful_callbacks[cb_name].append(cb_state)
|
|
else:
|
|
self.state.stateful_callbacks[cb_name] = cb_state
|
|
self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
|
|
# Save RNG state in non-distributed training
|
|
rng_states = {
|
|
'python': random.getstate(),
|
|
'numpy': np.random.get_state(),
|
|
'cpu': torch.random.get_rng_state(),
|
|
}
|
|
if torch.cuda.is_available():
|
|
if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
|
|
# In non distributed, we save the global
|
|
# CUDA RNG state (will take care of DataParallel)
|
|
rng_states['cuda'] = torch.cuda.random.get_rng_state_all()
|
|
else:
|
|
rng_states['cuda'] = torch.cuda.random.get_rng_state()
|
|
|
|
# A process can arrive here before the process 0 has a chance to
|
|
# save the model, in which case output_dir may not yet exist.
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
if self.args.world_size <= 1:
|
|
torch.save(rng_states, os.path.join(output_dir, 'rng_state.pth'))
|
|
else:
|
|
torch.save(
|
|
rng_states,
|
|
os.path.join(output_dir, f'rng_state_{self.args.process_index}.pth'),
|
|
)
|
|
if self.args.safe_serialization:
|
|
torch.save({'safe_serialization': True}, 'safe_serialization')
|
|
replace_index_file(output_dir)
|
|
|
|
torch.save = torch_native_save
|
|
if (self.state.global_step == self.state.max_steps):
|
|
success = self.flash_checkpointer.save_checkpoint_to_storage(self.state.global_step, True)
|
|
else:
|
|
success = self.flash_checkpointer.save_checkpoint_to_storage(self.state.global_step)
|
|
|
|
if not success:
|
|
logger.info(f'Skip saving the checkpoint of step {self.state.global_step} '
|
|
'because the latest checkpoint is not finished.')
|
|
shutil.rmtree(output_dir, ignore_errors=True)
|
|
|
|
if self.args.push_to_hub:
|
|
self._push_from_checkpoint(output_dir)
|
|
|
|
# Maybe delete some older checkpoints.
|
|
if self.args.should_save:
|
|
self._rotate_flash_checkpoints(use_mtime=True, output_dir=run_dir)
|
|
|
|
@staticmethod
|
|
@contextmanager
|
|
def _fix_grad_norm_nan():
|
|
from accelerate import Accelerator
|
|
origin_clip_grad_norm_ = Accelerator.clip_grad_norm_
|
|
|
|
def clip_grad_norm_(self, parameters, *args, **kwargs):
|
|
# If NaN occurs, ignore weight updates.
|
|
parameters = list(parameters)
|
|
grad_norm = origin_clip_grad_norm_(self, parameters, *args, **kwargs)
|
|
if isinstance(grad_norm, torch.Tensor) and grad_norm.isnan().item():
|
|
for p in parameters:
|
|
p.grad = None
|
|
return grad_norm
|
|
|
|
Accelerator.clip_grad_norm_ = clip_grad_norm_
|
|
try:
|
|
yield
|
|
finally:
|
|
Accelerator.clip_grad_norm_ = origin_clip_grad_norm_
|
|
|
|
def _patch_tasks(self):
|
|
if isinstance(self.model, PeftModel):
|
|
model = self.model.model
|
|
else:
|
|
model = self.model
|
|
task_type = self.task_type
|
|
sp_enabled = self.template.sequence_parallel_size > 1
|
|
pf_enabled = bool(self.template.padding_free)
|
|
padding_side = 'left' if pf_enabled else self.template.padding_side
|
|
|
|
if 'SentenceTransformer' in model.__class__.__name__:
|
|
|
|
def forward_transformer(transformer, features: Dict[str, torch.Tensor],
|
|
**kwargs) -> Dict[str, torch.Tensor]:
|
|
trans_features = {
|
|
key: value
|
|
for key, value in features.items()
|
|
if key in ['input_ids', 'attention_mask', 'token_type_ids', 'inputs_embeds', 'position_ids']
|
|
}
|
|
|
|
outputs = transformer.auto_model(**trans_features, **kwargs, return_dict=True)
|
|
token_embeddings = outputs[0]
|
|
features['token_embeddings'] = token_embeddings
|
|
|
|
if transformer.auto_model.config.output_hidden_states and 'hidden_states' in outputs:
|
|
features['all_layer_embeddings'] = outputs['hidden_states']
|
|
|
|
return features
|
|
|
|
from sentence_transformers.models import Transformer
|
|
if isinstance(model[0], Transformer):
|
|
model[0].forward = MethodType(forward_transformer, model[0])
|
|
|
|
def forward_sentence_transformer(sentence_transformer, **kwargs) -> Dict[str, torch.Tensor]:
|
|
input = kwargs
|
|
kwargs = {}
|
|
for idx, (module_name, module) in enumerate(sentence_transformer.named_children()):
|
|
from sentence_transformers.models import Router
|
|
if isinstance(module, Router):
|
|
module_kwargs = kwargs
|
|
else:
|
|
module_kwarg_keys = []
|
|
if sentence_transformer.module_kwargs is not None:
|
|
module_kwarg_keys = sentence_transformer.module_kwargs.get(module_name, [])
|
|
module_kwargs = {
|
|
key: value
|
|
for key, value in kwargs.items() if key in module_kwarg_keys or (
|
|
hasattr(module, 'forward_kwargs') and key in module.forward_kwargs)
|
|
}
|
|
output = module(input, **module_kwargs)
|
|
if idx == 0 and self.template.padding_free:
|
|
output = revert_padding_free(output, input, padding_side)
|
|
input = output
|
|
return {'last_hidden_state': input['sentence_embedding']}
|
|
|
|
model.forward = MethodType(forward_sentence_transformer, model)
|
|
else:
|
|
|
|
def _register_llm_hooks_in_order(llm_model: nn.Module, hooks: List[Callable]):
|
|
# hooks are provided in desired execution order.
|
|
# We use prepend=True and register in reverse to preserve the order.
|
|
for hook in reversed(hooks):
|
|
llm_model.register_forward_hook(hook, with_kwargs=True, prepend=True)
|
|
|
|
def _get_hook_target_model(task_type_: str) -> nn.Module:
|
|
# For embedding, we hook on the LM-head model because embedding outputs are typically
|
|
# produced from `output.logits` by `patch_output_normalizer` (registered on LM-head model).
|
|
if task_type_ == 'embedding':
|
|
return get_lm_head_model(self.model, model_meta=self.model.model_meta)
|
|
return get_llm_model(self.model, model_meta=self.model.model_meta)
|
|
|
|
# --- seq_cls / reranker / generative_reranker / embedding unified pipeline ---
|
|
if task_type in {'seq_cls', 'reranker', 'generative_reranker', 'embedding'}:
|
|
llm_model = _get_hook_target_model(task_type)
|
|
|
|
hooks: List[Callable] = []
|
|
|
|
if sp_enabled:
|
|
|
|
def sp_gather_hook(module, args, input, output):
|
|
return gather_sequence_parallel_outputs(output)
|
|
|
|
hooks.append(sp_gather_hook)
|
|
|
|
if pf_enabled:
|
|
if sp_enabled:
|
|
|
|
def revert_padding_free_hook(module, args, input, output):
|
|
# Use full packed position ids cached by sequence_parallel.prepare_inputs
|
|
position_ids = sequence_parallel.real_position_ids
|
|
tmp_input = {'position_ids': position_ids}
|
|
return revert_padding_free(output, tmp_input, padding_side)
|
|
else:
|
|
|
|
def revert_padding_free_hook(module, args, input, output):
|
|
return revert_padding_free(output, input, padding_side)
|
|
|
|
hooks.append(revert_padding_free_hook)
|
|
|
|
if hooks:
|
|
_register_llm_hooks_in_order(llm_model, hooks)
|
|
|
|
# wrappers for seq_cls / reranker (pooling/head must see gathered/reverted outputs)
|
|
if task_type in {'seq_cls', 'reranker'} and (sp_enabled or pf_enabled):
|
|
lm_head_model = get_lm_head_model(self.model, model_meta=self.model.model_meta)
|
|
|
|
if task_type == 'seq_cls':
|
|
|
|
@wraps(model.forward.__func__)
|
|
def seq_cls_forward(model, *args, **kwargs):
|
|
sp_kwargs = dict(kwargs)
|
|
|
|
def inner_forward(*args, **_kwargs):
|
|
return llm_model(*args, **_kwargs)
|
|
|
|
return transformers_seq_cls_forward(
|
|
lm_head_model,
|
|
*args,
|
|
origin_forward=inner_forward,
|
|
padding_side=padding_side,
|
|
**sp_kwargs,
|
|
)
|
|
|
|
model.forward = MethodType(seq_cls_forward, model)
|
|
else:
|
|
|
|
@wraps(model.forward.__func__)
|
|
def reranker_forward(model, *args, **kwargs):
|
|
sp_kwargs = dict(kwargs)
|
|
|
|
def inner_forward(*args, **_kwargs):
|
|
return llm_model(*args, **_kwargs)
|
|
|
|
padding_free_fn = getattr(model, 'padding_free_fn', None)
|
|
if callable(padding_free_fn):
|
|
output = inner_forward(*args, **sp_kwargs)
|
|
return padding_free_fn(output, sp_kwargs, padding_side)
|
|
|
|
return transformers_seq_cls_forward(
|
|
lm_head_model,
|
|
*args,
|
|
origin_forward=inner_forward,
|
|
padding_side=padding_side,
|
|
**sp_kwargs,
|
|
)
|
|
|
|
model.forward = MethodType(reranker_forward, model)
|
|
|
|
def _fix_gradient_checkpointing(self):
|
|
# fix use_reentrant
|
|
if hasattr(torch.utils.checkpoint, '_old_checkpoint'): # avoid double patching
|
|
return
|
|
args = self.args
|
|
if args.gradient_checkpointing_kwargs:
|
|
use_reentrant_ = args.gradient_checkpointing_kwargs.get('use_reentrant')
|
|
else:
|
|
use_reentrant_ = None
|
|
if use_reentrant_ is None:
|
|
if is_dist() and not self.is_deepspeed_enabled and not self.is_fsdp_enabled:
|
|
use_reentrant_ = False
|
|
else:
|
|
use_reentrant_ = True
|
|
logger.info(f'use_reentrant: {use_reentrant_}')
|
|
_old_checkpoint = torch.utils.checkpoint.checkpoint
|
|
|
|
@wraps(_old_checkpoint)
|
|
def _new_checkpoint(*args, use_reentrant=None, **kwargs):
|
|
return _old_checkpoint(*args, use_reentrant=use_reentrant_, **kwargs)
|
|
|
|
torch.utils.checkpoint._old_checkpoint = _old_checkpoint
|
|
torch.utils.checkpoint.checkpoint = _new_checkpoint
|
|
try:
|
|
# Fix the old version of transformers.
|
|
import transformers.modeling_utils
|
|
transformers.modeling_utils.checkpoint = _new_checkpoint
|
|
except (ImportError, AttributeError):
|
|
pass
|
|
|
|
def _prepare_gradient_checkpointing(self, model) -> None:
|
|
args = self.args
|
|
HfConfigFactory.set_config_attr(model.config, 'use_cache', False)
|
|
if args.gradient_checkpointing or args.vit_gradient_checkpointing:
|
|
dynamic_gradient_checkpointing(model, args.vit_gradient_checkpointing)
|
|
gc_kwargs = {}
|
|
parameters = inspect.signature(model.gradient_checkpointing_enable).parameters
|
|
if 'gradient_checkpointing_kwargs' in parameters:
|
|
gc_kwargs['gradient_checkpointing_kwargs'] = args.gradient_checkpointing_kwargs
|
|
if args.gradient_checkpointing:
|
|
model.gradient_checkpointing_enable(**gc_kwargs)
|
|
model.enable_input_require_grads()
|
|
|
|
model_meta = model.model_meta
|
|
model_arch = model_meta.model_arch
|
|
if model_meta.is_multimodal and model_arch:
|
|
for vision_tower_name in model_arch.vision_tower:
|
|
vision_tower = deep_getattr(model, vision_tower_name)
|
|
if hasattr(vision_tower, 'enable_input_require_grads'):
|
|
try:
|
|
if args.vit_gradient_checkpointing:
|
|
vision_tower.gradient_checkpointing_enable(**gc_kwargs)
|
|
vision_tower.enable_input_require_grads()
|
|
else:
|
|
vision_tower.gradient_checkpointing_disable()
|
|
vision_tower.disable_input_require_grads()
|
|
except (NotImplementedError, AttributeError, ValueError) as e:
|
|
logger.warning(f'prepare gradient_checkpointing failed: {e}')
|
|
# Avoid vit_gradient_checkpointing being overwritten by transformers.Trainer.gradient_checkpointing_enable.
|
|
self.args.gradient_checkpointing = False
|
|
|
|
def train(self, *args, **kwargs):
|
|
if self.model_meta.is_multimodal:
|
|
models = []
|
|
for model_name in ['model', 'ref_model', 'value_model', 'teacher_model']:
|
|
model = getattr(self, model_name, None)
|
|
if isinstance(model, nn.Module):
|
|
models.append(model)
|
|
|
|
reward_model = getattr(self, 'reward_model', None)
|
|
if reward_model is not None:
|
|
if isinstance(reward_model, list):
|
|
models.extend([m for m in reward_model if isinstance(m, nn.Module)])
|
|
elif isinstance(reward_model, nn.Module):
|
|
models.append(reward_model)
|
|
|
|
models = list(set(self.accelerator.unwrap_model(model) for model in models)) # Deduplicate
|
|
self.template.register_post_encode_hook(models)
|
|
logger.info(f'Successfully registered post_encode hook: {[model.__class__.__name__ for model in models]}.')
|
|
self._save_initial_model(self.args.output_dir)
|
|
|
|
# gradient_checkpointing
|
|
gradient_checkpointing = self.args.gradient_checkpointing
|
|
base_model = self.template.get_base_model(self.accelerator.unwrap_model(self.model)) # fix peftmodel
|
|
self._prepare_gradient_checkpointing(base_model)
|
|
with self.hub.patch_hub(), self._fix_grad_norm_nan(), self._patch_skip_first_batches(
|
|
), self._patch_deepspeed_load_checkpoint():
|
|
res = super().train(*args, **kwargs)
|
|
self.template.remove_post_encode_hook()
|
|
self.args.gradient_checkpointing = gradient_checkpointing # recover
|
|
return res
|
|
|
|
def push_to_hub(self, *args, **kwargs):
|
|
with self.hub.patch_hub():
|
|
return super().push_to_hub(*args, **kwargs)
|
|
|
|
@staticmethod
|
|
def compute_custom_metrics(metrics, key_prefix: str = ''):
|
|
logs = {}
|
|
# Synchronize keys to avoid getting stuck.
|
|
if dist.is_initialized():
|
|
all_keys = [None] * dist.get_world_size()
|
|
dist.all_gather_object(all_keys, list(metrics.keys()))
|
|
for key in set().union(*all_keys):
|
|
if key not in metrics:
|
|
metrics[key]
|
|
|
|
for k, metric in sorted(metrics.items()):
|
|
k = f'{key_prefix}{k}'
|
|
value = metric.compute()
|
|
metric.reset()
|
|
if isinstance(value, dict):
|
|
if len(value) == 1:
|
|
val = list(value.values())[0]
|
|
logs[k] = val
|
|
else:
|
|
for k_suffix, val in value.items():
|
|
new_k = f'{k}_{k_suffix}'
|
|
logs[new_k] = val
|
|
else:
|
|
logs[k] = value
|
|
for k in list(logs.keys()):
|
|
if logs[k] is None:
|
|
logs.pop(k)
|
|
return logs
|
|
|
|
def log(self, logs: Dict[str, float], *args, **kwargs) -> None:
|
|
mode = 'train' if self.model.training else 'eval'
|
|
metrics = self.custom_metrics[mode]
|
|
prefix = 'eval_' if mode == 'eval' else ''
|
|
logs.update(self.compute_custom_metrics(metrics, prefix))
|
|
return super().log(logs, *args, **kwargs)
|
|
|
|
def _maybe_log_save_evaluate(self, tr_loss, *args, **kwargs):
|
|
if self.control.should_log and self.state.global_step > self._globalstep_last_logged:
|
|
self.control.should_log = False
|
|
|
|
# all_gather + mean() to get average loss over all processes
|
|
if version.parse(transformers.__version__) >= version.parse('5.2.0'):
|
|
from transformers.trainer_pt_utils import nested_gather
|
|
tr_loss_scalar = nested_gather(tr_loss, self.args.parallel_mode).mean().item()
|
|
else:
|
|
tr_loss_scalar = self._nested_gather(tr_loss).mean().item()
|
|
loss = tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged)
|
|
logs: Dict[str, float] = {'loss': loss} # loss first
|
|
if version.parse(transformers.__version__) >= version.parse('4.38'):
|
|
grad_norm = args[0]
|
|
if grad_norm is not None:
|
|
logs['grad_norm'] = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
|
|
logs['learning_rate'] = self._get_learning_rate()
|
|
tr_loss -= tr_loss
|
|
self._total_loss_scalar += tr_loss_scalar
|
|
self._globalstep_last_logged = self.state.global_step
|
|
self.store_flos()
|
|
self.log(logs)
|
|
|
|
if self.args.eval_use_evalscope and self.control.should_evaluate:
|
|
try:
|
|
self._evalscope_eval()
|
|
except Exception as e:
|
|
logger.warning(f'Failed to call EvalScope evaluation function: {e}.')
|
|
|
|
if not self.eval_dataset:
|
|
self.control.should_evaluate = False
|
|
super()._maybe_log_save_evaluate(tr_loss, *args, **kwargs)
|
|
|
|
def create_loss_and_eval_metric(self, args):
|
|
res = {}
|
|
if args.eval_metric is not None:
|
|
eval_metric = eval_metrics_map[args.eval_metric](args, self)
|
|
res['compute_metrics'], res['preprocess_logits_for_metrics'] = (eval_metric.compute_metrics,
|
|
eval_metric.preprocess_logits_for_metrics)
|
|
if args.loss_type is not None:
|
|
res['compute_loss_func'] = loss_map[args.loss_type](args, self)
|
|
return res
|
|
|
|
def create_optimizer_and_scheduler(self, num_training_steps: int):
|
|
self.optimizer_callback.create_optimizer_and_scheduler(num_training_steps)
|
|
|
|
def create_optimizer(self, model=None):
|
|
self._optimizer_ori = self.optimizer = self.optimizer_callback.create_optimizer(model=model)
|
|
if self.optimizer is not None:
|
|
self.optimizer.param_groups = [pg for pg in self.optimizer.param_groups if len(pg['params']) > 0]
|
|
return self.optimizer
|
|
|
|
def create_scheduler(self, num_training_steps: int, optimizer=None):
|
|
if optimizer is None:
|
|
# fix deepspeed & cosine_with_min_lr (transformers 5.8.0)
|
|
optimizer = getattr(self, '_optimizer_ori', None)
|
|
self.lr_scheduler = self.optimizer_callback.create_scheduler(num_training_steps, optimizer)
|
|
return self.lr_scheduler
|
|
|
|
@staticmethod
|
|
def _get_listwise_reranker_preds(logits, labels):
|
|
positive_indices = torch.nonzero(labels == 1, as_tuple=False).squeeze(-1).tolist()
|
|
positive_indices.append(labels.shape[0])
|
|
preds = []
|
|
for i in range(len(positive_indices) - 1):
|
|
start, end = positive_indices[i], positive_indices[i + 1]
|
|
preds.append(logits[start:end].argmax())
|
|
preds = torch.tensor(preds)
|
|
labels = torch.tensor([0] * (len(positive_indices) - 1))
|
|
return preds, labels
|
|
|
|
def _compute_acc(self, outputs, labels, cu_seqlens=None) -> None:
|
|
args = self.args
|
|
logits = outputs.logits
|
|
metrics = None
|
|
task_type = self.task_type
|
|
problem_type = self.problem_type
|
|
if task_type == 'embedding':
|
|
return
|
|
elif task_type == 'seq_cls':
|
|
if problem_type == 'regression':
|
|
return
|
|
elif problem_type == 'multi_label_classification':
|
|
preds = logits.sigmoid() > 0.5
|
|
metrics = {'acc': (labels == preds).all(dim=-1)}
|
|
else:
|
|
preds = logits.argmax(dim=-1)
|
|
metrics = compute_acc(preds, labels)
|
|
elif task_type == 'causal_lm':
|
|
preds = logits.argmax(dim=-1)
|
|
if self.template.sequence_parallel_size > 1:
|
|
# Gather preds and labels across the sp group
|
|
if isinstance(preds, np.ndarray):
|
|
preds = torch.from_numpy(preds).to(get_current_device())
|
|
if isinstance(labels, np.ndarray):
|
|
labels = torch.from_numpy(labels).to(get_current_device())
|
|
assert labels.shape[1] == preds.shape[1]
|
|
|
|
if sequence_parallel.rp_world_size > 1:
|
|
position_ids = sequence_parallel.real_position_ids
|
|
position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
|
|
else:
|
|
position_ids = None
|
|
preds_output = sequence_parallel.gather(preds, dim=1, position_ids=position_ids)
|
|
labels_output = sequence_parallel.gather(labels, dim=1, position_ids=position_ids)
|
|
# roll back to fit compute_acc
|
|
labels_output = torch.roll(labels_output, shifts=1, dims=1)
|
|
preds = preds_output
|
|
labels = labels_output.int()
|
|
|
|
metrics = compute_acc(
|
|
preds,
|
|
labels,
|
|
acc_strategy=args.acc_strategy,
|
|
is_encoder_decoder=self.template.is_encoder_decoder,
|
|
cu_seqlens=cu_seqlens)
|
|
elif task_type in {'generative_reranker', 'reranker'}:
|
|
if logits.dim() == 2:
|
|
logits = logits.squeeze(-1)
|
|
if args.loss_type == 'listwise_reranker':
|
|
preds, labels = self._get_listwise_reranker_preds(logits, labels)
|
|
else:
|
|
preds = (logits > 0).long()
|
|
metrics = compute_acc(preds, labels.long())
|
|
if metrics:
|
|
mode = 'train' if self.model.training else 'eval'
|
|
for k, v in metrics.items():
|
|
self.custom_metrics[mode][k].update(v)
|
|
|
|
@torch.no_grad()
|
|
def _evalscope_eval(self):
|
|
from evalscope import TaskConfig, run_task
|
|
|
|
from ..pipelines.eval.utils import EvalModel
|
|
|
|
self.model.eval()
|
|
template = copy(self.template)
|
|
template.packing = False
|
|
template.padding_free = False
|
|
# prepare task config
|
|
task_config_kwargs = dict(
|
|
model=EvalModel(
|
|
model_name=f'model-step{self.state.global_step}',
|
|
model=self.model,
|
|
template=template,
|
|
max_batch_size=self.args.per_device_eval_batch_size,
|
|
),
|
|
eval_type='swift_custom',
|
|
datasets=self.args.eval_dataset,
|
|
dataset_args=self.args.eval_dataset_args,
|
|
limit=self.args.eval_limit,
|
|
work_dir=os.path.join(self.args.output_dir, 'eval'),
|
|
eval_batch_size=self.args.per_device_eval_batch_size,
|
|
generation_config=self.args.eval_generation_config or {'max_tokens': 512},
|
|
)
|
|
task_config_kwargs.update(self.args.extra_eval_args or {})
|
|
task_config = TaskConfig(**task_config_kwargs)
|
|
# start evaluation
|
|
eval_report = run_task(task_config)
|
|
# convert to dict
|
|
eval_dict = {f'test_{k}': v.score for k, v in eval_report.items()}
|
|
self.log(eval_dict)
|
|
|
|
self.model.train()
|
|
return eval_dict
|
|
|
|
def prepare_logits_to_keep(self, inputs):
|
|
labels = inputs['labels']
|
|
loss_scale = inputs.get('loss_scale')
|
|
if self.template.sequence_parallel_size > 1:
|
|
raise NotImplementedError()
|
|
if labels.shape[0] == 1 and not is_mp():
|
|
# device_map may encounter device mismatch issues.
|
|
loss_mask = (labels != -100)[0]
|
|
labels = labels[:, loss_mask]
|
|
labels = nn.functional.pad(labels, (1, 0), value=-100)
|
|
if loss_scale is not None:
|
|
loss_scale = loss_scale[:, loss_mask]
|
|
inputs['loss_scale'] = nn.functional.pad(loss_scale, (1, 0), value=0)
|
|
logits_to_keep = nn.functional.pad(loss_mask[1:], (0, 1), value=True)
|
|
else:
|
|
logits_to_keep = labels.shape[-1] - ((labels != -100).int().argmax(-1).min().item()) + 1
|
|
assert logits_to_keep > 0
|
|
labels = labels[:, -logits_to_keep:]
|
|
if loss_scale is not None:
|
|
inputs['loss_scale'] = loss_scale[:, -logits_to_keep:]
|
|
inputs['labels'] = labels
|
|
inputs['logits_to_keep'] = logits_to_keep
|
|
|
|
def get_cu_seqlens(self, position_ids, logits_to_keep) -> torch.Tensor:
|
|
cu_seqlens = get_packed_seq_params(position_ids)['cu_seq_lens_q']
|
|
res_cu_seqlens = cu_seqlens.clone()
|
|
if isinstance(logits_to_keep, torch.Tensor):
|
|
for i in range(cu_seqlens.shape[0] - 1):
|
|
start, end = cu_seqlens[i], cu_seqlens[i + 1]
|
|
res_cu_seqlens[i + 1:] -= (~logits_to_keep[start:end]).sum()
|
|
elif isinstance(logits_to_keep, int):
|
|
res_cu_seqlens[1:] -= position_ids.shape[-1] + 1 - logits_to_keep
|
|
return res_cu_seqlens
|
|
|
|
@contextmanager
|
|
def _patch_skip_first_batches(self):
|
|
from transformers import trainer
|
|
origin_skip_first_batches = trainer.skip_first_batches
|
|
|
|
def skip_first_batches(dataloader, num_batches=0):
|
|
if isinstance(dataloader, (DataLoaderShard, DataLoaderDispatcher)):
|
|
# DataLoaderMixin
|
|
return self.get_train_dataloader(skip_batches=num_batches)
|
|
else:
|
|
return origin_skip_first_batches(dataloader, num_batches)
|
|
|
|
trainer.skip_first_batches = skip_first_batches
|
|
try:
|
|
yield
|
|
finally:
|
|
trainer.skip_first_batches = origin_skip_first_batches
|
|
|
|
|
|
class DataLoaderMixin:
|
|
|
|
def get_sp_dataloader(self, dataset, batch_size, skip_batches=0):
|
|
|
|
data_collator = self.data_collator
|
|
if isinstance(dataset, datasets.Dataset):
|
|
dataset = self._remove_unused_columns(dataset, description='training')
|
|
else:
|
|
data_collator = self._get_collator_with_removed_columns(data_collator, description='training')
|
|
if hasattr(dataset, '__len__'):
|
|
sampler = SequenceParallelSampler(sequence_parallel, dataset, seed=42)
|
|
dataloader_params = {
|
|
'batch_size': batch_size,
|
|
'collate_fn': data_collator,
|
|
'num_workers': self.args.dataloader_num_workers,
|
|
'pin_memory': self.args.dataloader_pin_memory,
|
|
'persistent_workers': self.args.dataloader_persistent_workers,
|
|
}
|
|
|
|
if not isinstance(dataset, torch.utils.data.IterableDataset):
|
|
if skip_batches > 0:
|
|
from accelerate.data_loader import SkipBatchSampler
|
|
sampler = SkipBatchSampler(sampler, skip_batches=skip_batches * batch_size)
|
|
dataloader_params['sampler'] = sampler
|
|
dataloader_params['drop_last'] = self.args.dataloader_drop_last
|
|
dataloader_params['worker_init_fn'] = partial(
|
|
seed_worker, num_workers=self.args.dataloader_num_workers, rank=sequence_parallel.dp_rank)
|
|
|
|
return DataLoaderShard(dataset, device=self.accelerator.device, **dataloader_params)
|
|
else:
|
|
dataloader_params = {
|
|
'collate_fn': data_collator,
|
|
'num_workers': self.args.dataloader_num_workers,
|
|
'pin_memory': self.args.dataloader_pin_memory,
|
|
'persistent_workers': self.args.dataloader_persistent_workers,
|
|
'prefetch_factor': self.args.dataloader_prefetch_factor
|
|
}
|
|
if dist.is_initialized() and dataloader_params['prefetch_factor']:
|
|
dataloader_params['prefetch_factor'] = dataloader_params['prefetch_factor'] * dist.get_world_size()
|
|
dataloader = DataLoader(dataset, batch_size=batch_size, **dataloader_params)
|
|
dataloader = SequenceParallelDispatcher(
|
|
dataloader, sequence_parallel, self.accelerator.device, skip_batches=skip_batches)
|
|
return dataloader
|
|
|
|
def get_train_dataloader(self, skip_batches=0):
|
|
dataloader = None
|
|
if self.template.sequence_parallel_size > 1:
|
|
dataloader = self.get_sp_dataloader(self.train_dataset, self._train_batch_size, skip_batches=skip_batches)
|
|
if dataloader is None:
|
|
# Higher efficiency
|
|
if self.train_dataset is None:
|
|
raise ValueError('Trainer: training requires a train_dataset.')
|
|
args = self.args
|
|
train_dataset = self.train_dataset
|
|
|
|
dataloader_params = {
|
|
'collate_fn': self.data_collator,
|
|
'num_workers': args.dataloader_num_workers,
|
|
'pin_memory': args.dataloader_pin_memory,
|
|
'persistent_workers': args.dataloader_persistent_workers,
|
|
'prefetch_factor': args.dataloader_prefetch_factor
|
|
}
|
|
batch_sampler_params = {
|
|
'drop_last':
|
|
args.dataloader_drop_last,
|
|
'shuffle':
|
|
args.train_dataloader_shuffle,
|
|
'data_seed':
|
|
args.data_seed,
|
|
'tp_size':
|
|
args.deepspeed['tensor_parallel']['autotp_size']
|
|
if args.deepspeed and 'tensor_parallel' in args.deepspeed else 1,
|
|
}
|
|
|
|
if hasattr(train_dataset, '__len__'):
|
|
if args.group_by_length:
|
|
batch_sampler_params['group_by_length'] = args.group_by_length
|
|
batch_sampler_params['lengths'] = train_dataset['lengths']
|
|
batch_sampler = BatchSamplerShard(
|
|
len(train_dataset), batch_size=self._train_batch_size, **batch_sampler_params)
|
|
dataloader_params['worker_init_fn'] = partial(
|
|
seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index)
|
|
if skip_batches > 0:
|
|
from accelerate.data_loader import SkipBatchSampler
|
|
batch_sampler = SkipBatchSampler(batch_sampler, skip_batches=skip_batches)
|
|
dataloader_params['batch_sampler'] = batch_sampler
|
|
dataloader = DataLoaderShard(train_dataset, device=self.accelerator.device, **dataloader_params)
|
|
else:
|
|
# IterableDataset
|
|
if dist.is_initialized() and dataloader_params['prefetch_factor']:
|
|
dataloader_params['prefetch_factor'] = dataloader_params['prefetch_factor'] * dist.get_world_size()
|
|
dataloader = DataLoader(train_dataset, batch_size=self._train_batch_size, **dataloader_params)
|
|
dataloader = DataLoaderDispatcher(dataloader, self.accelerator.device, skip_batches=skip_batches)
|
|
return dataloader
|
|
|
|
@contextmanager
|
|
def _disable_group_by_length(self):
|
|
group_by_length = getattr(self.args, 'group_by_length', False)
|
|
self.args.group_by_length = False
|
|
try:
|
|
yield
|
|
finally:
|
|
self.args.group_by_length = group_by_length
|
|
|
|
def get_eval_dataloader(self, eval_dataset=None):
|
|
dataloader = None
|
|
if self.template.sequence_parallel_size > 1:
|
|
if eval_dataset is None and self.eval_dataset is None:
|
|
raise ValueError('Trainer: evaluation requires an eval_dataset.')
|
|
eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
|
|
dataloader = self.get_sp_dataloader(eval_dataset, self.args.eval_batch_size)
|
|
if dataloader is None:
|
|
with self._disable_group_by_length():
|
|
return super().get_eval_dataloader(eval_dataset=eval_dataset)
|
|
return dataloader
|