code: def run_grad_clip_norm(
      world_size: int, dtype: torch.dtype, device: str, norm_type: float
  ):
      print(f"{world_size}, {dtype}, {device}, {norm_type}")
      cuda_device = get_current_device()
      devices = [cuda_device] * 4
      if device == "cpu":
          devices = [torch.device("cpu")] * 4
      elif device == "mixed":
          devices = [cuda_device] * 2 + [torch.device("cpu")] * 2
      pg = ProcessGroup(tp_degree=world_size)
      params = [
          Parameter(torch.empty(4, 4, dtype=dtype, device=devices[i]))
          for i in range(4)
      ]
      colo_params = [
          ColoParameter(
              torch.empty(4, 4, dtype=dtype, device=devices[i]),
              spec=ColoTensorSpec(pg),
          )
          for i in range(4)
      ]
      for p, colo_p in zip(params, colo_params):
          grad = torch.rand_like(p)
          p.grad = grad
          colo_p.grad = grad.clone().detach()
      shard_param(colo_params[0])
      shard_param(colo_params[2])
      torch_norm = clip_grad_norm_(params, 1.0, norm_type=norm_type)
      colo_norm = clip_grad_norm(colo_params, 1.0, norm_type=norm_type)
      assert close(
          torch_norm, colo_norm
      ), f"diff: {abs(torch_norm-colo_norm)}"
      for p, colo_p in zip(params, colo_params):
          check_grad_equal(p, colo_p)
  
file_path: ColossalAI/tests/test_legacy/test_utils/test_norm_gradient_clipping.py
imports: import pytest
  import torch
  from torch.nn.parameter import Parameter
  from torch.nn.utils import clip_grad_norm_
  import colossalai
  from colossalai.legacy.tensor import ColoTensorSpec, ProcessGroup, distspec
  from colossalai.legacy.utils.common import clip_grad_norm
  from colossalai.logging import disable_existing_loggers
  from colossalai.tensor.colo_parameter import ColoParameter
  from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
  from colossalai.utils import get_current_device
label_desc: Use image node_id 4 for calling a global function with example usage: run_grad_clip_norm(world_size, dtype, device, norm_type) without return types