code: def run_test_row_parallel_linear(
      rank, model_parallel_size, filename, filename_rpc
  ):
      dist_init(rank, model_parallel_size, filename, filename_rpc)
  
      mpu.initialize_model_parallel(model_parallel_size)
      if torch.distributed.get_rank() == 0:
          print(
              "> testing RowParallelLinear with model parallel size: {}".format(
                  model_parallel_size
              )
          )
      model_parallel_size = mpu.get_model_parallel_world_size()
  
      seed = 12345
      set_random_seed(seed)
      input_size_coeff = 13
      input_size = input_size_coeff * model_parallel_size
      output_size_coeff = 17
      output_size = output_size_coeff * model_parallel_size
      batch_size = 7
  
      # Network
      identity_layer = IdentityLayer2D(batch_size, input_size).cuda()
      linear_layer = layers.RowParallelLinear(
          input_size, output_size, keep_master_weight_for_test=True
      ).cuda()
      loss_weight = torch.randn([batch_size, output_size]).cuda()
      # Forward
      input_ = identity_layer()
      output = linear_layer(input_)
      loss = torch.mul(output, loss_weight).sum()
      # Backward
      loss.backward()
  
      # Values.
      dLdY = loss_weight
      X = identity_layer.weight
      A = linear_layer.master_weight.cuda()
      dLdA = torch.matmul(dLdY.t(), X)
      dLdb = torch.matmul(
          torch.ones(batch_size, 1).cuda().t(), dLdY
      ).view(-1)
      dLdX = torch.matmul(dLdY, A)
  
      rank = mpu.get_model_parallel_rank()
      my_dLdA = (
          torch.split(dLdA, input_size_coeff, dim=1)[rank]
          .contiguous()
          .clone()
      )
      error = my_dLdA.sub(linear_layer.weight.grad).abs().max()
      torch.distributed.barrier()
      print(
          "   error in dLdA on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6, error
  
      error = dLdb.sub(linear_layer.bias.grad).abs().max()
      torch.distributed.barrier()
      print(
          "   error in dLdb on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6, error
  
      error = dLdX.sub(identity_layer.weight.grad).abs().max()
      torch.distributed.barrier()
      print(
          "   error in dLdX on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6, 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 4 for calling a global function with example usage: run_test_row_parallel_linear(rank, model_parallel_size, filename, filename_rpc) without return types