code: def run_test_checkpoint(strategy_name: str, shard: bool):
      if strategy_name == "ddp":
          strategy = DDPStrategy()
      elif strategy_name == "colossalai_gemini":
          strategy = GeminiStrategy(
              placement_policy="auto", initial_scale=2**5
          )
      elif strategy_name == "colossalai_zero2":
          strategy = LowLevelZeroStrategy(
              stage=2, placement_policy="cuda"
          )
      else:
          raise ValueError(f"Unsupported strategy '{strategy_name}'")
  
      with strategy.model_init_context():
          actor = GPTActor(config=GPT_CONFIG).cuda()
      actor_optim = HybridAdam(actor.parameters())
      actor, actor_optim = strategy.prepare((actor, actor_optim))
  
      train_step(strategy, actor, actor_optim)
  
      ctx = (
          tempfile.TemporaryDirectory()
          if dist.get_rank() == 0
          else nullcontext()
      )
  
      with ctx as dirname:
          rank0_dirname = [dirname]
          dist.broadcast_object_list(rank0_dirname)
          rank0_dirname = rank0_dirname[0]
  
          model_path = os.path.join(
              rank0_dirname, "model" if shard else f"model.pt"
          )
          strategy.save_model(actor, model_path)
          optim_path = os.path.join(
              rank0_dirname, "optim" if shard else "optim.pt"
          )
          strategy.save_optimizer(actor_optim, optim_path)
          dist.barrier()
  
          strategy.load_model(actor, model_path, strict=False)
          strategy.load_optimizer(actor_optim, optim_path)
          dist.barrier()
  
      train_step(strategy, actor, actor_optim)
  
file_path: colossalai/applications/Chat/tests/test_checkpoint.py
imports: import os
  import tempfile
  from contextlib import nullcontext
  import pytest
  import torch
  import torch.distributed
  from coati.models.gpt import GPTActor
  from coati.models.utils import calc_action_log_probs
  from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy, Strategy
  from transformers.models.gpt2.configuration_gpt2 import GPT2Config
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.testing import rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 3 for calling a global function with example usage: run_test_checkpoint(strategy_name, shard) without return types