code: def test_fsspec_with_dist(self):
      CHECKPOINT_DIR = self.temp_dir
  
      model = FSDP(MyTestModule().cuda())
      optim = torch.optim.Adam(model.parameters(), lr=0.1)
      model(torch.rand(8, 8, device=dist.get_rank())).sum().backward()
      optim.step()
  
      with FSDP.state_dict_type(
          model, StateDictType.SHARDED_STATE_DICT
      ):
          state_dict = {
              "model": model.state_dict(),
              "optim": FSDP.optim_state_dict(model, optim),
          }
  
          dcp.save_state_dict(
              state_dict=state_dict,
              storage_writer=FsspecWriter(CHECKPOINT_DIR),
              planner=dcp.DefaultSavePlanner(),
          )
  
      model_2 = FSDP(MyTestModule().cuda())
      optim_2 = torch.optim.Adam(model_2.parameters(), lr=0.1)
  
      with FSDP.summon_full_params(model):
          with FSDP.summon_full_params(model_2):
              for n_p1, n_p2 in zip(
                  model.named_parameters(), model_2.named_parameters()
              ):
                  self.assertNotEqual(n_p1[1], n_p2[1])
  
      # 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(),
          }
  
          dcp.load_state_dict(
              state_dict=state_dict,
              storage_reader=FsspecReader(CHECKPOINT_DIR),
              planner=dcp.DefaultLoadPlanner(),
          )
          model_2.load_state_dict(state_dict["model"])
  
          optim_state = load_sharded_optimizer_state_dict(
              model_state_dict=state_dict["model"],
              optimizer_key="optim",
              storage_reader=FsspecReader(CHECKPOINT_DIR),
          )
  
          flattened_osd = FSDP.optim_state_dict_to_load(
              model_2, optim_2, optim_state["optim"]
          )
          optim_2.load_state_dict(flattened_osd)
  
      with FSDP.summon_full_params(model):
          with FSDP.summon_full_params(model_2):
              for n_p1, n_p2 in zip(
                  model.named_parameters(), model_2.named_parameters()
              ):
                  self.assertEqual(n_p1[1], n_p2[1])
  
      def opt_at(opt, idx):
          return list(iter(opt.state.values()))[idx]
  
      # Adam lazily creates its state
      self.assertEqual(
          opt_at(optim, 0)["exp_avg"], opt_at(optim_2, 0)["exp_avg"]
      )
      self.assertEqual(
          opt_at(optim, 0)["exp_avg_sq"],
          opt_at(optim_2, 0)["exp_avg_sq"],
      )
  
imports: import shutil
  import tempfile
  from functools import wraps
  from typing import Any, Callable, Dict, Optional, Tuple
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.nn
  from torch.distributed.checkpoint._fsspec_filesystem import FsspecReader, FsspecWriter
  from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
  from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests, TestCase
  from torch.testing._internal.distributed._shard.sharded_tensor import ShardedTensorTestBase, with_comms
label_desc: Use image node_id 2 for calling the TestFSSpecWithDist obj's underlying member method code with example usage: obj.test_fsspec_with_dist() and returns: list