code: def run_test_multiple_groups(rank, world_size, tempfile_name):
      # Only work with the even ranks, to check that the global_rank indexing is properly used
      dist_init(
          rank=rank,
          world_size=world_size,
          tempfile_name=tempfile_name,
          backend="gloo",
      )
      sub_group_ranks = [0, 2, 4]
      process_group = torch.distributed.new_group(
          ranks=sub_group_ranks, backend="gloo"
      )
  
      # Make sure that all the ranks get different training data
      # So that the sync check in between their models is meaningful
      torch.manual_seed(rank)
      np.random.seed(rank)
  
      # Standard deep learning setup
      device = "cpu"
      epochs, batch, input_width, hidden, target_width = 5, 3, 20, 10, 5
      loss_fn = torch.nn.L1Loss().to(device)
  
      def check(optimizer):
          # Just run a couple of epochs, check that the model is properly updated
          for _ in range(epochs):
              target = torch.rand((batch, target_width), device=device)
              inputs = torch.rand((batch, input_width), device=device)
  
              def closure():
                  optimizer.zero_grad()
                  output = model(inputs)
                  loss = loss_fn(output, target)
                  loss /= world_size
                  loss.backward()
                  dist.all_reduce(
                      loss, group=process_group
                  )  # Not strictly needed for the test below
  
                  return loss
  
              _ = optimizer.step(closure=closure)
  
              # Check that all the params are the same on all ranks
              for pg in optimizer.param_groups:
                  for p in pg["params"]:
                      receptacle = (
                          [p.clone() for _ in sub_group_ranks]
                          if rank == 0
                          else []
                      )
                      dist.gather(
                          p, receptacle, dst=0, group=process_group
                      )
                      if rank == 0:
                          for sync_p in receptacle[1:]:
                              assert torch.all(
                                  torch.eq(receptacle[0], sync_p)
                              ), "Models differ in between ranks {} - {}".format(
                                  torch.norm(receptacle[0]),
                                  torch.norm(sync_p),
                              )
  
      if rank in sub_group_ranks:
          # Model fitting in the broadcast bucket
          model = torch.nn.Sequential(
              torch.nn.Linear(input_width, hidden),
              torch.nn.Linear(hidden, target_width),
          ).to(device)
  
          # With SGD, Momentum is required to get a state to shard
          optimizer = optim.OSS(
              model.parameters(),
              lr=0.1,
              momentum=0.99,
              group=process_group,
              broadcast_buffer_size=2**20,
          )
          check(optimizer)
  
          # Model not-fitting in the broadcast bucket
          model = torch.nn.Sequential(
              torch.nn.Linear(input_width, hidden),
              torch.nn.Linear(hidden, target_width),
          ).to(device)
  
          # With SGD, Momentum is required to get a state to shard
          optimizer = optim.OSS(
              model.parameters(),
              lr=0.1,
              momentum=0.99,
              group=process_group,
              broadcast_buffer_size=0,
          )
          check(optimizer)
  
      dist.destroy_process_group(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 19 for calling a global function with example usage: run_test_multiple_groups(rank, world_size, tempfile_name) and returns: loss