wehub-resource-sync
411 lines
15 KiB
Python
411 lines
15 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
# Part of the implementation is borrowed from huggingface/transformers.
|
|
import inspect
|
|
import math
|
|
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn.functional as F
|
|
from contextlib import contextmanager
|
|
from modelscope.hub.api import HubApi
|
|
from peft import PeftModel
|
|
from torch import nn
|
|
from torch.nn import CrossEntropyLoss, Module
|
|
from transformers import PreTrainedModel
|
|
from types import FunctionType, MethodType
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
|
|
|
from swift.model import ModelMeta
|
|
from swift.sequence_parallel import ChunkedCrossEntropyLoss, GatherLoss, sequence_parallel
|
|
from swift.utils import deep_getattr, get_dist_setting, get_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from .arguments import TrainingArguments
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def _get_deepspeed_elastic_world_size():
|
|
if dist.is_available() and dist.is_initialized():
|
|
return dist.get_world_size()
|
|
return get_dist_setting()[2]
|
|
|
|
|
|
def _enable_load_universal(ds_config):
|
|
if isinstance(ds_config, dict):
|
|
checkpoint = ds_config.get('checkpoint')
|
|
if not isinstance(checkpoint, dict):
|
|
checkpoint = {}
|
|
ds_config['checkpoint'] = checkpoint
|
|
checkpoint['load_universal'] = True
|
|
|
|
|
|
def enable_deepspeed_load_universal(args: 'TrainingArguments', trainer=None):
|
|
_enable_load_universal(getattr(args, 'deepspeed', None))
|
|
|
|
hf_ds_config = getattr(args, 'hf_deepspeed_config', None)
|
|
_enable_load_universal(getattr(hf_ds_config, 'config', None))
|
|
|
|
deepspeed_plugin = getattr(args, 'deepspeed_plugin', None)
|
|
if trainer is not None and deepspeed_plugin is None:
|
|
accelerator = getattr(trainer, 'accelerator', None)
|
|
state = getattr(accelerator, 'state', None)
|
|
deepspeed_plugin = getattr(state, 'deepspeed_plugin', None)
|
|
if deepspeed_plugin is not None:
|
|
_enable_load_universal(getattr(deepspeed_plugin, 'deepspeed_config', None))
|
|
plugin_hf_ds_config = getattr(deepspeed_plugin, 'hf_ds_config', None)
|
|
_enable_load_universal(getattr(plugin_hf_ds_config, 'config', None))
|
|
|
|
|
|
def prepare_deepspeed_elastic_config(args: 'TrainingArguments', state=None):
|
|
ds_config = args.deepspeed
|
|
if not ds_config:
|
|
return
|
|
if not isinstance(ds_config, dict):
|
|
logger.warning('DeepSpeed elastic expects args.deepspeed to be a dict, but got '
|
|
f'{type(ds_config).__name__}. Skip elastic config.')
|
|
return
|
|
|
|
from deepspeed.elasticity import compute_elastic_config
|
|
from deepspeed.git_version_info import version as __version__
|
|
|
|
enable_deepspeed_load_universal(args)
|
|
elasticity = ds_config.get('elasticity') or {}
|
|
if not elasticity:
|
|
logger.warning_once('DeepSpeed elastic callback is enabled, but no `elasticity` section is found in '
|
|
'the DeepSpeed config. Only `checkpoint.load_universal` is enabled.')
|
|
return
|
|
if elasticity.get('enabled') is False:
|
|
return
|
|
|
|
world_size = _get_deepspeed_elastic_world_size()
|
|
final_batch_size, _, micro_batch_size = compute_elastic_config(
|
|
ds_config=ds_config,
|
|
target_deepspeed_version=__version__,
|
|
world_size=world_size,
|
|
)
|
|
if world_size <= 0 or micro_batch_size <= 0:
|
|
raise ValueError('DeepSpeed elastic config produced invalid batch settings: '
|
|
f'world_size={world_size}, micro_batch_size={micro_batch_size}.')
|
|
gradient_accu_steps = max(1, final_batch_size // (micro_batch_size * world_size))
|
|
args.per_device_train_batch_size = micro_batch_size
|
|
args.gradient_accumulation_steps = gradient_accu_steps
|
|
if state is not None:
|
|
state.train_batch_size = args.per_device_train_batch_size * max(1, args.n_gpu)
|
|
logger.info_once('DeepSpeed elastic config is enabled. '
|
|
f'world_size: {world_size}, '
|
|
f'per_device_train_batch_size: {args.per_device_train_batch_size}, '
|
|
f'gradient_accumulation_steps: {args.gradient_accumulation_steps}')
|
|
|
|
|
|
def can_return_loss(model: Module) -> bool:
|
|
"""Check if a given model can return loss."""
|
|
if isinstance(model, PeftModel):
|
|
signature = inspect.signature(model.model.forward)
|
|
else:
|
|
signature = inspect.signature(model.forward)
|
|
for p in signature.parameters:
|
|
if p == 'return_loss' and signature.parameters[p].default is True:
|
|
return True
|
|
return False
|
|
|
|
|
|
def find_labels(model: Module) -> List[str]:
|
|
"""Find the labels used by a given model."""
|
|
model_name = model.__class__.__name__
|
|
if isinstance(model, PeftModel):
|
|
signature = inspect.signature(model.model.forward)
|
|
else:
|
|
signature = inspect.signature(model.forward)
|
|
if 'QuestionAnswering' in model_name:
|
|
return [p for p in signature.parameters if 'label' in p or p in ('start_positions', 'end_positions')]
|
|
else:
|
|
return [p for p in signature.parameters if 'label' in p]
|
|
|
|
|
|
def get_function(method_or_function: Union[MethodType, FunctionType]) -> FunctionType:
|
|
if isinstance(method_or_function, MethodType):
|
|
method_or_function = method_or_function.__func__
|
|
return method_or_function
|
|
|
|
|
|
def is_instance_of_ms_model(model: Module) -> bool:
|
|
"""avoid import modelscope: circular dependency problem"""
|
|
for m_cls in model.__class__.__mro__:
|
|
cls_name = m_cls.__name__
|
|
cls_module = m_cls.__module__
|
|
if cls_name == 'Model' and cls_module.startswith('modelscope'):
|
|
return True
|
|
return False
|
|
|
|
|
|
def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, **kwargs) -> torch.Tensor:
|
|
"""Common loss function for sequence parallel training"""
|
|
if hasattr(outputs, 'logits'):
|
|
logits = outputs.logits
|
|
else:
|
|
logits = outputs
|
|
device = logits.device
|
|
|
|
batch_size = logits.shape[0]
|
|
logits = logits.view(-1, logits.shape[-1])
|
|
labels = labels.flatten().to(device)
|
|
sploss_parallel_size = int(os.environ.get('CELOSS_PARALLEL_SIZE', '0'))
|
|
if sploss_parallel_size > 0:
|
|
loss = ChunkedCrossEntropyLoss.apply(logits, labels, sploss_parallel_size)
|
|
else:
|
|
loss_fct = CrossEntropyLoss(reduction='none')
|
|
loss = loss_fct(logits, labels)
|
|
if enable_dft_loss:
|
|
with torch.no_grad():
|
|
target_probs = torch.exp(-loss)
|
|
loss *= target_probs
|
|
position_ids = sequence_parallel.real_position_ids
|
|
if position_ids is not None:
|
|
position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
|
|
loss, labels = GatherLoss.apply(loss.reshape(batch_size, -1), labels.reshape(batch_size, -1), 1, position_ids)
|
|
if position_ids is not None and position_ids.min() == -1:
|
|
_pos_mask = position_ids >= 0
|
|
loss = loss[_pos_mask].contiguous()
|
|
|
|
return loss
|
|
|
|
|
|
def per_token_loss_func(outputs, labels, enable_dft_loss: bool = False, **kwargs):
|
|
logits = outputs.logits
|
|
# Upcast to float if we need to compute the loss to avoid potential precision issues
|
|
logits = logits.float()
|
|
labels = torch.roll(labels, shifts=-1, dims=-1).view(-1)
|
|
|
|
# Flatten the tokens
|
|
logits = logits.view(-1, logits.shape[-1])
|
|
# Enable model parallelism
|
|
labels = labels.to(logits.device)
|
|
loss = F.cross_entropy(logits, labels, ignore_index=-100, reduction='none')
|
|
if enable_dft_loss:
|
|
with torch.no_grad():
|
|
target_probs = torch.exp(-loss)
|
|
loss *= target_probs
|
|
return loss
|
|
|
|
|
|
def _kwargs_to_args(func, args, kwargs) -> Optional[List[Any]]:
|
|
parameters = inspect.signature(func).parameters
|
|
args = list(args)
|
|
parameters = list(parameters.items())[len(args):]
|
|
for key, param in parameters:
|
|
if key in kwargs:
|
|
args.append(kwargs[key])
|
|
elif param.default != param.empty:
|
|
args.append(param.default)
|
|
else:
|
|
return
|
|
return args
|
|
|
|
|
|
def _add_gradient_checkpointing(module_list):
|
|
|
|
requires_grad = None
|
|
|
|
def _new_forward(self, *args, **kwargs):
|
|
nonlocal requires_grad
|
|
if requires_grad is None:
|
|
requires_grad = any(p.requires_grad for p in self.parameters())
|
|
|
|
new_args = _kwargs_to_args(self.__old_forward, args, kwargs)
|
|
if new_args is not None and self.gradient_checkpointing and self.training:
|
|
if new_args and isinstance(new_args[0], torch.Tensor) and requires_grad and not new_args[0].requires_grad:
|
|
new_args[0].requires_grad_(True)
|
|
layer_ret = self._gradient_checkpointing_func(self.__old_forward, *new_args)
|
|
logger.info_once('Successfully using dynamic gradient checkpointing.')
|
|
else:
|
|
layer_ret = self.__old_forward(*args, **kwargs)
|
|
return layer_ret
|
|
|
|
for module in module_list:
|
|
module.gradient_checkpointing = False
|
|
if hasattr(module, '_old_forward'): # device_map
|
|
__old_forward = module._old_forward
|
|
module._old_forward = MethodType(_new_forward, module)
|
|
else:
|
|
__old_forward = module.forward
|
|
module.forward = MethodType(_new_forward, module)
|
|
module.__old_forward = __old_forward
|
|
|
|
|
|
def find_module_list(model) -> Optional[nn.ModuleList]:
|
|
module_lists = []
|
|
for m in model.modules():
|
|
if hasattr(m, 'gradient_checkpointing') or m.__class__.__name__ == 'CheckpointWrapper':
|
|
return
|
|
if (isinstance(m, (nn.ModuleList, nn.Sequential)) and len(m) >= 10
|
|
and 'mlp' not in m[0].__class__.__name__.lower()): # fix moe
|
|
module_lists.append(m)
|
|
if module_lists:
|
|
return max(module_lists, key=lambda x: len(x))
|
|
|
|
|
|
def dynamic_gradient_checkpointing(model, including_vit: bool = False) -> None:
|
|
if isinstance(model, PeftModel):
|
|
model = model.model
|
|
model_meta: ModelMeta = getattr(model, 'model_meta', None)
|
|
if model_meta is not None and model_meta.is_multimodal and model_meta.model_arch:
|
|
tower_names = model_meta.model_arch.language_model.copy()
|
|
if including_vit:
|
|
tower_names += model_meta.model_arch.vision_tower
|
|
else:
|
|
tower_names = [None]
|
|
|
|
model.supports_gradient_checkpointing = True
|
|
for tower_name in tower_names:
|
|
if tower_name is None:
|
|
model_tower = model
|
|
else:
|
|
model_tower = deep_getattr(model, tower_name)
|
|
if model_tower is None:
|
|
continue
|
|
model_tower.supports_gradient_checkpointing = True
|
|
module_list = find_module_list(model_tower)
|
|
if module_list is None:
|
|
continue
|
|
_add_gradient_checkpointing(module_list)
|
|
logger.info(f'Automatically add gradient_checkpointing to {model_tower.__class__}.')
|
|
|
|
|
|
@contextmanager
|
|
def disable_gradient_checkpointing(model: PreTrainedModel, gradient_checkpointing_kwargs: Optional[Dict] = None):
|
|
"""
|
|
Temporarily disable gradient checkpointing, restoring the previous state afterward.
|
|
|
|
When gradient checkpointing is enabled with use_reentrant=True (default), calling the model inside a
|
|
torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True").
|
|
Temporarily disable checkpointing to avoid this warning during inference.
|
|
|
|
Args:
|
|
model (`PreTrainedModel`):
|
|
Model for which to temporarily disable gradient checkpointing.
|
|
gradient_checkpointing_kwargs (`dict` or `None`, *optional*):
|
|
Additional kwargs for gradient checkpointing enabling.
|
|
"""
|
|
was_enabled = getattr(model, 'is_gradient_checkpointing', False)
|
|
if was_enabled:
|
|
model.gradient_checkpointing_disable()
|
|
try:
|
|
yield
|
|
finally:
|
|
if was_enabled:
|
|
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs)
|
|
|
|
|
|
def gather_for_unpadded_tensors(input_data, use_gather_object=False):
|
|
from accelerate.utils import gather_object
|
|
if getattr(sequence_parallel, 'dp_group', None) is not None:
|
|
input_data = sequence_parallel._gather_object_dp(input_data)
|
|
else:
|
|
input_data = gather_object(input_data)
|
|
output = []
|
|
for _data in input_data:
|
|
if len(_data.shape) == 0:
|
|
_data = _data.unsqueeze(0)
|
|
_data = _data.cpu()
|
|
output.append(_data)
|
|
if len(output[0].shape) == 1 and output[0].shape[0] > 1:
|
|
data = torch.stack(output, dim=0)
|
|
else:
|
|
data = torch.concat(output, dim=0)
|
|
return data
|
|
|
|
|
|
def calculate_max_steps(args: 'TrainingArguments', dataset) -> int:
|
|
if args.max_steps and args.max_steps > 0:
|
|
max_steps = args.max_steps
|
|
else:
|
|
len_dataset = len(dataset)
|
|
_, _, world_size, _ = get_dist_setting()
|
|
total_train_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps * world_size
|
|
num_update_steps_per_epoch = len_dataset // total_train_batch_size
|
|
num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
|
|
max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
|
|
return max_steps
|
|
|
|
|
|
def extract_version(name: str) -> Optional[int]:
|
|
if not name.startswith('v'):
|
|
return None
|
|
try:
|
|
num = name[1:].split('-', 1)[0]
|
|
return int(num)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def get_previous_version_from_path(current_path: str) -> Optional[str]:
|
|
from pathlib import Path
|
|
current = Path(current_path)
|
|
parent = current.parent
|
|
current_name = current.name
|
|
|
|
candidates = [d for d in parent.iterdir() if d.is_dir()]
|
|
|
|
valid = [(d.name, extract_version(d.name)) for d in candidates]
|
|
valid = [(name, ver) for name, ver in valid if ver is not None]
|
|
|
|
valid.sort(key=lambda x: x[1])
|
|
names = [name for name, _ in valid]
|
|
|
|
if current_name not in names:
|
|
return None
|
|
|
|
idx = names.index(current_name)
|
|
if idx == 0:
|
|
return None
|
|
|
|
prev_name = names[idx - 1]
|
|
return str(parent / prev_name)
|
|
|
|
|
|
def get_resume_dir(output_dir):
|
|
return get_previous_version_from_path(output_dir)
|
|
|
|
|
|
def replace_index_file(output_dir: str):
|
|
import json
|
|
import os
|
|
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, WEIGHTS_INDEX_NAME
|
|
index_file = os.path.join(output_dir, WEIGHTS_INDEX_NAME)
|
|
|
|
if not os.path.exists(index_file):
|
|
return
|
|
with open(index_file, 'r', encoding='utf-8') as f:
|
|
bin_data = json.load(f)
|
|
if 'weight_map' not in bin_data:
|
|
return
|
|
bin_data['weight_map'] = {
|
|
k: v.replace('pytorch_model', 'model').replace('.bin', '.safetensors')
|
|
for k, v in bin_data['weight_map'].items()
|
|
}
|
|
safe_path = os.path.join(output_dir, SAFE_WEIGHTS_INDEX_NAME)
|
|
with open(safe_path, 'w', encoding='utf-8') as f:
|
|
json.dump(bin_data, f, indent=2)
|
|
from contextlib import suppress
|
|
with suppress(FileNotFoundError):
|
|
os.remove(os.path.join(output_dir, WEIGHTS_INDEX_NAME))
|
|
|
|
|
|
@contextmanager
|
|
def patch_modelscope_hub_timeout():
|
|
__init__ = HubApi.__init__
|
|
|
|
def __new_init__(self, *args, **kwargs):
|
|
timeout = kwargs.get('timeout')
|
|
if timeout is not None and timeout > 5:
|
|
kwargs['timeout'] = 5
|
|
__init__(self, *args, **kwargs)
|
|
|
|
HubApi.__init__ = __new_init__
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
HubApi.__init__ = __init__
|