docstr: Check that ZeroRedundancyOptimizer works with a non-default process
  group consisting only of even ranks.
docstr_tok: ["Check","that","ZeroRedundancyOptimizer","works","with","a","non-default","process","group","consisting","only","of","even","ranks","."]
code: def test_nondefault_process_group(self):
      """Check that ZeroRedundancyOptimizer works with a non-default process
      group consisting only of even ranks."""
      # Skip the test if below the minimum world size since then the test is
      # trivial
      MIN_WORLD_SIZE = 4
      if self.world_size < MIN_WORLD_SIZE:
          common_distributed.logger.info(
              "Skipping `test_nondefault_process_group()` since world size "
              "of %s is less than %s",
              self.world_size,
              MIN_WORLD_SIZE,
          )
          return
      BACKEND = dist.Backend.GLOO
      self.dist_init(self.rank, self.world_size, BACKEND)
      # Use GPU if enough are available, or fall back to CPU otherwise, which
      # is fine since Gloo backend supports both
      if (
          torch.cuda.is_available()
          and torch.cuda.device_count() >= self.world_size
      ):
          device = torch.device(self.rank)
      else:
          device = torch.device("cpu")
      # Create a new process group consisting of the even ranks to exercise
      # the case where the global and local ranks do not necessarily match
      subgroup_ranks = [r for r in range(self.world_size) if r % 2 == 0]
      process_group = dist.new_group(
          ranks=subgroup_ranks,
          backend=BACKEND,
      )
      # Ranks not participating in the new process group are no longer needed
      if self.rank not in subgroup_ranks:
          return
  
      # Set different seeds across ranks so that each rank gets different
      # training data and hence the model sync check is meaningful
      torch.manual_seed(self.rank)
      np.random.seed(self.rank)
  
      EPOCHS, BATCH_SIZE, INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM = (
          5,
          3,
          20,
          10,
          5,
      )
      LR = 1e-3
      MOMENTUM = 0.99
      REFERENCE_RANK = 0
      assert (
          REFERENCE_RANK in subgroup_ranks
      ), "Reference rank must be in the new process group"
      loss_fn = torch.nn.L1Loss().to(device)
  
      def check(optimizer):
          for _ in range(EPOCHS):
              target = torch.rand(
                  (BATCH_SIZE, OUTPUT_DIM), device=device
              )
              inputs = torch.rand(
                  (BATCH_SIZE, INPUT_DIM), device=device
              )
  
              def closure():
                  optimizer.zero_grad()
                  output = model(inputs)
                  loss = loss_fn(output, target)
                  loss /= self.world_size
                  loss.backward()
                  dist.all_reduce(loss, group=process_group)
                  return loss
  
              _ = optimizer.step(closure=closure)
  
              # Check that the parameters match across ranks after a step
              for pg in optimizer.param_groups:
                  for p in pg["params"]:
                      receptacle = (
                          [p.clone() for _ in subgroup_ranks]
                          if self.rank == REFERENCE_RANK
                          else []
                      )
                      dist.gather(
                          p,
                          receptacle,
                          dst=REFERENCE_RANK,
                          group=process_group,
                      )
                      if self.rank == REFERENCE_RANK:
                          reference_param = receptacle[0]
                          for param in receptacle[1:]:
                              torch.testing.assert_close(
                                  reference_param,
                                  param,
                                  msg="Models differ between ranks",
                              )
  
      model = torch.nn.Sequential(
          torch.nn.Linear(INPUT_DIM, HIDDEN_DIM),
          torch.nn.Linear(HIDDEN_DIM, OUTPUT_DIM),
      ).to(device)
      optimizer = ZeroRedundancyOptimizer(
          model.parameters(),
          optimizer_class=SGD,
          lr=LR,
          momentum=MOMENTUM,  # ensure there exists state to shard
          process_group=process_group,
      )
      check(optimizer)
  
imports: import copy
  import os
  import sys
  import unittest
  from contextlib import nullcontext
  from typing import Any, cast, List
  import numpy
  import torch
  import torch.distributed
  from torch.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook import hook_with_zero_step, hook_with_zero_step_interleaved
  from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import allreduce_hook
  from torch.distributed.algorithms.join import Join, Joinable, JoinHook
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from torch.distributed.optim.zero_redundancy_optimizer import _broadcast_object
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import AdamW, SGD
  from torch.testing._internal import common_distributed
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, IS_WINDOWS, parametrize, run_tests, TEST_WITH_ASAN, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 12 for calling the TestZeroRedundancyOptimizerDistributed obj's underlying member method code with example usage: obj.test_nondefault_process_group() and returns: loss