add gpt-moe model for modelscope finetune

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11085918

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'

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',

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

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)

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

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:

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

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

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]', '<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]', '<sep>') + '\n'})
 
         max_epochs = 10

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]', '<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]', '<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()