code: def _test_fsdp_model_state(self, process_group) -> None:
      CHECKPOINT_DIR = self.temp_dir
  
      model = FSDP(torch.nn.Linear(8, 8, device="meta"))
      model(torch.rand(8, 8, device=dist.get_rank())).sum().backward()
  
      with FSDP.state_dict_type(
          model, StateDictType.SHARDED_STATE_DICT
      ):
          state_dict = {
              "model": model.state_dict(),
          }
  
          dist_cp.save_state_dict(
              state_dict=state_dict,
              storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
              planner=DefaultSavePlanner(),
          )
  
      model_2 = FSDP(
          torch.nn.Linear(8, 8, device="meta"),
          process_group=process_group,
      )
  
      with FSDP.summon_full_params(model):
          with FSDP.summon_full_params(model_2):
              self.assertNotEqual(model.weight, model_2.weight)
              self.assertNotEqual(model.bias, model_2.bias)
  
      # now load the model and ensure the values are the same
      with FSDP.state_dict_type(
          model_2, StateDictType.SHARDED_STATE_DICT
      ):
          state_dict = {
              "model": model_2.state_dict(),
          }
  
          dist_cp.load_state_dict(
              state_dict=state_dict,
              storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
              planner=DefaultLoadPlanner(),
          )
          model_2.load_state_dict(state_dict["model"])
  
      with FSDP.summon_full_params(model):
          with FSDP.summon_full_params(model_2):
              self.assertEqual(model.weight, model_2.weight)
              self.assertEqual(model.bias, model_2.bias)
  
file_path: pytorch/test/distributed/checkpoint/test_fsdp_model_state.py
imports: import torch
  import torch.distributed
  import torch.distributed.checkpoint
  from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import 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 FsdpModelStateCheckpoint obj's underlying member method code with example usage: obj._test_fsdp_model_state(process_group) without return types