code: def test_set_cuda_rng_state(self):
      for tensor_model_parallel_world_size in range(
          1, self.world_size + 1
      ):
          if self.world_size % tensor_model_parallel_world_size:
              continue
          msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}"
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_parallel_world_size
          )
  
          size, seed = 123, 1234
          torch.cuda.manual_seed(seed)
          tensor = torch.cuda.FloatTensor(size)
  
          rng_state = torch.cuda.get_rng_state()
          rng_state_clone = rng_state.clone()
  
          for _ in range(5):
              torch.randn(size, out=tensor)
          result_1 = tensor.clone()
  
          self.assertEqual(
              rng_state.sub(rng_state_clone).max(), 0, msg=msg
          )
          self.assertGreater(
              torch.cuda.get_rng_state().sub(rng_state_clone).max(),
              0,
              msg=msg,
          )
  
          new_rng_state = torch.cuda.get_rng_state()
          self.assertGreater(
              new_rng_state.sub(rng_state).max(), 0, msg=msg
          )
  
          tensor_parallel.random._set_cuda_rng_state(rng_state)
          for _ in range(5):
              torch.randn(size, out=tensor)
          tensor_parallel.random._set_cuda_rng_state(rng_state)
          for _ in range(5):
              torch.randn(size, out=tensor)
          result_2 = tensor.clone()
  
          self.assertEqual(result_2, result_1, msg=msg)
  
          self.assertEqual(
              rng_state.sub(rng_state_clone).max(), 0, msg=msg
          )
  
          parallel_state.destroy_model_parallel()
  
file_path: apex/tests/L0/run_transformer/test_random.py
imports: import logging
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer import parallel_state
  from apex.transformer import tensor_parallel
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
  from apex.transformer.testing.distributed_test_base import UccDistributedTestBase
label_desc: Use image node_id 1 for calling the TransformerRandomTestBase obj's underlying member method code with example usage: obj.test_set_cuda_rng_state() without return types