wehub-resource-sync
624 lines
23 KiB
Python
624 lines
23 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import accelerate
|
|
import copy
|
|
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import transformers
|
|
from accelerate.utils import find_device
|
|
from contextlib import contextmanager
|
|
from functools import wraps
|
|
from packaging import version
|
|
from peft import PeftModel
|
|
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
from transformers import PreTrainedModel, dynamic_module_utils, trainer
|
|
from transformers.modeling_outputs import SequenceClassifierOutputWithPast
|
|
from types import MethodType
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from swift.utils import (HfConfigFactory, deep_getattr, get_device_count, get_dist_setting, get_last_valid_indices,
|
|
get_logger, get_position_ids_from_cu_seqlens, is_mp, is_mp_ddp, safe_ddp_context, to_device,
|
|
to_float_dtype)
|
|
|
|
logger = get_logger()
|
|
|
|
transformers_version = version.parse(transformers.__version__)
|
|
transformers_5 = transformers_version >= version.parse('5.0.0')
|
|
|
|
|
|
def patch_fixed_float_dtype(module: torch.nn.Module, dtype):
|
|
"""Patch the module, to make sure the consisitent dtype."""
|
|
|
|
def get_float_dtype_hook(dtype):
|
|
|
|
def _float_dtype_hook(module, input, output):
|
|
return to_float_dtype(output, dtype)
|
|
|
|
return _float_dtype_hook
|
|
|
|
module.register_forward_hook(get_float_dtype_hook(dtype))
|
|
|
|
|
|
def patch_fixed_device(module: torch.nn.Module, device):
|
|
"""Move the output to the specific device"""
|
|
|
|
def get_device_hook(device):
|
|
|
|
def _device_hook(module, input, output):
|
|
return to_device(output, device)
|
|
|
|
return _device_hook
|
|
|
|
module.register_forward_hook(get_device_hook(device))
|
|
|
|
|
|
def patch_output_clone(module: torch.nn.Module):
|
|
"""Clone the output, to avoid the inplace problem"""
|
|
|
|
def _clone_hook(module, input, output):
|
|
return output.requires_grad_(True).clone()
|
|
|
|
module.register_forward_hook(_clone_hook)
|
|
|
|
|
|
def patch_get_input_embeddings(model, embedding_keys: str):
|
|
|
|
def get_input_embeddings(self) -> nn.Module:
|
|
return deep_getattr(model, embedding_keys)
|
|
|
|
model.get_input_embeddings = MethodType(get_input_embeddings, model)
|
|
|
|
|
|
def patch_output_normalizer(module: torch.nn.Module, model_meta):
|
|
|
|
def lm_head_forward(self, hidden_states):
|
|
return hidden_states
|
|
|
|
lm_heads = ['lm_head', 'output', 'embed_out', 'output_layer']
|
|
lm_head_model = get_lm_head_model(module, model_meta=model_meta, lm_heads=lm_heads)
|
|
|
|
found = False
|
|
for lm_head in lm_heads:
|
|
if hasattr(lm_head_model, lm_head):
|
|
getattr(lm_head_model, lm_head).forward = MethodType(lm_head_forward, getattr(lm_head_model, lm_head))
|
|
found = True
|
|
break
|
|
|
|
assert found, 'Cannot find the proper lm_head name'
|
|
|
|
def _output_embedding_hook(module, args, kwargs, output):
|
|
attention_mask = kwargs.get('attention_mask', None)
|
|
hidden_states = output.logits
|
|
sequence_lengths = -1 if attention_mask is None else get_last_valid_indices(attention_mask)
|
|
embeddings = hidden_states[torch.arange(hidden_states.shape[0], device=hidden_states.device), sequence_lengths]
|
|
embeddings = F.normalize(embeddings, p=2, dim=1)
|
|
return {
|
|
'last_hidden_state': embeddings.contiguous(),
|
|
}
|
|
|
|
lm_head_model.register_forward_hook(_output_embedding_hook, with_kwargs=True)
|
|
|
|
|
|
def patch_output_to_input_device(module: torch.nn.Module):
|
|
"""Patch the module, to make sure the output is in the same device with the input.
|
|
|
|
Args:
|
|
module: The module to be patched
|
|
"""
|
|
|
|
def _output_to_input_device_hook(module, args, kwargs, output):
|
|
device = find_device(args) or find_device(kwargs)
|
|
return to_device(output, device)
|
|
|
|
module.register_forward_hook(_output_to_input_device_hook, with_kwargs=True)
|
|
|
|
|
|
@contextmanager
|
|
def patch_device_map():
|
|
|
|
if not hasattr(PreTrainedModel, '_get_no_split_modules'):
|
|
yield
|
|
return
|
|
|
|
_get_no_split_modules = PreTrainedModel._get_no_split_modules
|
|
|
|
def _new_get_no_split_modules(self, device_map: str):
|
|
for module in self.modules():
|
|
if isinstance(module, PreTrainedModel) and module._no_split_modules is None:
|
|
module.__class__._no_split_modules = []
|
|
return _get_no_split_modules(self, device_map)
|
|
|
|
PreTrainedModel._get_no_split_modules = _new_get_no_split_modules
|
|
try:
|
|
yield
|
|
finally:
|
|
PreTrainedModel._get_no_split_modules = _get_no_split_modules
|
|
|
|
|
|
@contextmanager
|
|
def patch_ignore_check_imports():
|
|
import transformers.dynamic_module_utils as td
|
|
|
|
def _check_imports(filename) -> List[str]:
|
|
return td.get_relative_imports(filename)
|
|
|
|
_old_check_imports = td.check_imports
|
|
td.check_imports = _check_imports
|
|
try:
|
|
yield
|
|
finally:
|
|
td.check_imports = _old_check_imports
|
|
|
|
|
|
def get_lm_head_model(model, model_meta=None, lm_heads=None):
|
|
if isinstance(model, PeftModel):
|
|
model = model.model
|
|
model_meta = model_meta or model.model_meta
|
|
if lm_heads is None:
|
|
lm_heads = ['lm_head', 'output', 'embed_out', 'output_layer']
|
|
llm_prefix_list = getattr(model_meta.model_arch, 'language_model', None)
|
|
prefix_list = []
|
|
if llm_prefix_list:
|
|
prefix_list = llm_prefix_list[0].split('.')
|
|
|
|
current_model = model
|
|
for prefix in prefix_list:
|
|
current_model = getattr(current_model, prefix)
|
|
for lm_head in lm_heads:
|
|
if hasattr(current_model, lm_head):
|
|
return current_model
|
|
return model
|
|
|
|
|
|
def transformers_seq_cls_forward(self, *args, origin_forward, padding_side=None, **kwargs):
|
|
labels = kwargs.pop('labels', None)
|
|
return_dict = kwargs.pop('return_dict', None)
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
input_ids = kwargs.get('input_ids')
|
|
inputs_embeds = kwargs.get('inputs_embeds')
|
|
|
|
output = origin_forward(*args, **kwargs)
|
|
if hasattr(output, 'logits'):
|
|
output.logits = output.logits.to(self.score.weight.dtype)
|
|
elif 'last_hidden_state' in output:
|
|
output.logits = output['last_hidden_state'].to(self.score.weight.dtype)
|
|
logits = self.score(output.logits)
|
|
if input_ids is not None:
|
|
batch_size = input_ids.shape[0]
|
|
else:
|
|
batch_size = inputs_embeds.shape[0]
|
|
|
|
if padding_side == 'left':
|
|
pooled_logits = logits[:, -1]
|
|
else:
|
|
pad_token_id = HfConfigFactory.get_config_attr(self.config, 'pad_token_id')
|
|
if pad_token_id is None and batch_size != 1:
|
|
raise ValueError('Cannot handle batch sizes > 1 if no padding token is defined.')
|
|
if pad_token_id is None:
|
|
sequence_lengths = -1
|
|
else:
|
|
if input_ids is not None:
|
|
# if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
|
|
sequence_lengths = torch.eq(input_ids, pad_token_id).int().argmax(-1) - 1
|
|
sequence_lengths = sequence_lengths % input_ids.shape[-1]
|
|
elif kwargs.get('attention_mask') is not None:
|
|
sequence_lengths = get_last_valid_indices(kwargs['attention_mask'])
|
|
else:
|
|
sequence_lengths = -1
|
|
if isinstance(sequence_lengths, torch.Tensor):
|
|
sequence_lengths = sequence_lengths.to(logits.device)
|
|
|
|
pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
|
|
|
|
loss = None
|
|
if labels is not None:
|
|
labels = labels.to(logits.device)
|
|
if self.config.problem_type is None:
|
|
if self.num_labels == 1:
|
|
self.config.problem_type = 'regression'
|
|
elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
|
|
self.config.problem_type = 'single_label_classification'
|
|
else:
|
|
self.config.problem_type = 'multi_label_classification'
|
|
|
|
if self.config.problem_type == 'regression':
|
|
loss_fct = MSELoss()
|
|
if self.num_labels == 1:
|
|
loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
|
|
else:
|
|
loss = loss_fct(pooled_logits, labels)
|
|
elif self.config.problem_type == 'single_label_classification':
|
|
loss_fct = CrossEntropyLoss()
|
|
loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
|
|
elif self.config.problem_type == 'multi_label_classification':
|
|
loss_fct = BCEWithLogitsLoss()
|
|
loss = loss_fct(pooled_logits, labels)
|
|
if not return_dict:
|
|
output = (pooled_logits, ) + output[1:]
|
|
return ((loss, ) + output) if loss is not None else output
|
|
|
|
return SequenceClassifierOutputWithPast(
|
|
loss=loss,
|
|
logits=pooled_logits,
|
|
past_key_values=output.past_key_values,
|
|
hidden_states=output.hidden_states,
|
|
attentions=output.attentions,
|
|
)
|
|
|
|
|
|
def _patch_sequence_classification(model, model_meta):
|
|
hidden_size = HfConfigFactory.get_config_attr(model.config, 'hidden_size')
|
|
initializer_range = HfConfigFactory.get_config_attr(model.config, 'initializer_range')
|
|
|
|
lm_heads = ['lm_head', 'output', 'embed_out', 'output_layer']
|
|
lm_head_model = get_lm_head_model(model, model_meta, lm_heads)
|
|
lm_head_model.num_labels = model.config.num_labels
|
|
for lm_head in lm_heads:
|
|
if hasattr(lm_head_model, lm_head):
|
|
hidden_size = getattr(lm_head_model, lm_head).in_features
|
|
setattr(lm_head_model, lm_head, nn.Identity())
|
|
break
|
|
lm_head_model.score = nn.Linear(hidden_size, lm_head_model.num_labels, bias=False, dtype=lm_head_model.dtype)
|
|
if lm_head_model.score.weight.device == torch.device('meta'):
|
|
lm_head_model.score.to_empty(device='cpu')
|
|
lm_head_model.score.weight.data.normal_(mean=0.0, std=initializer_range)
|
|
|
|
origin_forward = lm_head_model.forward
|
|
|
|
@wraps(origin_forward.__func__)
|
|
def new_forward(self, *args, **kwargs):
|
|
return transformers_seq_cls_forward(self, *args, origin_forward=origin_forward, **kwargs)
|
|
|
|
lm_head_model.forward = MethodType(new_forward, lm_head_model)
|
|
|
|
|
|
@contextmanager
|
|
def patch_automodel_for_sequence_classification(model_info=None,
|
|
model_meta=None,
|
|
patch_from_pretrained=True,
|
|
patch_missing_init=True,
|
|
**kwargs):
|
|
"""
|
|
Context manager for patching AutoModel sequence classification.
|
|
|
|
Args:
|
|
model_info: Model information
|
|
model_meta: Model metadata
|
|
patch_from_pretrained (bool): Whether to patch PreTrainedModel.from_pretrained
|
|
patch_missing_init (bool): Whether to patch missing __init__ methods
|
|
**kwargs: Additional keyword arguments
|
|
"""
|
|
model_config = kwargs.get('model_config', None)
|
|
from_pretrained = PreTrainedModel.from_pretrained.__func__
|
|
|
|
# Patch 1: from_pretrained method
|
|
if patch_from_pretrained:
|
|
|
|
@classmethod
|
|
def _new_from_pretrained(cls, *args, **kwargs):
|
|
__init__ = cls.__init__
|
|
|
|
def __new_init__(self, *args, **kwargs):
|
|
__init__(self, *args, **kwargs)
|
|
_patch_sequence_classification(self, model_meta)
|
|
|
|
cls.__init__ = __new_init__
|
|
if hasattr(cls, '_tp_plan'): # fix tp_plan
|
|
cls._tp_plan = cls._tp_plan or {}
|
|
res = from_pretrained(cls, *args, **kwargs)
|
|
cls.__init__ = __init__
|
|
return res
|
|
else:
|
|
_new_from_pretrained = None
|
|
|
|
# Patch 2: missing __init__ methods
|
|
# https://github.com/modelscope/ms-swift/pull/5820
|
|
patched_classes = []
|
|
if patch_missing_init:
|
|
|
|
def get_all_subclasses(cls, include_root=True):
|
|
subclass_list = []
|
|
|
|
def recurse(cl):
|
|
for subclass in cl.__subclasses__():
|
|
subclass_list.append(subclass)
|
|
recurse(subclass)
|
|
|
|
recurse(cls)
|
|
|
|
ret = set(subclass_list)
|
|
if include_root:
|
|
ret.add(cls)
|
|
return ret
|
|
|
|
def create_default_init(cls):
|
|
"""Create a default __init__ method that calls super().__init__"""
|
|
|
|
def default_init(self, *args, **kwargs):
|
|
super(cls, self).__init__(*args, **kwargs)
|
|
|
|
return default_init
|
|
|
|
if model_config is not None:
|
|
# we should import in advance so that get_all_subclasses can find the class
|
|
archs = model_config.architectures
|
|
for arch in archs:
|
|
try:
|
|
getattr(transformers, arch)
|
|
except AttributeError:
|
|
continue
|
|
|
|
for subclass in get_all_subclasses(torch.nn.modules.module.Module):
|
|
if '__init__' not in subclass.__dict__:
|
|
subclass.__init__ = create_default_init(subclass)
|
|
patched_classes.append(subclass)
|
|
|
|
if patch_from_pretrained:
|
|
PreTrainedModel.from_pretrained = _new_from_pretrained
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
# Restore patches
|
|
if patch_from_pretrained:
|
|
PreTrainedModel.from_pretrained = classmethod(from_pretrained)
|
|
|
|
if patch_missing_init:
|
|
for subclass in patched_classes:
|
|
try:
|
|
if '__init__' in subclass.__dict__:
|
|
del subclass.__init__
|
|
except (AttributeError, TypeError):
|
|
pass
|
|
|
|
|
|
@contextmanager
|
|
def patch_automodel(model_info, model_meta, auto_model_cls, return_dummy_model, **kwargs):
|
|
from_pretrained = PreTrainedModel.from_pretrained.__func__
|
|
|
|
@classmethod
|
|
def _new_from_pretrained(cls, *args, **kwargs):
|
|
if 'AutoAWQFor' in auto_model_cls.__name__:
|
|
kwargs.pop('use_cache', None)
|
|
if model_info.quant_method == 'gptq':
|
|
cls.main_input_name = 'input_ids'
|
|
if hasattr(cls, '_tp_plan'): # fix tp_plan
|
|
cls._tp_plan = cls._tp_plan or {}
|
|
if return_dummy_model:
|
|
origin_torch_dtype = torch.get_default_dtype()
|
|
torch.set_default_dtype(kwargs['config'].torch_dtype)
|
|
model = cls(copy.deepcopy(kwargs['config']))
|
|
torch.set_default_dtype(origin_torch_dtype)
|
|
else:
|
|
model = from_pretrained(cls, *args, **kwargs)
|
|
return model
|
|
|
|
PreTrainedModel.from_pretrained = _new_from_pretrained
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
PreTrainedModel.from_pretrained = classmethod(from_pretrained)
|
|
|
|
|
|
def _get_max_memory(device_ids: List[int]) -> Dict[Union[int, str], int]:
|
|
"""add feat in accelerate to support MP + DDP"""
|
|
import psutil
|
|
|
|
# Make sure CUDA is initialized on each GPU to have the right memory info.
|
|
for i in device_ids:
|
|
_ = torch.tensor([0], device=i)
|
|
|
|
device_ids_set = set(device_ids)
|
|
max_memory = {}
|
|
for i in range(get_device_count()):
|
|
max_memory[i] = 0
|
|
if i in device_ids_set:
|
|
max_memory[i] = torch.cuda.mem_get_info(i)[0]
|
|
max_memory['cpu'] = psutil.virtual_memory().available
|
|
return max_memory
|
|
|
|
|
|
def _sync_max_memory(max_memory: Dict[Union[int, str], int]) -> Dict[Union[int, str], int]:
|
|
"""Make sure that the model structure of MP(device_map) is the same, when using DDP."""
|
|
max_memory_list = [v for k, v in max_memory.items() if (v > 0 and k != 'cpu')]
|
|
_, local_rank, world_size, _ = get_dist_setting()
|
|
src_tensor = torch.tensor(max_memory_list).to(local_rank)
|
|
tgt_tensor_list = [torch.zeros_like(src_tensor) for _ in range(world_size)]
|
|
dist.all_gather(tgt_tensor_list, src_tensor)
|
|
tgt_tensor = torch.stack(tgt_tensor_list, dim=0)
|
|
new_max_memory_iter = iter(tgt_tensor.min(dim=0)[0].tolist())
|
|
new_max_memory = {}
|
|
for k, v in max_memory.items():
|
|
new_max_memory[k] = v
|
|
if v > 0 and k != 'cpu':
|
|
new_max_memory[k] = next(new_max_memory_iter)
|
|
return new_max_memory
|
|
|
|
|
|
_mp_ddp_patched = False
|
|
|
|
|
|
def patch_mp_ddp():
|
|
"""Patch ddp with device_map.
|
|
After patching, the ddp can run with the device_map.
|
|
This should be called before any training starts.
|
|
"""
|
|
global _mp_ddp_patched
|
|
if _mp_ddp_patched:
|
|
return
|
|
_mp_ddp_patched = True
|
|
if is_mp_ddp():
|
|
if transformers_5:
|
|
from transformers.integrations import accelerate as tf_accelerate
|
|
get_balanced_memory = tf_accelerate.get_balanced_memory
|
|
infer_auto_device_map = tf_accelerate.infer_auto_device_map
|
|
else:
|
|
from accelerate.utils.modeling import get_balanced_memory, infer_auto_device_map
|
|
|
|
@wraps(infer_auto_device_map)
|
|
def _infer_auto_device_map_patch(model: nn.Module,
|
|
max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
|
|
**kwargs) -> Dict[str, Union[int, str, torch.device]]:
|
|
"""The auxiliary function for supports MP + DDP. Monkey Patching.
|
|
add feat in accelerate to support MP + DDP"""
|
|
verbose = kwargs.pop('verbose', False)
|
|
n_gpu = get_device_count()
|
|
_, local_rank, _, local_world_size = get_dist_setting()
|
|
device_ids = list(range(local_rank, n_gpu, local_world_size))
|
|
max_memory = _get_max_memory(device_ids)
|
|
max_memory = _sync_max_memory(max_memory)
|
|
max_memory = get_balanced_memory(model, max_memory, low_zero=False, **kwargs)
|
|
max_memory = {k: v for k, v in max_memory.items() if v > 0}
|
|
return infer_auto_device_map(model, max_memory, verbose=verbose, **kwargs)
|
|
|
|
_old_ddp_init = DDP.__init__
|
|
accelerate.accelerator.torch.nn.parallel.DistributedDataParallel.__init__ = (
|
|
lambda self, model, device_ids, output_device, *args, **kwargs: _old_ddp_init(self, model, *args, **kwargs))
|
|
if transformers_5:
|
|
tf_accelerate.infer_auto_device_map = _infer_auto_device_map_patch
|
|
else:
|
|
transformers.modeling_utils.infer_auto_device_map = _infer_auto_device_map_patch
|
|
transformers.modeling_utils.get_balanced_memory = lambda *args, **kwargs: {}
|
|
_old_accelerator_init = trainer.Accelerator.__init__
|
|
trainer.Accelerator.__init__ = (lambda self, device_placement=False, *args, **kwargs: _old_accelerator_init(
|
|
self, device_placement=device_placement, *args, **kwargs))
|
|
trainer.Accelerator.verify_device_map = lambda *args, **kwargs: False
|
|
|
|
|
|
@contextmanager
|
|
def patch_get_dynamic_module():
|
|
origin_get_cached_module_file = dynamic_module_utils.get_cached_module_file
|
|
|
|
def new_get_cached_module_file(pretrained_model_name_or_path, *args, **kwargs):
|
|
with safe_ddp_context(hash_id=str(pretrained_model_name_or_path)):
|
|
return origin_get_cached_module_file(pretrained_model_name_or_path, *args, **kwargs)
|
|
|
|
dynamic_module_utils.get_cached_module_file = new_get_cached_module_file
|
|
try:
|
|
yield
|
|
finally:
|
|
dynamic_module_utils.get_cached_module_file = origin_get_cached_module_file
|
|
|
|
|
|
@contextmanager
|
|
def patch_tp_plan(load_model: bool):
|
|
if not load_model or not is_mp() or transformers_version < version.parse('4.50') or 'WORLD_SIZE' not in os.environ:
|
|
yield
|
|
return
|
|
logger.info_once('Patch tp_plan.')
|
|
WORLD_SIZE = os.environ.get('WORLD_SIZE')
|
|
os.environ['_PATCH_WORLD_SIZE'] = WORLD_SIZE
|
|
os.environ.pop('WORLD_SIZE')
|
|
yield
|
|
os.environ['WORLD_SIZE'] = WORLD_SIZE
|
|
|
|
|
|
def revert_padding_free(outputs: Dict[str, Any], inputs: Dict[str, Any], padding_side='left'):
|
|
hidden_state_key = None
|
|
if 'last_hidden_state' in outputs:
|
|
hidden_state_key = 'last_hidden_state'
|
|
elif 'logits' in outputs:
|
|
hidden_state_key = 'logits'
|
|
elif 'token_embeddings' in outputs:
|
|
hidden_state_key = 'token_embeddings'
|
|
|
|
if hidden_state_key is None:
|
|
raise NotImplementedError()
|
|
last_hidden_state = outputs[hidden_state_key]
|
|
last_hidden_state = last_hidden_state.squeeze(dim=0)
|
|
if 'cu_seq_lens_q' in inputs:
|
|
position_ids = get_position_ids_from_cu_seqlens(inputs['cu_seq_lens_q'])
|
|
elif 'position_ids' in inputs and inputs['position_ids'].shape[0] == 1:
|
|
position_ids = inputs['position_ids']
|
|
else:
|
|
raise ValueError(
|
|
"revert_padding_free requires 'cu_seq_lens_q' or 'position_ids' in inputs, but neither was found.")
|
|
|
|
seq_lengths = []
|
|
pos = position_ids[0]
|
|
resets = torch.where(pos[1:] < pos[:-1])[0] + 1
|
|
|
|
if len(resets) == 0:
|
|
# Only one sequence in this batch item
|
|
seq_lengths = [pos.max().item() + 1]
|
|
else:
|
|
# Multiple sequences
|
|
start = 0
|
|
for end in resets:
|
|
seq_lengths.append(end - start)
|
|
start = end
|
|
seq_lengths.append(pos.shape[0] - start)
|
|
|
|
max_length = max(seq_lengths)
|
|
unpacked_logits = []
|
|
|
|
start = 0
|
|
for length in seq_lengths:
|
|
seq_state = last_hidden_state[start:start + length]
|
|
padding = torch.zeros(
|
|
(max_length - length, last_hidden_state.shape[-1])).to(last_hidden_state.dtype).to(last_hidden_state.device)
|
|
# re-padding
|
|
if padding_side == 'left':
|
|
seq_state = torch.cat((padding, seq_state), dim=0)
|
|
else:
|
|
seq_state = torch.cat((seq_state, padding), dim=0)
|
|
unpacked_logits.append(seq_state)
|
|
start += length
|
|
outputs[hidden_state_key] = torch.stack(unpacked_logits, dim=0)
|
|
return outputs
|
|
|
|
|
|
def gather_sequence_parallel_outputs(
|
|
outputs: Dict[str, Any],
|
|
tensor_keys: Optional[List[str]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Gather split tensors produced by sequence parallel training so that downstream
|
|
components (loss, metrics, etc.) can operate on full-length sequences.
|
|
"""
|
|
from swift.sequence_parallel import GatherTensor, sequence_parallel
|
|
|
|
tensor_keys = tensor_keys or ['logits', 'last_hidden_state', 'hidden_states']
|
|
|
|
position_ids = None
|
|
|
|
if sequence_parallel.rp_world_size and 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)
|
|
|
|
for key in tensor_keys:
|
|
if key in outputs:
|
|
outputs[key] = GatherTensor.apply(outputs[key], 1, position_ids)
|
|
|
|
return outputs
|
|
|
|
|
|
@contextmanager
|
|
def patch_attach_align_device_hook_on_blocks():
|
|
from accelerate import big_modeling
|
|
origin_attach_align_device_hook_on_blocks = big_modeling.attach_align_device_hook_on_blocks
|
|
|
|
def attach_align_device_hook_on_blocks(*args, **kwargs):
|
|
return
|
|
|
|
big_modeling.attach_align_device_hook_on_blocks = attach_align_device_hook_on_blocks
|
|
try:
|
|
yield
|
|
finally:
|
|
big_modeling.attach_align_device_hook_on_blocks = origin_attach_align_device_hook_on_blocks
|
|
|
|
|
|
def patch_module_forward(module, new_forward):
|
|
if getattr(module, '_patched', False):
|
|
return
|
|
module._patched = True
|
|
new_forward_wrapped = MethodType(new_forward, module)
|
|
if hasattr(module, '_old_forward'): # device_map
|
|
module._old_forward = new_forward_wrapped
|
|
else:
|
|
module.forward = new_forward_wrapped
|