code: def split_ddp_grad(grad, world_size):
      with torch.no_grad():
          grad = grad.clone().detach().flatten()
          padding_size = (
              world_size - grad.numel() % world_size
          ) % world_size
          if padding_size > 0:
              grad = torch.nn.functional.pad(grad, [0, padding_size])
          splited_grad = grad.split(grad.numel() // world_size)
      return splited_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 1 for calling a global function with example usage: split_ddp_grad(grad, world_size) and returns: splited_grad