From 6dd94ff2bcbc1d3b6e38e212427a388c5b6d9ae1 Mon Sep 17 00:00:00 2001
From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com>
Date: Sun, 9 Apr 2023 21:46:48 +0800
Subject: [PATCH] add first case for gpt3 test (#236)

---
 .../models/nlp/gpt3/distributed_gpt3.py       |  4 +-
 modelscope/utils/test_utils.py                |  7 +-
 .../trainers/test_finetune_gpt3_smoke_test.py | 99 +++++++++++++++++++
 3 files changed, 106 insertions(+), 4 deletions(-)
 create mode 100644 tests/trainers/test_finetune_gpt3_smoke_test.py

diff --git a/modelscope/models/nlp/gpt3/distributed_gpt3.py b/modelscope/models/nlp/gpt3/distributed_gpt3.py
index d0da9659..75bc6130 100644
--- a/modelscope/models/nlp/gpt3/distributed_gpt3.py
+++ b/modelscope/models/nlp/gpt3/distributed_gpt3.py
@@ -955,7 +955,6 @@ class DistributedGPT3(TorchModel):
                  megatron_cfg=None,
                  **kwargs):
         super().__init__(model_dir, *args, **kwargs)
-
         init_megatron_util(megatron_cfg, model_dir, rank=rank)
 
         self.config = GPT3Config.from_pretrained(model_dir)
@@ -981,7 +980,8 @@ class DistributedGPT3(TorchModel):
         load_model = pre_load(ckpt_rank, model_dir, tag=path_load_tag)
         load_model = split_state_dict(load_model, model, tensor_ws // ckpt_ws)
 
-        self.dist_model.load_state_dict(load_model)
+        self.dist_model.load_state_dict(
+            load_model, strict=kwargs.get('strict', True))
 
         self.inference_params = None
 
diff --git a/modelscope/utils/test_utils.py b/modelscope/utils/test_utils.py
index 291fa768..b4ce7299 100644
--- a/modelscope/utils/test_utils.py
+++ b/modelscope/utils/test_utils.py
@@ -365,8 +365,11 @@ class DistributedTestCase(unittest.TestCase):
               **kwargs):
         from .torch_utils import _find_free_port
         ip = socket.gethostbyname(socket.gethostname())
-        dist_start_cmd = '%s -m torch.distributed.launch --nproc_per_node=%d --master_addr=\'%s\' --master_port=%s' % (
-            sys.executable, num_gpus, ip, _find_free_port())
+        if 'dist_start_cmd' in kwargs:
+            dist_start_cmd = kwargs.pop('dist_start_cmd')
+        else:
+            dist_start_cmd = '%s -m torch.distributed.launch --nproc_per_node=%d ' \
+                             '--master_addr=\'%s\' --master_port=%s' % (sys.executable, num_gpus, ip, _find_free_port())
 
         return self._start(
             dist_start_cmd=dist_start_cmd,
diff --git a/tests/trainers/test_finetune_gpt3_smoke_test.py b/tests/trainers/test_finetune_gpt3_smoke_test.py
new file mode 100644
index 00000000..b3a9d43a
--- /dev/null
+++ b/tests/trainers/test_finetune_gpt3_smoke_test.py
@@ -0,0 +1,99 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import shutil
+import tempfile
+import unittest
+
+import torch
+
+from modelscope.metainfo import Trainers
+from modelscope.msdatasets import MsDataset
+from modelscope.trainers import build_trainer
+from modelscope.utils.hub import Config, read_config, snapshot_download
+from modelscope.utils.test_utils import DistributedTestCase, test_level
+
+
+@unittest.skipIf(not torch.cuda.is_available()
+                 or torch.cuda.device_count() <= 1, 'distributed unittest')
+class TestFinetuneGPT3Smoke(DistributedTestCase):
+
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+
+        self.tmp_dir = tempfile.TemporaryDirectory().name
+        if not os.path.exists(self.tmp_dir):
+            os.makedirs(self.tmp_dir)
+
+        self.model_dir = snapshot_download(
+            'damo/nlp_gpt3_text-generation_1.3B')
+        config: Config = read_config(
+            os.path.join(self.model_dir, 'configuration.json'))
+        config.megatron.world_size = 2
+        config.megatron.tensor_model_parallel_size = 2
+        config.dump(os.path.join(self.model_dir, 'configuration.json'))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+        super().tearDown()
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_multi_finetune_portry(self):
+        dist_start_cmd = 'torchrun --nproc_per_node 2'
+        self.start(finetune_poetry, num_gpus=2, dist_start_cmd=dist_start_cmd)
+
+    # TODO: add gpt3 trainer predict unittest
+
+
+def finetune_poetry(work_dir='./gpt3_poetry'):
+    dataset_dict = MsDataset.load('chinese-poetry-collection')
+    train_dataset = dataset_dict['train'].remap_columns({
+        'text1': 'src_txt'
+    }).select(range(20))
+    eval_dataset = dataset_dict['test'].remap_columns({
+        'text1': 'src_txt'
+    }).select(range(20))
+    max_epochs = 2
+    tmp_dir = './gpt3_poetry'
+
+    num_warmup_steps = 100
+
+    def noam_lambda(current_step: int):
+        current_step += 1
+        return min(current_step**(-0.5),
+                   current_step * num_warmup_steps**(-1.5))
+
+    def cfg_modify_fn(cfg):
+        cfg.train.lr_scheduler = {
+            'type': 'LambdaLR',
+            'lr_lambda': noam_lambda,
+            'options': {
+                'by_epoch': False
+            }
+        }
+        cfg.train.optimizer = {'type': 'AdamW', 'lr': 3e-4}
+        cfg.train.dataloader = {'batch_size_per_gpu': 2, 'workers_per_gpu': 1}
+        cfg.train.hooks.append({'type': 'MegatronHook'})
+        cfg.evaluation.dataloader = {
+            'batch_size_per_gpu': 2,
+            'workers_per_gpu': 1
+        }
+        cfg.evaluation.metrics = 'ppl'
+        cfg.num_hidden_layers = 1
+        cfg.model.strict = False
+        return cfg
+
+    kwargs = dict(
+        model='damo/nlp_gpt3_text-generation_1.3B',
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        max_epochs=max_epochs,
+        work_dir=tmp_dir,
+        cfg_modify_fn=cfg_modify_fn)
+
+    # Construct trainer and train
+    trainer = build_trainer(name=Trainers.gpt3_trainer, default_args=kwargs)
+    trainer.train()
+
+
+if __name__ == '__main__':
+    unittest.main()