class_docstr: Tests ``fully_shard`` model checkpointing.
class_docstr_tok: ["Tests","``","fully_shard","``","model","checkpointing","."]
code: def _test_save_dict_save_load_flow(
      self, ignore_modules: bool, sharded_state_dict: bool
  ):
      local_model = TransformerWithSharedParams.init(
          self.process_group,
          FSDPInitMode.NO_FSDP,
          CUDAInitMode.CUDA_BEFORE,
          deterministic=True,
      )
  
      # Force model parameters and buffers to be nonzero
      for tensor in itertools.chain(
          local_model.parameters(), local_model.buffers()
      ):
          if torch.count_nonzero(tensor) == 0:
              with torch.no_grad():
                  tensor.add_(torch.ones_like(tensor))
  
      save_model = copy.deepcopy(local_model)
      fully_shard(
          save_model,
          policy=ModuleWrapPolicy(
              {TransformerEncoderLayer, TransformerDecoderLayer}
          ),
          ignored_modules=(
              save_model.get_ignored_modules() if ignore_modules else []
          ),
      )
  
      # TODO: test state_dict_type after https://github.com/pytorch/pytorch/issues/90954 is resolved
      if not sharded_state_dict:
          FSDP.set_state_dict_type(
              save_model, StateDictType.FULL_STATE_DICT
          )
      else:
          FSDP.set_state_dict_type(
              save_model, StateDictType.SHARDED_STATE_DICT
          )
      state_dict = save_model.state_dict()
      local_state_dict = local_model.state_dict()
      self._check_state_dict_parity(
          local_state_dict, _gather_state_dict(state_dict)
      )
  
      load_model = TransformerWithSharedParams.init(
          self.process_group,
          FSDPInitMode.NO_FSDP,
          CUDAInitMode.CUDA_BEFORE,
      )
      _zero_model(load_model, zero_buffers=True, summon_full=False)
      fully_shard(
          load_model,
          policy=ModuleWrapPolicy(
              {TransformerDecoderLayer, TransformerEncoderLayer}
          ),
          ignored_modules=(
              load_model.get_ignored_modules() if ignore_modules else []
          ),
      )
      if not sharded_state_dict:
          FSDP.set_state_dict_type(
              load_model, StateDictType.FULL_STATE_DICT
          )
      else:
          FSDP.set_state_dict_type(
              load_model, StateDictType.SHARDED_STATE_DICT
          )
      load_model.load_state_dict(state_dict)
      self._check_model_parity(load_model, save_model)
  
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 5 for calling the TestModelCheckpointing obj's underlying member method code with example usage: obj._test_save_dict_save_load_flow(ignore_modules, sharded_state_dict) without return types