From 906fa673b44d4bf5249b33e8421875fbcc38b14a Mon Sep 17 00:00:00 2001 From: "jerry.lp" Date: Sat, 17 Dec 2022 05:52:57 +0800 Subject: [PATCH] add gpt-moe model for modelscope finetune Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11085918 --- modelscope/metainfo.py | 1 + modelscope/models/nlp/__init__.py | 2 + modelscope/models/nlp/gpt_moe/__init__.py | 2 + .../models/nlp/gpt_moe/distributed_gpt_moe.py | 115 +++++++++++---- .../models/nlp/gpt_moe/moe/sharded_moe.py | 10 +- .../models/nlp/gpt_moe/text_generation.py | 20 ++- .../nlp/distributed_gpt_moe_pipeline.py | 3 +- modelscope/trainers/nlp/gpt_moe_trainer.py | 61 ++++++++ tests/trainers/test_finetune_gpt3.py | 19 +-- tests/trainers/test_finetune_gpt_moe.py | 131 ++++++++++++++++++ 10 files changed, 320 insertions(+), 44 deletions(-) create mode 100644 modelscope/trainers/nlp/gpt_moe_trainer.py create mode 100644 tests/trainers/test_finetune_gpt_moe.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index bdee20fd..bcc6dd2f 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -358,6 +358,7 @@ class Trainers(object): text_generation_trainer = 'text-generation-trainer' nlp_plug_trainer = 'nlp-plug-trainer' gpt3_trainer = 'nlp-gpt3-trainer' + gpt_moe_trainer = 'nlp-gpt-moe-trainer' # audio trainers speech_frcrn_ans_cirm_16k = 'speech_frcrn_ans_cirm_16k' diff --git a/modelscope/models/nlp/__init__.py b/modelscope/models/nlp/__init__.py index 44aa813a..3a32e44f 100644 --- a/modelscope/models/nlp/__init__.py +++ b/modelscope/models/nlp/__init__.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: from .deberta_v2 import DebertaV2ForMaskedLM, DebertaV2Model from .gpt_neo import GPTNeoModel from .gpt3 import GPT3ForTextGeneration, DistributedGPT3 + from .gpt_moe import GPTMoEForTextGeneration, DistributedGPTMoE from .heads import SequenceClassificationHead from .palm_v2 import PalmForTextGeneration from .ponet import PoNetForMaskedLM, PoNetModel, PoNetConfig @@ -60,6 +61,7 @@ else: 'csanmt': ['CsanmtForTranslation'], 'heads': ['SequenceClassificationHead'], 'gpt3': ['GPT3ForTextGeneration', 'DistributedGPT3'], + 'gpt_moe': ['GPTMoEForTextGeneration', 'DistributedGPTMoE'], 'structbert': [ 'SbertForFaqQuestionAnswering', 'SbertForMaskedLM', diff --git a/modelscope/models/nlp/gpt_moe/__init__.py b/modelscope/models/nlp/gpt_moe/__init__.py index 3010e64f..5508b9a3 100644 --- a/modelscope/models/nlp/gpt_moe/__init__.py +++ b/modelscope/models/nlp/gpt_moe/__init__.py @@ -8,12 +8,14 @@ if TYPE_CHECKING: from .backbone import GPTMoEModel from .text_generation import GPTMoEForTextGeneration from .tokenizer import JiebaBPETokenizer + from .distributed_gpt_moe import DistributedGPTMoE else: _import_structure = { 'configuration': ['GPTMoEConfig'], 'backbone': ['GPTMoEModel'], 'text_generation': ['GPTMoEForTextGeneration'], 'tokenizer': ['JiebaBPETokenizer'], + 'distributed_gpt_moe': ['DistributedGPTMoE'], } import sys diff --git a/modelscope/models/nlp/gpt_moe/distributed_gpt_moe.py b/modelscope/models/nlp/gpt_moe/distributed_gpt_moe.py index 9adf332c..31ca48f9 100644 --- a/modelscope/models/nlp/gpt_moe/distributed_gpt_moe.py +++ b/modelscope/models/nlp/gpt_moe/distributed_gpt_moe.py @@ -27,6 +27,7 @@ from transformers.modeling_utils import PreTrainedModel from modelscope.models import TorchModel from modelscope.models.nlp.gpt_moe import GPTMoEConfig +from modelscope.outputs import TextGenerationModelOutput, TokenGeneratorOutput from modelscope.utils.nlp.distributed import initialize_distributed from modelscope.utils.torch_utils import set_random_seed_mpu from .checkpointing import load_checkpoint @@ -42,7 +43,6 @@ class GPTMoEParallelMLP(nn.Module): moe=False, enable_expert_tensor_parallelism=False): super().__init__() - # Project to 4h. self.dense_h_to_4h = mpu.ColumnParallelLinearV3( config, @@ -606,6 +606,8 @@ class GPTMoEParallelTransformerLayer(nn.Module): # Layer norm post the self attention. layernorm_output = self.post_attention_layernorm(layernorm_input) + moe_loss = torch.tensor( + 0.0, device=layernorm_output.device, dtype=layernorm_output.dtype) mlp_bias = torch.tensor( 0.0, device=layernorm_output.device, dtype=layernorm_output.dtype) @@ -635,7 +637,7 @@ class GPTMoEParallelTransformerLayer(nn.Module): output = mpu.make_viewless_tensor( inp=output, requires_grad=output.requires_grad, keep_graph=True) - return output + return output, moe_loss class GPTMoEParallelTransformer(nn.Module): @@ -743,18 +745,19 @@ class GPTMoEParallelTransformer(nn.Module): with rng_context: # Forward pass. + moe_losses = [] for index in range(self.num_layers): layer = self._get_layer(index) - hidden_states = layer( + hidden_states, moe_loss = layer( hidden_states, attention_mask, inference_params=inference_params) + moe_losses.append(moe_loss) # Final layer norm. if self.post_process and self.post_layer_norm: hidden_states = self.final_layernorm(hidden_states) - - return hidden_states + return (hidden_states, *moe_losses) class GPTMoETransformerLanguageModel(nn.Module): @@ -805,7 +808,7 @@ class GPTMoETransformerLanguageModel(nn.Module): # Run encoder. if enc_hidden_states is None: if self.encoder is not None: - encoder_output = self.encoder( + encoder_output, *moe_losses = self.encoder( encoder_input, enc_attn_mask, inference_params=inference_params) @@ -814,7 +817,7 @@ class GPTMoETransformerLanguageModel(nn.Module): else: encoder_output = enc_hidden_states.to(encoder_input.dtype) - return encoder_output + return (encoder_output, *moe_losses) def load_state_dict(self, state_dict, strict=True): """Customized load.""" @@ -929,31 +932,53 @@ class GPTMoEModel(PreTrainedModel): return attention_mask, position_ids + @staticmethod + def post_language_model_processing(input_, labels, word_embeddings_weight, + sequence_parallel): + + # Output. Format [s b h] + # Parallel logits. + input_parallel = input_ + # Matrix multiply. + logits_parallel = mpu.LinearWithGradAccumulationAndAsyncCommunication.apply( + input_parallel, word_embeddings_weight, None, False, False, + sequence_parallel) + + output = logits_parallel + + if labels is None: + # [s b h] => [b s h] + return output.transpose(0, 1).contiguous() + else: + # [b s] => [s b] + labels = labels.transpose(0, 1).contiguous() + loss = mpu.vocab_parallel_cross_entropy(output.float(), labels) + # [s b] => [b, s] + loss = loss.transpose(0, 1).contiguous() + return loss + def forward(self, input_ids, attention_mask=None, position_ids=None, inference_params=None, + labels=None, **kwargs): if attention_mask is None and position_ids is None: attention_mask, position_ids = \ self.build_attention_mask_and_position_ids(input_ids) - lm_output = self.language_model( + lm_output, *moe_losses = self.language_model( input_ids, position_ids, attention_mask, inference_params=inference_params) - logits_parallel = mpu.LinearWithGradAccumulationAndAsyncCommunication.apply( - lm_output, self.word_embeddings_weight(), None, False, True, + lm_output = self.post_language_model_processing( + lm_output, labels, self.word_embeddings_weight(), self.config.sequence_parallel) - # Gather if needed. - output = logits_parallel - if not self.parallel_output: - output = mpu.gather_from_model_parallel_region(logits_parallel) - return output.transpose(0, 1).contiguous() + return (lm_output, *moe_losses) def load_state_dict(self, state_dict, strict=True): """Customized load.""" @@ -1126,28 +1151,63 @@ class DistributedGPTMoE(TorchModel): load_ds_ckpts=self.config.load_ds_ckpts) self.inference_params = None - def forward_step(self, tokens, attention_mask, position_ids): - logits = self.dist_model( + def train(self, mode: bool = True): + if mode: + self.inference_params = None + return super().train(mode) + + def forward(self, + tokens, + attention_mask=None, + position_ids=None, + labels=None, + prompt_length=None): + + outputs, *other_losses = self.dist_model( tokens, attention_mask, position_ids, - inference_params=self.inference_params) - self.inference_params.sequence_len_offset += tokens.size(1) - return logits + inference_params=self.inference_params, + labels=labels) + + if labels is None: + self.inference_params.sequence_len_offset += tokens.size(1) + return TextGenerationModelOutput(logits=outputs) + else: + + moe_losses = [] + for moe_loss in other_losses: + if moe_loss is not None: + moe_losses.append(moe_loss) + moe_loss = sum(moe_losses) * 0.01 + + loss_mask = torch.ones( + tokens.size(), dtype=torch.float, device=tokens.device) + + losses = outputs.float() + loss_mask = loss_mask.view(-1).float() + loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() + + loss = loss + moe_loss + + return TextGenerationModelOutput(loss=loss) def generate(self, tokens, temperature=1.0, use_eod_token_for_early_termination=True, stop_on_double_eol=False, - stop_on_eol=False): - lengths = torch.tensor([tokens.size(1)], device=tokens.device) + stop_on_eol=False, + **kwargs): + batch_size = tokens.size(0) + lengths = kwargs.pop( + 'prompt_length', + torch.tensor([tokens.size(1)], device=tokens.device)) pads = torch.ones( - 1, self.config.tokens_to_generate, + batch_size, self.config.tokens_to_generate, device=tokens.device).long() * self.config.eod_id tokens = torch.cat((tokens, pads), dim=-1) - batch_size = tokens.size(0) min_prompt_length = lengths.min().item() max_sequence_length = tokens.size(1) max_sequence_length = min(max_sequence_length, @@ -1176,6 +1236,7 @@ class DistributedGPTMoE(TorchModel): with torch.no_grad(): attention_mask, position_ids = \ GPTMoEModel.build_attention_mask_and_position_ids(tokens) + prev_context_length = 0 for context_length in range(min_prompt_length, max_sequence_length): @@ -1188,8 +1249,8 @@ class DistributedGPTMoE(TorchModel): ..., prev_context_length:context_length, :context_length] # logits will be meanigful only in the last pipeline stage. - logits = self.forward_step(tokens2use, attention_mask2use, - positions2use) + logits = self(tokens2use, attention_mask2use, + positions2use).logits # Sample. last_token_logits = logits[:, -1, :] @@ -1233,4 +1294,4 @@ class DistributedGPTMoE(TorchModel): break tokens = tokens[:, :(context_length + 1)] - return tokens + return TokenGeneratorOutput(sequences=tokens) diff --git a/modelscope/models/nlp/gpt_moe/moe/sharded_moe.py b/modelscope/models/nlp/gpt_moe/moe/sharded_moe.py index a7d73d5d..86c591c9 100644 --- a/modelscope/models/nlp/gpt_moe/moe/sharded_moe.py +++ b/modelscope/models/nlp/gpt_moe/moe/sharded_moe.py @@ -26,6 +26,15 @@ from torch.nn import Module from ..configuration import logger from .mappings import drop_tokens, gather_tokens +try: + # To enable Tutel MoE optimizations: + # python3 -m pip install --user --upgrade git+https://github.com/microsoft/tutel@v0.1.x + from tutel import moe as tutel_moe + TUTEL_INSTALLED = True +except ImportError: + # Fail silently so we don't spam logs unnecessarily if user isn't using tutel + TUTEL_INSTALLED = False + try: from apex.normalization import FusedLayerNorm as _FusedLayerNorm @@ -444,7 +453,6 @@ class MOELayer(Base): # Implement Algorithm 2 from GShard paper. d_model = input[0].shape[-1] - # Initial implementation -> Reshape into S tokens by dropping sequence dimension. # Reshape into G groups so that each group can distribute tokens equally # group_size = kwargs['group_size'] if 'group_size' in kwargs.keys() else 1 diff --git a/modelscope/models/nlp/gpt_moe/text_generation.py b/modelscope/models/nlp/gpt_moe/text_generation.py index 59245917..917ac5b2 100644 --- a/modelscope/models/nlp/gpt_moe/text_generation.py +++ b/modelscope/models/nlp/gpt_moe/text_generation.py @@ -1,9 +1,12 @@ # Copyright (c) Alibaba, Inc. and its affiliates. from typing import Dict +from transformers import BertTokenizer + from modelscope.metainfo import Models from modelscope.models.base import Tensor, TorchModel from modelscope.models.builder import MODELS +from modelscope.models.nlp.gpt_moe import GPTMoEModel from modelscope.utils.constant import Tasks __all__ = ['GPTMoEForTextGeneration'] @@ -20,12 +23,15 @@ class GPTMoEForTextGeneration(TorchModel): """ super().__init__(model_dir, *args, **kwargs) - from modelscope.models.nlp.gpt_moe import GPTMoEModel - from transformers import BertTokenizer - print('****') - print(model_dir) - self.model = GPTMoEModel.from_pretrained(model_dir) - self.tokenizer = BertTokenizer.from_pretrained(model_dir) + # Temporarily compatible with DistributedGPT3 and GPT3Model, + # the base/large model based on GPT3Model will be replaced in the future, + # and GPT3Model will be deprecated + if 'model_parallel_size' in kwargs: + from modelscope.models.nlp import DistributedGPTMoE + self.model = DistributedGPTMoE(model_dir, **kwargs) + else: + self.model = GPTMoEModel.from_pretrained(model_dir) + self.tokenizer = BertTokenizer.from_pretrained(model_dir) def forward(self, input: Dict[str, Tensor]) -> Dict[str, Tensor]: """return the result by the model @@ -43,6 +49,8 @@ class GPTMoEForTextGeneration(TorchModel): return self.model(**input) def generate(self, input: Dict[str, Tensor]) -> Dict[str, Tensor]: + if not isinstance(self.model, GPTMoEModel): + return self.model.generate(**input) assert 'input_ids' in input, "generate function must accept 'input_ids' key" input_ids = input['input_ids'] if 'attention_mask' in input: diff --git a/modelscope/pipelines/nlp/distributed_gpt_moe_pipeline.py b/modelscope/pipelines/nlp/distributed_gpt_moe_pipeline.py index 71e48a11..bf4672ac 100644 --- a/modelscope/pipelines/nlp/distributed_gpt_moe_pipeline.py +++ b/modelscope/pipelines/nlp/distributed_gpt_moe_pipeline.py @@ -50,5 +50,6 @@ class DistributedGPTMoEPipeline(DistributedPipeline): from modelscope.outputs import OutputKeys return { OutputKeys.TEXT: - self.preprocessor.tokenizer.detokenize(inputs[0].tolist()) + self.preprocessor.tokenizer.detokenize( + inputs.sequences[0].tolist()) } diff --git a/modelscope/trainers/nlp/gpt_moe_trainer.py b/modelscope/trainers/nlp/gpt_moe_trainer.py new file mode 100644 index 00000000..8d431881 --- /dev/null +++ b/modelscope/trainers/nlp/gpt_moe_trainer.py @@ -0,0 +1,61 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os +from collections.abc import Mapping +from typing import List + +import torch +from megatron import mpu + +from modelscope.metainfo import Trainers +from modelscope.models import TorchModel +from modelscope.trainers.builder import TRAINERS +from modelscope.trainers.nlp_trainer import NlpEpochBasedTrainer +from modelscope.utils.config import Config +from modelscope.utils.file_utils import func_receive_dict_inputs + + +@TRAINERS.register_module(module_name=Trainers.gpt_moe_trainer) +class GPTMoETrainer(NlpEpochBasedTrainer): + + def rebuild_config(self, cfg: Config): + super().rebuild_config(cfg) + cfg.model.rank = int(os.environ.get('LOCAL_RANK', -1)) + cfg.model.master_ip = os.environ.get('MASTER_ADDR', '127.0.0.1') + cfg.model.master_port = os.environ.get('MASTER_PORT', '29500') + return cfg + + def train_step(self, model: TorchModel, inputs: Mapping): + keys = list(inputs.keys()) + datatype = torch.int64 + inputs = mpu.broadcast_data(keys, inputs, datatype) + return super().train_step(model, inputs) + + def _decode(self, tokens): + tokenizer = self.eval_preprocessor.tokenizer + return tokenizer.detokenize(tokens.tolist()) + + def evaluation_step(self, data): + model = self.model.module if self._dist else self.model + 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) + + prompt_length: List[int] = data['prompt_length'] + result['preds'] = [ + self._decode(seq[skip_len:]) + for seq, skip_len in zip(result['sequences'], prompt_length) + ] + data['tgts'] = [ + self._decode(seq[skip_len - 1:]) + for seq, skip_len in zip(data['labels'], prompt_length) + ] + assert len(result['preds']) == len(data['tgts']) + + return result diff --git a/tests/trainers/test_finetune_gpt3.py b/tests/trainers/test_finetune_gpt3.py index e2110cfa..7a9e03d0 100644 --- a/tests/trainers/test_finetune_gpt3.py +++ b/tests/trainers/test_finetune_gpt3.py @@ -21,13 +21,14 @@ class TestFinetuneTextGeneration(unittest.TestCase): shutil.rmtree(self.tmp_dir) super().tearDown() - @unittest.skip + @unittest.skip( + 'skip since the test requires multiple GPU and takes a long time to run' + ) def test_finetune_poetry(self): dataset_dict = MsDataset.load('chinese-poetry-collection') - train_dataset = dataset_dict['train'].to_hf_dataset().rename_columns( - {'text1': 'src_txt'}) - eval_dataset = dataset_dict['test'].to_hf_dataset().rename_columns( + train_dataset = dataset_dict['train'].remap_columns( {'text1': 'src_txt'}) + eval_dataset = dataset_dict['test'].remap_columns({'text1': 'src_txt'}) max_epochs = 10 tmp_dir = './gpt3_poetry' @@ -66,17 +67,17 @@ class TestFinetuneTextGeneration(unittest.TestCase): name=Trainers.gpt3_trainer, default_args=kwargs) trainer.train() - @unittest.skip + @unittest.skip( + 'skip since the test requires multiple GPU and takes a long time to run' + ) def test_finetune_dureader(self): # DuReader_robust-QG is an example data set, # users can also use their own data set for training dataset_dict = MsDataset.load('DuReader_robust-QG') - train_dataset = dataset_dict['train'].to_hf_dataset() \ - .rename_columns({'text1': 'src_txt', 'text2': 'tgt_txt'}) \ + train_dataset = dataset_dict['train'].remap_columns({'text1': 'src_txt', 'text2': 'tgt_txt'}) \ .map(lambda example: {'src_txt': example['src_txt'].replace('[SEP]', '') + '\n'}) - eval_dataset = dataset_dict['validation'].to_hf_dataset() \ - .rename_columns({'text1': 'src_txt', 'text2': 'tgt_txt'}) \ + eval_dataset = dataset_dict['validation'].remap_columns({'text1': 'src_txt', 'text2': 'tgt_txt'}) \ .map(lambda example: {'src_txt': example['src_txt'].replace('[SEP]', '') + '\n'}) max_epochs = 10 diff --git a/tests/trainers/test_finetune_gpt_moe.py b/tests/trainers/test_finetune_gpt_moe.py new file mode 100644 index 00000000..3bbcae14 --- /dev/null +++ b/tests/trainers/test_finetune_gpt_moe.py @@ -0,0 +1,131 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +import shutil +import tempfile +import unittest + +from modelscope.metainfo import Trainers +from modelscope.msdatasets import MsDataset +from modelscope.trainers import build_trainer + + +class TestFinetuneTextGeneration(unittest.TestCase): + + test_model_id = 'PAI/nlp_gpt3_text-generation_0.35B_MoE-64' + + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(self.tmp_dir): + os.makedirs(self.tmp_dir) + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + super().tearDown() + + @unittest.skip( + 'skip since the test requires multiple GPU and takes a long time to run' + ) + def test_finetune_poetry(self): + dataset_dict = MsDataset.load('chinese-poetry-collection') + train_dataset = dataset_dict['train'].remap_columns( + {'text1': 'src_txt'}) + eval_dataset = dataset_dict['test'].remap_columns({'text1': 'src_txt'}) + max_epochs = 10 + tmp_dir = './gpt_moe_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': 1, + 'workers_per_gpu': 1 + } + return cfg + + kwargs = dict( + model=self.test_model_id, + 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.gpt_moe_trainer, default_args=kwargs) + trainer.train() + + @unittest.skip( + 'skip since the test requires multiple GPU and takes a long time to run' + ) + def test_finetune_dureader(self): + # DuReader_robust-QG is an example data set, + # users can also use their own data set for training + dataset_dict = MsDataset.load('DuReader_robust-QG') + + train_dataset = dataset_dict['train'].remap_columns({'text1': 'src_txt', 'text2': 'tgt_txt'}) \ + .map(lambda example: {'src_txt': example['src_txt'].replace('[SEP]', '') + '\n'}) + eval_dataset = dataset_dict['validation'].remap_columns({'text1': 'src_txt', 'text2': 'tgt_txt'}) \ + .map(lambda example: {'src_txt': example['src_txt'].replace('[SEP]', '') + '\n'}) + + max_epochs = 10 + tmp_dir = './gpt_moe_dureader' + + num_warmup_steps = 200 + + 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': 16, + 'workers_per_gpu': 1 + } + cfg.train.hooks.append({ + 'type': 'EvaluationHook', + 'by_epoch': True, + 'interval': 1 + }) + cfg.preprocessor.sequence_length = 512 + cfg.model.checkpoint_model_parallel_size = 1 + return cfg + + kwargs = dict( + model=self.test_model_id, + 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.gpt_moe_trainer, default_args=kwargs) + trainer.train() + + +if __name__ == '__main__': + unittest.main()