code: def run_test_reproducibility(
      rank, world_size, tempfile_name, broadcast_fp16
  ):
      dist_init(rank, world_size, tempfile_name)
      device = (
          torch.device(rank)
          if torch.cuda.device_count() > 1
          else DEVICE
      )
      torch.cuda.set_device(rank)
  
      # Run a dummy step so that the optimizer state dict exists
      batch, input_width, hidden, target_width = 3, 3, 3, 5
      target = torch.rand((batch, target_width), device=device)
      inputs = torch.rand((batch, input_width), device=device)
  
      model = torch.nn.Sequential(
          torch.nn.Linear(input_width, hidden),
          torch.nn.Linear(hidden, target_width),
      )
      model.to(device)
      model = DDP(model, device_ids=[device])
  
      loss_fn = torch.nn.L1Loss()
      loss_fn.to(device)
  
      optimizer = optim.OSS(
          model.parameters(),
          optim=torch.optim.RMSprop,
          lr=0.1,
          broadcast_fp16=broadcast_fp16,
      )
  
      def closure():
          optimizer.zero_grad()
          output = model(inputs)
          loss = loss_fn(output, target)
          loss.backward()
          return loss
  
      _ = optimizer.step(closure=closure)
  
      # Get a snapshot of the state at this point
      optimizer_state_dict = copy.deepcopy(
          optimizer.state_dict(all_ranks=True)
      )
      model_state_dict = copy.deepcopy(model.state_dict())
  
      # Run two steps, log the loss
      _ = optimizer.step(closure=closure)
      reference_loss = optimizer.step(closure=closure)
  
      # Load the optimizer state dict, rewind the state two steps back
      optimizer.load_state_dict(optimizer_state_dict)
      model.load_state_dict(model_state_dict)
  
      # Run two new steps, log the loss again and check that we get the same
      _ = optimizer.step(closure=closure)
      test_loss = optimizer.step(closure=closure)
  
      assert torch.allclose(
          reference_loss, test_loss
      ), f"{reference_loss} vs {test_loss}. Reproducibility is broken"
  
      # Check that no matter what the buffer is back to fp32
      for device in optimizer.buckets.keys():
          for bucket in optimizer.buckets[device].values():
              assert bucket.buffer.dtype == torch.float32
      dist.destroy_process_group()
  
imports: import copy
  from math import inf
  import tempfile
  from typing import Any, Dict, Type, cast
  import unittest
  import numpy
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torch.nn.parallel import DistributedDataParallel
  from fairscale.fair_dev.testing.testing import check_same_model_params, check_same_models_across_ranks, skip_if_no_cuda, skip_if_py39_no_cuda, skip_if_single_gpu
  from fairscale.internal import torch_version
  import fairscale.optim
label_desc: Use image node_id 17 for calling a global function with example usage: run_test_reproducibility(rank, world_size, tempfile_name, broadcast_fp16) and returns: loss