diff --git a/modelscope/models/nlp/gpt3/configuration.py b/modelscope/models/nlp/gpt3/configuration.py
index 9c98cae8..0a1d16d7 100644
--- a/modelscope/models/nlp/gpt3/configuration.py
+++ b/modelscope/models/nlp/gpt3/configuration.py
@@ -21,6 +21,83 @@ logger = logging.get_logger()
 
 
 class GPT3Config(PretrainedConfig):
+    r"""
+    Configuration classes for GPT-3 model.
+
+    Class attributes:
+
+    - **model_type** (`str`) -- An identifier for the model type, serialized into the JSON file, can be used to recreate
+      the correct object in [`~transformers.AutoConfig`].
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 25600):
+            Vocabulary size of the GPT model. Defines the number of different
+            tokens that can be represented by the `inputs_ids` passed when
+            calling [`GPT3Model`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the decoder layers and the pooler layer.
+        ffn_hidden_size (`int`, *optional*, defaults to None):
+            Dimensionality of the ffn layer, None defaults to four times the hidden_size.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the
+            Transformer decoder.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (often named feed-forward)
+            layer in the Transformer decoder.
+        hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the
+            decoder and pooler. If string, `"gelu"`, `"relu"`, `"silu"` and
+            `"gelu_new"` are supported.
+        hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the
+            embeddings, decoder, and pooler.
+        attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (`int`, *optional*, defaults to 512):
+            The maximum sequence length that this model might ever be used with.
+            Typically set this to something large just in case (e.g., 512 or
+            1024 or 2048).
+        type_vocab_size (`int`, *optional*, defaults to 2):
+            The vocabulary size of the `token_type_ids` passed when calling
+            [`GPT3Model`].
+        layernorm_epsilon (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        bias_gelu_fusion (`bool`, *optional*, defaults to True):
+            Whether to use gelu activation function when mixing bias.
+        fp32_residual_connection (`bool`, *optional*, defaults to False):
+            Whether to use fp32 for residual connection
+            between layers to improve accuracy.
+        sequence_parallel (`bool`, *optional*, defaults to False):
+            Whether to use sequence parallel during training.
+        bf16 (`bool`, *optional*, defaults to `False`):
+            Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training.
+            Requires Ampere or higher NVIDIA architecture or using CPU (no_cuda).
+            This is an experimental API and it may change.
+        fp16 (`bool`, *optional*, defaults to `False`):
+            Whether to use fp16 16-bit (mixed) precision training instead of 32-bit training.
+        apply_query_key_layer_scaling (`bool`, *optional*, defaults to `True`):
+            Whether to scale query and key layer parameters during training.
+        init_method_std (`float`, *optional*, defaults to `0.02`):
+            The standard deviation of the normal distribution for initialization process.
+        eod_id (`int`, *optional*, defaults to `1`):
+            The end of text label for tokenizer, also indicates the end of the generation.
+        tokens_to_generate (`int`, *optional*, defaults to 100):
+            Number of tokens to generate.
+        top_k (`int`, *optional*, defaults to 0):
+            Number of highest probability vocabulary tokens to keep for
+            top-k-filtering that will be used by default in
+            the `generate` method of the model.
+        top_p (`float`, *optional*, defaults to 0.9):
+            Value that will be used by default in the `generate` method of the model
+            for `top_p`. If set to float < 1,
+            only the most probable tokens with probabilities that add up to `top_p`
+            or higher are kept for generation.
+        temperature (`float`, *optional*, defaults to 1.0):
+            The value used to module the next token probabilities that will be used
+            by default in the `generate` method of the model. Must be strictly positive.
+    """
 
     model_type = 'gpt3'
 
@@ -53,10 +130,11 @@ class GPT3Config(PretrainedConfig):
             hidden_dropout=0.1,
             init_method_std=0.02,
             # generate
-            eod_id=7,
+            eod_id=1,
             tokens_to_generate=100,
             top_k=0,
             top_p=0.9,
+            temperature=1.0,
             **kwargs):
         super().__init__(layer_norm_eps=layernorm_epsilon, **kwargs)
 
@@ -95,6 +173,7 @@ class GPT3Config(PretrainedConfig):
         self.tokens_to_generate = tokens_to_generate
         self.top_k = top_k
         self.top_p = top_p
+        self.temperature = temperature
 
         TORCH_MAJOR = int(torch.__version__.split('.')[0])
         TORCH_MINOR = int(torch.__version__.split('.')[1])
diff --git a/modelscope/models/nlp/gpt3/distributed_gpt3.py b/modelscope/models/nlp/gpt3/distributed_gpt3.py
index e469f866..1c4505a0 100644
--- a/modelscope/models/nlp/gpt3/distributed_gpt3.py
+++ b/modelscope/models/nlp/gpt3/distributed_gpt3.py
@@ -13,14 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import io
 import math
 import os
