code: def run_zero_test(local_rank, world_size, stage=1):
      criterion = torch.nn.CrossEntropyLoss()
  
      zero_model = MoeModel()
      optimizer = torch.optim.Adam(zero_model.parameters())
      plugin = LowLevelZeroPlugin(stage=stage, precision="fp32")
      booster = Booster(plugin=plugin)
      zero_model, optimizer, _, _, _ = booster.boost(
          zero_model, 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_model = torch_model.cuda()
      grad_handler = MoeGradientHandler(torch_model)
  
      # assert zero model
      for (torch_name, torch_param), (zero_name, zero_param) in zip(
          torch_model.named_parameters(),
          zero_model.module.named_parameters(),
      ):
          assert zero_name == torch_name
          assert torch.allclose(zero_param.data, torch_param.data)
  
      data = torch.randn(16, 4).cuda()
      label = torch.randint(0, 4, (16,)).cuda()
  
      torch_out = run_fwd_bwd(torch_model, data, label, criterion, None)
      zero_out = run_fwd_bwd(
          zero_model, data, label, criterion, optimizer
      )
      assert torch.allclose(torch_out, zero_out)
      grad_handler.handle_gradient()
  
      for (zero_name, zero_param), (torch_name, torch_param) in zip(
          zero_model.module.named_parameters(),
          torch_model.named_parameters(),
      ):
          assert zero_name == torch_name
          zero_grad_list = optimizer._grad_store.get_partitioned_gradients_by_param_id(
              0, id(zero_param)
          )
          if hasattr(zero_param, "moe_info"):
              assert len(zero_grad_list) == 0
              assert torch.allclose(zero_param.grad, torch_param.grad)
          else:
              assert len(zero_grad_list) > 0
              torch_grad_list = split_ddp_grad(
                  torch_param.grad, world_size
              )
              if stage == 2:
                  torch_grad_list = torch_grad_list[
                      local_rank : local_rank + 1
                  ]
              assert len(zero_grad_list) == len(torch_grad_list)
              for zero_grad, torch_grad in zip(
                  zero_grad_list, torch_grad_list
              ):
                  assert torch.allclose(zero_grad, torch_grad)
  
file_path: ColossalAI/tests/test_moe/test_moe_zero_fwd_bwd.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 colossalai.testing.random import seed_all
  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_test(local_rank, world_size, stage) without return types