code: def test_parallel_embedding(self) -> None:
      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,
          )
          set_random_seed(self.SEED + 1)
          input_tensor = torch.randint(
              0,
              self.VOCAB_SIZE,
              (
                  self.BATCH_SIZE,
                  self.SEQUENCE_LENGTH,
              ),
              device="cuda",
          )
          loss_weight = torch.randn(
              (
                  self.BATCH_SIZE,
                  self.SEQUENCE_LENGTH,
                  self.HIDDEN_SIZE,
              ),
              device="cuda",
          )
  
          set_random_seed(self.SEED)
          embedding_torch = nn.Embedding(
              self.VOCAB_SIZE,
              self.HIDDEN_SIZE,
          ).cuda()
          output_torch = embedding_torch(input_tensor)
          loss_torch = torch.mul(output_torch, loss_weight).sum()
          loss_torch.backward()
  
          # N.B.(mkozuki): With affine weight initialization on GPU,
          # it's super difficult to keep the consistency with nn.Embedding.
          # Thus, turning on `use_cpu_initialization`.
          set_random_seed(self.SEED)
          embedding_vocab_parallel = layers.VocabParallelEmbedding(
              self.VOCAB_SIZE,
              self.HIDDEN_SIZE,
              init_method=nn.init.normal_,
              use_cpu_initialization=True,
          ).cuda()
          output_vocab_parallel = embedding_vocab_parallel(input_tensor)
          loss_vocab_parallel = torch.mul(
              output_vocab_parallel, loss_weight
          ).sum()
          loss_vocab_parallel.backward()
  
          msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}"
          self.assertEqual(output_torch, output_vocab_parallel, msg=msg)
          self.assertEqual(loss_torch, loss_vocab_parallel, msg=msg)
  
          splitted_weight_torch = torch.split(
              embedding_torch.weight.grad,
              self.VOCAB_SIZE // tensor_model_parallel_world_size,
              0,
          )[parallel_state.get_tensor_model_parallel_rank()]
          self.assertEqual(
              splitted_weight_torch,
              embedding_vocab_parallel.weight.grad,
              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 4 for calling the TensorParallelLayerTestBase obj's underlying member method code with example usage: obj.test_parallel_embedding() without return types