code: def exam_zero_1_2():
      """
      In this test, we want to test whether zero stage 1 and 2
      deliver the same numerical results despite different communication
      pattern
  
      we use these prefixes to differentiate the zero stage
      oss: partition optimizer states
      pg: partition gradients and optimizer states
  
      """
      local_rank = torch.distributed.get_rank()
      seed_all(2001)
  
      # create model
      zero1_model = MlpModel().cuda()
      zero2_model = copy.deepcopy(zero1_model)
  
      # create optimizer
      zero1_optimizer = torch.optim.Adam(zero1_model.parameters(), lr=1)
      zero2_optimizer = torch.optim.Adam(zero2_model.parameters(), lr=1)
      zero1_optimizer = LowLevelZeroOptimizer(
          zero1_optimizer,
          overlap_communication=True,
          initial_scale=128,
          verbose=True,
      )
      zero2_optimizer = LowLevelZeroOptimizer(
          zero2_optimizer,
          overlap_communication=True,
          partition_grad=True,
          initial_scale=128,
      )
      # create data
      seed_all(2001 + local_rank)
      input_data = torch.randn(32, 123).cuda()
  
      zero1_output = zero1_model(input_data)
      zero2_output = zero2_model(input_data)
      assert torch.equal(zero1_output, zero2_output)
  
      # zero-dp backward
      zero1_optimizer.backward(zero1_output.mean().float())
      zero2_optimizer.backward(zero2_output.mean().float())
  
      # check grad
      z1g_list = (
          zero1_optimizer._grad_store.get_working_grads_by_group_id(0)
      )
      z2g_list = (
          zero2_optimizer._grad_store.get_working_grads_by_group_id(0)
      )
      for z1g, z2g in zip(z1g_list, z2g_list):
          assert torch.equal(z1g, z2g)
  
      # step
      zero1_optimizer.step()
      zero2_optimizer.step()
  
      # check updated param
      for z1p, z2p in zip(
          zero1_model.parameters(), zero2_model.parameters()
      ):
          assert torch.equal(z1p.data, z2p.data)
  
file_path: colossalai/tests/test_zero/test_low_level/test_zero1_2.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 parameterize, 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 3 for calling a global function with example usage: exam_zero_1_2() without return types