code: def test_2d_optim_state_dict(self, is_even_sharded_model):
      simple_model = (
          SimpleModel if is_even_sharded_model else SimpleModelUneven
      )
  
      # Create a model without wrapper
      torch.manual_seed(0)
      no_wrap_model = simple_model().cuda(self.rank)
      no_wrap_state_dict = no_wrap_model.state_dict()
      no_wrap_optim = torch.optim.Adam(
          no_wrap_model.parameters(), lr=0.01
      )
      no_wrap_model(
          no_wrap_model.get_input().cuda(self.rank)
      ).sum().backward()
      no_wrap_optim.step()
      no_wrap_osd = get_optimizer_state_dict(
          no_wrap_model, optimizers=no_wrap_optim
      )
  
      # Create a model and sharded it with 2D FSDP + TP
      torch.manual_seed(0)
      mesh_2d = init_device_mesh(
          self.device_type,
          (2, self.world_size // 2),
          mesh_dim_names=("dp", "tp"),
      )
      parallelize_plan = {
          "net1": ColwiseParallel(),
          "net2": RowwiseParallel(),
      }
      model_2d = parallelize_module(
          simple_model().cuda(), mesh_2d["tp"], parallelize_plan
      )
      model_2d = FSDP(
          model_2d, device_mesh=mesh_2d["dp"], use_orig_params=True
      )
      FSDP.set_state_dict_type(
          model_2d,
          StateDictType.SHARDED_STATE_DICT,
      )
      optim_2d = torch.optim.Adam(model_2d.parameters(), lr=0.01)
      model_2d(model_2d.get_input().cuda(self.rank)).sum().backward()
      optim_2d.step()
      optim_2d_osd = get_optimizer_state_dict(
          model_2d, optimizers=optim_2d
      )
      ref_optim_2d_osd = deepcopy(optim_2d_osd)
  
      no_wrap_osd_states = no_wrap_osd["state"]
      optim_2d_osd_states = optim_2d_osd["state"]
  
      self.assertEqual(
          len(no_wrap_osd_states), len(optim_2d_osd_states)
      )
      self.assertEqual(
          no_wrap_osd_states.keys(), optim_2d_osd_states.keys()
      )
      for fqn, states in no_wrap_osd_states.items():
          dist_states = optim_2d_osd_states.get(fqn)
  
          for state_name, state in states.items():
              dist_state = dist_states.get(state_name)
              # If a state  is DTensor, we all gather it in both DP and TP dimension to
              # compare with no_wrap state.
              if isinstance(dist_state, DT):
                  dist_state = (
                      dist_state.cuda()
                      .redistribute(
                          placements=(Replicate(), Replicate())
                      )
                      .to_local()
                  )
              self.assertTrue(isinstance(dist_state, torch.Tensor))
              self.assertTrue(torch.allclose(state, dist_state))
  
      # Update the parameters 2d optim states will be different from ref_optim_state_dict.
      model_2d(model_2d.get_input().cuda(self.rank)).sum().backward()
      optim_2d.step()
  
      set_optimizer_state_dict(
          model_2d,
          optimizers=optim_2d,
          optim_state_dict=ref_optim_2d_osd,
      )
      new_optim_2d_osd = get_optimizer_state_dict(
          model_2d, optimizers=optim_2d
      )
  
      ref_optim_2d_osd_states = ref_optim_2d_osd["state"]
      new_optim_2d_osd_states = optim_2d_osd["state"]
  
      # Compare the new optim state dict after load with the reference one
      self.assertEqual(
          len(ref_optim_2d_osd_states), len(new_optim_2d_osd_states)
      )
      self.assertEqual(
          ref_optim_2d_osd_states.keys(), new_optim_2d_osd_states.keys()
      )
      for fqn, states in ref_optim_2d_osd_states.items():
          new_states = new_optim_2d_osd_states.get(fqn)
  
          for state_name, state in states.items():
              new_state = new_states.get(state_name)
  
              if isinstance(new_state, DT):
                  self.assertEqual(
                      new_state.placements, state.placements
                  )
                  self.assertEqual(
                      new_state.device_mesh, state.device_mesh
                  )
                  self.assertTrue(
                      torch.allclose(
                          new_state.to_local(), state.to_local()
                      )
                  )
              else:
                  self.assertEqual(new_state, state)
  
imports: import io
  from copy import deepcopy
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from torch.distributed._tensor import DTensor, init_device_mesh, Replicate, Shard
  from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict, set_optimizer_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp._common_utils import _get_module_fsdp_state, clean_tensor_name
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  from torch.distributed.tensor.parallel.fsdp import DTensorExtensions
  from torch.distributed.tensor.parallel.input_reshard import input_reshard
  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
label_desc: Use image node_id 4 for calling the TestNew2dParallelStateDict obj's underlying member method code with example usage: obj.test_2d_optim_state_dict(is_even_sharded_model) without return types