code: def run_test_cuda_rng_tracker(
      rank, model_parallel_size, filename, filename_rpc
  ):
      dist_init(rank, model_parallel_size, filename, filename_rpc)
  
      if torch.distributed.get_rank() == 0:
          print(
              "> testing cuda rng tracker with size {} ...".format(
                  model_parallel_size
              )
          )
  
      mpu.initialize_model_parallel(model_parallel_size)
      model_parallel_size = mpu.get_model_parallel_world_size()
  
      seed_1 = 1234
      seed_2 = 4321
      size = [12, 21]
      tensor = torch.cuda.FloatTensor(size)
  
      # Set to seed_1 and generate two tensors.
      torch.cuda.manual_seed(seed_1)
      torch.randn(size, out=tensor)
      target_11 = tensor.clone()
      torch.randn(size, out=tensor)
      target_12 = tensor.clone()
  
      # Set to seed_2 and generate two tensors.
      torch.cuda.manual_seed(seed_2)
      torch.randn(size, out=tensor)
      target_21 = tensor.clone()
      torch.randn(size, out=tensor)
      target_22 = tensor.clone()
  
      # Now if we interleave seed_1 and seed_2,
      # we should still get the same tensors
      torch.cuda.manual_seed(seed_1)
      get_cuda_rng_tracker().add("test", seed_2)
  
      torch.randn(size, out=tensor)
      result_11 = tensor.clone()
  
      with get_cuda_rng_tracker().fork("test"):
          torch.randn(size, out=tensor)
          result_21 = tensor.clone()
  
      torch.randn(size, out=tensor)
      result_12 = tensor.clone()
  
      with get_cuda_rng_tracker().fork("test"):
          torch.randn(size, out=tensor)
          result_22 = tensor.clone()
  
      diff = result_11.sub(result_21).abs().max()
      diff = min(diff, result_12.sub(result_22).abs().max())
      print(
          "   max diff in generated tensors (should be non-zero) on global rank {}: {}".format(
              torch.distributed.get_rank(), diff
          )
      )
      assert diff > 1.0e-6
      error = max(
          result_11.sub(target_11).abs().max(),
          result_12.sub(target_12).abs().max(),
      )
      error = max(error, result_21.sub(target_21).abs().max())
      error = max(error, result_22.sub(target_22).abs().max())
      print(
          "   max error in generated tensors (should be zero) on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error < 1.0e-6
  
      # Reset the tracker
      get_cuda_rng_tracker().reset()
  
      # 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_random.py
imports: import torch
  from fairscale.fair_dev.testing.testing import dist_init, spawn_for_all_world_sizes
  from fairscale.nn.model_parallel import initialize
  from fairscale.nn.model_parallel import random
  from fairscale.nn.model_parallel.random import get_cuda_rng_tracker, model_parallel_cuda_manual_seed
label_desc: Use image node_id 2 for calling a global function with example usage: run_test_cuda_rng_tracker(rank, model_parallel_size, filename, filename_rpc) without return types