add finetune_text_generation

1. Add TrainingArgs and cli call form for text generation task
2. Fix dp+tp finetune bug
Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11929345

examples/pytorch/text_generation/finetune_text_generation.py

@@ -0,0 +1,102 @@
+from dataclasses import dataclass, field
+
+from modelscope.metainfo import Trainers
+from modelscope.msdatasets import MsDataset
+from modelscope.trainers import EpochBasedTrainer, build_trainer
+from modelscope.trainers.training_args import TrainingArgs
+
+
+@dataclass
+class TextGenerationArguments(TrainingArgs):
+
+    trainer: str = field(
+        default=Trainers.default, metadata={
+            'help': 'The trainer used',
+        })
+
+    work_dir: str = field(
+        default='./tmp',
+        metadata={
+            'help': 'The working path for saving checkpoint',
+        })
+
+    src_txt: str = field(
+        default=None,
+        metadata={
+            'help': 'The source text key of preprocessor',
+            'cfg_node': 'preprocessor.src_txt'
+        })
+
+    tgt_txt: str = field(
+        default=None,
+        metadata={
+            'help': 'The target text key of preprocessor',
+            'cfg_node': 'preprocessor.tgt_txt'
+        })
+
+    preprocessor: str = field(
+        default=None,
+        metadata={
+            'help': 'The preprocessor type',
+            'cfg_node': 'preprocessor.type'
+        })
+
+    lr_scheduler: str = field(
+        default=None,
+        metadata={
+            'help': 'The lr scheduler type',
+            'cfg_node': 'train.lr_scheduler.type'
+        })
+
+    world_size: int = field(
+        default=None,
+        metadata={
+            'help': 'The parallel world size',
+            'cfg_node': 'megatron.world_size'
+        })
+
+    tensor_model_parallel_size: int = field(
+        default=None,
+        metadata={
+            'help': 'The tensor model parallel size',
+            'cfg_node': 'megatron.tensor_model_parallel_size'
+        })
+
+    def __call__(self, config):
+        config = super().__call__(config)
+        if config.train.lr_scheduler.type == 'noam':
+            config.train.lr_scheduler = {
+                'type': 'LambdaLR',
+                'lr_lambda': noam_lambda,
+                'options': {
+                    'by_epoch': False
+                }
+            }
+            config.train.hooks.append({'type': 'MegatronHook'})
+        return config
+
+
+def noam_lambda(current_step: int):
+    current_step += 1
+    return min(current_step**(-0.5), current_step * 100**(-1.5))
+
+
+args = TextGenerationArguments.from_cli(task='text-generation')
+
+print(args)
+
+dataset = MsDataset.load(args.dataset_name)
+train_dataset = dataset['train']
+eval_dataset = dataset['validation' if 'validation' in dataset else 'test']
+
+kwargs = dict(
+    model=args.model,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    seed=args.seed,
+    work_dir=args.work_dir,
+    cfg_modify_fn=args)
+
+trainer: EpochBasedTrainer = build_trainer(
+    name=args.trainer, default_args=kwargs)
+trainer.train()

examples/pytorch/text_generation/run_train.sh

@@ -0,0 +1,22 @@
+DATA_PARALLEL_SIZE=2
+TENSOR_MODEL_PARALLEL_SIZE=2
+
+WORLD_SIZE=$(($DATA_PARALLEL_SIZE * $TENSOR_MODEL_PARALLEL_SIZE))
+
+
+PYTHONPATH=. torchrun --nproc_per_node $WORLD_SIZE examples/pytorch/text_generation/finetune_text_generation.py \
+    --trainer 'nlp-gpt3-trainer' \
+    --work_dir './tmp' \
+    --model 'damo/nlp_gpt3_text-generation_1.3B' \
+    --dataset_name 'chinese-poetry-collection' \
+    --preprocessor 'text-gen-jieba-tokenizer' \
+    --src_txt 'text1' \
+    --tgt_txt 'text2' \
+    --max_epochs 3 \
+    --per_device_train_batch_size 16 \
+    --lr 3e-4 \
+    --lr_scheduler 'noam' \
+    --eval_metrics 'ppl' \
+    --world_size $WORLD_SIZE \
+    --tensor_model_parallel_size $TENSOR_MODEL_PARALLEL_SIZE \
+    # --dataset_name 'DuReader_robust-QG' \ # input&output