-from os import path as osp
+from collections import OrderedDict
 from typing import Callable, Dict, List, Optional, Union
 
 import torch
-from megatron_util import mpu
+from megatron_util import get_args, mpu
 from megatron_util.global_vars import get_global_memory_buffer
 from megatron_util.model import (AttnMaskType, Float16Module, LayerNorm,
                                  bias_gelu_impl)
@@ -29,11 +28,9 @@ from torch import nn
 from torch.nn import functional as F
 from transformers.modeling_utils import PreTrainedModel
 
-from modelscope.fileio import File
 from modelscope.models import TorchModel
 from modelscope.models.nlp.gpt3 import GPT3Config
 from modelscope.outputs import TextGenerationModelOutput, TokenGeneratorOutput
-from modelscope.utils.checkpoint import weights_to_cpu
 from modelscope.utils.megatron_utils import init_megatron_util
 from modelscope.utils.nlp.load_checkpoint import pre_load
 
@@ -948,21 +945,6 @@ def split_state_dict(state_dict: Dict[str, torch.Tensor], model: GPT3Model,
     return state_dict
 
 
-def save_checkpoint(model: torch.nn.Module, filename: str, **kwargs) -> None:
-    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
-        model = model.module
-
-    checkpoint = {'module': weights_to_cpu(model.state_dict())}
-    mp_rank = mpu.get_tensor_model_parallel_rank()
-    filename = osp.join(
-        osp.dirname(filename), 'model',
-        'mp_rank_{:02d}'.format(mp_rank) + '_model_states.pt')
-
-    with io.BytesIO() as f:
-        torch.save(checkpoint, f)
-        File.write(f.getvalue(), filename)
-
-
 class DistributedGPT3(TorchModel):
 
     def __init__(self,
@@ -992,7 +974,8 @@ class DistributedGPT3(TorchModel):
         self.dist_model = model
 
         tensor_ws = mpu.get_tensor_model_parallel_world_size()
-        ckpt_ws = kwargs.pop('checkpoint_model_parallel_size', tensor_ws)
+        ckpt_ws = get_args().get('checkpoint_tensor_model_parallel_size',
+                                 tensor_ws)
         ckpt_rank = mpu.get_tensor_model_parallel_rank() * ckpt_ws // tensor_ws
         load_model = pre_load(ckpt_rank, model_dir, tag=path_load_tag)
         load_model = split_state_dict(load_model, model, tensor_ws // ckpt_ws)
@@ -1011,8 +994,8 @@ class DistributedGPT3(TorchModel):
                 attention_mask=None,
                 position_ids=None,
                 labels=None,
-                prompt_length=None,
-                is_pair=(False, )):
+                prompts_len=None,
+                inputs_len=None):
 
         logits, losses = self.dist_model(
             tokens,
@@ -1026,10 +1009,15 @@ class DistributedGPT3(TorchModel):
             self.inference_params.sequence_len_offset += tokens.size(1)
         else:
             loss_mask = torch.ones(
-                tokens.size(), dtype=torch.float, device=tokens.device)
-            if is_pair[0]:
-                for i, length in enumerate(prompt_length):
-                    loss_mask[i, :length] = 0
+                labels.size(), dtype=torch.float, device=tokens.device)
+            if inputs_len is None:
+                for i, l in enumerate(prompts_len):
+                    loss_mask[i, l:] = 0
+            else:
+                for i, l in enumerate(inputs_len):
+                    loss_mask[i, l - 1:] = 0
+                for i, l in enumerate(prompts_len):
+                    loss_mask[i, :l - 1] = 0
 
             losses = losses.float()
             loss_mask = loss_mask.view(-1).float()
@@ -1039,15 +1027,15 @@ class DistributedGPT3(TorchModel):
 
     def sample(self,
                tokens,
-               temperature=1.0,
+               prompts_len=None,
                use_eod_token_for_early_termination=True,
                stop_on_double_eol=False,
                stop_on_eol=False,
                **kwargs):
         batch_size = tokens.size(0)
-        lengths = kwargs.pop(
-            'prompt_length',
-            torch.tensor([tokens.size(1)], device=tokens.device))
+        lengths = prompts_len
+        if lengths is None:
+            lengths = torch.tensor([tokens.size(1)], device=tokens.device)
         pads = torch.ones(
             batch_size, self.config.tokens_to_generate,
             device=tokens.device).long() * self.config.eod_id
@@ -1096,9 +1084,9 @@ class DistributedGPT3(TorchModel):
             last_token_logits = logits[:, -1, :]
             new_sample = sample(
                 last_token_logits,
-                top_k=self.config.top_k,
-                top_p=self.config.top_p,
-                temperature=temperature,
+                top_k=kwargs.pop('top_k', self.config.top_k),
+                top_p=kwargs.pop('top_p', self.config.top_p),
+                temperature=kwargs.pop('temperature', self.config.temperature),
                 vocab_size=self.config.vocab_size)
 
             # If a prompt length is smaller or equal th current context
@@ -1257,6 +1245,11 @@ class DistributedGPT3(TorchModel):
     def state_dict(self, destination=None, prefix='', keep_vars=False):
         return self.dist_model.state_dict(destination, prefix, keep_vars)
 
+    def load_state_dict(self,
+                        state_dict: 'OrderedDict[str, torch.Tensor]',
+                        strict: bool = True):
+        return self.dist_model.load_state_dict(state_dict, strict)
+
     def save_pretrained(self,
                         target_folder: Union[str, os.PathLike],
                         save_checkpoint_names: Union[str, List[str]] = None,
@@ -1265,13 +1258,15 @@ class DistributedGPT3(TorchModel):
                         **kwargs):
         # DistributedPipeline type is different from task name
         config['pipeline']['type'] = 'gpt3-generation'
-        # a temp fix for master_ip, master_port and rank
-        # can be removed after refactoring megatron_util
-        for unused_key in ('master_ip', 'master_port', 'rank'):
-            config['model'].pop(unused_key, None)
+
+        config['model'].pop('rank', None)
+        config['megatron'].pop('checkpoint_tensor_model_parallel_size', None)
+        tp_size = get_args().tensor_model_parallel_size
+        pp_size = get_args().pipeline_model_parallel_size
+        config['megatron']['world_size'] = tp_size * pp_size
 
         return super().save_pretrained(target_folder, save_checkpoint_names,
-                                       save_checkpoint, config, **kwargs)
+                                       save_function, config, **kwargs)
 
 
 class BeamHypotheses:
