code: def run_zero_optim_test(local_rank, world_size, stage=1):
      criterion = torch.nn.CrossEntropyLoss()
  
      MOE_MANAGER.__init__()
      MOE_MANAGER.setup(
          parallel="EP",
      )
      zero_model = MoeModel(enable_load_balance=True)
      zero_optimizer = torch.optim.Adam(zero_model.parameters())
      plugin = LowLevelZeroPlugin(
          stage=stage, precision="bf16", verbose=True
      )
      booster = Booster(plugin=plugin)
      zero_model, zero_optimizer, _, _, _ = booster.boost(
          zero_model, zero_optimizer
      )
  
      MOE_MANAGER.__init__()
      MOE_MANAGER.setup(parallel="EP")
      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().bfloat16()
      grad_handler = MoeGradientHandler(torch_model)
  
      # run to update expert load
      data = (
          torch.randn(16, 4).cuda().bfloat16() / 1000 / (local_rank + 1)
      )
      label = torch.randint(0, 4, (16,)).cuda()
  
      # run torch model twice
      run_fwd_bwd(torch_model, data, label, criterion, None)
      grad_handler.handle_gradient()
      torch_optimizer.step()
      torch_optimizer.zero_grad()
      run_fwd_bwd(torch_model, data, label, criterion, None)
      grad_handler.handle_gradient()
  
      # get optim and load status in zero model
      run_fwd_bwd(zero_model, data, label, criterion, zero_optimizer)
      zero_optimizer.step()
      zero_optimizer.zero_grad()
      with torch.no_grad():
          origin_out = zero_model(data)
  
      # load balance
      apply_load_balance(zero_model, zero_optimizer)
  
      # run again to test
      zero_out = run_fwd_bwd(
          zero_model, data, label, criterion, zero_optimizer
      )
      torch.allclose(origin_out, zero_out)
  
      # assert optim
      torch_optimizer.step()
      torch_out = run_fwd_bwd(torch_model, data, label, criterion, None)
      zero_optimizer.step()
      zero_out = run_fwd_bwd(
          zero_model, data, label, criterion, zero_optimizer
      )
      assert torch.allclose(
          zero_out, torch_out, atol=3e-5
      ), f"zero_out:{zero_out}\ntorch_out{torch_out}"
  
file_path: ColossalAI/tests/test_moe/test_moe_load_balance.py
imports: import pytest
  import torch
  import torch.distributed
  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.layers import apply_load_balance
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.tensor.moe_tensor.api import is_moe_tensor
  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