code: def _row_parallel_linear_test_impl(
      self,
      gradient_accumulation_fusion: bool,
      accumulation_in_fp16: bool,
      sequence_parallel_enabled: bool,
  ) -> None:
      tensor_shape = (
          self.SEQUENCE_LENGTH,
          self.BATCH_SIZE,
          self.HIDDEN_SIZE,
      )
      for tensor_model_parallel_world_size in range(
          1 + int(sequence_parallel_enabled), 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,
          )
          set_random_seed(self.SEED)
  
          linear = layers.RowParallelLinear(
              self.HIDDEN_SIZE,
              self.HIDDEN_SIZE,
              keep_master_weight_for_test=True,
              params_dtype=torch.float32,
              use_cpu_initialization=True,
              gradient_accumulation_fusion=gradient_accumulation_fusion,
              accumulation_in_fp16=accumulation_in_fp16,
              sequence_parallel_enabled=sequence_parallel_enabled,
              # n.b.(mkozuki): RowParallelLinear is constructed with `input_is_parallel=True`
              # by default, e.g. https://github.com/NVIDIA/NeMo/blob/782b4e1652aaa43c8be390d9\
              # db0dc89544afa080/nemo/collections/nlp/modules/common/megatron/transformer.py#L204
              input_is_parallel=True,
          ).cuda()
          if accumulation_in_fp16:
              linear = linear.half()
          # Simulate the situation where fusion of weight grad calculation and gradient accumulation is enabled.
          if gradient_accumulation_fusion:
              with torch.no_grad():
                  linear.weight.main_grad = torch.zeros_like(
                      linear.weight
                  )
  
          msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}"
  
          with torch.no_grad():
              orig_input_tensor = torch.randn(
                  tensor_shape, requires_grad=True, device="cuda"
              )
              orig_loss_weight = torch.randn(
                  tensor_shape, device="cuda"
              )
              input_tensor = orig_input_tensor.chunk(
                  chunks=tensor_model_parallel_world_size,
                  dim=2,
              )[
                  parallel_state.get_tensor_model_parallel_rank()
              ].contiguous()
              if sequence_parallel_enabled:
                  loss_weight = orig_loss_weight.chunk(
                      chunks=tensor_model_parallel_world_size,
                      dim=0,
                  )[parallel_state.get_tensor_model_parallel_rank()]
              else:
                  loss_weight = orig_loss_weight
              if accumulation_in_fp16:
                  orig_input_tensor = orig_input_tensor.half()
                  input_tensor = input_tensor.half()
                  loss_weight = loss_weight.half()
          input_tensor.requires_grad_()
          output, _ = linear(input_tensor)
          loss = torch.mul(output, loss_weight).sum()
          loss.backward()
          self.assertIsNotNone(input_tensor.grad, msg=msg)
  
          ref_linear = nn.Linear(
              in_features=self.HIDDEN_SIZE,
              out_features=self.HIDDEN_SIZE,
              bias=False,
              device="cuda",
          )
          with torch.no_grad():
              dldy = orig_loss_weight.clone()
              x = orig_input_tensor.clone()
              ref_linear.weight.copy_(linear.master_weight)
              if accumulation_in_fp16:
                  ref_linear = ref_linear.half()
          x.requires_grad_()
          expected_output = ref_linear(x)
          expected_loss = torch.mul(expected_output, dldy).sum()
          expected_loss.backward()
  
          if not accumulation_in_fp16:
              if sequence_parallel_enabled:
                  self.assertEqual(
                      x=output,
                      y=expected_output.chunk(
                          chunks=tensor_model_parallel_world_size,
                          dim=0,
                      )[
                          parallel_state.get_tensor_model_parallel_rank()
                      ],
                      msg=msg,
                  )
              else:
                  self.assertEqual(
                      x=output,
                      y=expected_output,
                      msg=msg,
                  )
  
          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 14 for calling the TensorParallelLayerTestBase obj's underlying member method code with example usage: obj._row_parallel_linear_test_impl(gradient_accumulation_fusion, accumulation_in_fp16, sequence_parallel_enabled) without return types