class_docstr: Tests ``fully_shard`` model checkpointing.
class_docstr_tok: ["Tests","``","fully_shard","``","model","checkpointing","."]
docstr: E2E test of save + load with rank0_only + CPU offload for TransformerWithSharedParams
  on the composable path.
docstr_tok: ["E2E","test","of","save","+","load","with","rank0_only","+","CPU","offload","for","TransformerWithSharedParams","on","the","composable","path","."]
code: def test_state_dict_save_load_flow(self):
      """
      E2E test of save + load with rank0_only + CPU offload for TransformerWithSharedParams
      on the composable path.
      """
      self.run_subtests(
          {
              "ignore_modules": [False, True],
              "sharded_state_dict": [False, True],
          },
          self._test_save_dict_save_load_flow,
      )
  
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 4 for calling the TestModelCheckpointing obj's underlying member method code with example usage: obj.test_state_dict_save_load_flow() without return types