docstr: Tests that the order of keys in the state dict does not matter when loading
  If order was not accounted for, the following test would cause a deadlock.
docstr_tok: ["Tests","that","the","order","of","keys","in","the","state","dict","does","not","matter","when","loading","If","order","was","not","accounted","for",",","the","following","test","would","cause","a","deadlock","."]
code: def test_different_ordered_state_dict_keys(self):
      """Tests that the order of keys in the state dict does not matter when loading
      If order was not accounted for, the following test would cause a deadlock.
      """
  
      world_size = self.world_size
  
      class Foo:
          def state_dict(self):
              return {}
  
          def load_state_dict(self, state_dict):
              tl = [
                  torch.ones(2, dtype=torch.int64, device="cuda")
                  for _ in range(world_size)
              ]
              t = (
                  torch.arange(2, dtype=torch.int64, device="cuda")
                  + 1
                  + 2 * dist.get_rank()
              )
              dist.all_gather(tl, t, async_op=False)
  
      class Bar:
          def state_dict(self):
              return {}
  
          def load_state_dict(self, state_dict):
              tensor = (
                  torch.arange(2, dtype=torch.int64, device="cuda")
                  + 1
                  + 2 * dist.get_rank()
              )
              dist.all_reduce(tensor, op=ReduceOp.SUM)
  
      if self.rank == 0:
          sd = {
              "A": Foo(),
              "B": Bar(),
          }
      else:
          sd = {
              "B": Bar(),
              "A": Foo(),
          }
  
      DCP.save(sd, DCP.FileSystemWriter(self.temp_dir))
      DCP.load(sd, DCP.FileSystemReader(self.temp_dir))
  
imports: import time
  from enum import auto, Enum
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.nn
  import torch.nn.functional
  from torch.distributed._tensor.device_mesh import init_device_mesh
  from torch.distributed.checkpoint.filesystem import _FileSystemCheckpointer
  from torch.distributed.checkpoint.state_dict import _patch_model_state_dict, _patch_optimizer_state_dict, get_state_dict
  from torch.distributed.distributed_c10d import ReduceOp
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.api import ShardingStrategy
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  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
  from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
  from torch.testing._internal.distributed.common_state_dict import VerifyStateDictMixin
label_desc: Use image node_id 7 for calling the TestE2ESaveAndLoad obj's underlying member method code with example usage: obj.test_different_ordered_state_dict_keys() and returns: dict, dict