From eecdd90412092ab435be41227bf9d5b5f4ba3a1f Mon Sep 17 00:00:00 2001 From: "hemu.zp" Date: Fri, 10 Mar 2023 15:08:56 +0800 Subject: [PATCH] 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 --- .../finetune_text_generation.py | 102 ++++++++++++++++++ examples/pytorch/text_generation/run_train.sh | 22 ++++ .../models/nlp/gpt3/distributed_gpt3.py | 3 +- modelscope/trainers/nlp/gpt3_trainer.py | 6 +- .../trainers/nlp/text_generation_trainer.py | 7 +- modelscope/utils/megatron_utils.py | 22 ++-- 6 files changed, 143 insertions(+), 19 deletions(-) create mode 100644 examples/pytorch/text_generation/finetune_text_generation.py create mode 100644 examples/pytorch/text_generation/run_train.sh diff --git a/examples/pytorch/text_generation/finetune_text_generation.py b/examples/pytorch/text_generation/finetune_text_generation.py new file mode 100644 index 00000000..5168e00e --- /dev/null +++ b/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() diff --git a/examples/pytorch/text_generation/run_train.sh b/examples/pytorch/text_generation/run_train.sh new file mode 100644 index 00000000..cbecd11a --- /dev/null +++ b/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 diff --git a/modelscope/models/nlp/gpt3/distributed_gpt3.py b/modelscope/models/nlp/gpt3/distributed_gpt3.py index be5d5f4c..d0da9659 100644 --- a/modelscope/models/nlp/gpt3/distributed_gpt3.py +++ b/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. diff --git a/modelscope/trainers/nlp/gpt3_trainer.py b/modelscope/trainers/nlp/gpt3_trainer.py index c6b6a17f..ee5fbfae 100644 --- a/modelscope/trainers/nlp/gpt3_trainer.py +++ b/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) diff --git a/modelscope/trainers/nlp/text_generation_trainer.py b/modelscope/trainers/nlp/text_generation_trainer.py index fa6a448f..0021f7fc 100644 --- a/modelscope/trainers/nlp/text_generation_trainer.py +++ b/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']] diff --git a/modelscope/utils/megatron_utils.py b/modelscope/utils/megatron_utils.py index 9f2b2c09..11e79831 100644 --- a/modelscope/utils/megatron_utils.py +++ b/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