code: def run_test_multiple_groups(
      rank, world_size, tempfile_name, backend, reduce_buffer_size
  ):
      # Only work with the even ranks, to check that the global_rank indexing is properly used
      dist.init_process_group(
          init_method="file://" + tempfile_name,
          backend=backend,
          rank=rank,
          world_size=world_size,
      )
  
      sub_group_ranks = [0, 2]
      process_group = torch.distributed.new_group(
          ranks=sub_group_ranks, backend=backend
      )
  
      # 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 = "cuda"
      torch.cuda.set_device(rank)
  
      epochs, batch, input_width, hidden, target_width = 5, 3, 20, 10, 5
      loss_fn = torch.nn.L1Loss().to(device)
  
      def check(optimizer, model):
          # 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.backward()
                  return loss
  
              _ = optimizer.step(closure=closure)
  
              # Check that all the params are the same on all ranks
              check_same_models_across_ranks(
                  model,
                  process_group,
                  params_should_be_equal=True,
                  check_broadcast_buffers=True,
              )
  
      if rank in sub_group_ranks:
          # 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 = OSS(
              model.parameters(),
              group=process_group,
              lr=1e-3,
              momentum=0.99,
          )
          model = ShardedDataParallel(
              model,
              optimizer,
              process_group=process_group,
              reduce_buffer_size=reduce_buffer_size,
          )
          check(optimizer, model)
  
      dist.destroy_process_group(process_group)
  
imports: from contextlib import suppress
  import numpy
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torch.nn import Linear, Sequential
  from fairscale.fair_dev.testing.testing import GPT2, SGDWithPausingCompute, available_devices, check_same_models_across_ranks, skip_if_less_than_four_gpu, skip_if_no_cuda, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn.data_parallel import ShardedDataParallel
  from fairscale.optim import OSS
label_desc: Use image node_id 23 for calling a global function with example usage: run_test_multiple_groups(rank, world_size, tempfile_name, backend, reduce_buffer_size) and returns: loss