add first case for gpt3 test (#236)

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
 

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,

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()