wehub-resource-sync
304 lines
13 KiB
Python
304 lines
13 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
|
|
import math
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from contextlib import contextmanager, nullcontext
|
|
from megatron.core import mpu
|
|
from megatron.core.extensions.transformer_engine import TEDotProductAttention
|
|
from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
|
|
from megatron.core.tensor_parallel import VocabParallelEmbedding
|
|
from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region,
|
|
gather_from_tensor_model_parallel_region)
|
|
from typing import Any, Dict
|
|
|
|
from swift.utils import HfConfigFactory, get_logger, to_device, to_float_dtype
|
|
from .megatron_lm_utils import get_batch_on_this_cp_rank
|
|
from .utils import forward_step_helper, get_packed_seq_params, get_padding_to
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def _test_params_sum(model):
|
|
total_sum = 0
|
|
zero_count = 0
|
|
n_parameter = 0
|
|
for n, p in model.named_parameters():
|
|
n_parameter += 1
|
|
sum_ = p.to(device='cuda', dtype=torch.float32).abs().sum().cpu().item()
|
|
if sum_ == 0 and '.lora_B.' not in n:
|
|
zero_count += 1
|
|
logger.warning(f'n: {n}, sum: {sum_}')
|
|
elif math.isnan(sum_) or math.isinf(sum_) or sum_ > 1e10:
|
|
logger.warning(f'n: {n}, sum: {sum_}')
|
|
else:
|
|
total_sum += sum_
|
|
cond = mpu.get_data_parallel_rank() == 0
|
|
logger.info_if(f'n_parameter: {n_parameter}', cond=cond)
|
|
logger.info_if(f'total_sum: {total_sum}', cond=cond)
|
|
logger.info_if(f'zero_count: {zero_count}', cond=cond)
|
|
|
|
|
|
def _find_modules(model, recurse: bool = True, prefix='', ignore_modules=None):
|
|
ignore_modules = ignore_modules or []
|
|
for k in ignore_modules:
|
|
if prefix.startswith(k):
|
|
return []
|
|
else:
|
|
named_children = list(model.named_children())
|
|
|
|
modules = []
|
|
for n, module in named_children:
|
|
if module.__class__ is nn.ModuleList:
|
|
modules += _find_modules(module, False, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules)
|
|
elif recurse:
|
|
modules += _find_modules(module, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules)
|
|
else:
|
|
modules.append(module)
|
|
if not named_children:
|
|
modules.append(model)
|
|
return modules
|
|
|
|
|
|
@contextmanager
|
|
def _model_cpu_forward_context(modules,
|
|
torch_dtype=None,
|
|
compute_device=None,
|
|
share_embedding: bool = False,
|
|
target_device='cpu'):
|
|
for module in modules:
|
|
try:
|
|
origin_torch_dtype = next(module.parameters()).dtype
|
|
except StopIteration:
|
|
pass
|
|
else:
|
|
break
|
|
embeddings = None
|
|
if share_embedding:
|
|
embeddings = [module for module in modules if isinstance(module, (nn.Embedding, VocabParallelEmbedding))]
|
|
|
|
def _to_cuda_hook(module, args):
|
|
if compute_device is not None or torch_dtype is not None:
|
|
module.to(device=compute_device, dtype=torch_dtype)
|
|
args = to_float_dtype(args, dtype=torch_dtype)
|
|
return args
|
|
|
|
def _to_cpu_hook(module, args, output):
|
|
if share_embedding and module in embeddings or 'rotaryemb' in module.__class__.__name__.lower():
|
|
return
|
|
module.to(device=target_device, dtype=origin_torch_dtype)
|
|
|
|
hooks = []
|
|
for module in modules:
|
|
hooks.append(module.register_forward_pre_hook(_to_cuda_hook))
|
|
hooks.append(module.register_forward_hook(_to_cpu_hook))
|
|
try:
|
|
yield
|
|
finally:
|
|
for hook in hooks:
|
|
hook.remove()
|
|
|
|
|
|
def get_examples(mm_type: str) -> Dict[str, Any]:
|
|
if mm_type == 'image':
|
|
data = {
|
|
'messages': [{
|
|
'role': 'user',
|
|
'content': '<image>describe the image.'
|
|
}, {
|
|
'role':
|
|
'assistant',
|
|
'content':
|
|
'The image depicts a close-up of a kitten with striking features. '
|
|
'The kitten has a white and gray coat with distinct black stripes, '
|
|
'particularly noticeable on its face and ears. Its eyes are large '
|
|
'and expressive, with a captivating blue hue that stands out against '
|
|
"the darker fur around them. The kitten's nose is small and pink, "
|
|
'and it has long, delicate whiskers extending from either side of its mouth. '
|
|
"The background is blurred, drawing attention to the kitten's face and "
|
|
'making it the focal point of the image. The overall impression is '
|
|
'one of cuteness and charm.'
|
|
}],
|
|
'images': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png']
|
|
}
|
|
elif mm_type == 'audio':
|
|
data = {
|
|
'messages': [{
|
|
'role': 'user',
|
|
'content': '<audio>Caption the audio.'
|
|
}, {
|
|
'role': 'assistant',
|
|
'content': "The audio contains a male voice speaking the phrase '今天天气真好呀' in Mandarin."
|
|
}],
|
|
'audios': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/weather.wav']
|
|
}
|
|
else: # text
|
|
data = {
|
|
'messages': [
|
|
{
|
|
'role': 'user',
|
|
'content': 'Introduction to ms-swift.'
|
|
},
|
|
{
|
|
'role':
|
|
'assistant',
|
|
'content':
|
|
'ms-swift is an official framework provided by the ModelScope community for fine-tuning '
|
|
'and deploying large language models and multi-modal large models.'
|
|
},
|
|
]
|
|
}
|
|
return data
|
|
|
|
|
|
def broadcast_mg_logits(mg_logits=None, src_rank=None):
|
|
if not dist.is_initialized():
|
|
return
|
|
rank = dist.get_rank()
|
|
if src_rank is None:
|
|
src_rank = dist.get_world_size() - 1
|
|
if rank == src_rank:
|
|
meta = [tuple(mg_logits.shape), str(mg_logits.dtype).split('.', 1)[1]]
|
|
else:
|
|
meta = [None, None]
|
|
|
|
dist.broadcast_object_list(meta, src=src_rank)
|
|
shape, dtype = meta
|
|
dtype = getattr(torch, dtype)
|
|
|
|
if rank != src_rank:
|
|
mg_logits = torch.empty(shape, dtype=dtype, device='cuda')
|
|
|
|
dist.broadcast(mg_logits, src=src_rank)
|
|
|
|
return mg_logits
|
|
|
|
|
|
@contextmanager
|
|
def _patch_attention_fp32(compute_dtype):
|
|
forward = TEDotProductAttention.forward
|
|
|
|
def new_forward(self, query_layer, key_layer, value_layer, *args, **kwargs):
|
|
torch_dtype = query_layer.dtype
|
|
query_layer = query_layer.to(compute_dtype)
|
|
key_layer = key_layer.to(compute_dtype)
|
|
value_layer = value_layer.to(compute_dtype)
|
|
res = forward(self, query_layer, key_layer, value_layer, *args, **kwargs)
|
|
res = res.to(dtype=torch_dtype)
|
|
return res
|
|
|
|
TEDotProductAttention.forward = new_forward
|
|
try:
|
|
yield
|
|
finally:
|
|
TEDotProductAttention.forward = forward
|
|
|
|
|
|
def test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=None):
|
|
if test_convert_dtype is None:
|
|
test_convert_dtype = getattr(args, 'test_convert_dtype', torch.float32)
|
|
template.set_mode('train')
|
|
_test_params_sum(mg_model)
|
|
|
|
config = mg_model.config
|
|
is_multimodal = config.is_multimodal
|
|
if is_multimodal:
|
|
test_mm_type = getattr(config, 'test_mm_type', 'image')
|
|
else:
|
|
test_mm_type = 'text'
|
|
mg_language_model = mg_model.language_model if is_multimodal else mg_model
|
|
if mg_language_model.config.fp8 is not None:
|
|
raise ValueError('fp8 models currently do not support testing convert_precision. '
|
|
'Please set `--test_convert_precision false`.')
|
|
share_embedding = mg_language_model.share_embeddings_and_output_weights
|
|
if hf_model is not None:
|
|
hf_model.eval()
|
|
if dist.get_world_size() == 1:
|
|
_test_params_sum(hf_model)
|
|
inputs = template.encode(get_examples(test_mm_type), return_length=True)
|
|
hf_inputs = to_device(template.data_collator([inputs]), 'cuda')
|
|
template.register_post_encode_hook([hf_model])
|
|
HfConfigFactory.set_config_attr(hf_model.config, 'use_cache', False)
|
|
model_arch = hf_model.model_meta.model_arch
|
|
ignore_modules = (model_arch.vision_tower + model_arch.aligner) if is_multimodal else []
|
|
hf_modules = _find_modules(hf_model, ignore_modules=ignore_modules)
|
|
with torch.inference_mode(), _model_cpu_forward_context(
|
|
hf_modules, test_convert_dtype, share_embedding=share_embedding):
|
|
hf_inputs.pop('text_position_ids', None)
|
|
hf_logits = hf_model(**hf_inputs).logits
|
|
hf_logits = hf_logits.to('cuda')
|
|
hf_model.to('cpu')
|
|
|
|
template.use_megatron = True
|
|
inputs = [
|
|
template.encode(get_examples(test_mm_type), return_length=True) for _ in range(2 if args.padding_free else 1)
|
|
]
|
|
mg_inputs = to_device(template.data_collator(inputs, padding_to=get_padding_to(args)), 'cuda')
|
|
mg_model.eval()
|
|
# thd
|
|
text_position_ids = mg_inputs.pop('text_position_ids', None)
|
|
if text_position_ids is None:
|
|
text_position_ids = mg_inputs.get('position_ids')
|
|
if args.padding_free:
|
|
mg_inputs['packed_seq_params'] = get_packed_seq_params(args, text_position_ids)
|
|
mg_language_model.config.fp8 = None # compat fp8
|
|
mg_modules = _find_modules(mg_language_model, ignore_modules=['visual'])
|
|
for key in ['labels', 'seq_lens', 'attention_mask_2d']:
|
|
mg_inputs.pop(key, None)
|
|
mg_inputs = get_batch_on_this_cp_rank(args, mg_inputs)
|
|
_param = next(mg_language_model.parameters())
|
|
mg_dtype = _param.dtype
|
|
mg_device = _param.device
|
|
if args.model_type == 'minimax_m2':
|
|
# router to bfloat16 (expert_bias). No need to do this when actually training.
|
|
for n, m in mg_language_model.named_modules():
|
|
if n.endswith('router'):
|
|
m.to(mg_dtype)
|
|
if getattr(config, 'enable_hyper_connections', False):
|
|
for param in mg_language_model.decoder.parameters(recurse=False):
|
|
param.data = param.data.cuda()
|
|
attention_context = (
|
|
_patch_attention_fp32(mg_dtype) if args.attention_backend.name in {'flash', 'fused'} else nullcontext())
|
|
with torch.inference_mode(), _model_cpu_forward_context(
|
|
mg_modules, test_convert_dtype, 'cuda', share_embedding=share_embedding,
|
|
target_device=mg_device), attention_context:
|
|
mg_logits = forward_step_helper(mg_model, mg_inputs, dtype=test_convert_dtype)
|
|
if args.tensor_model_parallel_size > 1 and args.task_type != 'seq_cls':
|
|
if mg_logits is not None:
|
|
mg_logits = gather_from_tensor_model_parallel_region(mg_logits)
|
|
if args.context_parallel_size > 1:
|
|
if mg_logits is not None:
|
|
mg_logits = gather_from_sequence_parallel_region(
|
|
mg_logits.transpose(0, 1), group=mpu.get_context_parallel_group())
|
|
# Contiguous CP already gathers ranks in original token order, so the
|
|
# zigzag un-balancing must be skipped (it is only correct for zigzag split).
|
|
if getattr(args, 'cp_partition_mode', 'zigzag') != 'contiguous':
|
|
mg_logits = _undo_attention_load_balancing(mg_logits, args.context_parallel_size)
|
|
mg_logits = mg_logits.transpose(0, 1)
|
|
mg_logits = broadcast_mg_logits(mg_logits)
|
|
if hf_model is None:
|
|
return
|
|
if args.task_type == 'seq_cls':
|
|
mg_logits = mg_logits[:, -1]
|
|
mean_diff = (mg_logits - hf_logits).abs().mean().item()
|
|
max_diff = (mg_logits - hf_logits).abs().max().item()
|
|
print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
|
|
else:
|
|
mg_logits = mg_logits[:, :hf_logits.shape[1]]
|
|
token_mean_diff = (mg_logits - hf_logits).abs().mean(dim=-1)
|
|
mean_diff = token_mean_diff.mean().item()
|
|
max_diff = (mg_logits - hf_logits).abs().max().item()
|
|
loss_mask = (torch.roll(hf_inputs['labels'], -1) != -100)
|
|
mean_diff_with_loss = token_mean_diff[loss_mask].mean().item()
|
|
max_diff_with_loss = (mg_logits - hf_logits)[loss_mask].abs().max().item()
|
|
print(f'token_mean_diff: {token_mean_diff}')
|
|
print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
|
|
print(f'mean_diff (with loss): {mean_diff_with_loss}, max_diff (with loss): {max_diff_with_loss} '
|
|
'(Please check that mean_diff (with loss) is less than 0.1).')
|
|
hf_tokens = hf_logits.argmax(-1)
|
|
mg_tokens = mg_logits.argmax(-1)
|
|
print(f'hf_tokens: {hf_tokens[0].tolist()}\nmg_tokens: {mg_tokens[0].tolist()}')
|
|
print(f'token_diff: {(hf_tokens != mg_tokens).sum().item()}')
|
|
print(f'token_diff (with loss): {(hf_tokens[loss_mask] != mg_tokens[loss_mask]).sum().item()}')
|