class_docstr: Tests ``fully_shard`` model checkpointing.
class_docstr_tok: ["Tests","``","fully_shard","``","model","checkpointing","."]
docstr: Tests that the full state dict saved from a module with ``fully_shard``
  applied to the global root matches that of an equivalent local module. Also
  ensure that this state_dict can be reloaded into a composable module and
  is equivalent to the original composable module.
docstr_tok: ["Tests","that","the","full","state","dict","saved","from","a","module","with","``","fully_shard","``","applied","to","the","global","root","matches","that","of","an","equivalent","local","module",".","Also","ensure","that","this","state_dict","can","be","reloaded","into","a","composable","module","and","is","equivalent","to","the","original","composable","module","."]
code: def test_state_dict_save_load_root_fully_shard(self):
      """
      Tests that the full state dict saved from a module with ``fully_shard``
      applied to the global root matches that of an equivalent local module. Also
      ensure that this state_dict can be reloaded into a composable module and
      is equivalent to the original composable module.
      """
      local_model = CompositeParamModel(device=torch.device("cuda"))
      save_composable = copy.deepcopy(local_model)
      fully_shard(
          save_composable, policy=ModuleWrapPolicy({UnitModule})
      )
      local_sd = local_model.state_dict()
      composable_sd = save_composable.state_dict()
      self._check_state_dict_parity(local_sd, composable_sd)
  
      # Validate load
      load_composable = fully_shard(
          copy.deepcopy(local_model),
          policy=ModuleWrapPolicy({UnitModule}),
      )
      _zero_model(load_composable, summon_full=False)
      for p in load_composable.parameters():
          self.assertEqual(p.sum(), 0)
  
      sd = {k: v.clone() for k, v in composable_sd.items()}
      load_composable.load_state_dict(sd)
      self._check_model_parity(load_composable, save_composable)
  
imports: import copy
  import itertools
  import sys
  from typing import Dict
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed._composable import fully_shard
  from torch.distributed._state_dict_utils import _gather_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel, StateDictType
  from torch.distributed.fsdp.api import ShardingStrategy
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
  from torch.testing._internal.common_dist_composable import CompositeParamModel, UnitModule
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import _zero_model, CUDAInitMode, FSDPInitMode, FSDPTest, TransformerWithSharedParams
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 2 for calling the TestModelCheckpointing obj's underlying member method code with example usage: obj.test_state_dict_save_load_root_fully_shard() without return types