diff --git a/modelscope/models/nlp/gpt3/text_generation.py b/modelscope/models/nlp/gpt3/text_generation.py
index b4a1182c..368cd2b5 100644
--- a/modelscope/models/nlp/gpt3/text_generation.py
+++ b/modelscope/models/nlp/gpt3/text_generation.py
@@ -1,4 +1,5 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from collections import OrderedDict
 from typing import Dict
 
 import torch
@@ -9,6 +10,7 @@ from modelscope.models.base import Tensor, TorchModel
 from modelscope.models.builder import MODELS
 from modelscope.models.nlp.gpt3 import GPT3Model
 from modelscope.utils.constant import Tasks
+from modelscope.utils.hub import read_config
 
 __all__ = ['GPT3ForTextGeneration']
 
@@ -27,7 +29,7 @@ class GPT3ForTextGeneration(TorchModel):
         # 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 'world_size' in kwargs:
+        if 'megatron' in read_config(model_dir):
             from modelscope.models.nlp import DistributedGPT3
             self.model = DistributedGPT3(model_dir, **kwargs)
         else:
@@ -66,3 +68,11 @@ class GPT3ForTextGeneration(TorchModel):
         if not isinstance(self.model, GPT3Model):
             return self.model.save_pretrained(*args, **kwargs)
         return super().save_pretrained(*args, **kwargs)
+
+    def state_dict(self, destination=None, prefix='', keep_vars=False):
+        return self.model.state_dict(destination, prefix, keep_vars)
+
+    def load_state_dict(self,
+                        state_dict: 'OrderedDict[str, Tensor]',
+                        strict: bool = True):
+        return self.model.load_state_dict(state_dict, strict)
diff --git a/modelscope/models/nlp/gpt3/tokenizer.py b/modelscope/models/nlp/gpt3/tokenizer.py
index ba29891e..eed3f80f 100644
--- a/modelscope/models/nlp/gpt3/tokenizer.py
+++ b/modelscope/models/nlp/gpt3/tokenizer.py
@@ -32,7 +32,7 @@ class JiebaBPETokenizer:
         except ImportError:
             raise ImportError(
                 'You need to install jieba to use JiebaTokenizer. '
-                'See https://pypi.org/project/rjieba/ for installation.')
+                'See https://pypi.org/project/jieba/ for installation.')
         self.jieba = jieba
         self.new_line = self.vocab['\n']
         self.sep_token = self.vocab['<sep>']
@@ -64,7 +64,9 @@ class JiebaBPETokenizer:
             return self.tokenizer.encode(
                 text, is_pretokenized=False, add_special_tokens=True).ids
 
-    def detokenize(self, token_ids):
+    def detokenize(self, token_ids: List[int], early_stop: bool = True) -> str:
+        if early_stop and self.sep_token in token_ids:
+            token_ids = token_ids[:token_ids.index(self.sep_token)]
         text = self.tokenizer.decode(token_ids, skip_special_tokens=True)
         return text
 
