code: def test_hsdp_checkpoint(self, is_even_sharded_model) -> None:
      CHECKPOINT_DIR = self.temp_dir
      simple_model = (
          SimpleModel if is_even_sharded_model else SimpleModelUneven
      )
  
      mesh_2d = init_device_mesh(
          self.device_type, (2, self.world_size // 2)
      )
      model = FSDP(
          simple_model().cuda(),
          sharding_strategy=ShardingStrategy.HYBRID_SHARD,
          device_mesh=mesh_2d,
      )
      optim = torch.optim.Adam(model.parameters(), lr=0.1)
  
      FSDP.set_state_dict_type(
          model,
          StateDictType.SHARDED_STATE_DICT,
      )
      state_dict = {"model": model.state_dict()}
      state_dict_to_save = deepcopy(state_dict)
  
      dist_cp.save_state_dict(
          state_dict=state_dict_to_save,
          storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
          planner=DefaultSavePlanner(),
      )
  
      # Update the parameters so current model state_dict now be different from state_dict_to_save.
      model(model.get_input()).sum().backward()
      optim.step()
  
      # At this point, the current state dict is different from state_dict_to_save.
      for (k1, v1), (k2, v2) in zip(
          state_dict_to_save["model"].items(),
          model.state_dict().items(),
      ):
          self.assertEqual(k1, k2)
          self.assertEqual(v1.device_mesh, v2.device_mesh)
          self.assertEqual(v1.placements, v2.placements)
          self.assertNotEqual(v1.to_local(), v2.to_local())
  
      dist_cp.load_state_dict(
          state_dict=state_dict_to_save,
          storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
          planner=DefaultLoadPlanner(),
      )
      model.load_state_dict(state_dict_to_save["model"])
  
      state_dict_after_load = model.state_dict()
      # After loading, the current model state dict should be the same as state_dict_to_save.
      for (k1, v1), (k2, v2) in zip(
          state_dict_to_save["model"].items(),
          model.state_dict().items(),
      ):
          self.assertEqual(k1, k2)
          self.assertEqual(v1.device_mesh, v2.device_mesh)
          self.assertEqual(v1.placements, v2.placements)
          self.assertEqual(v1.to_local(), v2.to_local())
  
imports: from copy import deepcopy
  import torch
  import torch.distributed.checkpoint
  import torch.nn
  import torch.nn.functional
  from torch.distributed._tensor import init_device_mesh, Replicate
  from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy, StateDictType
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
  from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
label_desc: Use image node_id 2 for calling the TestHSDPCheckpoint obj's underlying member method code with example usage: obj.test_hsdp_checkpoint(is_even_sharded_model) without return types