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_consolidate_weights(
      self, config, rank, group, paths=None, transformer=False
  ):
      """FSDP.gather_full_optim_state_dict() should return something very similar to optimizer.state_dict()"""
      # Establish reference behavior.
  
      if transformer:
          fsdp = self.get_wrapped_model(group, config=config).cuda()
      else:
          fsdp = FullyShardedDataParallel(
              MixtureOfExperts(group, wrapper_config=config)
          ).cuda()
  
      optim = Adam(
          fsdp.parameters(),
          lr=0.01,
      )
      optim.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)
          optim.step()
  
      # each worker saves a checkpoint with local_state_dict
      cp_data = {
          "weights": {
              k: v.cpu() for k, v in fsdp.local_state_dict().items()
          },
          "meta": fsdp.local_metadata_dict(),
      }
      torch.save(cp_data, paths[fsdp.rank])
      full_model_state_dict = fsdp.state_dict()
      torch.distributed.barrier()
      if fsdp.rank > 0:
          return
      all_checkpoints = [torch.load(p) for p in paths]
      consolidated_checkpoint = (
          FullyShardedDataParallel.consolidate_shard_weights(
              shard_weights=[c["weights"] for c in all_checkpoints],
              shard_metadata=[c["meta"] for c in all_checkpoints],
          )
      )
      full_model_extra = set(full_model_state_dict).difference(
          set(consolidated_checkpoint)
      )
      consolidated_extra = set(consolidated_checkpoint).difference(
          set(full_model_state_dict)
      )
      msg = f"full model extra keys: {full_model_extra}, consolidated extra {consolidated_extra}"
      for k in full_model_state_dict.keys():
          assert (
              consolidated_checkpoint[k].shape
              == full_model_state_dict[k].shape
          )
      assert set(full_model_state_dict.keys()) == set(
          consolidated_checkpoint.keys()
      ), msg
  
imports: import functools
  import os
  import tempfile
  from parameterized import parameterized
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.optim import Adam
  from fairscale.fair_dev.testing.testing import in_temporary_directory, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn import FullyShardedDataParallel
  from tests.nn.data_parallel.test_fsdp import DistributedTest, MixtureOfExperts, rename_test, spawn_and_init
label_desc: Use image node_id 2 for calling the TestConsolidatedWeights obj's underlying member method code with example usage: obj._test_consolidate_weights(config, rank, group, paths, transformer) without return types