diff --git a/modelscope/pipelines/base.py b/modelscope/pipelines/base.py
index 374b05d0..a3c15695 100644
--- a/modelscope/pipelines/base.py
+++ b/modelscope/pipelines/base.py
@@ -411,7 +411,7 @@ class DistributedPipeline(Pipeline):
         else:
             self.model_dir = snapshot_download(model)
         self.cfg = read_config(self.model_dir)
-        self.world_size = self.cfg.model.world_size
+        self.world_size = self._get_world_size(self.cfg)
         self.model_pool = None
         self.device_name = 'cpu'
         self.device = create_device(self.device_name)
@@ -423,17 +423,17 @@ class DistributedPipeline(Pipeline):
         self.model_pool = Pool(self.world_size)
         master_ip = '127.0.0.1' if 'master_ip' not in kwargs else kwargs[
             'master_ip']
+        os.environ['MASTER_ADDR'] = master_ip
         master_port = '29500' if 'master_port' not in kwargs else kwargs[
             'master_port']
         from modelscope.utils.torch_utils import _find_free_port, _is_free_port
         if not _is_free_port(int(master_port)):
             master_port = str(_find_free_port())
+        os.environ['MASTER_PORT'] = master_port
         self.model_pool.map(
             partial(
                 self.__class__._instantiate_one,
                 model_dir=self.model_dir,
-                master_ip=master_ip,
-                master_port=master_port,
                 **self.cfg.model,
                 **kwargs), ranks)
         self.models = []
@@ -487,6 +487,12 @@ class DistributedPipeline(Pipeline):
         """
         pass
 
+    def _get_world_size(self, cfg: Config) -> int:
+        m_world_size = cfg.safe_get('megatron.world_size')
+        if m_world_size is None:
+            return cfg.safe_get('model.world_size')
+        return m_world_size
+
 
 def collate_fn(data, device):
     """Prepare the input just before the forward function.
diff --git a/modelscope/pipelines/nlp/mglm_text_summarization_pipeline.py b/modelscope/pipelines/nlp/mglm_text_summarization_pipeline.py
index 330331a5..de55ce3d 100644
--- a/modelscope/pipelines/nlp/mglm_text_summarization_pipeline.py
+++ b/modelscope/pipelines/nlp/mglm_text_summarization_pipeline.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2022 Zhipu.AI
 
+import os
 from typing import Any, Dict, Optional, Union
 
 from modelscope.metainfo import Pipelines
@@ -30,6 +31,8 @@ class MGLMTextSummarizationPipeline(Pipeline):
         self.model.eval()
         if preprocessor is None:
             preprocessor = MGLMSummarizationPreprocessor()
+        from modelscope.utils.torch_utils import _find_free_port
+        os.environ['MASTER_PORT'] = str(_find_free_port())
         super().__init__(model=model, preprocessor=preprocessor, **kwargs)
 
     # define the forward pass
diff --git a/modelscope/preprocessors/nlp/text_generation_preprocessor.py b/modelscope/preprocessors/nlp/text_generation_preprocessor.py
index 76e1a4f9..c3b91485 100644
--- a/modelscope/preprocessors/nlp/text_generation_preprocessor.py
+++ b/modelscope/preprocessors/nlp/text_generation_preprocessor.py
@@ -2,7 +2,7 @@
 
 import os
 import os.path as osp
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 import torch
@@ -223,14 +223,6 @@ class TextGenerationJiebaPreprocessor(TextGenerationPreprocessorBase):
         """
         return self.tokenizer.detokenize(tokens)
 
-    def _truncate(self, array: np.ndarray) -> np.ndarray:
-        if len(array) < self.max_length:
-            return np.pad(
-                array, (0, self.max_length - len(array)),
-                constant_values=self.tokenizer.eod)
-        else:
-            return array[:self.max_length]
-
     def _tokenize_text(self, sequence1, sequence2=None, **kwargs):
         """Tokenize the text.
 
@@ -246,18 +238,46 @@ class TextGenerationJiebaPreprocessor(TextGenerationPreprocessorBase):
                 'input_ids':
                 torch.tensor(self.tokenizer.tokenize(sequence1)).unsqueeze_(0)
             }
+        # continue write train: | inputs | <sep> |
+        # input & output train: | inputs | outputs | <sep> |
         else:
