class_docstr: Tests ``fully_shard`` model checkpointing.
class_docstr_tok: ["Tests","``","fully_shard","``","model","checkpointing","."]
docstr: Checks that ``local_sd`` and ``composable_sd`` are the same.
docstr_tok: ["Checks","that","``","local_sd","``","and","``","composable_sd","``","are","the","same","."]
code: def _check_state_dict_parity(
      self, local_sd: Dict, composable_sd: Dict
  ):
      """Checks that ``local_sd`` and ``composable_sd`` are the same."""
      # Check that all keys match
      self.assertEqual(set(composable_sd.keys()), set(local_sd.keys()))
      # Check value shapes
      for k in composable_sd.keys():
          v1 = composable_sd[k]
          v2 = local_sd[k]
          self.assertEqual(
              v1.shape,
              v2.shape,
              f"Shape mismatch for {k} {v1.shape} vs {v2.shape}",
          )
      # Check actual values
      for k in composable_sd.keys():
          v1 = composable_sd[k]
          v2 = local_sd[k]
          self.assertEqual(
              v1, v2, f"Param mismatch for {k}: {v1} vs {v2}"
          )
  
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 7 for calling the TestModelCheckpointing obj's underlying member method code with example usage: obj._check_state_dict_parity(local_sd, composable_sd) without return types