From 49339941dfedc0f1cc0b00c242d4d7e12f928181 Mon Sep 17 00:00:00 2001
From: "hemu.zp" <hemu.zp@alibaba-inc.com>
Date: Thu, 11 May 2023 15:50:30 +0800
Subject: [PATCH] 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 |  3 ++-
 tools/convert_megatron_ckpt.py     | 31 ++++++++++++++++++++++++++++++
 tools/convert_megatron_ckpt.sh     |  7 +++++++
 3 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 tools/convert_megatron_ckpt.py
 create mode 100644 tools/convert_megatron_ckpt.sh

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 \