code: def test_all_gather_parity(self) -> None:
      if self.DISTRIBUTED_BACKEND == "ucc":
          self.skipTest(
              "torch_ucc does NOT support `torch.distributed._all_gather_base` as of 2022/06/15"
          )
      from torch.distributed.distributed_c10d import (
          all_gather,
          _all_gather_base,
      )  # NOQA
  
      for tensor_model_parallel_world_size in range(
          1, self.world_size + 1
      ):
          if self.world_size % tensor_model_parallel_world_size:
              continue
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_parallel_world_size,
          )
          tensor_model_parallel_rank = (
              parallel_state.get_tensor_model_parallel_rank()
          )
          cur_tensor_model_device = torch.device(
              f"cuda:{tensor_model_parallel_rank}"
          )
          with torch.no_grad():
              tensor = tensor_model_parallel_rank * torch.ones(
                  self.tensor_shape,
                  dtype=torch.float32,
                  device=cur_tensor_model_device,
              )
          numel = tensor.numel()
          numel_gathered = tensor_model_parallel_world_size * numel
          gathered = torch.empty(
              torch.Size((numel_gathered,)),
              device=cur_tensor_model_device,
              dtype=torch.float32,
              requires_grad=False,
          )
          chunks = [
              gathered[i * numel : (i + 1) * numel]
              for i in range(tensor_model_parallel_world_size)
          ]
          all_gather(
              chunks,
              tensor,
              group=parallel_state.get_tensor_model_parallel_group(),
          )
  
          gathered_for_base = torch.empty(
              torch.Size((numel_gathered,)),
              device=cur_tensor_model_device,
              dtype=torch.float32,
              requires_grad=False,
          )
          _all_gather_base(
              gathered_for_base,
              tensor,
              group=parallel_state.get_tensor_model_parallel_group(),
          )
  
          msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}"
          self.assertEqual(gathered, gathered_for_base, msg=msg)
          parallel_state.destroy_model_parallel()
  
imports: import logging
  import unittest
  import typing
  import torch
  import torch.nn
  from torch.testing._internal import common_utils
  from apex.transformer import parallel_state
  from apex.transformer.tensor_parallel import layers
  from apex.transformer.testing.commons import set_random_seed
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
  from apex.transformer.testing.distributed_test_base import UccDistributedTestBase
label_desc: Use image node_id 2 for calling the TensorParallelLayerTestBase obj's underlying member method code with example usage: obj.test_all_gather_parity() without return types