code: def test(self):
      param1 = torch.nn.Parameter(torch.Tensor([0]))
      param1.grad = torch.Tensor([1])
      param2 = torch.nn.Parameter(torch.Tensor([0]))
      param2.grad = torch.Tensor([dist.get_rank() + 1])
      # param2 is now MoE parameter
      param2.allreduce = False
  
      parameters = [param1, param2]
  
      groups._create_expert_and_data_parallel(2)
  
      norm = ds_utils.clip_grad_norm_(parameters, max_norm=0.1)
      norm = torch.Tensor([norm]).to(dist.get_rank())
  
      world_size = dist.get_world_size()
      gathered_norm = [torch.zeros(1).cuda() for i in range(world_size)]
  
      dist.all_gather(gathered_norm, norm)
  
      assert (
          gathered_norm[0] == gathered_norm[1]
      ), "norm at rank 0 does not match the norm at rank 1"
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/runtime/test_runtime_utils.py
imports: import torch
  from torch._utils import _flatten_dense_tensors
  import deepspeed.comm
  import pytest
  import deepspeed.runtime.utils
  import deepspeed.utils.groups
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling the TestClibGradNorm obj's underlying member method code with example usage: obj.test() without return types