docstr: FSDP.gather_full_optim_state_dict() should return something very similar to optimizer.state_dict()
docstr_tok: ["FSDP.gather_full_optim_state_dict","(",")","should","return","something","very","similar","to","optimizer.state_dict","(",")"]
code: def _test_consolidated_optimizer(
      self,
      config,
      rank,
      group,
      optim_fn=torch.optim.SGD,
      transformer=False,
      expert_group=None,
  ):
      """FSDP.gather_full_optim_state_dict() should return something very similar to optimizer.state_dict()"""
      # Establish reference behavior.
      if transformer:
          unwrapped_model = TransformerWithSharedParams(
              group, wrapper_config=config
          ).cuda()
          fsdp = self.get_wrapped_model(group, config=config).cuda()
      else:
          unwrapped_model = MixtureOfExperts(
              group, wrapper_config=None, expert_group=expert_group
          ).cuda()
          fsdp = FullyShardedDataParallel(
              MixtureOfExperts(
                  group,
                  wrapper_config=config,
                  expert_group=expert_group,
              )
          ).cuda()
  
      try:
          fsdp_optim = optim_fn(
              fsdp.parameters(),
              lr=0.01,
          )
          optim_unwrapped = optim_fn(
              unwrapped_model.parameters(), lr=0.01
          )
      except TypeError:  # Adadelta
          fsdp_optim = optim_fn(fsdp.parameters())
          optim_unwrapped = optim_fn(unwrapped_model.parameters())
  
      fsdp_optim.zero_grad()
      optim_unwrapped.zero_grad()
      with torch.cuda.amp.autocast(enabled=True):
          x = fsdp.module.get_input(torch.device("cuda"))
          output = fsdp(*x)
          loss = fsdp.module.get_loss(x, output).to("cuda")
          fsdp.module.run_backward(loss)
          fsdp_optim.step()
  
          output = unwrapped_model(*x)
          loss = unwrapped_model.get_loss(x, output)
          unwrapped_model.run_backward(loss)
          optim_unwrapped.step()
      unwrapped_sd = optim_unwrapped.state_dict()
  
      if not transformer and not expert_group:
          no_broadcast_children = [
              x
              for x in get_fsdp_instances(fsdp)
              if x.no_broadcast_optim_state
          ]
          assert (
              len(no_broadcast_children) == 1
          ), f"Length of non shared params {len(no_broadcast_children)}"
          assert get_fsdp_instances(fsdp)[-1].no_broadcast_optim_state
      torch.cuda.empty_cache()
      cuda_gb_before = (
          torch.cuda.memory_stats(fsdp.rank)[
              "allocated_bytes.all.current"
          ]
          / 1024**3
      )
      tstart = time()
      sd = fsdp.gather_full_optim_state_dict(
          fsdp_optim, recipient_rank=0
      )
      duration = time() - tstart
      assert (
          duration < fsdp.world_size
      ), f"gather optim state took {duration} seconds, suspect change in _consolidate"
  
      cuda_gb_after = (
          torch.cuda.memory_stats(fsdp.rank)[
              "allocated_bytes.all.current"
          ]
          / 1024**3
      )
      mem_usg_gb = cuda_gb_after - cuda_gb_before
      assert (
          mem_usg_gb == 0
      ), f"gather_full_optim_state_dict used {mem_usg_gb:.2f} CUDA GB, max allowed is 0"
      assert cuda_gb_after > 0, "got 0 memory usage, logging is broken"
  
      if fsdp.rank > 0:
          assert sd is None
          return
  
      # assert whole state dict on CPU
      for k, v in sd["state"].items():
          for buffer_name, t in v.items():
              if torch.is_tensor(t):
                  msg = f"got device {t.device} for {k}: {buffer_name}. expected CPU"
                  assert t.device == torch.device("cpu"), msg
  
      if expert_group:
          sd_state = recursive_copy_to_device(
              sd["state"], non_blocking=False, device="cpu"
          )
          orig_state = recursive_copy_to_device(
              unwrapped_sd["state"], non_blocking=False, device="cpu"
          )
  
          assert_equal(len(sd_state.keys()), len(orig_state.keys()))
  
          assert_equal(
              sum(
                  [
                      all_tensors_numel_except_for_step(v)
                      for k, v in sd_state.items()
                  ]
              ),
              sum(
                  [
                      all_tensors_numel_except_for_step(v)
                      for k, v in orig_state.items()
                  ]
              ),
          )
          return
  
      assert "uncollected_local_ids" in sd
      sd_copy = copy.deepcopy(sd)
      unflat_state = sd_copy["state"]
      shard_sd = fsdp.get_shard_from_optim_state_dict(sd_copy)
      shard_sd = recursive_copy_to_device(
          shard_sd, non_blocking=False, device="cpu"
      )
      state_after_get_shard = sd_copy["state"]
      # sd is changed in-place in case there are extra states.
      assert not objects_are_equal(unflat_state, state_after_get_shard)
      del sd_copy
  
      assert_equal(len(sd["state"]), len(unwrapped_sd["state"]))
      assert_equal(
          len(sd["param_groups"][0]["params"]),
          len(unwrapped_sd["param_groups"][0]["params"]),
      )
      assert_equal(
          sum(
              [
                  all_tensors_numel_except_for_step(v)
                  for k, v in sd["state"].items()
              ]
          ),
          sum(
              [
                  all_tensors_numel_except_for_step(v)
                  for k, v in unwrapped_sd["state"].items()
              ]
          ),
      )
  
      original_shard_sd = fsdp_optim.state_dict()
      assert_equal(
          len(shard_sd["state"]), len(original_shard_sd["state"])
      )
      assert_equal(shard_sd.keys(), original_shard_sd.keys())
      original_shard_sd = recursive_copy_to_device(
          original_shard_sd, non_blocking=False, device="cpu"
      )
      # Before asserting that the dicts are equal, we check keys individually to allow nice tracebacks.
      assert_equal(
          [
              all_tensors_numel_except_for_step(v)
              for k, v in shard_sd["state"].items()
          ],
          [
              all_tensors_numel_except_for_step(v)
              for k, v in original_shard_sd["state"].items()
          ],
      )
      assert_equal(
          [v for k, v in shard_sd["param_groups"][0].items()],
          [v for k, v in original_shard_sd["param_groups"][0].items()],
      )
      objects_are_equal(
          shard_sd["state"],
          original_shard_sd["state"],
          raise_exception=True,
      )
      objects_are_equal(
          {k: shard_sd[k] for k in original_shard_sd},
          original_shard_sd,
          raise_exception=True,
      )
  
imports: import copy
  import functools
  from time import time
  import unittest
  from parameterized import parameterized
  import torch
  from torch import nn
  from torch.optim import SGD, Adadelta, Adam
  from fairscale.fair_dev.testing.testing import dist_init, objects_are_equal, spawn_for_all_world_sizes
  from fairscale.internal.params import recursive_copy_to_device
  from fairscale.nn.data_parallel import FullyShardedDataParallel, get_fsdp_instances
  from fairscale.nn.data_parallel.fsdp_optim_utils import is_singleton_tensor
  from .test_fsdp import DistributedTest, DummyProcessGroup, MixtureOfExperts, TransformerWithSharedParams, rename_test, spawn_and_init
label_desc: Use image node_id 3 for calling the TestOptimizerUtils obj's underlying member method code with example usage: obj._test_consolidated_optimizer(config, rank, group, optim_fn, transformer, expert_group) without return types