code: def run_test_set_cuda_rng_state(
      rank, model_parallel_size, filename, filename_rpc
  ):
      dist_init(rank, model_parallel_size, filename, filename_rpc)
  
      if torch.distributed.get_rank() == 0:
          print(
              "> testing set_rng_state with size {} ...".format(
                  model_parallel_size
              )
          )
  
      mpu.initialize_model_parallel(model_parallel_size)
      model_parallel_size = mpu.get_model_parallel_world_size()
  
      size = 123
      seed = 1234
      torch.cuda.manual_seed(1234)
      tensor = torch.cuda.FloatTensor(size)
  
      # Get the state
      rng_state = torch.cuda.get_rng_state()
      rng_state_copy = rng_state.clone()
  
      # Do some stuff.
      for _ in range(5):
          torch.randn(size, out=tensor)
      result_1 = tensor.clone()
  
      assert rng_state.sub(rng_state_copy).max() == 0
      assert torch.cuda.get_rng_state().sub(rng_state_copy).max() > 0
  
      # State should be different.
      new_rng_state = torch.cuda.get_rng_state()
      max_diff = new_rng_state.sub(rng_state).max()
      print(
          "   max diff in rng state (should be non-zero) on global rank {}: {}".format(
              torch.distributed.get_rank(), max_diff
          )
      )
      assert max_diff > 0
  
      # Reset the rng state and do the same stuff.
      random._set_cuda_rng_state(rng_state)
      for _ in range(5):
          torch.randn(size, out=tensor)
      random._set_cuda_rng_state(rng_state)
      for _ in range(5):
          torch.randn(size, out=tensor)
      result_2 = tensor.clone()
  
      # Results should be the same
      error = result_2.sub(result_1).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
  
      # Input state should have remained intact.
      error = rng_state.sub(rng_state_copy).max()
      print(
          "   max error in rng state (should be zero) on global rank {}: {}".format(
              torch.distributed.get_rank(), error
          )
      )
      assert error == 0
  
      # 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 1 for calling a global function with example usage: run_test_set_cuda_rng_state(rank, model_parallel_size, filename, filename_rpc) without return types