code: def finetune(self, pretrain_dir: str, finetune_dir: str) -> None:
      device_mesh = init_device_mesh(
          self.device_type, (self.world_size,)
      )
  
      model = FineTuningModel().cuda()
      # TODO: make the parallelism more complicated, e.g., using 2D + DDP.
      model = FSDP(model, use_orig_params=True, device_mesh=device_mesh)
      optim = torch.optim.Adam(model.parameters(), lr=1e-3)
  
      # Simulate that the fine tuning restart after 3 iterations
      for i in range(2):
          # Load pretrain submodules checkpoint
          pretrain_state_dict, _ = get_state_dict(
              model,
              submodules={model.pretrain},
              options=StateDictOptions(keep_submodule_prefixes=False),
          )
          dist_cp.load_state_dict(
              {"model": pretrain_state_dict},
              storage_reader=dist_cp.FileSystemReader(pretrain_dir),
          )
          set_state_dict(
              model,
              model_state_dict={model.pretrain: pretrain_state_dict},
              options=StateDictOptions(strict=False),
          )
  
          try:
              # Load training submodules checkpoint
              model_state_dict, optim_state_dict = get_state_dict(
                  model,
                  optimizers=optim,
                  options=StateDictOptions(ignore_frozen_params=True),
              )
              dist_cp.load_state_dict(
                  {
                      "model": model_state_dict,
                      "optim": optim_state_dict,
                  },
                  storage_reader=dist_cp.FileSystemReader(pretrain_dir),
              )
              set_state_dict(
                  model,
                  optimizers=optim,
                  model_state_dict=model_state_dict,
                  optim_state_dict=optim_state_dict,
                  options=StateDictOptions(strict=False),
              )
          except KeyError:
              # If this is the first round of the fine tuning, then nothing is saved.
              # If this is the restart of the fine tuning, then checkpoint should exit.
              self.assertEqual(i, 0)
  
          # Training
          for j in range(3):
              batch = torch.rand(32, DIM, device="cuda")
              loss = model(batch).sum()
              loss.backward()
              optim.step()
              optim.zero_grad()
  
          # Save state_dict
          model_state_dict, optim_state_dict = get_state_dict(
              model,
              optimizers=optim,
              options=StateDictOptions(ignore_frozen_params=True),
          )
          saved_state_dict = {
              "model": model_state_dict,
              "optim": optim_state_dict,
          }
          dist_cp.save_state_dict(
              state_dict=saved_state_dict,
              storage_writer=dist_cp.FileSystemWriter(finetune_dir),
          )
  
imports: import os
  import sys
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.nn
  from torch.distributed._tensor import init_device_mesh
  from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict, StateDictOptions
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
  from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
label_desc: Use image node_id 4 for calling the TestFineTuning obj's underlying member method code with example usage: obj.finetune(pretrain_dir, finetune_dir) without return types