code: def test_hsdp_fsdp_checkpoint_conversion(
      self, is_even_sharded_model
  ) -> None:
      CHECKPOINT_DIR = self.temp_dir
      simple_model = (
          SimpleModel if is_even_sharded_model else SimpleModelUneven
      )
  
      # save the hsdp model state_dict
      mesh_2d = init_device_mesh(
          self.device_type, (2, self.world_size // 2)
      )
      hsdp_model = FSDP(
          simple_model().cuda(),
          sharding_strategy=ShardingStrategy.HYBRID_SHARD,
          device_mesh=mesh_2d,
      )
      FSDP.set_state_dict_type(
          hsdp_model,
          StateDictType.SHARDED_STATE_DICT,
      )
      hsdp_state_dict = {"model": hsdp_model.state_dict()}
      dist_cp.save_state_dict(
          state_dict=hsdp_state_dict,
          storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
          planner=DefaultSavePlanner(),
      )
  
      # initialize a fsdp model to load checkpoint into
      mesh_1d = init_device_mesh(self.device_type, (self.world_size,))
      fsdp_model = FSDP(
          simple_model().cuda(),
          device_mesh=mesh_1d,
      )
      FSDP.set_state_dict_type(
          fsdp_model,
          StateDictType.SHARDED_STATE_DICT,
      )
      fsdp_state_dict = {"model": fsdp_model.state_dict()}
  
      # at this point, the hsdp model parameters are different from fsdp model parameters.
      for (k1, v1), (k2, v2) in zip(
          hsdp_state_dict["model"].items(),
          fsdp_state_dict["model"].items(),
      ):
          self.assertEqual(k1, k2)
          self.assertNotEqual(v1.device_mesh, v2.device_mesh)
          self.assertNotEqual(v1.placements, v2.placements)
          v1_all_gather = v1.redistribute(
              mesh_2d, placements=(Replicate(), Replicate())
          )
          v2_all_gather = v2.redistribute(
              mesh_1d, placements=(Replicate(),)
          )
          self.assertNotEqual(
              v1_all_gather.to_local(), v2_all_gather.to_local()
          )
  
      # load the fsdp state_dict from storage
      dist_cp.load_state_dict(
          state_dict=fsdp_state_dict,
          storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
          planner=DefaultLoadPlanner(),
      )
      fsdp_model.load_state_dict(fsdp_state_dict["model"])
  
      state_dict_after_load = fsdp_model.state_dict()
      # After loading, the current model state dict should be the same as hsdp_state_dict.
      for (k1, v1), (k2, v2) in zip(
          hsdp_state_dict["model"].items(),
          state_dict_after_load.items(),
      ):
          self.assertEqual(k1, k2)
          self.assertNotEqual(v1.device_mesh, v2.device_mesh)
          self.assertNotEqual(v1.placements, v2.placements)
          v1_all_gather = v1.redistribute(
              mesh_2d, placements=(Replicate(), Replicate())
          )
          v2_all_gather = v2.redistribute(
              mesh_1d, placements=(Replicate(),)
          )
          self.assertEqual(
              v1_all_gather.to_local(), v2_all_gather.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 3 for calling the TestHSDPCheckpoint obj's underlying member method code with example usage: obj.test_hsdp_fsdp_checkpoint_conversion(is_even_sharded_model) without return types