-            tokens = self.tokenizer.tokenize(sequence1)
-            prompt_length = min(len(tokens), self.max_length - 1)
-            if sequence2 is not None:
-                tokens += self.tokenizer.tokenize(sequence2)
-            tokens = self._truncate(np.array(tokens))
-            return {
-                'tokens': tokens[:-1],
-                'labels': tokens[1:],
-                'prompt_length': prompt_length,
-                'is_pair': int(sequence2 is not None),
-            }
+            input_tokens = self.tokenizer.tokenize(sequence1)
+            if sequence2 is None:
+                return self._only_input(input_tokens)
+            else:
+                return self._input_and_output(
+                    input_tokens, self.tokenizer.tokenize(sequence2))
+
+    def _only_input(self, input_tokens: List[int]) -> Dict[str, Any]:
+        prompts_len = len(input_tokens)
+        input_tokens.append(self.tokenizer.sep_token)
+        tokens = self._truncate(np.asarray(input_tokens))
+        return {
+            'tokens': tokens[:-1],
+            'labels': tokens[1:],
+            'prompts_len': min(prompts_len, self.max_length),
+        }
+
+    def _input_and_output(self, input_tokens: List[int],
+                          output_tokens: List[int]) -> Dict[str, Any]:
+        tokens = input_tokens[:]
+        tokens.extend(output_tokens)
+        tokens.append(self.tokenizer.sep_token)
+        inputs_len = len(tokens)
+        tokens = self._truncate(np.asarray(tokens))
+        return {
+            'tokens': tokens[:-1],
+            'labels': tokens[1:],
+            'prompts_len': min(len(input_tokens), self.max_length),
+            'inputs_len': min(inputs_len, self.max_length),
+        }
+
+    def _truncate(self, array: np.ndarray) -> np.ndarray:
+        if len(array) < self.max_length:
+            return np.pad(
+                array, (0, self.max_length - len(array)), constant_values=0)
+        else:
+            return array[:self.max_length]
 
 
 @PREPROCESSORS.register_module(
diff --git a/modelscope/trainers/hooks/checkpoint_hook.py b/modelscope/trainers/hooks/checkpoint_hook.py
index eefcbfa3..57698f3a 100644
--- a/modelscope/trainers/hooks/checkpoint_hook.py
+++ b/modelscope/trainers/hooks/checkpoint_hook.py
@@ -14,6 +14,7 @@ from modelscope.utils.checkpoint import (load_checkpoint, save_checkpoint,
                                          save_configuration)
 from modelscope.utils.constant import LogKeys, ModelFile
 from modelscope.utils.logger import get_logger
+from modelscope.utils.megatron_utils import is_megatron_initialized
 from modelscope.utils.torch_utils import is_master
 from .builder import HOOKS
 from .hook import Hook
diff --git a/modelscope/trainers/hooks/logger/tensorboard_hook.py b/modelscope/trainers/hooks/logger/tensorboard_hook.py
index d7ce0dac..efc4d833 100644
--- a/modelscope/trainers/hooks/logger/tensorboard_hook.py
+++ b/modelscope/trainers/hooks/logger/tensorboard_hook.py
@@ -41,7 +41,7 @@ class TensorboardHook(LoggerHook):
         self.out_dir = out_dir
         self.skip_keys = skip_keys
 
-    @master_only
+    @master_only()
     def before_run(self, trainer):
         super(TensorboardHook, self).before_run(trainer)
         try:
@@ -59,7 +59,7 @@ class TensorboardHook(LoggerHook):
             f'tensorboard files will be saved to {self.out_dir}')
         self.writer = SummaryWriter(self.out_dir)
 
-    @master_only
+    @master_only()
     def log(self, trainer):
         if len(trainer.visualization_buffer.output) > 0:
             self.visualization_log(trainer)
@@ -112,6 +112,6 @@ class TensorboardHook(LoggerHook):
         # avoiding repeated writing of the same image buffer every self.interval
         trainer.visualization_buffer.clear_output()
 
-    @master_only
+    @master_only()
     def after_run(self, trainer):
         self.writer.close()
diff --git a/modelscope/trainers/nlp/gpt3_trainer.py b/modelscope/trainers/nlp/gpt3_trainer.py
index afda6424..22f244eb 100644
--- a/modelscope/trainers/nlp/gpt3_trainer.py
+++ b/modelscope/trainers/nlp/gpt3_trainer.py
@@ -1,14 +1,9 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
 import os
-from collections.abc import Mapping
 from typing import Any, Dict, List
 
-import torch
-from megatron_util import mpu
-
 from modelscope.metainfo import Trainers
-from modelscope.models import TorchModel
 from modelscope.models.nlp import GPT3ForTextGeneration
 from modelscope.trainers.builder import TRAINERS
 from modelscope.trainers.nlp_trainer import NlpEpochBasedTrainer
@@ -23,12 +18,6 @@ class GPT3Trainer(NlpEpochBasedTrainer):
         cfg.model.rank = int(os.environ.get('RANK', 0))
         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())
@@ -37,28 +26,25 @@ class GPT3Trainer(NlpEpochBasedTrainer):
         model = self.model.module if self._dist else self.model
         model.eval()
 
-        if self._is_pair(data):
+        if 'inputs_len' in data:
             return self._generate_eval(model, data)
         else:
             return self._forward_eval(model, data)
 
