code: def run_zero_optim_test(local_rank, world_size, stage=1):
      criterion = torch.nn.CrossEntropyLoss()
  
      zero_model = MoeModel()
      zero_optimizer = torch.optim.Adam(zero_model.parameters())
      plugin = LowLevelZeroPlugin(stage=stage, precision="fp32")
      booster = Booster(plugin=plugin)
      zero_model, zero_optimizer, _, _, _ = booster.boost(
          zero_model, zero_optimizer
      )
  
      torch_model = MoeModel()
      for zero_param, torch_param in zip(
          zero_model.parameters(), torch_model.parameters()
      ):
          torch_param.data.copy_(zero_param.data)
      torch_optimizer = torch.optim.Adam(torch_model.parameters())
      torch_model = torch_model.cuda()
      grad_handler = MoeGradientHandler(torch_model)
  
      for _ in range(2):
          data = torch.randn(16, 4).cuda() / (local_rank + 1)
          label = torch.randint(0, 4, (16,)).cuda()
          run_fwd_bwd(torch_model, data, label, criterion, None)
          run_fwd_bwd(
              zero_model, data, label, criterion, zero_optimizer
          )
          grad_handler.handle_gradient()
  
          torch_optimizer.step()
          zero_optimizer.step()
  
          for (torch_name, torch_param), (zero_name, zero_param) in zip(
              torch_model.named_parameters(),
              zero_model.named_parameters(),
          ):
              assert torch.allclose(
                  torch_param.data, zero_param.data
              ), f"{torch_name}\ntorch_param {torch_param.data}\nzero_param {zero_param.data}"
  
          torch_optimizer.zero_grad()
          zero_optimizer.zero_grad()
  
file_path: ColossalAI/tests/test_moe/test_moe_zero_optim.py
imports: import pytest
  import torch
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import LowLevelZeroPlugin
  from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.testing import rerun_if_address_is_in_use, spawn
  from tests.test_moe.moe_utils import MoeGradientHandler, MoeModel
label_desc: Use image node_id 3 for calling a global function with example usage: run_zero_optim_test(local_rank, world_size, stage) without return types