modelscope/models/nlp/gpt3/distributed_gpt3.py

@@ -952,10 +952,11 @@ class DistributedGPT3(TorchModel):
                  rank,
                  path_load_tag='model',
                  *args,
+                 megatron_cfg=None,
                  **kwargs):
         super().__init__(model_dir, *args, **kwargs)
 
-        init_megatron_util(model_dir=model_dir, rank=rank)
+        init_megatron_util(megatron_cfg, model_dir, rank=rank)
 
         self.config = GPT3Config.from_pretrained(model_dir)
         # Build model.

modelscope/trainers/nlp/gpt3_trainer.py

@@ -8,7 +8,7 @@ import torch
 from torch import nn
 
 from modelscope.metainfo import Trainers
-from modelscope.models.base import TorchModel
+from modelscope.models.base import Model, TorchModel
 from modelscope.models.nlp import GPT3ForTextGeneration
 from modelscope.trainers.builder import TRAINERS
 from modelscope.trainers.nlp_trainer import NlpEpochBasedTrainer
@@ -81,3 +81,7 @@ class GPT3Trainer(NlpEpochBasedTrainer):
     def _forward_eval(self, model: GPT3ForTextGeneration,
                       data: Dict[str, Any]) -> Dict[str, Any]:
         return model.forward(data)
+
+    def build_model(self) -> TorchModel:
+        return Model.from_pretrained(
+            self.model_dir, cfg_dict=self.cfg, megatron_cfg=self.cfg.megatron)

modelscope/trainers/nlp/text_generation_trainer.py

@@ -22,12 +22,7 @@ class TextGenerationTrainer(NlpEpochBasedTrainer):
         model.eval()
 
         with torch.no_grad():
-            if isinstance(
-                    data,
-                    Mapping) and not func_receive_dict_inputs(model.generate):
-                result = model.generate(**data)
-            else:
-                result = model.generate(data)
+            result = model.generate(data)
 
         result['preds'] = [self._decode(seq) for seq in result['sequences']]
         data['tgts'] = [self._decode(seq) for seq in data['labels']]

modelscope/utils/megatron_utils.py

@@ -20,24 +20,24 @@ _DEFAULT_CFG_WITH_MODEL_TYPE = {
 _IS_MEGATRON_INITIALIZED = False
 
 
-def init_megatron_util(cfg=None, model_dir=None, **kwargs):
+def init_megatron_util(megatron_cfg=None, model_dir=None, **kwargs):
     from modelscope.utils.hub import read_config
     from megatron_util import initialize_megatron
 
-    assert not (cfg is None and model_dir is None), \
+    assert not (megatron_cfg is None and model_dir is None), \
         'cfg and model_dir cannot both be None when initializing megatron_util'
-    if cfg is None:
+    if megatron_cfg is None:
         cfg = read_config(model_dir)
-    try:
-        megatron_cfg = cfg.megatron
-    except AttributeError:
         try:
-            model_type = cfg.model.type
+            megatron_cfg = cfg.megatron
         except AttributeError:
-            # Fit models without model type, such as mglm
-            model_type = cfg.pipeline.type
-        megatron_cfg = _DEFAULT_CFG_WITH_MODEL_TYPE[model_type] \
-            if model_type in _DEFAULT_CFG_WITH_MODEL_TYPE else {}
+            try:
+                model_type = cfg.model.type
+            except AttributeError:
+                # Fit models without model type, such as mglm
+                model_type = cfg.pipeline.type
+            megatron_cfg = _DEFAULT_CFG_WITH_MODEL_TYPE[model_type] \
+                if model_type in _DEFAULT_CFG_WITH_MODEL_TYPE else {}
     megatron_cfg.update(kwargs)
     initialize_megatron(megatron_cfg)
     global _IS_MEGATRON_INITIALIZED