code: def run_state_dict_distributed(rank, world_size, tempfile_name):
      dist_init(rank, world_size, tempfile_name, backend="gloo")
  
      device = torch.device(rank)
      torch.manual_seed(
          rank
      )  # make sure that the different rank get different data
  
      # Setup two problems in parallel, we'll make sure that the second track (with save/load) follows the first one(untouched)
      # We split the model in two to test the multiple param groups support
      batch, input_width, hidden, target_width = 3, 20, 10, 5
      target = torch.rand((batch, target_width), device=device)
      inputs = torch.rand((batch, input_width), device=device)
  
      model_oss1 = torch.nn.Sequential(
          torch.nn.Linear(input_width, hidden),
          torch.nn.Linear(hidden, hidden),
      ).to(device)
      head_oss1 = torch.nn.Linear(hidden, target_width).to(device)
  
      model_oss2 = copy.deepcopy(model_oss1)
      head_oss2 = copy.deepcopy(head_oss1)
  
      # For this test the gradients are (all) reduced in the same way in between the torch reference and fairscale.
      # Normally OSS would use ShardedDDP and only reduce to the proper rank, but this does not change the
      # gradient norm computation from OSS and adds a dependency.
      # to keep the comparison apples-to-apples DDP is used in both cases
      model_oss1 = DDP(
          module=model_oss1,
          device_ids=[rank],
      )
      sharded_optimizer1 = optim.OSS(
          model_oss1.parameters(), lr=0.1, momentum=0.99
      )
      sharded_optimizer1.add_param_group(
          {"params": head_oss1.parameters()}
      )
  
      model_oss2 = DDP(
          module=model_oss2,
          device_ids=[rank],
      )
      sharded_optimizer2 = optim.OSS(
          model_oss2.parameters(), lr=0.1, momentum=0.99
      )
      sharded_optimizer2.add_param_group(
          {"params": head_oss2.parameters()}
      )
  
      loss_fn = torch.nn.L1Loss().to(device)
  
      def run_grad_step(model, head, optimizer):
          model.zero_grad()
          outputs = head(model(inputs))
  
      # pull the current state, broadcast it to all ranks
      sharded_optimizer2.consolidate_state_dict(
          recipient_rank=RECIPIENT_RANK
      )  # all ranks
      state_dict2 = (
          sharded_optimizer2.state_dict()
          if rank == RECIPIENT_RANK
          else {}
      )
      state_dict2 = sync_object_ranks(
          state_dict2, RECIPIENT_RANK, device
      )
  
      # re-create a new optimizer from scratch with absurd values, load the previous state
      sharded_optimizer2 = optim.OSS(
          model_oss2.parameters(), lr=1e6, momentum=0.0001
      )
      sharded_optimizer2.add_param_group(
          {"params": head_oss2.parameters()}
      )
      sharded_optimizer2.load_state_dict(state_dict2)
      check_same_model_params(
          model_oss1,
          model_oss2,
          "parameters of the two identical models have diverged (before any steps)",
      )
  
      # now take a step and check that parameters are equal
      run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
      run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
      check_same_model_params(
          model_oss1,
          model_oss2,
          "parameters of the two identical models have diverged (after stepping)",
      )
  
      # save the state dict for one model only, then distribute to the other ranks
      sharded_optimizer2.consolidate_state_dict(
          recipient_rank=RECIPIENT_RANK
      )  # all ranks
      state_dict2 = (
          sharded_optimizer2.state_dict()
          if rank == RECIPIENT_RANK
          else {}
      )
      state_dict2 = sync_object_ranks(
          state_dict2, RECIPIENT_RANK, device
      )
  
      # Check that the pulled state and the .param_groups attribute are in sync
      for replica in range(len(state_dict2["param_groups"])):
          for k in state_dict2["param_groups"][replica].keys():
              if k != "params":
                  assert (
                      state_dict2["param_groups"][replica][k]
                      == sharded_optimizer2.param_groups[0][k]
                  )
  
      # take a step
      run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
      run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
      check_same_model_params(
          model_oss1,
          model_oss2,
          "parameters of the two identical models have diverged (after consolidating)",
      )
  
      # save again for one rank, then distribute to the others
      sharded_optimizer2.consolidate_state_dict(
          recipient_rank=RECIPIENT_RANK
      )  # all ranks
      state_dict2 = (
          sharded_optimizer2.state_dict()
          if rank == RECIPIENT_RANK
          else {}
      )
      state_dict2 = sync_object_ranks(
          state_dict2, RECIPIENT_RANK, device
      )
  
      # reload the state_dict
      sharded_optimizer2 = optim.OSS(
          model_oss2.parameters(), lr=0.1, momentum=0.99
      )
      sharded_optimizer2.add_param_group(
          {"params": head_oss2.parameters()}
      )
      sharded_optimizer2.load_state_dict(state_dict2)
  
      # take a step
      run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
      run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
      check_same_model_params(
          model_oss1,
          model_oss2,
          "parameters of the two identical models have diverged (after reloading)",
      )
  
      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 23 for calling a global function with example usage: run_state_dict_distributed(rank, world_size, tempfile_name) without return types