-    @staticmethod
-    def _is_pair(data: Dict[str, Any]) -> bool:
-        return 'is_pair' in data and bool(data['is_pair'][0])
-
     def _generate_eval(self, model: GPT3ForTextGeneration,
                        data: Dict[str, Any]) -> Dict[str, Any]:
-        data['do_sample'] = False
+        # Force greedy decoding in non-open tasks
+        data.update(top_k=1, top_p=0.)
         result = model.generate(data)
 
-        prompt_length: List[int] = data['prompt_length']
+        prompts_len: List[int] = data['prompts_len']
         result['preds'] = [
             self._decode(seq[skip_len:])
-            for seq, skip_len in zip(result['sequences'], prompt_length)
+            for seq, skip_len in zip(result['sequences'], prompts_len)
         ]
         data['tgts'] = [
             self._decode(seq[skip_len - 1:])
-            for seq, skip_len in zip(data['labels'], prompt_length)
+            for seq, skip_len in zip(data['labels'], prompts_len)
         ]
         return result
 
diff --git a/modelscope/trainers/trainer.py b/modelscope/trainers/trainer.py
index 46138ec7..843b1c2f 100644
--- a/modelscope/trainers/trainer.py
+++ b/modelscope/trainers/trainer.py
@@ -38,6 +38,7 @@ from modelscope.utils.data_utils import to_device
 from modelscope.utils.device import create_device
 from modelscope.utils.file_utils import func_receive_dict_inputs
 from modelscope.utils.logger import get_logger
+from modelscope.utils.megatron_utils import is_megatron_initialized
 from modelscope.utils.registry import build_from_cfg
 from modelscope.utils.torch_utils import (get_dist_info, get_local_rank,
                                           init_dist, is_dist, is_master,
@@ -683,6 +684,13 @@ class EpochBasedTrainer(BaseTrainer):
             find_unused_parameters=True,
             device_ids=[torch.cuda.current_device()])
 
+        if is_megatron_initialized():
+            from megatron_util import mpu
+            dp_cfg.update({
+                'output_device': torch.cuda.current_device(),
+                'process_group': mpu.get_data_parallel_group()
+            })
+
         return build_parallel(dp_cfg)
 
     def unwrap_module(self, model) -> Union[nn.Module, TorchModel]:
diff --git a/modelscope/trainers/utils/inference.py b/modelscope/trainers/utils/inference.py
index 90201d72..b0dbe4bf 100644
--- a/modelscope/trainers/utils/inference.py
+++ b/modelscope/trainers/utils/inference.py
@@ -10,7 +10,9 @@ import torch
 from torch import distributed as dist
 from tqdm import tqdm
 
+from modelscope.utils.constant import DistributedParallelType
 from modelscope.utils.data_utils import to_device
+from modelscope.utils.megatron_utils import is_megatron_initialized
 from modelscope.utils.torch_utils import (broadcast, get_dist_info, is_master,
                                           make_tmp_dir)
 
@@ -187,9 +189,10 @@ def evaluate_batch(trainer, data, metric_classes, vis_closure):
 
 
 def get_metric_values(metric_classes):
-    rank, world_size = get_dist_info()
+    _, world_size = get_dist_info()
     metric_values = {}
-    if rank == 0:
+    if is_master(
+            DistributedParallelType.DP if is_megatron_initialized() else None):
         for metric_cls in metric_classes:
             metric_values.update(metric_cls.evaluate())
     if world_size > 1:
@@ -218,16 +221,17 @@ def collect_results_cpu(result_part, tmpdir=None):
     rank, world_size = get_dist_info()
     if tmpdir is None:
         tmpdir = make_tmp_dir()
-    if not os.path.exists(tmpdir) and is_master():
+    if not os.path.exists(tmpdir) and is_master(DistributedParallelType.TP):
         os.makedirs(tmpdir)
     dist.barrier()
 
     # dump the part result to the dir
