code: def run_test_cross_entropy(
      rank, model_parallel_size, filename, filename_rpc
  ):
      dist_init(rank, model_parallel_size, filename, filename_rpc)
  
      if torch.distributed.get_rank() == 0:
          print(
              "> testing cross entropy 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 = 13
      seq_length = 17
      vocab_size_per_partition = 11
      logits_scale = 1000.0
      vocab_size = vocab_size_per_partition * model_parallel_size
      seed = 1234
  
      loss_torch, grad_torch = torch_cross_entropy(
          batch_size, seq_length, vocab_size, logits_scale, seed
      )
      loss_mpu, grad_mpu = mpu_cross_entropy(
          batch_size, seq_length, vocab_size, logits_scale, seed
      )
  
      error = loss_torch.sub_(loss_mpu).abs().max()
      print(
          "   max error in loss on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6
  
      error = grad_torch.sub_(grad_mpu).abs().max()
      print(
          "   max error in grad on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6
  
      # 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_cross_entropy.py
imports: import torch
  import torch.nn.functional
  from fairscale.fair_dev.testing.testing import IdentityLayer, dist_init, set_random_seed, spawn_for_all_world_sizes
  from fairscale.nn.model_parallel import initialize
  from fairscale.nn.model_parallel.cross_entropy import vocab_parallel_cross_entropy
  from fairscale.nn.model_parallel.mappings import scatter_to_model_parallel_region
label_desc: Use image node_id 3 for calling a global function with example usage: run_test_cross_entropy(rank, model_parallel_size, filename, filename_rpc) without return types