code: def run_ddp_parity(
      rank,
      world_size,
      backend,
      temp_file_name,
      change_train_graph,
      broadcast_fp16,
  ):
      url = "file://" + temp_file_name
      dist.init_process_group(
          init_method=url,
          backend=backend,
          rank=rank,
          world_size=world_size,
      )
  
      device = torch.device("cuda")
      torch.cuda.set_device(rank)
      torch.manual_seed(rank)
      np.random.seed(rank)
      hidden = 5
      in_channels = 3
      out_channels = 3
      batch = 64
  
      def check_optimizer_equivalence(
          optimizer: Type[torch.optim.Optimizer],
          change_train_graph: bool = False,
      ):
          # Any model works. Add one different buffer per rank
          trunk = torch.nn.Sequential(
              torch.nn.Linear(in_channels, hidden),
              torch.nn.Linear(hidden, hidden),
              torch.nn.Linear(hidden, hidden),
          )
          trunk.register_buffer("test_buffer", torch.ones(1) * rank)
          trunk.to(device)
  
          head = torch.nn.Linear(hidden, out_channels).to(device)
  
          # Define a model to be trained by OSS
          oss_module = torch.nn.Sequential(trunk, head)
  
          # Make sure that the param groups are interleaved, to catch an ordering bug in the state dict
          oss_trainable_params = [
              {
                  "params": list(trunk.parameters())[:-1]
                  + list(head.parameters()),
                  "lr": 1e-5,
              },
              {"params": list(trunk.parameters())[-1], "lr": 1e-4},
          ]
  
          optimizer_settings: Dict[Any, Any] = {}
          if isinstance(optimizer, torch.optim.SGD):
              optimizer_settings["momentum"] = 0.9
  
          sharded_optimizer = optim.OSS(
              params=oss_trainable_params,
              optim=optimizer,
              group=None,
              broadcast_buffer_size=2**10,
              **optimizer_settings,
          )
  
          oss_ddp_model = DDP(
              module=oss_module,
              device_ids=[rank],
              broadcast_buffers=True,
              find_unused_parameters=True,
          )
  
          # Define a model to be trained by normal pytorch + DDP
          ddp_trunk = copy.deepcopy(trunk)
          ddp_head = copy.deepcopy(head)
          ddp_module = torch.nn.Sequential(ddp_trunk, ddp_head)
  
          ddp_trainable_params = [
              {
                  "params": list(ddp_trunk.parameters())[:-1]
                  + list(ddp_head.parameters()),
                  "lr": 1e-5,
              },
              {"params": list(ddp_trunk.parameters())[-1], "lr": 1e-4},
          ]
          ddp_optimizer = optimizer(ddp_trainable_params, **optimizer_settings)  # type: ignore
          ddp_model = DDP(
              module=ddp_module,
              device_ids=[rank],
              broadcast_buffers=True,
              find_unused_parameters=True,
          )
  
          def check_step():
              input_tensor = torch.rand((batch, in_channels)).to(device)
  
              def closure_ddp(input_tensor=input_tensor):
                  ddp_optimizer.zero_grad()
                  ddp_loss = ddp_model(input_tensor).abs().sum()
                  ddp_loss.backward()
                  return ddp_loss
  
              def closure_sharded(input_tensor=input_tensor):
                  sharded_optimizer.zero_grad()
                  sharded_loss = oss_ddp_model(input_tensor).abs().sum()
                  sharded_loss.backward()
                  return sharded_loss
  
              loss_ddp = cast(
                  torch.Tensor, ddp_optimizer.step(closure=closure_ddp)
              )
              loss_sharded_optim = cast(
                  torch.Tensor,
                  sharded_optimizer.step(closure=closure_sharded),
              )
  
              assert torch.allclose(
                  loss_ddp, loss_sharded_optim, rtol=1e-3
              ), f"Losses differ in between Pytorch optim and OSS\n {loss_ddp.item()} - {loss_sharded_optim.item()} - world size {world_size}"
  
              check_same_model_params(oss_ddp_model, ddp_model)
  
          # The model should be synchronized in between the ranks at construction time, check that
          check_same_model_params(oss_ddp_model, ddp_model)
  
          # The models should stay the same in between ddp and sharded optimizer
          for i in range(5):
              check_step()
  
              # Check that altering the trainable parameters does not cause DDP and OSS to diverge
              if change_train_graph:
                  # Flip the first parameter from trainable to non-trainable and vice-versa
                  next(
                      ddp_module.parameters()
                  ).requires_grad = not next(
                      ddp_module.parameters()
                  ).requires_grad
                  next(
                      oss_module.parameters()
                  ).requires_grad = not next(
                      oss_module.parameters()
                  ).requires_grad
                  # sharded_optimizer.refresh_trainable()
  
          # Check that the checkpoints are compatible (post pytorch 1.5)
          if torch_version()[1] > 5:
              # - get states
              ddp_state_dict = ddp_optimizer.state_dict()
              sharded_optimizer.consolidate_state_dict(
                  recipient_rank=RECIPIENT_RANK
              )
              sharded_optim_state_dict = (
                  sharded_optimizer.state_dict()
                  if rank == RECIPIENT_RANK
                  else {}
              )
              sharded_optim_state_dict = sync_object_ranks(
                  sharded_optim_state_dict, RECIPIENT_RANK, device
              )
  
              # - cross load the states
              # run one step and check that the models are still the same
              ddp_state_dict_ref = copy.deepcopy(
                  ddp_state_dict
              )  # OSS will remove some states
              ddp_optimizer.load_state_dict(
                  sharded_optim_state_dict
              )  # mixup on purpose !
              sharded_optimizer.load_state_dict(ddp_state_dict)
              check_step()
  
              #  - self load, rewind, check no problem
              # run one step and check that the models are still the same
              ddp_optimizer.load_state_dict(ddp_state_dict_ref)
              sharded_optimizer.load_state_dict(
                  sharded_optim_state_dict
              )
              check_step()
  
      for opt in [torch.optim.Adam, torch.optim.SGD]:
          check_optimizer_equivalence(
              opt, change_train_graph=change_train_graph
          )
  
      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 25 for calling a global function with example usage: run_ddp_parity(rank, world_size, backend, temp_file_name, change_train_graph, broadcast_fp16) and returns: ddp_loss, sharded_loss