code: def exam_zero_1_torch_ddp_ckpt():
      """
      We examine the state_dict of zero and DDP.
      Moreover, we examine the zero's loading checkpoint of a torch ckpt.
      """
      local_rank = torch.distributed.get_rank()
      seed_all(1453)
  
      # create models
      torch_model = MlpModel().cuda()
      zero_model = copy.deepcopy(torch_model)
  
      torch_model = DDP(torch_model.cuda(), static_graph=True).cuda()
  
      # create optimizer
      zero_optimizer = torch.optim.Adam(zero_model.parameters(), lr=1)
  
      # we only test stage 1 here
      # the state dicts of stage 1 and stage 2 are the same
      zero_optimizer = LowLevelZeroOptimizer(
          zero_optimizer,
          overlap_communication=True,
          initial_scale=1,
          reduce_bucket_size=262144,
      )
  
      torch_optimizer = torch.optim.Adam(torch_model.parameters(), lr=1)
  
      seed_all(1453 + local_rank)
      # create
      input_data = torch.rand(4, 12).cuda()
  
      # forward
      zero_output = zero_model(input_data)
      torch_output = torch_model(input_data)
  
      # backward
      zero_optimizer.backward(zero_output.mean().float())
      torch_output.mean().backward()
  
      # step
      zero_optimizer.step()
      torch_optimizer.step()
  
      torch_state_dict = torch_optimizer.state_dict()
      zero_state_dict = zero_optimizer.state_dict()
  
      # examine the original state dict
      for torch_state, zero_state in zip(
          torch_state_dict["state"].values(),
          zero_state_dict["state"].values(),
      ):
          for t_v, z_v in zip(
              torch_state.values(), zero_state.values()
          ):
              loose_close(t_v, z_v)
  
      # empty the optimzer state
      zero_optimizer.optim.state = []
  
      # zero load a torch checkpoint
      zero_optimizer.load_state_dict(copy.deepcopy(torch_state_dict))
      zero_state_dict = zero_optimizer.state_dict()
  
      # examine the loaded state dict
      for torch_state, zero_state in zip(
          torch_state_dict["state"].values(),
          zero_state_dict["state"].values(),
      ):
          for t_v, z_v in zip(
              torch_state.values(), zero_state.values()
          ):
              loose_close(t_v, z_v)
  
file_path: ColossalAI/tests/test_zero/test_low_level/test_zero_ckpt.py
imports: import copy
  import pytest
  import torch
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing import assert_close
  import colossalai
  from colossalai.testing import rerun_if_address_is_in_use, spawn
  from colossalai.testing.random import seed_all
  from colossalai.zero import LowLevelZeroOptimizer
label_desc: Use image node_id 2 for calling a global function with example usage: exam_zero_1_torch_ddp_ckpt() without return types