docstr: Check that ZeroRedundancyOptimizer exposes the expected state dict
  interface, irrespective of the sharding.
docstr_tok: ["Check","that","ZeroRedundancyOptimizer","exposes","the","expected","state","dict","interface",",","irrespective","of","the","sharding","."]
code: def test_state_dict(self):
      """Check that ZeroRedundancyOptimizer exposes the expected state dict
      interface, irrespective of the sharding."""
      self.dist_init(self.rank)
      LR1 = 0.1
      LR2 = 0.01
      MOMENTUM = 0.9
      RECIPIENT_RANK = (
          0  # rank 0 is the only rank since the world size is 1
      )
      x = torch.tensor([1.0], device=self.device, requires_grad=True)
      o = ZeroRedundancyOptimizer(
          [x],
          optimizer_class=SGD,
          lr=LR1,
          momentum=MOMENTUM,
      )
      x.backward()
      o.step()
      self.assertEqual(x, torch.tensor([0.9], device=self.device))
      self.assertEqual(
          o.optim.state[x]["momentum_buffer"],
          torch.tensor([1.0], device=self.device),
      )
  
      o.zero_grad()
      o.consolidate_state_dict(to=RECIPIENT_RANK)
      state_dict = o.state_dict()
  
      # Check that the state dict has keys compliant with PyTorch
      self.assertIn("param_groups", state_dict.keys())
      self.assertIn("state", state_dict.keys())
  
      # Check that the state has the expected keys
      self.assertEqual(state_dict["param_groups"][0]["lr"], 0.1)
      self.assertEqual(state_dict["param_groups"][0]["momentum"], 0.9)
      self.assertFalse(state_dict["param_groups"][0]["nesterov"])
      self.assertEqual(
          state_dict["param_groups"][0]["weight_decay"], 0.0
      )
      self.assertEqual(state_dict["param_groups"][0]["dampening"], 0.0)
  
      # Check that the state and the `param_groups` attribute are in sync
      for k in state_dict["param_groups"][0]:
          if k != "params":
              self.assertEqual(
                  state_dict["param_groups"][0][k],
                  o.param_groups[0][k],
              )
  
      # Check that the state is reloaded with the correct values and device
      o = ZeroRedundancyOptimizer([x], optimizer_class=SGD, lr=LR2)
      o.load_state_dict(state_dict)
      self.assertEqual(
          o.optim.state[x]["momentum_buffer"],
          torch.tensor([1.0], device=self.device),
      )
  
      # We should we using `LR1` and not `LR2` after reloading, both within
      # the optimizer and as exposed by the `param_groups` attribute
      self.assertEqual(o.param_groups[0]["lr"], LR1)
      x.backward()
      o.step()
      self.assertEqual(x, torch.tensor([0.71], device=self.device))
      self.assertEqual(
          o.optim.state[x]["momentum_buffer"],
          torch.tensor([1.9], device=self.device),
      )
  
      # Check that the exposed `param_groups`` are on the proper device
      self.assertEqual(o.param_groups[0]["params"][0].device, x.device)
  
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 1 for calling the TestZeroRedundancyOptimizerSingleRank obj's underlying member method code with example usage: obj.test_state_dict() without return types