code: def _column_parallel_linear_test_impl(
      self,
      async_tensor_model_parallel_allreduce: bool,
      gradient_accumulation_fusion: bool,
      accumulation_in_fp16: bool,
      sequence_parallel_enabled: bool,
  ):
      for tensor_model_parallel_world_size in range(
          1, self.world_size + 1
      ):
          if (
              async_tensor_model_parallel_allreduce
              and sequence_parallel_enabled
          ):
              if tensor_model_parallel_world_size == 1:
                  continue
          if self.world_size % tensor_model_parallel_world_size:
              continue
          msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}"
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_parallel_world_size,
          )
  
          input_tensor_shape = self.tensor_shape
          expected_output_shape = self.tensor_shape
          # When sequence parallel, `gather_output` is disabled, i.e.,
          # output of matmul isn't gathered in dimension of feature/hidden (last dim).
          if sequence_parallel_enabled:
              expected_output_shape[
                  -1
              ] //= tensor_model_parallel_world_size
  
          # tensor's shape is [sequence length, batch size, hidden size]
          set_random_seed(self.SEED)
          linear = layers.ColumnParallelLinear(
              self.HIDDEN_SIZE,
              self.HIDDEN_SIZE,
              bias=False,
              keep_master_weight_for_test=True,
              params_dtype=torch.float32,
              use_cpu_initialization=True,
              gather_output=not sequence_parallel_enabled,
              no_async_tensor_model_parallel_allreduce=not async_tensor_model_parallel_allreduce,
              gradient_accumulation_fusion=gradient_accumulation_fusion,
              accumulation_in_fp16=accumulation_in_fp16,
              sequence_parallel_enabled=sequence_parallel_enabled,
          ).cuda()
          if accumulation_in_fp16:
              linear = linear.half()
  
          # Simulate the situation where fusion of weight grad calculation and gradient accumulation happens.
          if gradient_accumulation_fusion:
              with torch.no_grad():
                  linear.weight.main_grad = torch.zeros_like(
                      linear.weight
                  )
  
          orig_input_tensor = torch.randn(
              input_tensor_shape, device="cuda", requires_grad=True
          )
          if accumulation_in_fp16:
              orig_input_tensor = orig_input_tensor.half()
          if sequence_parallel_enabled:
              input_tensor = list(
                  orig_input_tensor.chunk(
                      tensor_model_parallel_world_size, dim=0
                  )
              )[parallel_state.get_tensor_model_parallel_rank()]
          else:
              input_tensor = orig_input_tensor
          output, _ = linear(input_tensor)
          # The order of dimension is expected to be (sequence, batch, hidden)
          self.assertEqual(output.shape, expected_output_shape, msg=msg)
  
          orig_loss_weight = torch.randn(
              input_tensor_shape, device="cuda"
          )
          if accumulation_in_fp16:
              orig_loss_weight = orig_loss_weight.half()
          if sequence_parallel_enabled:
              loss_weight = orig_loss_weight.chunk(
                  tensor_model_parallel_world_size,
                  dim=2,
              )[parallel_state.get_tensor_model_parallel_rank()]
          else:
              loss_weight = orig_loss_weight
          loss = torch.mul(output, loss_weight).sum()
          loss.backward()
  
          with torch.no_grad():
              dldy = orig_loss_weight.clone()
              x = orig_input_tensor.clone()
              ref_linear = nn.Linear(
                  in_features=self.HIDDEN_SIZE,
                  out_features=self.HIDDEN_SIZE,
                  bias=False,
                  device="cuda",
              )
              if accumulation_in_fp16:
                  ref_linear = ref_linear.half()
              # NOTE(mkozuki): `master_weight` is available because `keep_master_weight_for_test` is set.
              ref_linear.weight.copy_(linear.master_weight)
          x.requires_grad_()
          expected_output = ref_linear(x)
          if sequence_parallel_enabled:
              chunk = expected_output.chunk(
                  tensor_model_parallel_world_size,
                  dim=2,
              )[parallel_state.get_tensor_model_parallel_rank()]
              self.assertEqual(
                  x=output,
                  y=chunk,
                  msg=msg,
              )
          else:
              self.assertEqual(
                  x=output,
                  y=expected_output,
                  msg=msg,
              )
  
          expected_loss = torch.mul(expected_output, dldy).sum()
          expected_loss.backward()
          grad_attr_name = (
              "main_grad" if gradient_accumulation_fusion else "grad"
          )
          # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel.
          if tensor_model_parallel_world_size == 1:
              self.assertEqual(
                  x=getattr(linear.weight, grad_attr_name),
                  y=ref_linear.weight.grad.chunk(
                      chunks=tensor_model_parallel_world_size,
                      dim=0,
                  )[parallel_state.get_tensor_model_parallel_rank()],
                  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 21 for calling the TensorParallelLayerTestBase obj's underlying member method code with example usage: obj._column_parallel_linear_test_impl(async_tensor_model_parallel_allreduce, gradient_accumulation_fusion, accumulation_in_fp16, sequence_parallel_enabled) without return types