code: def run_fsdp_checkpoint_example(rank, world_size):
      # Set up world pg
      os.environ["MASTER_ADDR"] = "localhost"
      os.environ["MASTER_PORT"] = "12355"
  
      # Initialize the process group
      dist.init_process_group(
          "cpu:gloo,cuda:nccl", rank=rank, world_size=world_size
      )
      torch.cuda.set_device(rank)
  
      # Create a model
      model_1, optim_1 = init_model()
  
      # Save the model to CHECKPOINT_DIR
      with FSDP.state_dict_type(
          model_1, StateDictType.SHARDED_STATE_DICT
      ):
          state_dict = {
              "model": model_1.state_dict(),
              "optim": FSDP.optim_state_dict(model_1, optim_1),
          }
  
          dist_cp.save_state_dict(
              state_dict=state_dict,
              storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
          )
  
      # Create a second model
      model_2, optim_2 = init_model()
  
      # Print the model parameters for both models.
      # Before loading, the parameters should be different.
      print_params("Before loading", model_1, model_2, optim_1, optim_2)
  
      # Load model_2 with parameters saved in CHECKPOINT_DIR
      with FSDP.state_dict_type(
          model_2, StateDictType.SHARDED_STATE_DICT
      ):
          state_dict = {
              "model": model_2.state_dict(),
              # cannot load the optimizer state_dict together with the model state_dict
          }
  
          dist_cp.load_state_dict(
              state_dict=state_dict,
              storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
          )
          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=dist_cp.FileSystemReader(CHECKPOINT_DIR),
          )
  
          flattened_osd = FSDP.optim_state_dict_to_load(
              model_2, optim_2, optim_state["optim"]
          )
          optim_2.load_state_dict(flattened_osd)
  
      # Print the model parameters for both models.
      # After loading, the parameters should be the same.
      print_params("After loading", model_1, model_2, optim_1, optim_2)
  
      # Shut down world pg
      dist.destroy_process_group()
  
file_path: pytorch/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py
imports: import os
  import shutil
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.multiprocessing
  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
label_desc: Use image node_id 4 for calling a global function with example usage: run_fsdp_checkpoint_example(rank, world_size) without return types