code: def __init__(self, num_layers, mp_size, args_others, topo, **kwargs):
      from megatron.initialize import initialize_megatron
  
      args_defaults = {
          "vocab_file": get_test_path("gpt2-vocab.json"),
          "merge_file": get_test_path("gpt2-merges.txt"),
          "tokenizer_type": "GPT2BPETokenizer",
      }
  
      args_defaults.update(args_others)
  
      # setting "make-vocab-size-divisible-by" to avoid word-embedding size change in resizing testing.
      sys.argv.extend(
          [
              "--model-parallel-size",
              str(mp_size),
              "--make-vocab-size-divisible-by",
              str(1),
          ]
      )
  
      initialize_megatron(
          args_defaults=args_defaults, ignore_unknown_args=True
      )
  
      from megatron.model.transformer import ParallelTransformerLayer
  
      class ParallelTransformerLayerPipe(ParallelTransformerLayer):
          def forward(self, args):
              # hardcode attn mask for testing, PP requires the attn_mask to be stashed
              attention_mask = torch.tensor(
                  [[True]], device=torch.cuda.current_device()
              )
              return super().forward(args, attention_mask)
  
      layers = []
      for x in range(num_layers):
          layers.append(
              LayerSpec(
                  ParallelTransformerLayerPipe,
                  self.gpt2_attention_mask_func,
                  self.init_method_normal(0.02),
                  self.scaled_init_method_normal(0.02, num_layers),
                  x,
              )
          )
      super().__init__(
          layers=layers,
          loss_fn=torch.nn.CrossEntropyLoss(),
          topology=topo,
          **kwargs,
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/megatron_model.py
imports: import torch
  import os
  import sys
  import math
  from .common import get_test_path
  from deepspeed.pipe import PipelineModule, LayerSpec
label_desc: Use image node_id 1 to create a new MockGPT2ModelPipe object from inherited base classes: PipelineModule with example: obj = MockGPT2ModelPipe(num_layers, mp_size, args_others, topo)