code: def build_model(
      model_provider_func: Callable[
          [Any, Dict[str, Any]], torch.nn.Module
      ],
      wrap_with_ddp: bool = True,
      virtual_pipeline_model_parallel_size: Optional[int] = None,
      model_type: ModelType = ModelType.encoder_or_decoder,
      *args: Any,
      **kwargs: Any,
  ) -> List[torch.nn.Module]:
      """Build the model satisfying pipeline model parallel requirements.
  
      This function sets `pre_process` and `post_process` to `**kwargs` and pass `*args` and `**kwargs` to
      `model_provider_func`.
  
      Args:
          model_provider_func: A function which takes `*args` and `**kwargs` and returns a `nn.Module`.
          wrap_with_ddp: If :obj:`True`, wrap the instantiated model
              with `torch.nn.parallel.distributed.DistributedDataParallel`, a.k.a. `DDP`.
          virtual_pipeline_model_parallel_size: Specify when using interleaving scheduling pipeline model parallel.
          model_type:
          *args: arguments for model provider func
          **kwargs: Keyword arguments for model provider func
  
      Returns:
          a list of `nn.Module`(s). If `virtual_pipeline_model_parallel_size` is not None,
          the list has multiple models, otherwise one.
      """
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and virtual_pipeline_model_parallel_size is not None
      ):
          model = []
          for i in range(virtual_pipeline_model_parallel_size):
              cur_args = args
              cur_kwargs = kwargs
              parallel_state.set_virtual_pipeline_model_parallel_rank(i)
              # Set pre_process and post_process only after virtual rank is set.
              pre_process = parallel_state.is_pipeline_first_stage()
              post_process = parallel_state.is_pipeline_last_stage()
              cur_kwargs.update(
                  {
                      "pre_process": pre_process,
                      "post_process": post_process,
                  }
              )
              this_model = model_provider_func(*cur_args, **cur_kwargs)
              model.append(this_model)
      else:
          cur_args = args
          cur_kwargs = kwargs
          if model_type == ModelType.encoder_or_decoder:
              pre_process = parallel_state.is_pipeline_first_stage()
              post_process = parallel_state.is_pipeline_last_stage()
              cur_kwargs.update(
                  {
                      "pre_process": pre_process,
                      "post_process": post_process,
                  }
              )
              model = model_provider_func(*cur_args, **cur_kwargs)
          elif model_type == ModelType.encoder_and_decoder:
              pre_process = parallel_state.is_pipeline_first_stage()
              post_process = parallel_state.is_pipeline_last_stage()
              # `add_encoder` & `add_decoder` logic.
              add_encoder, add_decoder = True, True
              if (
                  parallel_state.get_pipeline_model_parallel_world_size()
                  > 1
              ):
                  split_rank = (
                      parallel_state.get_pipeline_model_parallel_split_rank()
                  )
                  if split_rank is None:
                      raise RuntimeError(
                          "Split rank needs to be specified for model with both encoder and decoder."
                      )
                  rank = (
                      parallel_state.get_pipeline_model_parallel_rank()
                  )
                  world_size = (
                      parallel_state.get_pipeline_model_parallel_world_size()
                  )
                  pre_process = rank == 0 or rank == split_rank
                  post_process = rank == (split_rank - 1) or rank == (
                      world_size - 1
                  )
                  add_encoder = (
                      parallel_state.is_pipeline_stage_before_split()
                  )
                  add_decoder = (
                      parallel_state.is_pipeline_stage_after_split()
                  )
              cur_kwargs.update(
                  {
                      "pre_process": pre_process,
                      "post_process": post_process,
                      "add_encoder": add_encoder,
                      "add_decoder": add_decoder,
                  }
              )
              model = model_provider_func(*cur_args, **cur_kwargs)
          model.model_type = model_type
  
      if not isinstance(model, list):
          model = [model]
  
      # Set tensor model parallel attributes if not set.
      # Only parameters that are already tensor model parallel have these
      # attributes set for them. We should make sure the default attributes
      # are set for all params so the optimizer can use them.
      for model_module in model:
          for param in model_module.parameters():
              set_defaults_if_not_set_tensor_model_parallel_attributes(
                  param
              )
  
      # Print number of parameters.
      if (
          parallel_state.model_parallel_is_initialized()
          and parallel_state.get_data_parallel_rank() == 0
      ):
          msg = " > number of parameters on (tensor, pipeline) model parallel rank ({}, {}): {}".format(
              parallel_state.get_tensor_model_parallel_rank(),
              parallel_state.get_pipeline_model_parallel_rank(),
              _calc_number_of_params(model),
          )
          print(msg, flush=True)
  
      # GPU allocation.
      for model_module in model:
          model_module.cuda(torch.cuda.current_device())
  
      if wrap_with_ddp:
          i = torch.cuda.current_device()
          model = [
              torch.nn.parallel.distributed.DistributedDataParallel(
                  model_module,
                  device_ids=[i],
                  output_device=i,
                  process_group=parallel_state.get_data_parallel_group(),
              )
              for model_module in model
          ]
      return model
  
file_path: apex/apex/transformer/pipeline_parallel/schedules/common.py
imports: from typing import Any, Callable, Dict, List, Tuple, Union, Optional, Sequence
  import torch
  from torch.autograd.variable import Variable
  from apex.normalization.fused_layer_norm import FusedLayerNorm
  from apex.transformer import parallel_state
  from apex.transformer.enums import ModelType
  from apex.transformer.pipeline_parallel.p2p_communication import FutureTensor
  from apex.transformer.pipeline_parallel.utils import get_num_microbatches
  from apex.transformer.pipeline_parallel.utils import listify_model
  from apex.transformer.pipeline_parallel.utils import unwrap_model
  from apex.transformer.pipeline_parallel.utils import get_model_type
  from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 1 for calling a global function with example usage: build_model(model_provider_func, wrap_with_ddp, virtual_pipeline_model_parallel_size, model_type) and returns: model