code: def test_dtensor_sharded_tensor_state_dict_identical(
      self, offload_to_cpu
  ):
      mesh_2d = init_device_mesh(
          self.device_type, (2, self.world_size // 2)
      )
      model, optim = self._create_model(mesh_2d)
  
      FSDP.set_state_dict_type(
          model,
          StateDictType.SHARDED_STATE_DICT,
          state_dict_config=ShardedStateDictConfig(
              offload_to_cpu=offload_to_cpu
          ),
          optim_state_dict_config=ShardedOptimStateDictConfig(
              offload_to_cpu=offload_to_cpu
          ),
      )
      dtensor_sd = model.state_dict()
      dtensor_osd = FSDP.optim_state_dict(model, optim)
  
      ref_model, ref_optim = self._create_model()
      FSDP.set_state_dict_type(
          ref_model,
          StateDictType.SHARDED_STATE_DICT,
          state_dict_config=ShardedStateDictConfig(
              offload_to_cpu=offload_to_cpu
          ),
          optim_state_dict_config=ShardedOptimStateDictConfig(
              offload_to_cpu=offload_to_cpu
          ),
      )
      sharded_tensor_sd = ref_model.state_dict()
      sharded_tensor_osd = FSDP.optim_state_dict(ref_model, ref_optim)
  
      # Check dtensor and sharded_tensor model state dict values are identical
      for dtensor_sd_item, sharded_tensor_sd_item in zip(
          dtensor_sd.items(), sharded_tensor_sd.items()
      ):
          k1, v1 = dtensor_sd_item
          k2, v2 = sharded_tensor_sd_item
          self.assertEqual(k1, k2)
  
          self.assertEqual(type(v1), DTensor)
          self.assertEqual(type(v2), ShardedTensor)
          # check whether local_tensor are the same
          self.assertEqual(v1.to_local(), v2.local_tensor())
          # check whether device are the same
          self.assertEqual(
              v1.to_local().device, v2.local_tensor().device
          )
  
      # Check dtensor and sharde_tensor optim state dict values are identical
      for dtensor_osd_state, sharded_tensor_osd_state in zip(
          dtensor_osd["state"].items(),
          sharded_tensor_osd["state"].items(),
      ):
          # check FQN are the same
          self.assertEqual(
              dtensor_osd_state[0], sharded_tensor_osd_state[0]
          )
          for dtensor_hyper_param, sharded_tensor_hyper_param in zip(
              dtensor_osd_state[1].items(),
              sharded_tensor_osd_state[1].items(),
          ):
              k1, v1 = dtensor_hyper_param
              k2, v2 = sharded_tensor_hyper_param
              self.assertEqual(k1, k2)
  
              if k1 != "step":
                  self.assertEqual(type(v1), DTensor)
                  self.assertEqual(type(v2), ShardedTensor)
                  # check whether local_tensor are the same
                  self.assertEqual(v1.to_local(), v2.local_tensor())
                  # check whether device are the same
                  self.assertEqual(
                      v1.to_local().device, v2.local_tensor().device
                  )
              else:
                  self.assertEqual(v1, v2)
  
imports: import io
  from copy import deepcopy
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._tensor import DTensor, Replicate, Shard
  from torch.distributed.device_mesh import _mesh_resources, init_device_mesh
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.api import ShardedOptimStateDictConfig, ShardedStateDictConfig, ShardingStrategy, StateDictType
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, skip_if_lt_x_gpu, with_comms
label_desc: Use image node_id 4 for calling the TestHSDPWithDeviceMeshAndDTensor obj's underlying member method code with example usage: obj.test_dtensor_sharded_tensor_state_dict_identical(offload_to_cpu) without return types