Add convert megatron ckpt script

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12555222

* add convert megatron ckpt script

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)

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)

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 \