docstr: Check that ZeroRedundancyOptimizer properly handles adding a new
  parameter group a posteriori and that all ranks get a shard of the
  contained parameters.
  
  NOTE: The correctness of this test depends on the ZeRO implementation
  using the sorted-greedy partitioning algorithm. For details, see
  ``ZeroRedundancyOptimizer._partition_parameters()`` in
  zero_redundancy_optimizer.py.
docstr_tok: ["Check","that","ZeroRedundancyOptimizer","properly","handles","adding","a","new","parameter","group","a","posteriori","and","that","all","ranks","get","a","shard","of","the","contained","parameters",".","NOTE",":","The","correctness","of","this","test","depends","on","the","ZeRO","implementation","using","the","sorted-greedy","partitioning","algorithm",".","For","details",",","see","``","ZeroRedundancyOptimizer._partition_parameters","(",")","``","in","zero_redundancy_optimizer.py","."]
code: def test_add_param_group(self):
      """Check that ZeroRedundancyOptimizer properly handles adding a new
      parameter group a posteriori and that all ranks get a shard of the
      contained parameters.
  
      NOTE: The correctness of this test depends on the ZeRO implementation
      using the sorted-greedy partitioning algorithm. For details, see
      ``ZeroRedundancyOptimizer._partition_parameters()`` in
      zero_redundancy_optimizer.py.
      """
      self.dist_init(self.rank)
      LR = 0.01
  
      # Test with all parameters trainable to begin with
      def all_trainable():
          params = []
          sizes = [9, 7, 5, 3]
          sizes_world = sizes * self.world_size
          for size in sizes_world[:-1]:
              params.append(torch.rand(size, 1))
  
          # Make sure that the params are trainable so that they are factored
          # into the size-based parameter partitioning
          for p in params:
              p.requires_grad = True
  
          o = ZeroRedundancyOptimizer(
              params, optimizer_class=SGD, lr=LR
          )
          self.assertEqual(len(o.param_groups), 1)
          o.add_param_group({"params": [torch.rand(3, 1)]})
          # Verify that new group is added to the correct partition, making
          # all partitions have the same elements
          self.assertEqual(len(o.param_groups), 2)
          self.assertEqual(
              sum(
                  [
                      x.numel()
                      for g in o.optim.param_groups
                      for x in g["params"]
                  ]
              ),
              sum(sizes),
          )
          self.assertEqual(len(o.optim.param_groups), 2)
  
      # Test a pathological config with a first big non-trainable param
      def some_trainable():
          params = []
          for size in [100, 3, 5, 2, 6, 4]:
              params.append(torch.rand(size, 1))
  
          # Make sure that all but the first param are trainable so that they
          # are factored into the size-based parameter partitioning
          for p in params[1:]:
              p.requires_grad = True
  
          o = ZeroRedundancyOptimizer(
              params, optimizer_class=SGD, lr=LR
          )
          self.assertEqual(len(o.param_groups), 1)
          o.add_param_group({"params": [torch.rand(3, 1)]})
          self.assertEqual(len(o.param_groups), 2)
          self.assertEqual(len(o.optim.param_groups), 2)
  
      all_trainable()
      some_trainable()
  
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 9 for calling the TestZeroRedundancyOptimizerDistributed obj's underlying member method code with example usage: obj.test_add_param_group() without return types