code: def initialize_model_parallel(
      tensor_model_parallel_size_: int = 1,
      pipeline_model_parallel_size_: int = 1,
      virtual_pipeline_model_parallel_size_: Optional[int] = None,
      pipeline_model_parallel_split_rank_: Optional[int] = None,
      use_fp8_: bool = False,
      init_mpi_proc_group: bool = False,
      *,
      default_backend: Optional[str] = None,
      p2p_backend: Optional[str] = None,
  ) -> None:
      """
      Initialize model data parallel groups.
  
      Arguments:
          tensor_model_parallel_size: number of GPUs used to parallelize model tensor.
          pipeline_model_parallel_size: number of GPUs used to parallelize model pipeline.
          virtual_pipeline_model_parallel_size: number of virtual stages (interleaved pipeline).
          pipeline_model_parallel_split_rank: for models with both encoder and decoder, rank in pipeline with split point.
          use_fp8_: FP8 training that needs AMAX reduction across data-parallel ranks.
          init_mpi_proc_group: Create a MPI process group, which is used for UCX-based communication APIs.
      Keyword Arguments:
          default_backend: Backend of process groups except for pipeline parallel ones.
              If :obj:`None`, the backend specified in `torch.distributed.init_process_group` will be used.
          p2p_backend: Backend of process groups for pipeline model parallel.
              If :obj:`None`, the backend specified in `torch.distributed.init_process_group` will be used.
  
      .. note::
          `torch_ucc <https://github.com/facebookresearch/torch_ucc>`_ is
          necessary for "ucc" backend.
  
      Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
      use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
      the model pipeline. The present function will
      create 8 tensor model-parallel groups, 4 pipeline model-parallel groups
      and 8 data-parallel groups as:
          8 data_parallel groups:
              [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
          8 tensor model-parallel groups:
              [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
          4 pipeline model-parallel groups:
              [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]
      Note that for efficiency, the caller should make sure adjacent ranks
      are on the same DGX box. For example if we are using 2 DGX-1 boxes
      with a total of 16 GPUs, rank 0 to 7 belong to the first box and
      ranks 8 to 15 belong to the second box.
      """
      # Get world size and rank. Ensure some consistencies.
      assert torch.distributed.is_initialized()
      assert default_backend is None or default_backend in (
          "nccl",
          "ucc",
      )
      assert p2p_backend is None or p2p_backend in ("nccl", "ucc")
      if "ucc" in (default_backend, p2p_backend):
          if not HAS_UCC:
              raise ImportError(
                  "UCC backend requires pytorch source build with UCC installed and enabled"
              )
          warnings.warn(
              "`ucc` backend support is experimental",
              ExperimentalWarning,
          )
      if default_backend == "ucc":
          warnings.warn(
              "The UCC's functionality as `default_backend` is not well verified",
              ExperimentalWarning,
          )
  
      # Saving the NCCL_NET type for reusing it at the epilogue
      default_nccl_net = os.getenv("NCCL_NET")
  
      world_size: int = torch.distributed.get_world_size()
      tensor_model_parallel_size: int = min(
          tensor_model_parallel_size_, world_size
      )
      pipeline_model_parallel_size: int = min(
          pipeline_model_parallel_size_, world_size
      )
      if (
          world_size
          % (tensor_model_parallel_size * pipeline_model_parallel_size)
          != 0
      ):
          raise RuntimeError(
              f"`world_size` ({world_size}) is not divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) x pipeline_model_parallel_size ({pipeline_model_parallel_size})"
          )
      data_parallel_size: int = world_size // (
          tensor_model_parallel_size * pipeline_model_parallel_size
      )
      if torch.distributed.get_rank() == 0:
          _logger.info(
              "> initializing tensor model parallel with size {}".format(
                  tensor_model_parallel_size
              )
          )
          _logger.info(
              "> initializing pipeline model parallel with size {}".format(
                  pipeline_model_parallel_size
              )
          )
          _logger.info(
              "> initializing data parallel with size {}".format(
                  data_parallel_size
              )
          )
  
      num_tensor_model_parallel_groups: int = (
          world_size // tensor_model_parallel_size
      )
      num_pipeline_model_parallel_groups: int = (
          world_size // pipeline_model_parallel_size
      )
      num_data_parallel_groups: int = world_size // data_parallel_size
  
      if virtual_pipeline_model_parallel_size_ is not None:
          # n.b. (eqy) This check was inherited from Megatron-LM, need to revisit
          # the root cause as we do see numerical mismatches with 2 stages and
          # the interleaved schedule
          assert pipeline_model_parallel_size_ > 2, (
              "pipeline-model-parallel size should be greater than 2 with "
              "interleaved schedule"
          )
          global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
          global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
          _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0
          _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = (
              virtual_pipeline_model_parallel_size_
          )
  
      if pipeline_model_parallel_split_rank_ is not None:
          global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
          _PIPELINE_MODEL_PARALLEL_SPLIT_RANK = (
              pipeline_model_parallel_split_rank_
          )
  
      rank = torch.distributed.get_rank()
  
      # Build the data-parallel groups.
      global _DATA_PARALLEL_GROUP
      assert (
          _DATA_PARALLEL_GROUP is None
      ), "data parallel group is already initialized"
      all_data_parallel_group_ranks = []
      for i in range(pipeline_model_parallel_size):
          start_rank = i * num_pipeline_model_parallel_groups
          end_rank = (i + 1) * num_pipeline_model_parallel_groups
          for j in range(tensor_model_parallel_size):
              ranks = range(
                  start_rank + j, end_rank, tensor_model_parallel_size
              )
              all_data_parallel_group_ranks.append(list(ranks))
              group = new_process_group(ranks, backend=default_backend)
              if rank in ranks:
                  _DATA_PARALLEL_GROUP = group
  
      # Build the amax-reduction groups for fp8 precision conversion.
      if use_fp8_:
          global _AMAX_REDUCTION_GROUP
          assert (
              _AMAX_REDUCTION_GROUP is None
          ), "amax reduction group is already initialized"
          amax_group_size: int = (
              tensor_model_parallel_size * data_parallel_size
          )
          num_amax_groups: int = world_size // amax_group_size
          for i in range(num_amax_groups):
              start_rank = i * amax_group_size
              end_rank = (i + 1) * amax_group_size
              ranks = range(start_rank, end_rank)
              group = torch.distributed.new_group(
                  ranks, backend=default_backend
              )
              if rank in ranks:
                  _AMAX_REDUCTION_GROUP = group
  
      # Build the model-parallel groups.
      global _MODEL_PARALLEL_GROUP
      assert (
          _MODEL_PARALLEL_GROUP is None
      ), "model parallel group is already initialized"
      for i in range(data_parallel_size):
          ranks = [
              data_parallel_group_ranks[i]
              for data_parallel_group_ranks in all_data_parallel_group_ranks
          ]
          group = new_process_group(ranks, backend=default_backend)
          if rank in ranks:
              _MODEL_PARALLEL_GROUP = group
  
      # Build the tensor model-parallel groups.
      global _TENSOR_MODEL_PARALLEL_GROUP
      assert (
          _TENSOR_MODEL_PARALLEL_GROUP is None
      ), "tensor model parallel group is already initialized"
      for i in range(num_tensor_model_parallel_groups):
          ranks = list(
              range(
                  i * tensor_model_parallel_size,
                  (i + 1) * tensor_model_parallel_size,
              )
          )
          group = new_process_group(ranks, backend=default_backend)
          if rank in ranks:
              _TENSOR_MODEL_PARALLEL_GROUP = group
  
      # Build the pipeline model-parallel groups and embedding groups
      # (first and last rank in each pipeline model-parallel group).
      global _PIPELINE_MODEL_PARALLEL_GROUP
      global _PIPELINE_GLOBAL_RANKS
      assert (
          _PIPELINE_MODEL_PARALLEL_GROUP is None
      ), "pipeline model parallel group is already initialized"
      global _EMBEDDING_GROUP
      global _EMBEDDING_GLOBAL_RANKS
      assert (
          _EMBEDDING_GROUP is None
      ), "embedding group is already initialized"
      global _POSITION_EMBEDDING_GROUP
      global _POSITION_EMBEDDING_GLOBAL_RANKS
      assert (
          _POSITION_EMBEDDING_GROUP is None
      ), "position embedding group is already initialized"
      global _ENCODER_RELATIVE_POSITION_EMBEDDING_GROUP
      global _DECODER_RELATIVE_POSITION_EMBEDDING_GROUP
      global _ENCODER_RELATIVE_POSITION_EMBEDDING_GLOBAL_RANKS
      global _DECODER_RELATIVE_POSITION_EMBEDDING_GLOBAL_RANKS
      assert (
          _ENCODER_RELATIVE_POSITION_EMBEDDING_GROUP is None
          or _DECODER_RELATIVE_POSITION_EMBEDDING_GROUP is None
      ), "relative position embedding group is already initialized"
      for i in range(num_pipeline_model_parallel_groups):
          ranks = range(
              i, world_size, num_pipeline_model_parallel_groups
          )
          group = new_process_group(ranks, backend=p2p_backend)
          if rank in ranks:
              _PIPELINE_MODEL_PARALLEL_GROUP = group
              _PIPELINE_GLOBAL_RANKS = ranks
          # Setup embedding group (to exchange gradients between
          # first and last stages).
          encoder_relative_position_embedding_ranks = None
          decoder_relative_position_embedding_ranks = None
          if len(ranks) > 1:
              embedding_ranks = [ranks[0], ranks[-1]]
              position_embedding_ranks = [ranks[0]]
              encoder_relative_position_embedding_ranks = [ranks[0]]
              decoder_relative_position_embedding_ranks = [ranks[0]]
              if pipeline_model_parallel_split_rank_ is not None:
                  encoder_relative_position_embedding_ranks = ranks[
                      :pipeline_model_parallel_split_rank_
                  ]
                  decoder_relative_position_embedding_ranks = ranks[
                      pipeline_model_parallel_split_rank_:
                  ]
                  if (
                      ranks[pipeline_model_parallel_split_rank_]
                      not in embedding_ranks
                  ):
                      embedding_ranks = [
                          ranks[0],
                          ranks[pipeline_model_parallel_split_rank_],
                          ranks[-1],
                      ]
                  if (
                      ranks[pipeline_model_parallel_split_rank_]
                      not in position_embedding_ranks
                  ):
                      position_embedding_ranks = [
                          ranks[0],
                          ranks[pipeline_model_parallel_split_rank_],
                      ]
          else:
              embedding_ranks = ranks
              position_embedding_ranks = ranks
              encoder_relative_position_embedding_ranks = ranks
              decoder_relative_position_embedding_ranks = ranks
  
          group = new_process_group(
              embedding_ranks, backend=p2p_backend
          )
          if rank in embedding_ranks:
              _EMBEDDING_GROUP = group
          if rank in ranks:
              _EMBEDDING_GLOBAL_RANKS = embedding_ranks
  
          group = new_process_group(
              position_embedding_ranks, backend=p2p_backend
          )
          if rank in position_embedding_ranks:
              _POSITION_EMBEDDING_GROUP = group
          if rank in ranks:
              _POSITION_EMBEDDING_GLOBAL_RANKS = (
                  position_embedding_ranks
              )
  
          if encoder_relative_position_embedding_ranks:
              group = new_process_group(
                  encoder_relative_position_embedding_ranks,
                  backend=p2p_backend,
              )
          if rank in encoder_relative_position_embedding_ranks:
              _ENCODER_RELATIVE_POSITION_EMBEDDING_GROUP = group
          if rank in ranks:
              _ENCODER_RELATIVE_POSITION_EMBEDDING_GLOBAL_RANKS = (
                  encoder_relative_position_embedding_ranks
              )
  
          if decoder_relative_position_embedding_ranks:
              group = new_process_group(
                  decoder_relative_position_embedding_ranks,
                  backend=p2p_backend,
              )
          if rank in decoder_relative_position_embedding_ranks:
              _DECODER_RELATIVE_POSITION_EMBEDDING_GROUP = group
          if rank in ranks:
              _DECODER_RELATIVE_POSITION_EMBEDDING_GLOBAL_RANKS = (
                  decoder_relative_position_embedding_ranks
              )
  
      if init_mpi_proc_group:
          torch.distributed.new_group(backend="mpi")
  
      if default_nccl_net == "Socket":
          set_nccl_socket_envs()
      elif default_nccl_net == "IB":
          set_nccl_ib_envs()
      elif default_nccl_net is None:
          os.unsetenv("NCCL_NET")
      else:
          os.environ["NCCL_NET"] = default_nccl_net
  
imports: from typing import Tuple, Optional
  import warnings
  import os
  import torch
  from apex.transformer.log_util import get_transformer_logger
  from apex.transformer._ucc_util import HAS_UCC
label_desc: Use image node_id 8 for calling a global function with example usage: initialize_model_parallel(tensor_model_parallel_size_, pipeline_model_parallel_size_, virtual_pipeline_model_parallel_size_, pipeline_model_parallel_split_rank_, use_fp8_, init_mpi_proc_group) without return types