diff --git a/modelscope/utils/megatron_utils.py b/modelscope/utils/megatron_utils.py index 922cb53d..53b5aacb 100644 --- a/modelscope/utils/megatron_utils.py +++ b/modelscope/utils/megatron_utils.py @@ -96,15 +96,16 @@ def convert_megatron_checkpoint( log_master( f'origin_num_partitions: {origin_num_partitions}, target_num_partitions: {target_num_partitions}' ) - os.makedirs(target_dir, exist_ok=True) if origin_num_partitions < target_num_partitions: + os.makedirs(target_dir, exist_ok=True) state_dict = _split_checkpoint( model, checkpoint_dir, target_num_partitions // origin_num_partitions) _save_converted_checkpoint(state_dict, target_dir) log_master('Split checkpoints succeeded.') elif origin_num_partitions > target_num_partitions: + os.makedirs(target_dir, exist_ok=True) state_dict = _merge_checkpoint( model, checkpoint_dir, origin_num_partitions // target_num_partitions) diff --git a/tools/convert_megatron_ckpt.py b/tools/convert_megatron_ckpt.py new file mode 100644 index 00000000..f9b8f8f3 --- /dev/null +++ b/tools/convert_megatron_ckpt.py @@ -0,0 +1,31 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import argparse +import os + +from modelscope.models import Model +from modelscope.utils.megatron_utils import convert_megatron_checkpoint + + +def unwrap_model(model): + for name in ('model', 'module', 'dist_model'): + while hasattr(model, name): + model = getattr(model, name) + return model + + +parser = argparse.ArgumentParser( + description='Split or merge your megatron_based checkpoint.') +parser.add_argument( + '--model_dir', type=str, required=True, help='Checkpoint to be converted.') +parser.add_argument( + '--target_dir', type=str, required=True, help='Target save path.') +args = parser.parse_args() + +model = Model.from_pretrained( + args.model_dir, + rank=int(os.getenv('RANK')), + megatron_cfg={'tensor_model_parallel_size': int(os.getenv('WORLD_SIZE'))}) +unwrapped_model = unwrap_model(model) + +convert_megatron_checkpoint(unwrapped_model, model.model_dir, args.target_dir) diff --git a/tools/convert_megatron_ckpt.sh b/tools/convert_megatron_ckpt.sh new file mode 100644 index 00000000..86e94877 --- /dev/null +++ b/tools/convert_megatron_ckpt.sh @@ -0,0 +1,7 @@ +TARGET_TENSOR_MODEL_PARALLEL_SIZE=1 +ORIGIN_MODEL='damo/nlp_gpt3_text-generation_1.3B' +TARGET_DIR='./target' + +torchrun --nproc_per_node $TARGET_TENSOR_MODEL_PARALLEL_SIZE tools/convert_megatron_ckpt.py \ + --model_dir $ORIGIN_MODEL \ + --target_dir $TARGET_DIR \