docstr: Check ZeroRedundancyOptimizer's parameter sharding at construction
  time.
  
  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","ZeroRedundancyOptimizer","'s","parameter","sharding","at","construction","time",".","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_sharding(self):
      """
      Check ZeroRedundancyOptimizer's parameter sharding at construction
      time.
  
      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
      sizes = [9, 7, 5, 3]
      params = []
      for size in sizes * self.world_size:
          params.append(torch.rand(size, 1))
      o = ZeroRedundancyOptimizer(params, optimizer_class=SGD, lr=LR)
      self.assertEqual(
          sum([x.numel() for x in o.optim.param_groups[0]["params"]]),
          sum(sizes),
      )
  
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 8 for calling the TestZeroRedundancyOptimizerDistributed obj's underlying member method code with example usage: obj.test_sharding() without return types