This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
|
||||
from .convert_utils import test_convert_precision
|
||||
from .megatron_lm_utils import (disable_forward_pre_hook, enable_forward_pre_hook, get_batch_on_this_cp_rank,
|
||||
get_optimizer_param_scheduler, init_persistent_async_worker, initialize_megatron,
|
||||
initialize_tp_communicators, load_mcore_checkpoint, maybe_finalize_async_save,
|
||||
save_mcore_checkpoint, should_disable_forward_pre_hook, warmup_jit_function, wrap_model)
|
||||
from .parallel_utils import logical_and_across_model_parallel_group, reduce_max_stat_across_model_parallel_group
|
||||
from .patcher import patch_merge_fn, patch_torch_dist_shard
|
||||
from .router_replay_utils import (RouterReplayHelper, apply_router_replay_patch, get_local_topk_idx_for_current_rank,
|
||||
get_router_replay_data, set_router_replay_data)
|
||||
from .utils import (forward_step_helper, get_packed_seq_params, get_padding_to, prepare_mcore_model,
|
||||
reconstruct_tensor_cp)
|
||||
@@ -0,0 +1,303 @@
|
||||
# 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()}')
|
||||
@@ -0,0 +1,743 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
# Parts of the functions in this file are code borrowed from NVIDIA/Megatron-LM
|
||||
import copy
|
||||
import dataclasses
|
||||
import megatron.core
|
||||
import numpy as np
|
||||
import os
|
||||
import random
|
||||
import torch
|
||||
from argparse import Namespace
|
||||
from contextlib import contextmanager
|
||||
from datetime import timedelta
|
||||
from mcore_bridge import set_random_seed, split_cp_inputs, unwrap_model
|
||||
from megatron.core import dist_checkpointing, mpu, parallel_state, tensor_parallel
|
||||
from megatron.core.dist_checkpointing.mapping import ShardedObject
|
||||
from megatron.core.dist_checkpointing.serialization import (get_default_load_sharded_strategy,
|
||||
get_default_save_sharded_strategy)
|
||||
from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest
|
||||
from megatron.core.dist_checkpointing.strategies.fully_parallel import (FullyParallelLoadStrategyWrapper,
|
||||
FullyParallelSaveStrategyWrapper)
|
||||
from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy
|
||||
from megatron.core.distributed import DistributedDataParallel as DDP
|
||||
from megatron.core.distributed import DistributedDataParallelConfig
|
||||
from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train
|
||||
from megatron.core.fusions.fused_bias_gelu import bias_gelu
|
||||
from megatron.core.fusions.fused_bias_swiglu import bias_swiglu
|
||||
from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
|
||||
from megatron.core.transformer.module import Float16Module
|
||||
from megatron.core.utils import get_torch_version, is_te_min_version, is_torch_min_version
|
||||
from packaging import version
|
||||
from transformers.utils import is_torch_npu_available
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from swift.utils import check_json_format, get_logger, init_process_group, is_master, set_device
|
||||
from .patcher import patch_merge_fn
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
mcore_017 = version.parse(megatron.core.__version__) >= version.parse('0.17.0rc0')
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_megatron_timeout(distributed_timeout_minutes):
|
||||
origin_create_group = parallel_state.create_group
|
||||
|
||||
def create_group(ranks=None, timeout=None, *_args, **kwargs):
|
||||
if timeout is None:
|
||||
timeout = timedelta(minutes=distributed_timeout_minutes)
|
||||
return origin_create_group(ranks, timeout, *_args, **kwargs)
|
||||
|
||||
parallel_state.create_group = create_group
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
parallel_state.create_group = origin_create_group
|
||||
|
||||
|
||||
def _initialize_mpu(args):
|
||||
"""Initialize torch.distributed and core model parallel."""
|
||||
if not torch.distributed.is_initialized():
|
||||
set_device()
|
||||
init_process_group(args.ddp_backend, args.ddp_timeout)
|
||||
args.rank = torch.distributed.get_rank()
|
||||
args.world_size = torch.distributed.get_world_size()
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
logger.info('model parallel is already initialized')
|
||||
else:
|
||||
distributed_timeout_minutes = args.ddp_timeout // 60
|
||||
with _patch_megatron_timeout(distributed_timeout_minutes):
|
||||
mpu.initialize_model_parallel(
|
||||
args.tensor_model_parallel_size,
|
||||
args.pipeline_model_parallel_size,
|
||||
args.virtual_pipeline_model_parallel_size,
|
||||
context_parallel_size=args.context_parallel_size,
|
||||
expert_model_parallel_size=args.expert_model_parallel_size,
|
||||
expert_tensor_parallel_size=args.expert_tensor_parallel_size,
|
||||
distributed_timeout_minutes=distributed_timeout_minutes,
|
||||
)
|
||||
if is_master():
|
||||
logger.info(f'TP: {args.tensor_model_parallel_size}, PP: {args.pipeline_model_parallel_size}, '
|
||||
f'VPP: {args.virtual_pipeline_model_parallel_size}, CP: {args.context_parallel_size}, '
|
||||
f'EP: {args.expert_model_parallel_size}, ETP: {args.expert_tensor_parallel_size}')
|
||||
|
||||
|
||||
def initialize_megatron(args):
|
||||
# Pytorch distributed.
|
||||
_initialize_mpu(args)
|
||||
|
||||
# Random seeds for reproducibility.
|
||||
logger.info(f'Setting random seeds to {args.seed}.')
|
||||
set_random_seed(args.seed, args.data_parallel_random_init, args.te_rng_tracker)
|
||||
|
||||
# Setup MoE aux loss scale value.
|
||||
if args.model_info.is_moe_model:
|
||||
from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler
|
||||
MoEAuxLossAutoScaler.set_loss_scale(torch.ones(1, device=torch.cuda.current_device()))
|
||||
|
||||
|
||||
def _get_rng_state():
|
||||
"""Collect rng state across data parallel ranks."""
|
||||
rng_state = {
|
||||
'random_rng_state': random.getstate(),
|
||||
'np_rng_state': np.random.get_state(),
|
||||
'torch_rng_state': torch.get_rng_state(),
|
||||
'cuda_rng_state': torch.cuda.get_rng_state(),
|
||||
'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states()
|
||||
}
|
||||
|
||||
# data_parallel_random_init False
|
||||
rng_state_list = [rng_state]
|
||||
|
||||
pp_rank = mpu.get_pipeline_model_parallel_rank()
|
||||
pp_size = mpu.get_pipeline_model_parallel_world_size()
|
||||
tp_rank = mpu.get_tensor_model_parallel_rank()
|
||||
tp_size = mpu.get_tensor_model_parallel_world_size()
|
||||
rng_state_list = ShardedObject(
|
||||
'rng_state',
|
||||
rng_state_list, (pp_size, tp_size), (pp_rank, tp_rank),
|
||||
replica_id=mpu.get_data_parallel_rank(with_context_parallel=True))
|
||||
return rng_state_list
|
||||
|
||||
|
||||
def _generate_state_dict(args,
|
||||
models,
|
||||
optimizer=None,
|
||||
opt_param_scheduler=None,
|
||||
rng_state=None,
|
||||
iteration=None,
|
||||
model_sd_kwargs=None,
|
||||
optim_sd_kwargs=None):
|
||||
model_sd_kwargs = model_sd_kwargs or {}
|
||||
state_dict = {
|
||||
'args': Namespace(**check_json_format(vars(args))),
|
||||
'checkpoint_version': 3.0,
|
||||
}
|
||||
if iteration is not None:
|
||||
state_dict['iteration'] = iteration
|
||||
for i, m in enumerate(models):
|
||||
key = 'model'
|
||||
if len(models) > 1:
|
||||
key = f'model{i}'
|
||||
model_sd = models[i].sharded_state_dict(**model_sd_kwargs)
|
||||
state_dict[key] = model_sd
|
||||
|
||||
if not args.no_save_optim:
|
||||
if optimizer is not None:
|
||||
state_dict['optimizer'] = _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs or {})
|
||||
if opt_param_scheduler is not None:
|
||||
state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict()
|
||||
|
||||
if not args.no_save_rng and rng_state is not None:
|
||||
state_dict['rng_state'] = rng_state
|
||||
return state_dict
|
||||
|
||||
|
||||
def _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs):
|
||||
if is_torch_npu_available():
|
||||
from swift.model.npu_patch.megatron_checkpoint import optimizer_sharded_state_dict
|
||||
return optimizer_sharded_state_dict(optimizer, state_dict, **optim_sd_kwargs)
|
||||
return optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs)
|
||||
|
||||
|
||||
def _load_optimizer_state_dict(optimizer, state_dict):
|
||||
if is_torch_npu_available():
|
||||
from swift.model.npu_patch.megatron_checkpoint import load_optimizer_state_dict
|
||||
load_optimizer_state_dict(optimizer, state_dict)
|
||||
return
|
||||
optimizer.load_state_dict(state_dict)
|
||||
|
||||
|
||||
def _filter_adapter_state_dict(state_dict, peft_format: bool, adapter_name: str = 'default'):
|
||||
"""
|
||||
When peft_format is True, keep only the PEFT format state_dict;
|
||||
when False, remove the PEFT format state_dict.
|
||||
|
||||
This function ensures it is called when tuner_type != 'full'.
|
||||
"""
|
||||
if 'model' in state_dict:
|
||||
n_models = 1
|
||||
else:
|
||||
n_models = 0
|
||||
while f'model{n_models}' in state_dict:
|
||||
n_models += 1
|
||||
|
||||
for i in range(n_models):
|
||||
if i == 0 and n_models == 1:
|
||||
model_key = 'model'
|
||||
else:
|
||||
model_key = f'model{i}'
|
||||
new_state_dict = {}
|
||||
state_dict_model = state_dict[model_key]
|
||||
for k, v in state_dict_model.items():
|
||||
if peft_format:
|
||||
if '.lora_A.' in k or '.lora_B.' in k or '.modules_to_save.' in k:
|
||||
new_state_dict[k] = v
|
||||
else:
|
||||
if '.lora_A.' in k or '.lora_B.' in k or 'original_module.' in k:
|
||||
continue
|
||||
k = k.replace('base_layer.', '')
|
||||
k = k.replace(f'modules_to_save.{adapter_name}.', '')
|
||||
v.key = v.key.replace('base_layer.', '')
|
||||
v.key = v.key.replace(f'modules_to_save.{adapter_name}.', '')
|
||||
new_state_dict[k] = v
|
||||
state_dict[model_key] = new_state_dict
|
||||
|
||||
|
||||
def _preprocess_common_before_consistancy_check(common_state_dict):
|
||||
# Convert args key of type namespace to dictionary
|
||||
preprocessed_common_state_dict = copy.deepcopy(common_state_dict)
|
||||
preprocessed_common_state_dict['args'] = vars(preprocessed_common_state_dict['args'])
|
||||
# Remove rank and local rank from state dict if it exists, since they are expected to be different
|
||||
preprocessed_common_state_dict['args'].pop('local_rank', None)
|
||||
preprocessed_common_state_dict['args'].pop('rank', None)
|
||||
return preprocessed_common_state_dict
|
||||
|
||||
|
||||
def get_sharded_sd_metadata(args):
|
||||
sharded_sd_metadata = {'singleton_local_shards': False, 'chained_optim_avoid_prefix': True}
|
||||
force_pre_mcore_014 = not is_torch_min_version('2.6a0')
|
||||
if force_pre_mcore_014 and not args.dist_ckpt_save_pre_mcore_014:
|
||||
args.dist_ckpt_save_pre_mcore_014 = True
|
||||
logger.warning(f'PyTorch version {get_torch_version()} below 2.6 detected.'
|
||||
f' Forcing dist_ckpt_save_pre_mcore_014 behavior.')
|
||||
|
||||
if args.dist_ckpt_save_pre_mcore_014:
|
||||
sharded_sd_metadata['distrib_optim_sharding_type'] = 'fully_sharded_model_space'
|
||||
else:
|
||||
if args.dist_ckpt_optim_fully_reshardable:
|
||||
sharded_sd_metadata['distrib_optim_sharding_type'] = 'fully_reshardable'
|
||||
sharded_sd_metadata[
|
||||
'distrib_optim_fully_reshardable_mem_efficient'] = args.distrib_optim_fully_reshardable_mem_efficient
|
||||
else:
|
||||
sharded_sd_metadata['distrib_optim_sharding_type'] = 'dp_reshardable'
|
||||
return sharded_sd_metadata
|
||||
|
||||
|
||||
def save_mcore_checkpoint(
|
||||
args,
|
||||
models,
|
||||
optimizer=None,
|
||||
opt_param_scheduler=None,
|
||||
iteration=1,
|
||||
output_dir: Optional[str] = None,
|
||||
peft_format: bool = False,
|
||||
):
|
||||
if output_dir is None:
|
||||
output_dir = args.output_dir
|
||||
models = unwrap_model(models)
|
||||
rng_state = _get_rng_state() if models else None
|
||||
checkpoint_dir = os.path.join(output_dir, f'iter_{iteration:07d}')
|
||||
sharded_sd_metadata = get_sharded_sd_metadata(args)
|
||||
os.makedirs(checkpoint_dir, exist_ok=True)
|
||||
|
||||
state_dict = _generate_state_dict(
|
||||
args,
|
||||
models,
|
||||
optimizer,
|
||||
opt_param_scheduler,
|
||||
rng_state,
|
||||
iteration=iteration,
|
||||
model_sd_kwargs={'metadata': sharded_sd_metadata},
|
||||
optim_sd_kwargs={'metadata': sharded_sd_metadata},
|
||||
)
|
||||
_filter_adapter_state_dict(state_dict, peft_format)
|
||||
if mcore_017:
|
||||
save_strategy = TorchDistSaveShardedStrategy()
|
||||
else:
|
||||
save_strategy = get_default_save_sharded_strategy()
|
||||
save_strategy = FullyParallelSaveStrategyWrapper(
|
||||
save_strategy,
|
||||
mpu.get_data_parallel_group(with_context_parallel=True),
|
||||
)
|
||||
kwargs = {'content_metadata': sharded_sd_metadata}
|
||||
async_save = args.async_save
|
||||
if not models: # save GPU memory
|
||||
assert 'optimizer' not in state_dict
|
||||
async_save = False
|
||||
common_path = os.path.join(checkpoint_dir, 'common.pt')
|
||||
if is_master():
|
||||
state_dict.update(kwargs)
|
||||
torch.save(state_dict, common_path)
|
||||
async_save_request = None
|
||||
else:
|
||||
async_save_request = dist_checkpointing.save(
|
||||
state_dict,
|
||||
checkpoint_dir,
|
||||
save_strategy,
|
||||
async_sharded_save=async_save,
|
||||
validate_access_integrity=True,
|
||||
preprocess_common_before_consistancy_check=_preprocess_common_before_consistancy_check,
|
||||
**kwargs)
|
||||
|
||||
if not async_save:
|
||||
assert async_save_request is None
|
||||
# Wait so everyone is done (necessary)
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
|
||||
if is_master():
|
||||
|
||||
tracker_path = os.path.join(output_dir, 'latest_checkpointed_iteration.txt')
|
||||
try:
|
||||
from megatron.core.msc_utils import open_file
|
||||
except ImportError:
|
||||
open_file = open
|
||||
with open_file(tracker_path, 'w') as f:
|
||||
f.write(str(iteration))
|
||||
|
||||
def iter_finalize_fn():
|
||||
logger.info(f'Successfully saved Megatron model weights in `{output_dir}`.')
|
||||
|
||||
if async_save:
|
||||
assert async_save_request is not None
|
||||
async_save_request.add_finalize_fn(iter_finalize_fn)
|
||||
else:
|
||||
iter_finalize_fn()
|
||||
|
||||
if async_save:
|
||||
schedule_async_save(async_save_request)
|
||||
|
||||
|
||||
# Singleton manager of async calls
|
||||
# The default is `TemporalAsyncCaller`
|
||||
_async_calls_queue = AsyncCallsQueue()
|
||||
|
||||
|
||||
def init_persistent_async_worker():
|
||||
global _async_calls_queue
|
||||
# Recreate the async_calls_queue for persistent worker
|
||||
# This duplicate step is for backward compatiblity
|
||||
_async_calls_queue = AsyncCallsQueue(persistent=True)
|
||||
|
||||
|
||||
def schedule_async_save(async_request: AsyncRequest):
|
||||
"""Schedule the async save request.
|
||||
|
||||
Args:
|
||||
async_request (AsyncRequest): the async save request.
|
||||
"""
|
||||
_async_calls_queue.schedule_async_request(async_request)
|
||||
|
||||
|
||||
def maybe_finalize_async_save(args, blocking: bool = False, terminate=False):
|
||||
"""Finalizes active async save calls.
|
||||
|
||||
Args:
|
||||
blocking (bool, optional): if True, will wait until all active requests
|
||||
are done. Otherwise, finalizes only the async request that already
|
||||
finished. Defaults to False.
|
||||
terminate (bool, optional): if True, the asynchronous queue will
|
||||
be closed as the last action of this function.
|
||||
"""
|
||||
if not args.async_save:
|
||||
return
|
||||
|
||||
_async_calls_queue.maybe_finalize_async_calls(blocking, no_dist=False)
|
||||
|
||||
if terminate:
|
||||
_async_calls_queue.close()
|
||||
|
||||
|
||||
def is_empty_async_queue() -> bool:
|
||||
"""Check if async calls queue is empty. This result is consistent across ranks."""
|
||||
return _async_calls_queue.get_num_unfinalized_calls() == 0
|
||||
|
||||
|
||||
def _load_iteration(tracker_path: str):
|
||||
if not os.path.exists(tracker_path):
|
||||
return 0
|
||||
with open(tracker_path, 'r') as f:
|
||||
iteration = int(f.read())
|
||||
# Get the max iteration retrieved across the ranks.
|
||||
if torch.distributed.is_initialized():
|
||||
iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda')
|
||||
torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
|
||||
iteration = iters_cuda[0].item()
|
||||
return iteration
|
||||
|
||||
|
||||
def load_mcore_checkpoint(args,
|
||||
ddp_models: list,
|
||||
optimizer=None,
|
||||
opt_param_scheduler=None,
|
||||
load_arg: str = 'mcore_model',
|
||||
adapter_name: str = 'default'):
|
||||
if load_arg in {'mcore_adapter', 'mcore_ref_adapter'}:
|
||||
peft_format = True
|
||||
else:
|
||||
# 'mcore_model', 'mcore_ref_model'
|
||||
peft_format = False
|
||||
load_dir = getattr(args, load_arg)
|
||||
|
||||
no_load_optim = args.no_load_optim
|
||||
no_load_rng = args.no_load_rng
|
||||
finetune = args.finetune
|
||||
if not peft_format and args.tuner_type != 'full':
|
||||
# When training with LoRA and loading the base model
|
||||
no_load_optim = True
|
||||
no_load_rng = True
|
||||
finetune = True
|
||||
models = unwrap_model(ddp_models)
|
||||
tracker_path = os.path.join(load_dir, 'latest_checkpointed_iteration.txt')
|
||||
iteration = _load_iteration(tracker_path)
|
||||
checkpoint_dir = os.path.join(load_dir, f'iter_{iteration:07d}')
|
||||
state_dict = dist_checkpointing.load_common_state_dict(checkpoint_dir)
|
||||
|
||||
ckpt_tp_pp = (
|
||||
state_dict['args'].tensor_model_parallel_size,
|
||||
state_dict['args'].pipeline_model_parallel_size,
|
||||
)
|
||||
run_tp_pp = (
|
||||
args.tensor_model_parallel_size,
|
||||
args.pipeline_model_parallel_size,
|
||||
)
|
||||
mismatch_msg = f'(TP, PP) mismatch after resume ({run_tp_pp} vs {ckpt_tp_pp} from checkpoint)'
|
||||
# Determine if RNG state will be loaded
|
||||
if (ckpt_tp_pp == run_tp_pp and not finetune and not no_load_rng
|
||||
and not getattr(state_dict['args'], 'no_save_rng', False)):
|
||||
gen_sd_rng_state = _get_rng_state() # we can load the rng state
|
||||
else:
|
||||
gen_sd_rng_state = None
|
||||
if ckpt_tp_pp != run_tp_pp:
|
||||
logger.info(f'{mismatch_msg}: RNG state will be ignored')
|
||||
sharded_sd_metadata = state_dict.get('content_metadata')
|
||||
if (not finetune and not no_load_optim and not getattr(state_dict['args'], 'no_save_optim', False)):
|
||||
gen_sd_optim = optimizer
|
||||
gen_sd_opt_param_scheduler = opt_param_scheduler
|
||||
|
||||
if (args.use_distributed_optimizer and ckpt_tp_pp != run_tp_pp
|
||||
and (sharded_sd_metadata or {}).get('distrib_optim_sharding_type') not in {
|
||||
'fully_reshardable',
|
||||
'fully_sharded_model_space',
|
||||
'fsdp_dtensor',
|
||||
}):
|
||||
raise RuntimeError(f'{mismatch_msg}: not supported for DistributedOptimizer')
|
||||
else:
|
||||
gen_sd_optim, gen_sd_opt_param_scheduler = None, None
|
||||
optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True)
|
||||
model_sd_kwargs = dict(metadata=sharded_sd_metadata)
|
||||
# TODO: check no_save_optim
|
||||
sharded_state_dict = _generate_state_dict(
|
||||
args,
|
||||
models,
|
||||
gen_sd_optim,
|
||||
gen_sd_opt_param_scheduler,
|
||||
gen_sd_rng_state,
|
||||
iteration=iteration,
|
||||
model_sd_kwargs=model_sd_kwargs,
|
||||
optim_sd_kwargs=optim_sd_kwargs)
|
||||
_filter_adapter_state_dict(sharded_state_dict, peft_format, adapter_name=adapter_name)
|
||||
model_keys = [k for k in sharded_state_dict.keys() if k.startswith('model')] # compat vpp
|
||||
for k in model_keys:
|
||||
patch_merge_fn(sharded_state_dict[k])
|
||||
if mcore_017:
|
||||
load_strategy = TorchDistLoadShardedStrategy()
|
||||
else:
|
||||
load_strategy = get_default_load_sharded_strategy(checkpoint_dir)
|
||||
|
||||
load_strategy = FullyParallelLoadStrategyWrapper(load_strategy,
|
||||
mpu.get_data_parallel_group(with_context_parallel=True))
|
||||
state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_dir, load_strategy)
|
||||
|
||||
if finetune:
|
||||
iteration = 0
|
||||
if 'args' in state_dict and not finetune:
|
||||
args.consumed_train_samples = getattr(state_dict['args'], 'consumed_train_samples', 0)
|
||||
|
||||
if len(ddp_models) == 1:
|
||||
ddp_models[0].load_state_dict(state_dict['model'], strict=False)
|
||||
else:
|
||||
for i, m in enumerate(ddp_models):
|
||||
if f'model{i}' not in state_dict:
|
||||
continue
|
||||
m.load_state_dict(state_dict[f'model{i}'])
|
||||
|
||||
if not finetune and not no_load_optim:
|
||||
if optimizer is not None:
|
||||
_load_optimizer_state_dict(optimizer, state_dict['optimizer'])
|
||||
if opt_param_scheduler is not None:
|
||||
opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler'])
|
||||
elif (args.fp16 or args.bf16) and optimizer is not None:
|
||||
optimizer.reload_model_params()
|
||||
|
||||
if not finetune and not no_load_rng:
|
||||
if 'rng_state' in state_dict:
|
||||
rng_state = state_dict['rng_state']
|
||||
if args.data_parallel_random_init:
|
||||
rng_state = rng_state[mpu.get_data_parallel_rank()]
|
||||
else:
|
||||
rng_state = rng_state[0]
|
||||
random.setstate(rng_state['random_rng_state'])
|
||||
np.random.set_state(rng_state['np_rng_state'])
|
||||
torch.set_rng_state(rng_state['torch_rng_state'])
|
||||
torch.cuda.set_rng_state(rng_state['cuda_rng_state'])
|
||||
tensor_parallel.get_cuda_rng_tracker().set_states(rng_state['rng_tracker_states'])
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
|
||||
logger.info(f'Successfully loaded Megatron model weights from: {load_dir}')
|
||||
return iteration
|
||||
|
||||
|
||||
def wrap_model(args, models, wrap_with_ddp: bool = True):
|
||||
# Set tensor model parallel attributes if not set.
|
||||
# Only parameters that are already tensor model parallel have these
|
||||
# attributes set for them. We should make sure the default attributes
|
||||
# are set for all params so the optimizer can use them.
|
||||
for m in models:
|
||||
for param in m.parameters():
|
||||
tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param)
|
||||
m.cuda(torch.cuda.current_device())
|
||||
# Fp16
|
||||
config = models[0].config
|
||||
if args.fp16 or args.bf16:
|
||||
models = [Float16Module(config, model_module) for model_module in models]
|
||||
|
||||
# DDP
|
||||
if not wrap_with_ddp:
|
||||
return models
|
||||
|
||||
kwargs = {}
|
||||
for f in dataclasses.fields(DistributedDataParallelConfig):
|
||||
if hasattr(args, f.name):
|
||||
kwargs[f.name] = getattr(args, f.name)
|
||||
|
||||
# compat: SWIFT keeps the user-facing Megatron-LM arg name, while MCore
|
||||
# DistributedDataParallelConfig expects grad_reduce_in_fp32.
|
||||
if hasattr(args, 'accumulate_allreduce_grads_in_fp32'):
|
||||
kwargs['grad_reduce_in_fp32'] = args.accumulate_allreduce_grads_in_fp32
|
||||
kwargs['check_for_nan_in_grad'] = True
|
||||
ddp_config = DistributedDataParallelConfig(**kwargs)
|
||||
|
||||
# In the Megatron FSDP and DDP use path, we need to initialize the bucket size.
|
||||
# If bucket_size is not provided as an input, use sane default.
|
||||
# If using very large dp_sizes, make buckets larger to ensure that chunks used in NCCL
|
||||
# ring-reduce implementations are large enough to remain bandwidth-bound rather than
|
||||
# latency-bound.
|
||||
if ddp_config.bucket_size is None:
|
||||
ddp_config.bucket_size = max(40000000, 1000000 * mpu.get_data_parallel_world_size(with_context_parallel=True))
|
||||
# Set bucket_size to infinity if overlap_grad_reduce is False.
|
||||
if not ddp_config.overlap_grad_reduce:
|
||||
ddp_config.bucket_size = None
|
||||
|
||||
# Setup stream for DDP initialization with proper synchronization.
|
||||
ddp_stream = torch.cuda.Stream()
|
||||
ddp_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(ddp_stream):
|
||||
models = [
|
||||
DDP(
|
||||
config=config,
|
||||
ddp_config=ddp_config,
|
||||
module=model_chunk,
|
||||
# Turn off bucketing for model_chunk 2 onwards, since communication for these
|
||||
# model chunks is overlapped with compute anyway.
|
||||
disable_bucketing=(model_chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step,
|
||||
) for (model_chunk_idx, model_chunk) in enumerate(models)
|
||||
]
|
||||
# Ensure DDP initialization completes before proceeding on the default stream.
|
||||
torch.cuda.current_stream().wait_stream(ddp_stream)
|
||||
|
||||
# Broadcast params from data parallel src rank to other data parallel ranks.
|
||||
if args.data_parallel_random_init:
|
||||
for m in models:
|
||||
m.broadcast_params()
|
||||
|
||||
return models
|
||||
|
||||
|
||||
def get_optimizer_param_scheduler(args, optimizer):
|
||||
# Iteration-based training.
|
||||
if args.lr_decay_iters is None:
|
||||
args.lr_decay_iters = args.train_iters
|
||||
lr_decay_steps = args.lr_decay_iters * args.global_batch_size
|
||||
wd_incr_steps = args.train_iters * args.global_batch_size
|
||||
wsd_decay_steps = None
|
||||
if args.lr_wsd_decay_iters is not None:
|
||||
wsd_decay_steps = args.lr_wsd_decay_iters * args.global_batch_size
|
||||
if args.lr_warmup_fraction is not None:
|
||||
lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
|
||||
else:
|
||||
lr_warmup_steps = args.lr_warmup_iters * args.global_batch_size
|
||||
|
||||
opt_param_scheduler = OptimizerParamScheduler(
|
||||
optimizer,
|
||||
init_lr=args.lr_warmup_init,
|
||||
max_lr=args.lr,
|
||||
min_lr=args.min_lr,
|
||||
lr_warmup_steps=lr_warmup_steps,
|
||||
lr_decay_steps=lr_decay_steps,
|
||||
lr_decay_style=args.lr_decay_style,
|
||||
start_wd=args.start_weight_decay,
|
||||
end_wd=args.end_weight_decay,
|
||||
wd_incr_steps=wd_incr_steps,
|
||||
wd_incr_style=args.weight_decay_incr_style,
|
||||
wsd_decay_steps=wsd_decay_steps,
|
||||
lr_wsd_decay_style=args.lr_wsd_decay_style,
|
||||
)
|
||||
|
||||
return opt_param_scheduler
|
||||
|
||||
|
||||
def should_disable_forward_pre_hook(args):
|
||||
"""Block forward pre-hook for certain configurations."""
|
||||
return args.use_distributed_optimizer and args.overlap_param_gather
|
||||
|
||||
|
||||
def enable_forward_pre_hook(model_chunks):
|
||||
for model_chunk in model_chunks:
|
||||
assert isinstance(model_chunk, DDP)
|
||||
model_chunk.enable_forward_pre_hook()
|
||||
|
||||
|
||||
def disable_forward_pre_hook(model_chunks, param_sync=True):
|
||||
for model_chunk in model_chunks:
|
||||
assert isinstance(model_chunk, DDP)
|
||||
model_chunk.disable_forward_pre_hook(param_sync=param_sync)
|
||||
|
||||
|
||||
def initialize_tp_communicators(args, config):
|
||||
"""initializing the communicators with user buffers for high-performance tensor-model-parallel
|
||||
communication overlap"""
|
||||
from transformer_engine.pytorch import module as te_module
|
||||
input_shape = [
|
||||
(args.seq_length * args.micro_batch_size) // args.context_parallel_size,
|
||||
config.hidden_size,
|
||||
]
|
||||
|
||||
if is_te_min_version('2.7.0'):
|
||||
UserBufferQuantizationMode = te_module.base.UserBufferQuantizationMode
|
||||
quantization_modes = [UserBufferQuantizationMode.FP8 if args.fp8 else UserBufferQuantizationMode.NONE]
|
||||
if args.fp8 is not None and args.first_last_layers_bf16 and (args.num_layers_at_start_in_bf16 > 0
|
||||
or args.num_layers_at_end_in_bf16 > 0):
|
||||
quantization_modes.append(UserBufferQuantizationMode.NONE)
|
||||
# The process group with the target bootstrap backend is created in Transformer Engine.
|
||||
te_module.base.initialize_ub(
|
||||
shape=input_shape,
|
||||
tp_size=args.tensor_model_parallel_size,
|
||||
quantization_modes=quantization_modes,
|
||||
)
|
||||
elif is_te_min_version('1.9.0'):
|
||||
# The process group with the target bootstrap backend is created in Transformer Engine.
|
||||
te_module.base.initialize_ub(
|
||||
shape=input_shape,
|
||||
tp_size=args.tensor_model_parallel_size,
|
||||
use_fp8=(args.fp8 is not None),
|
||||
)
|
||||
|
||||
|
||||
def warmup_jit_function(config, args):
|
||||
if args.bf16:
|
||||
dtype = torch.bfloat16
|
||||
elif args.fp16:
|
||||
dtype = torch.float16
|
||||
else:
|
||||
dtype = torch.float32
|
||||
|
||||
bias = torch.rand(config.ffn_hidden_size // config.tensor_model_parallel_size, dtype=dtype, device='cuda')
|
||||
input_tensor = torch.rand(
|
||||
(
|
||||
args.seq_length // config.context_parallel_size,
|
||||
args.micro_batch_size,
|
||||
config.ffn_hidden_size // config.tensor_model_parallel_size,
|
||||
),
|
||||
dtype=dtype,
|
||||
device='cuda',
|
||||
)
|
||||
# Warmup JIT fusions with the input_tensor grad_enable state of both forward
|
||||
# prop and recomputation
|
||||
for bias_grad, input_grad in zip([True, True], [False, True]):
|
||||
bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad
|
||||
for _ in range(5):
|
||||
if config.swiglu:
|
||||
output = bias_swiglu(input_tensor, bias)
|
||||
else:
|
||||
output = bias_gelu(bias, input_tensor)
|
||||
del bias, input_tensor, output
|
||||
|
||||
# Warmup fused bias+dropout+add
|
||||
if config.sequence_parallel:
|
||||
seq_length = args.seq_length // mpu.get_tensor_model_parallel_world_size()
|
||||
else:
|
||||
seq_length = args.seq_length
|
||||
input_tensor = torch.rand(
|
||||
(seq_length // config.context_parallel_size, args.micro_batch_size, config.hidden_size),
|
||||
dtype=dtype,
|
||||
device='cuda',
|
||||
)
|
||||
residual = torch.rand(
|
||||
(seq_length // config.context_parallel_size, args.micro_batch_size, config.hidden_size),
|
||||
dtype=dtype,
|
||||
device='cuda',
|
||||
)
|
||||
bias = torch.rand((config.hidden_size), dtype=dtype, device='cuda').expand_as(residual)
|
||||
dropout_rate = 0.1
|
||||
# Warmup JIT fusions with the input_tensor grad_enable state of both forward
|
||||
# prop and recomputation
|
||||
for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]):
|
||||
input_tensor.requires_grad = input_grad
|
||||
bias.requires_grad = bias_grad
|
||||
residual.requires_grad = residual_grad
|
||||
for _ in range(5):
|
||||
output = bias_dropout_add_fused_train([input_tensor, bias], residual, dropout_rate)
|
||||
del bias, input_tensor, residual, output
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def get_batch_on_this_cp_rank(args, batch: Dict[str, Any]):
|
||||
"""Slice batch input along sequence dimension into multiple chunks,
|
||||
which are parallelized across GPUs in a context parallel group.
|
||||
"""
|
||||
|
||||
# With causal masking, each token only attends to its prior tokens. Simply split
|
||||
# sequence into CP chunks can result in severe load imbalance. That's to say, chunks
|
||||
# at the end of sequence have bigger workload than others. To address this issue,
|
||||
# we split sequence into 2*CP ranks. Assuming CP=2, we then get 4 chunks, chunk_0
|
||||
# and chunk_3 are assigned to GPU0, chunk_1 and chunk_2 are assigned to GPU1, so
|
||||
# that we can get balanced workload among GPUs in a context parallel group.
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
if cp_size > 1:
|
||||
keys = ['labels', 'position_ids', 'loss_scale']
|
||||
if not args.is_multimodal:
|
||||
# Multimodal models will handle CP in input_embeds.
|
||||
keys.append('input_ids')
|
||||
|
||||
packed_seq_params = batch.get('packed_seq_params')
|
||||
cp_partition_mode = getattr(args, 'cp_partition_mode', 'zigzag')
|
||||
kwargs = {}
|
||||
if cp_partition_mode == 'contiguous':
|
||||
kwargs['cp_partition_mode'] = 'contiguous'
|
||||
for key, val in batch.items():
|
||||
if key not in keys:
|
||||
continue
|
||||
if args.task_type in ('seq_cls', 'embedding', 'generative_reranker') and key == 'labels':
|
||||
continue
|
||||
if val is not None:
|
||||
batch[key] = split_cp_inputs(val, getattr(packed_seq_params, 'cu_seqlens_q', None), -1, **kwargs)
|
||||
attention_mask = batch.get('attention_mask')
|
||||
if is_torch_npu_available() and attention_mask is not None and attention_mask.ndim >= 4:
|
||||
batch['attention_mask'] = split_cp_inputs(attention_mask, getattr(packed_seq_params, 'cu_seqlens_q', None),
|
||||
-2, **kwargs)
|
||||
|
||||
return batch
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import torch
|
||||
from megatron.core import mpu
|
||||
|
||||
|
||||
def reduce_max_stat_across_model_parallel_group(stat: float) -> float:
|
||||
"""
|
||||
Ranks without an optimizer will have no grad_norm or num_zeros_in_grad stats.
|
||||
We need to ensure the logging and writer rank has those values.
|
||||
This function reduces a stat tensor across the model parallel group.
|
||||
|
||||
We use an all_reduce max since the values have already been summed across optimizer ranks where possible
|
||||
"""
|
||||
stat = torch.tensor([stat], dtype=torch.float32, device=torch.cuda.current_device())
|
||||
torch.distributed.all_reduce(stat, op=torch.distributed.ReduceOp.MAX, group=mpu.get_model_parallel_group())
|
||||
return stat.item()
|
||||
|
||||
|
||||
def logical_and_across_model_parallel_group(input: bool) -> bool:
|
||||
"""
|
||||
This function gathers a bool value across the model parallel group
|
||||
"""
|
||||
input = int(bool(input))
|
||||
input = torch.tensor([input], dtype=torch.int, device=torch.cuda.current_device())
|
||||
torch.distributed.all_reduce(input, op=torch.distributed.ReduceOp.MIN, group=mpu.get_model_parallel_group())
|
||||
return bool(input.item())
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import torch
|
||||
from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory
|
||||
from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy
|
||||
|
||||
from swift.utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def patch_torch_dist_shard(thread_count):
|
||||
__init__ = TorchDistSaveShardedStrategy.__init__
|
||||
|
||||
def __new_init__(*args, **kwargs):
|
||||
kwargs['thread_count'] = thread_count
|
||||
return __init__(*args, **kwargs)
|
||||
|
||||
TorchDistSaveShardedStrategy.__init__ = __new_init__
|
||||
|
||||
|
||||
def patch_merge_fn(state_dict_model):
|
||||
# https://github.com/NVIDIA/Megatron-LM/issues/1380
|
||||
|
||||
def sh_ten_merge_fn(sub_state_dict):
|
||||
with torch.no_grad():
|
||||
shared_storage = sub_state_dict[0].untyped_storage()
|
||||
if all(shared_storage.data_ptr() == tensor.untyped_storage().data_ptr() for tensor in sub_state_dict):
|
||||
element_size = sub_state_dict[0].element_size()
|
||||
total_numel = sum(tensor.numel() for tensor in sub_state_dict)
|
||||
if shared_storage.nbytes() == total_numel * element_size:
|
||||
dim_0 = sum(tensor.shape[0] for tensor in sub_state_dict)
|
||||
shape = (dim_0, ) + sub_state_dict[0].shape[1:]
|
||||
combined_tensor = torch.empty(
|
||||
shape, dtype=sub_state_dict[0].dtype,
|
||||
device=sub_state_dict[0].device).set_(shared_storage, 0, shape)
|
||||
return combined_tensor
|
||||
return torch.cat(sub_state_dict)
|
||||
|
||||
for v in state_dict_model.values():
|
||||
if isinstance(v, ShardedTensorFactory) and 'apply_swiglu_sharded_factory' in v.merge_fn.__qualname__:
|
||||
v.merge_fn = sh_ten_merge_fn
|
||||
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
Router Replay Utilities
|
||||
Utilities for handling router replay functionality in Megatron models.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from mcore_bridge import split_cp_inputs
|
||||
from megatron.core import mpu
|
||||
from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region
|
||||
from megatron.core.transformer.transformer_block import get_num_layers_to_build
|
||||
from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
|
||||
|
||||
from swift.utils import get_logger
|
||||
from swift.utils.torch_utils import get_current_device
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
try:
|
||||
from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
|
||||
from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher
|
||||
ROUTER_REPLAY_AVAILABLE = True
|
||||
except ImportError:
|
||||
logger.warning('RouterReplay not available in current megatron-core version')
|
||||
RouterReplay = None
|
||||
RouterReplayAction = None
|
||||
MoEAlltoAllTokenDispatcher = None
|
||||
ROUTER_REPLAY_AVAILABLE = False
|
||||
|
||||
device_name = get_current_device()
|
||||
|
||||
|
||||
def is_moe_layer(tf_config, layer_idx):
|
||||
moe_layer_freq = getattr(tf_config, 'moe_layer_freq', None)
|
||||
if isinstance(moe_layer_freq, int):
|
||||
return layer_idx % moe_layer_freq == 0
|
||||
elif isinstance(moe_layer_freq, list):
|
||||
return moe_layer_freq[layer_idx] == 1
|
||||
else:
|
||||
raise ValueError(f'Unsupported moe_layer_freq type: {type(moe_layer_freq)}')
|
||||
|
||||
|
||||
def get_moe_num_layers_to_build(tf_config, vp_stage=None, pp_rank=None):
|
||||
"""Count the number of MoE layers assigned to the current rank.
|
||||
When ``moe_layer_freq`` is 1 or unset, every transformer layer is an MoE
|
||||
layer, so the count equals the total layer count. Otherwise only layers
|
||||
whose global index satisfies the frequency predicate are counted.
|
||||
Args:
|
||||
config: Megatron TransformerConfig providing layer layout information.
|
||||
vp_stage: Virtual-pipeline stage index (None defaults to current).
|
||||
pp_rank: Pipeline-parallel rank (None defaults to current).
|
||||
Returns:
|
||||
Number of MoE layers on the specified rank/stage.
|
||||
"""
|
||||
total_layers = get_num_layers_to_build(tf_config, vp_stage=vp_stage, pp_rank=pp_rank)
|
||||
|
||||
layer_offset = get_transformer_layer_offset(tf_config, vp_stage=vp_stage)
|
||||
local_global_indices = range(layer_offset, layer_offset + total_layers)
|
||||
|
||||
num_moe_layers = sum(1 for idx in local_global_indices if is_moe_layer(tf_config, idx))
|
||||
|
||||
return num_moe_layers
|
||||
|
||||
|
||||
def get_local_layer_range(tf_config, vp_rank=None, only_moe_layer=True):
|
||||
vp_size = tf_config.virtual_pipeline_model_parallel_size
|
||||
if vp_size is not None:
|
||||
vp_rank = 0 if vp_rank is None else vp_rank
|
||||
offset = 0
|
||||
for pre_vp_stage in range(vp_size):
|
||||
if pre_vp_stage == vp_rank:
|
||||
break
|
||||
num_layers_to_build = get_moe_num_layers_to_build(
|
||||
tf_config, pre_vp_stage) if only_moe_layer else get_num_layers_to_build(tf_config, pre_vp_stage)
|
||||
offset += num_layers_to_build
|
||||
else:
|
||||
offset = 0
|
||||
count = get_moe_num_layers_to_build(tf_config, vp_rank) if only_moe_layer else get_num_layers_to_build(
|
||||
tf_config, vp_rank)
|
||||
return offset, count
|
||||
|
||||
|
||||
def get_local_topk_idx_for_current_rank(global_topk_idx, tf_config, packed_seq_params=None):
|
||||
if global_topk_idx is None:
|
||||
return None
|
||||
# 1. pp slice
|
||||
# For the hybrid model, global_topk_idx contains data from all layers
|
||||
# because vLLM reports routed_experts across all transformer layers(including dense).
|
||||
# However megatron only has routers for MoE layers.
|
||||
# So local_topk_idx should filter only data from the MoE layer.
|
||||
layer_offset = get_transformer_layer_offset(tf_config, vp_stage=0)
|
||||
offset, count = get_local_layer_range(
|
||||
tf_config, tf_config.virtual_pipeline_model_parallel_size, only_moe_layer=False)
|
||||
num_layers = offset + count
|
||||
moe_layer_idx = torch.tensor([
|
||||
layer_idx for layer_idx in range(layer_offset, layer_offset + num_layers) if is_moe_layer(tf_config, layer_idx)
|
||||
],
|
||||
dtype=torch.long,
|
||||
device=global_topk_idx.device)
|
||||
local_topk_idx = torch.index_select(global_topk_idx, dim=2, index=moe_layer_idx)
|
||||
# 2. cp slice
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
if cp_size > 1:
|
||||
local_topk_idx = split_cp_inputs(local_topk_idx, getattr(packed_seq_params, 'cu_seqlens_q', None), 1)
|
||||
# 3. sp slice
|
||||
local_topk_idx = scatter_to_sequence_parallel_region(local_topk_idx.transpose(0, 1)).transpose(0, 1)
|
||||
return local_topk_idx
|
||||
|
||||
|
||||
def get_router_replay_data(tf_config, vp_rank=None):
|
||||
router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
|
||||
layers_topk_idx = []
|
||||
for router in router_instances_list:
|
||||
layers_topk_idx.append(router.recorded_topk_idx.to(torch.uint8))
|
||||
# layer_num, seq_len, topk -> 1, seq_len, layer_num, topk
|
||||
layers_topk_idx = torch.stack(layers_topk_idx).transpose(0, 1).unsqueeze(0).to(device_name)
|
||||
return layers_topk_idx
|
||||
|
||||
|
||||
def set_router_replay_data(layers_topk_idx, tf_config, vp_rank=None):
|
||||
# bs, seq_len, layer_num, topk -> layer_num, total_seq_len, topk
|
||||
layers_topk_idx_reshape = layers_topk_idx.flatten(0, 1).transpose(0, 1).to(device_name)
|
||||
offset, count = get_local_layer_range(tf_config, vp_rank)
|
||||
router_instances_list = RouterReplay.global_router_replay_instances[offset:offset + count]
|
||||
for i, router in enumerate(router_instances_list):
|
||||
router.set_target_indices(layers_topk_idx_reshape[i + offset].to(torch.int64))
|
||||
|
||||
|
||||
class RouterReplayHelper:
|
||||
"""Helper class to query router replay state and locate local RouterReplay instances."""
|
||||
|
||||
@staticmethod
|
||||
def get_micro_batch_router_list(tf_config, vp_rank=None):
|
||||
"""
|
||||
Return the list of RouterReplay instances corresponding to the current micro-batch and local
|
||||
(pp_rank, vp_stage) layer range.
|
||||
|
||||
When virtual pipeline (VPP) is enabled, the local range for the PP rank is expanded to include
|
||||
all VP stages by multiplying the per-VP count by vp_size. The returned slice is taken from the
|
||||
global RouterReplay.global_router_replay_instances list.
|
||||
|
||||
Args:
|
||||
tf_config: Configuration object used to compute layer assignments.
|
||||
vp_rank (Optional[int]): Explicit virtual pipeline stage to query. If None, the current VP
|
||||
rank from Megatron parallel state is used when available.
|
||||
Returns:
|
||||
list: A contiguous sublist of RouterReplay.router_instances for the local layer range.
|
||||
"""
|
||||
offset, count = get_local_layer_range(tf_config, vp_rank)
|
||||
router_instances_list = RouterReplay.global_router_replay_instances[offset:offset + count]
|
||||
return router_instances_list
|
||||
|
||||
@staticmethod
|
||||
def is_r2_record_action(tf_config, vp_rank=None) -> bool:
|
||||
"""Return True if the current router_replay_action is RECORD (R2) for the local router instances.
|
||||
|
||||
This inspects the first local RouterReplay instance's router_replay_action and compares it to
|
||||
RouterReplayAction.RECORD.
|
||||
"""
|
||||
router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
|
||||
return (router_instances_list and router_instances_list[0].router_replay_action == RouterReplayAction.RECORD)
|
||||
|
||||
@staticmethod
|
||||
def is_replay_forward_action(tf_config, vp_rank=None) -> bool:
|
||||
"""Return True if the current router_replay_action is REPLAY_FORWARD for the local router instances.
|
||||
|
||||
This inspects the first local RouterReplay instance's router_replay_action and compares it to
|
||||
RouterReplayAction.REPLAY_FORWARD.
|
||||
"""
|
||||
router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
|
||||
return (router_instances_list
|
||||
and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_FORWARD)
|
||||
|
||||
@staticmethod
|
||||
def is_replay_backward_action(tf_config, vp_rank=None) -> bool:
|
||||
"""Return True if the current router_replay_action is REPLAY_BACKWARD for the local router instances.
|
||||
|
||||
This inspects the first local RouterReplay instance's router_replay_action and compares it to
|
||||
RouterReplayAction.REPLAY_BACKWARD.
|
||||
"""
|
||||
router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
|
||||
return (router_instances_list
|
||||
and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_BACKWARD)
|
||||
|
||||
|
||||
def apply_router_replay_patch():
|
||||
"""
|
||||
Applies the monkey patch for MoE Router Replay functionality.
|
||||
"""
|
||||
logger.info('Applying Router Replay Patch...')
|
||||
|
||||
assert ROUTER_REPLAY_AVAILABLE, \
|
||||
'The routing replay is not supported. Please upgrade megatron-core to 0.16.0 or higher'
|
||||
|
||||
# Patch MoEAlltoAllTokenDispatcher.preprocess to handle router replay
|
||||
# When router replay is enabled, duplicate indices in top_indices can cause
|
||||
# routing_map.sum() < num_tokens * topk, leading to split size mismatch in alltoall.
|
||||
if MoEAlltoAllTokenDispatcher is not None and not hasattr(MoEAlltoAllTokenDispatcher, '_preprocess_patched'):
|
||||
original_preprocess = MoEAlltoAllTokenDispatcher.preprocess
|
||||
|
||||
def patched_preprocess(self, routing_map):
|
||||
"""Patched preprocess that handles router replay correctly for alltoall dispatcher."""
|
||||
# Call original preprocess
|
||||
result = original_preprocess(self, routing_map)
|
||||
# Fix num_out_tokens when router replay is enabled
|
||||
if (getattr(self.config, 'moe_enable_routing_replay', False) and not self.drop_and_pad
|
||||
and self.config.moe_expert_capacity_factor is None
|
||||
and not (getattr(self.config, 'moe_router_padding_for_quantization', None)
|
||||
or getattr(self.config, 'moe_router_padding_for_fp8', None))):
|
||||
# With router replay, duplicate indices can reduce the actual routed
|
||||
# token count, so derive it from the routing map instead.
|
||||
self.num_out_tokens = int(routing_map.sum().item())
|
||||
return result
|
||||
|
||||
MoEAlltoAllTokenDispatcher.preprocess = patched_preprocess
|
||||
MoEAlltoAllTokenDispatcher._preprocess_patched = True
|
||||
@@ -0,0 +1,286 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import re
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from megatron.core import mpu
|
||||
from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear
|
||||
from megatron.core.inference.communication_utils import recv_from_prev_pipeline_rank_, send_to_next_pipeline_rank
|
||||
from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
|
||||
from megatron.core.packed_seq_params import PackedSeqParams
|
||||
from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
|
||||
from megatron.core.transformer.moe.router import TopKRouter
|
||||
from torch import nn
|
||||
from transformers.utils import is_torch_npu_available
|
||||
|
||||
from swift.tuners import LoraConfig, Swift
|
||||
from swift.utils import (activate_parameters, deep_getattr, find_layers, freeze_parameters, get_logger,
|
||||
get_model_parameter_info)
|
||||
from swift.utils import get_packed_seq_params as _get_packed_seq_params
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def find_all_linears(model, extra_layers=None):
|
||||
|
||||
def _cond(name, module):
|
||||
if (extra_layers and isinstance(module, tuple(extra_layers))) or name != 'output_layer' and isinstance(
|
||||
module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, nn.Linear)):
|
||||
return True
|
||||
return False
|
||||
|
||||
return find_layers(model, _cond)
|
||||
|
||||
|
||||
def find_router(model):
|
||||
return find_layers(model, lambda name, module: isinstance(module, TopKRouter))
|
||||
|
||||
|
||||
def find_embedding(model):
|
||||
return find_layers(model, lambda name, module: isinstance(module, LanguageModelEmbedding))
|
||||
|
||||
|
||||
def get_multimodal_target_regex(
|
||||
args,
|
||||
model,
|
||||
*,
|
||||
freeze_llm: bool = False,
|
||||
freeze_vit: bool = True,
|
||||
freeze_aligner: bool = True,
|
||||
include_embedding: bool = False,
|
||||
include_router: bool = False,
|
||||
) -> str:
|
||||
megatron_model_meta = args.megatron_model_meta
|
||||
modules = []
|
||||
visual_cls = megatron_model_meta.visual_cls
|
||||
vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower]
|
||||
aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner]
|
||||
if not freeze_llm:
|
||||
modules.append('language_model')
|
||||
if not freeze_vit:
|
||||
modules += vision_tower
|
||||
if not freeze_aligner:
|
||||
modules += aligner
|
||||
assert len(modules) > 0, f'modules: {modules}'
|
||||
extra_layers = []
|
||||
if include_embedding:
|
||||
extra_layers.append(LanguageModelEmbedding)
|
||||
if include_router:
|
||||
extra_layers.append(TopKRouter)
|
||||
|
||||
res = []
|
||||
for module in modules:
|
||||
rejected_modules = []
|
||||
if not freeze_vit:
|
||||
for _aligner in aligner:
|
||||
if _aligner.startswith(f'{module}.'):
|
||||
rejected_modules.append(_aligner)
|
||||
|
||||
sub_module = deep_getattr(model, module)
|
||||
if sub_module is None:
|
||||
continue
|
||||
target_modules = find_all_linears(sub_module, extra_layers)
|
||||
if not target_modules:
|
||||
continue
|
||||
target_modules = [tm for tm in target_modules if tm]
|
||||
target_pattern = rf'.*\.({"|".join(target_modules)})' if target_modules else ''
|
||||
rejected_pattern = rf'(?!({"|".join(rejected_modules)}))' if rejected_modules else ''
|
||||
res.append(rf'{rejected_pattern}{re.escape(module)}(?=\.){target_pattern}')
|
||||
|
||||
return rf'^({"|".join(res)})$'
|
||||
|
||||
|
||||
def get_target_modules(args, model):
|
||||
if isinstance(args.target_modules, str):
|
||||
return args.target_modules
|
||||
target_modules = args.target_modules.copy()
|
||||
if 'all-linear' in target_modules:
|
||||
if args.is_multimodal and not args.language_model_only:
|
||||
if args.tuner_type == 'lora_llm':
|
||||
kwargs = {
|
||||
'freeze_llm': False,
|
||||
'freeze_vit': True,
|
||||
'freeze_aligner': True,
|
||||
}
|
||||
else: # lora
|
||||
kwargs = {
|
||||
'freeze_llm': args.freeze_llm,
|
||||
'freeze_vit': args.freeze_vit,
|
||||
'freeze_aligner': args.freeze_aligner,
|
||||
}
|
||||
return get_multimodal_target_regex(
|
||||
args,
|
||||
model,
|
||||
include_embedding='all-embedding' in target_modules,
|
||||
include_router='all-router' in target_modules,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
target_modules.remove('all-linear')
|
||||
target_modules += find_all_linears(model)
|
||||
if 'all-embedding' in target_modules:
|
||||
target_modules.remove('all-embedding')
|
||||
target_modules += find_embedding(model)
|
||||
if 'all-router' in target_modules:
|
||||
target_modules.remove('all-router')
|
||||
target_modules += find_router(model)
|
||||
return target_modules
|
||||
|
||||
|
||||
def get_modules_to_save(args, model):
|
||||
if args.task_type == 'seq_cls':
|
||||
args.modules_to_save.append('output_layer')
|
||||
modules_to_save = args.modules_to_save.copy()
|
||||
if 'all-embedding' in args.modules_to_save:
|
||||
modules_to_save.remove('all-embedding')
|
||||
modules_to_save += find_embedding(model)
|
||||
return modules_to_save
|
||||
|
||||
|
||||
def prepare_adapter(args, model):
|
||||
target_modules = get_target_modules(args, model)
|
||||
modules_to_save = get_modules_to_save(args, model)
|
||||
lora_kwargs = {
|
||||
'r': args.lora_rank,
|
||||
'target_modules': target_modules,
|
||||
'lora_alpha': args.lora_alpha,
|
||||
'lora_dropout': args.lora_dropout,
|
||||
'bias': args.lora_bias,
|
||||
'modules_to_save': modules_to_save,
|
||||
'use_rslora': args.use_rslora,
|
||||
}
|
||||
lora_config = LoraConfig(task_type='CAUSAL_LM', lora_dtype=args.lora_dtype, **lora_kwargs)
|
||||
logger.info(f'lora_config: {lora_config}')
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
if args.mcore_ref_adapter or args.ref_adapters:
|
||||
model.add_adapter('ref_adapter', lora_config)
|
||||
model.base_model._cast_adapter_dtype(adapter_name='ref_adapter', autocast_adapter_dtype=True)
|
||||
for n, p in model.named_parameters():
|
||||
if '.ref_adapter.' in n:
|
||||
p.requires_grad = False
|
||||
return model
|
||||
|
||||
|
||||
def _prepare_full_vit(args, model):
|
||||
megatron_model_meta = args.megatron_model_meta
|
||||
visual_cls = megatron_model_meta.visual_cls
|
||||
vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower]
|
||||
aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner]
|
||||
for module_prefix in vision_tower + aligner:
|
||||
module = deep_getattr(model, module_prefix)
|
||||
if module is not None:
|
||||
module.requires_grad_(True)
|
||||
|
||||
|
||||
def prepare_mcore_model(args, model):
|
||||
if args.tuner_type == 'full':
|
||||
freeze_parameters(model, args.freeze_parameters_ratio, args.freeze_parameters, args.freeze_parameters_regex)
|
||||
if args.trainable_parameters or args.trainable_parameters_regex:
|
||||
activate_parameters(model, args.trainable_parameters, args.trainable_parameters_regex)
|
||||
elif args.tuner_type in {'lora', 'lora_llm'}:
|
||||
model = prepare_adapter(args, model)
|
||||
if args.tuner_type == 'lora_llm':
|
||||
_prepare_full_vit(args, model)
|
||||
logger.info(f'model: {model}')
|
||||
logger.info_if(
|
||||
f'[rank{dist.get_rank()}] model_parameter_info: {get_model_parameter_info(model)}',
|
||||
cond=mpu.get_data_parallel_rank() == 0)
|
||||
return model
|
||||
|
||||
|
||||
def forward_step_helper(model, inputs, dtype=None):
|
||||
config = model.config
|
||||
dtype = dtype or config.params_dtype
|
||||
if not mpu.is_pipeline_first_stage():
|
||||
recv_shape_buffer = torch.empty((3, ), device=torch.cuda.current_device(), dtype=torch.int64)
|
||||
recv_from_prev_pipeline_rank_(recv_shape_buffer)
|
||||
recv_buffer = torch.empty(recv_shape_buffer.tolist(), device=torch.cuda.current_device(), dtype=dtype)
|
||||
recv_from_prev_pipeline_rank_(recv_buffer)
|
||||
model.set_input_tensor(recv_buffer)
|
||||
output_tensor = model(**inputs)
|
||||
if not mpu.is_pipeline_last_stage():
|
||||
recv_shape_buffer = torch.tensor(output_tensor.shape, device=torch.cuda.current_device(), dtype=torch.int64)
|
||||
send_to_next_pipeline_rank(recv_shape_buffer)
|
||||
send_to_next_pipeline_rank(output_tensor)
|
||||
output_tensor = None
|
||||
|
||||
return output_tensor
|
||||
|
||||
|
||||
def get_padding_to(args):
|
||||
padding_to = None
|
||||
if args.tensor_model_parallel_size > 1 and args.sequence_parallel:
|
||||
padding_to = args.tensor_model_parallel_size
|
||||
if args.context_parallel_size > 1:
|
||||
padding_to = (padding_to or 1) * args.context_parallel_size
|
||||
origin_padding_to = padding_to
|
||||
fp8_format = getattr(args, 'fp8_format', None) or getattr(args, 'fp8', None)
|
||||
fp4_format = getattr(args, 'fp4_format', None) or getattr(args, 'fp4', None)
|
||||
if args.fp8_recipe == 'blockwise':
|
||||
padding_to = (padding_to or 1) * 128
|
||||
elif args.fp8_recipe == 'mxfp8':
|
||||
# MXFP8 uses a block size of 32. Under sequence parallel, the sequence is
|
||||
# split across TP ranks, so each per-rank shard (seq_len / TP) must itself
|
||||
# be divisible by 32. Pad the total length to TP * 32 to guarantee this.
|
||||
padding_to = (padding_to or 1) * 32
|
||||
elif fp8_format is not None or fp4_format is not None:
|
||||
padding_to = (padding_to or 1) * 16
|
||||
if args.attention_backend == 'fused':
|
||||
padding_to = max(padding_to or 1, ((origin_padding_to) or 1) * 64)
|
||||
return padding_to
|
||||
|
||||
|
||||
def get_packed_seq_params(args, position_ids: torch.Tensor) -> PackedSeqParams:
|
||||
params = _get_packed_seq_params(position_ids)
|
||||
packed = PackedSeqParams(
|
||||
cu_seqlens_q=params['cu_seq_lens_q'],
|
||||
cu_seqlens_kv=params['cu_seq_lens_k'],
|
||||
max_seqlen_q=params['max_length_q'],
|
||||
max_seqlen_kv=params['max_length_k'],
|
||||
qkv_format='thd',
|
||||
)
|
||||
if hasattr(packed, 'cp_partition_mode'):
|
||||
packed.cp_partition_mode = args.cp_partition_mode
|
||||
|
||||
if is_torch_npu_available():
|
||||
packed.cu_seqlens_q_padded = params['cu_seq_lens_q']
|
||||
packed.cu_seqlens_kv_padded = params['cu_seq_lens_k']
|
||||
|
||||
return packed
|
||||
|
||||
|
||||
def reconstruct_tensor_cp(tensor, packed_seq_params, dim=1) -> torch.Tensor:
|
||||
"""In CP mode, all-gather and undo the load-balanced (zigzag) chunking
|
||||
produced by ``split_cp_inputs``, restoring the full sequence in original
|
||||
token order along ``dim``.
|
||||
|
||||
Args:
|
||||
tensor: CP-sharded local tensor whose sequence dim is at ``dim``.
|
||||
packed_seq_params: ``PackedSeqParams`` for THD inputs, or ``None`` for
|
||||
regular ``[B, S, ...]`` inputs.
|
||||
dim: Sequence dimension index of ``tensor`` (default: 1).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Full-sequence tensor with the same shape as ``tensor``
|
||||
except the size at ``dim`` is multiplied by ``cp_size``.
|
||||
"""
|
||||
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
if cp_size <= 1:
|
||||
return tensor
|
||||
|
||||
cp_rank = mpu.get_context_parallel_rank()
|
||||
cp_group = mpu.get_context_parallel_group()
|
||||
|
||||
# All-gather across CP ranks (preserve local autograd graph for `tensor`).
|
||||
output_list = [torch.empty_like(tensor) for _ in range(cp_size)]
|
||||
torch.distributed.all_gather(output_list, tensor.contiguous(), group=cp_group)
|
||||
output_list[cp_rank] = tensor
|
||||
gathered = torch.cat(output_list, dim=dim)
|
||||
|
||||
# `_undo_attention_load_balancing` assumes sequence dim is 0; transpose if needed.
|
||||
if dim != 0:
|
||||
gathered = gathered.transpose(0, dim).contiguous()
|
||||
out = _undo_attention_load_balancing(gathered, cp_size, packed_seq_params)
|
||||
if dim != 0:
|
||||
out = out.transpose(0, dim).contiguous()
|
||||
return out
|
||||
Reference in New Issue
Block a user