-    with open(os.path.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f:
-        pickle.dump(result_part, f)
+    if is_master(DistributedParallelType.TP):
+        with open(os.path.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f:
+            pickle.dump(result_part, f)
     dist.barrier()
     # collect all parts
-    if rank != 0:
+    if not is_master():
         return None
     else:
         # load results of all parts from tmp dir
@@ -261,14 +265,19 @@ def collect_results_gpu(result_part):
     Returns:
         list: The collected results.
     """
-    rank, world_size = get_dist_info()
+    _, world_size = get_dist_info()
+    group = None
+    if is_megatron_initialized():
+        from megatron_util import mpu
+        group = mpu.get_data_parallel_group()
+
     # dump result part to tensor with pickle
     part_tensor = torch.tensor(
         bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
     # gather all result part tensor shape
     shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
     shape_list = [shape_tensor.clone() for _ in range(world_size)]
-    dist.all_gather(shape_list, shape_tensor)
+    dist.all_gather(shape_list, shape_tensor, group)
     # padding result part tensor to max length
     shape_max = torch.tensor(shape_list).max()
     part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
@@ -277,9 +286,9 @@ def collect_results_gpu(result_part):
         part_tensor.new_zeros(shape_max) for _ in range(world_size)
     ]
     # gather all result part
-    dist.all_gather(part_recv_list, part_send)
+    dist.all_gather(part_recv_list, part_send, group)
 
-    if rank == 0:
+    if is_master():
         part_list = []
         for recv, shape in zip(part_recv_list, shape_list):
             part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes())
diff --git a/modelscope/utils/checkpoint.py b/modelscope/utils/checkpoint.py
index f9a36b5f..ef1fa003 100644
--- a/modelscope/utils/checkpoint.py
+++ b/modelscope/utils/checkpoint.py
@@ -20,6 +20,7 @@ from modelscope.fileio import File, LocalStorage
 from modelscope.utils.config import Config, JSONIteratorEncoder
 from modelscope.utils.constant import ConfigFields, ModelFile
 from modelscope.utils.logger import get_logger
+from modelscope.utils.torch_utils import is_master
 
 logger = get_logger()
 
@@ -585,7 +586,8 @@ def save_pretrained(model,
     ignore_file_set = set(origin_file_to_be_ignored)
     ignore_file_set.add(ModelFile.CONFIGURATION)
     ignore_file_set.add('.*')
-    if hasattr(model, 'model_dir') and model.model_dir is not None:
+    if hasattr(model,
+               'model_dir') and model.model_dir is not None and is_master():
         copytree(
             model.model_dir,
             target_folder,
diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py
index 4f340015..667575ff 100644
--- a/modelscope/utils/constant.py
+++ b/modelscope/utils/constant.py
@@ -509,3 +509,10 @@ class MetaDataFields:
 
 
 DatasetVisibilityMap = {1: 'private', 3: 'internal', 5: 'public'}
+
+
+class DistributedParallelType(object):
+    """Parallel Strategies for Distributed Models"""
+    DP = 'data_parallel'
+    TP = 'tensor_model_parallel'
+    PP = 'pipeline_model_parallel'
diff --git a/modelscope/utils/megatron_utils.py b/modelscope/utils/megatron_utils.py
index 240572d1..9f2b2c09 100644
--- a/modelscope/utils/megatron_utils.py
+++ b/modelscope/utils/megatron_utils.py
@@ -1,12 +1,5 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
-from typing import Optional
-
-from megatron_util import initialize_megatron
-
-from modelscope.utils.config import Config
-from modelscope.utils.hub import read_config
-
 _DEFAULT_CFG_WITH_MODEL_TYPE = {
     'gpt-moe': {
         'version': 'moe',
@@ -24,10 +17,13 @@ _DEFAULT_CFG_WITH_MODEL_TYPE = {
     },
 }
 
+_IS_MEGATRON_INITIALIZED = False
+
+
+def init_megatron_util(cfg=None, model_dir=None, **kwargs):
+    from modelscope.utils.hub import read_config
+    from megatron_util import initialize_megatron
 
-def init_megatron_util(cfg: Optional[Config] = None,
-                       model_dir: Optional[str] = None,
-                       **kwargs):
     assert not (cfg is None and model_dir is None), \
         'cfg and model_dir cannot both be None when initializing megatron_util'
     if cfg is None:
@@ -44,3 +40,9 @@ def init_megatron_util(cfg: Optional[Config] = None,
             if model_type in _DEFAULT_CFG_WITH_MODEL_TYPE else {}
     megatron_cfg.update(kwargs)
     initialize_megatron(megatron_cfg)
+    global _IS_MEGATRON_INITIALIZED
+    _IS_MEGATRON_INITIALIZED = True
+
+
+def is_megatron_initialized() -> bool:
+    return _IS_MEGATRON_INITIALIZED
diff --git a/modelscope/utils/nlp/load_checkpoint.py b/modelscope/utils/nlp/load_checkpoint.py
index 920097a0..c32f4d8b 100755
--- a/modelscope/utils/nlp/load_checkpoint.py
+++ b/modelscope/utils/nlp/load_checkpoint.py
@@ -66,7 +66,7 @@ def pre_load(mp_rank, load_dir, tag=''):
     load_path = _get_ckpt_name(mp_rank, load_dir, tag)
     checkpoint = torch.load(
         load_path, map_location=lambda storage, loc: storage)
-    return checkpoint['module']
+    return checkpoint['module'] if 'module' in checkpoint else checkpoint
 
 
 def _load_checkpoint(model,
diff --git a/modelscope/utils/torch_utils.py b/modelscope/utils/torch_utils.py
index 7e0a9129..3d315716 100644
--- a/modelscope/utils/torch_utils.py
+++ b/modelscope/utils/torch_utils.py
@@ -14,6 +14,9 @@ import torch
 import torch.multiprocessing as mp
 from torch import distributed as dist
 
+from modelscope.utils.constant import DistributedParallelType
+from modelscope.utils.megatron_utils import is_megatron_initialized
+
 
 def _find_free_port() -> str:
     # Copied from https://github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py # noqa: E501
@@ -107,14 +110,12 @@ def _init_dist_slurm(backend: str, port: Optional[int] = None) -> None:
 
 def get_dist_info() -> Tuple[int, int]:
     if is_dist():
-        try:
+        group = None
+        if is_megatron_initialized():
             from megatron_util import mpu
-            assert mpu.model_parallel_is_initialized()
-            rank = mpu.get_data_parallel_rank()
-            world_size = mpu.get_data_parallel_world_size()
-        except (ImportError, AssertionError):
-            rank = dist.get_rank()
-            world_size = dist.get_world_size()
+            group = mpu.get_data_parallel_group()
+        rank = dist.get_rank(group)
+        world_size = dist.get_world_size(group)
     else:
         rank = 0
         world_size = 1
@@ -160,18 +161,37 @@ def is_dist():
     return dist.is_available() and dist.is_initialized()
 
 
-def is_master():
-    return dist.get_rank() == 0 if is_dist() else True
+def is_master(group=None):
+    if isinstance(group, str):
+        group = _parse_parallel_group(group)
+    return dist.get_rank(group) == 0 if is_dist() else True
 
 
-def master_only(func: Callable) -> Callable:
+def _parse_parallel_group(group: str):
+    from megatron_util import mpu
+    if group == DistributedParallelType.DP:
+        return mpu.get_data_parallel_group()
+    if group == DistributedParallelType.TP:
+        return mpu.get_tensor_model_parallel_group()
+    if group == DistributedParallelType.PP:
+        return mpu.get_pipeline_model_parallel_group()
+    raise ValueError(
+        f"Wrong group '{group}'. Supported groups are '{DistributedParallelType.DP}', "
+        f"'{DistributedParallelType.TP}' or '{DistributedParallelType.PP}'")
 
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        if is_master():
-            return func(*args, **kwargs)
 
-    return wrapper
+def master_only(group=None):
+
+    def decorate(func: Callable) -> Callable:
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if is_master(group):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorate
 
 
 def make_tmp_dir():
diff --git a/tests/trainers/test_finetune_gpt3.py b/tests/trainers/test_finetune_gpt3.py
index 563d271c..546a515d 100644
--- a/tests/trainers/test_finetune_gpt3.py
+++ b/tests/trainers/test_finetune_gpt3.py
@@ -7,9 +7,10 @@ import unittest
 from modelscope.metainfo import Trainers
 from modelscope.msdatasets import MsDataset
 from modelscope.trainers import build_trainer
+from modelscope.utils.test_utils import DistributedTestCase, test_level
 
 
-class TestFinetuneTextGeneration(unittest.TestCase):
+class TestFinetuneTextGeneration(DistributedTestCase):
 
     def setUp(self):
         print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
@@ -21,62 +22,6 @@ class TestFinetuneTextGeneration(unittest.TestCase):
         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 = './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': 16,
-                'workers_per_gpu': 1
-            }
-            cfg.train.hooks.append({
-                'type': 'EvaluationHook',
-                'by_epoch': True,
-                'interval': 1
-            })
-            cfg.evaluation.dataloader = {
-                'batch_size_per_gpu': 8,
-                'workers_per_gpu': 1
-            }
-            cfg.evaluation.metrics = 'ppl'
-            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()
-
     @unittest.skip(
         'skip since the test requires multiple GPU and takes a long time to run'
     )
@@ -137,6 +82,69 @@ class TestFinetuneTextGeneration(unittest.TestCase):
             name=Trainers.gpt3_trainer, default_args=kwargs)
         trainer.train()
 
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_single_finetune_portry(self):
+        finetune_poetry()
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_multi_finetune_portry(self):
+        self.start(
+            finetune_poetry, num_gpus=4, work_dir=self.tmp_dir, dp_tp=True)
+
+    # TODO: add gpt3 trainer predict unittest
+
+
+def finetune_poetry(dp_tp=False):
+    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 = 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': 'EvaluationHook',
+            'by_epoch': True,
+            'interval': 1
+        })
+        cfg.evaluation.dataloader = {
+            'batch_size_per_gpu': 8,
+            'workers_per_gpu': 1
+        }
+        cfg.evaluation.metrics = 'ppl'
+        cfg.train.train_iters_per_epoch = 10
+        if dp_tp:
+            cfg.megatron = {'world_size': 4, 'tensor_model_parallel_size': 2}
+        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()