code: def run_test_parallel_embedding(
      rank, model_parallel_size, filename, filename_rpc
  ):
      dist_init(rank, model_parallel_size, filename, filename_rpc)
  
      if torch.distributed.get_rank() == 0:
          print(
              "> testing parallel embedding with model parallel size {} ...".format(
                  model_parallel_size
              )
          )
  
      mpu.initialize_model_parallel(model_parallel_size)
      model_parallel_size = mpu.get_model_parallel_world_size()
  
      batch_size = 17
      seq_length = 23
      vocab_size = 48
      hidden_size = 16
      seed = 1236
  
      set_random_seed(123)
      input_data = (
          torch.LongTensor(size=(batch_size, seq_length))
          .random_(0, vocab_size)
          .cuda()
      )
      loss_weight = torch.randn(
          [batch_size, seq_length, hidden_size]
      ).cuda()
  
      set_random_seed(seed)
      embedding_original = torch.nn.Embedding(
          vocab_size, hidden_size
      ).cuda()
  
      output = embedding_original(input_data)
      loss_original = torch.mul(output, loss_weight).sum()
      loss_original.backward()
  
      set_random_seed(seed)
      embedding_parallel = layers.ParallelEmbedding(
          vocab_size, hidden_size, init_method=init.normal_
      ).cuda()
      output = embedding_parallel(input_data)
      loss_parallel = torch.mul(output, loss_weight).sum()
      loss_parallel.backward()
  
      set_random_seed(seed)
      embedding_vocab_parallel = layers.VocabParallelEmbedding(
          vocab_size, hidden_size, init_method=init.normal_
      ).cuda()
      output = embedding_vocab_parallel(input_data)
      loss_vocab_parallel = torch.mul(output, loss_weight).sum()
      loss_vocab_parallel.backward()
  
      torch.distributed.barrier()
      error = loss_parallel.sub(loss_original).abs()
      print(
          "   error in loss (parallel) on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-12, "error: {}".format(error)
  
      torch.distributed.barrier()
      error = loss_vocab_parallel.sub(loss_original).abs()
      print(
          "   error in loss (vocab parallel) on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-12, "error: {}".format(error)
  
      weight_grad_orig = torch.split(
          embedding_original.weight.grad,
          hidden_size // model_parallel_size,
          1,
      )[mpu.get_model_parallel_rank()]
      error = (
          embedding_parallel.weight.grad.sub(weight_grad_orig)
          .abs()
          .max()
      )
      print(
          "   error in grad (parallel) on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-12, "error: {}".format(error)
  
      weight_grad_orig = torch.split(
          embedding_original.weight.grad,
          vocab_size // model_parallel_size,
          0,
      )[mpu.get_model_parallel_rank()]
      error = (
          embedding_vocab_parallel.weight.grad.sub(weight_grad_orig)
          .abs()
          .max()
      )
      print(
          "   error in grad (vocab parallel) on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-12, "error: {}".format(error)
  
      # Reset groups
      mpu.destroy_model_parallel()
  
      torch.distributed.barrier()
      if torch.distributed.get_rank() == 0:
          print(">> passed the test :-)")
  
file_path: fairscale/tests/nn/model_parallel/test_layers.py
imports: import os
  import pytest
  import torch
  import torch.nn.init
  from torch.nn.parameter import Parameter
  from fairscale.fair_dev.testing.testing import dist_init, set_random_seed, spawn_for_all_world_sizes
  from fairscale.nn.model_parallel import initialize
  from fairscale.nn.model_parallel import layers
label_desc: Use image node_id 1 for calling a global function with example usage: run_test_parallel_embedding(rank, model_parallel_size, filename, filename_rpc) without return types