code: def test(self, check_using_norm):
      groups._create_expert_and_data_parallel(2)
  
      param1 = torch.nn.Parameter(torch.Tensor([0]))
      param1.grad = torch.Tensor([1])
      param2 = torch.nn.Parameter(torch.Tensor([0]))
      if dist.get_rank() == 0:
          param2.grad = torch.Tensor([1])
      else:
          param2.grad = torch.Tensor([float("inf")])
      param2.allreduce = False
      # param2 is now MoE parameter
      parameters = [param1, param2]
      if check_using_norm:
          grads_group_flat = [
              _flatten_dense_tensors([p.grad for p in parameters])
          ]
          norm = ds_utils.get_weight_norm(grads_group_flat)
          overflow_checker = ds_utils.CheckOverflow([parameters])
          overflow = overflow_checker.check_using_norm(
              [norm], reduce_overflow=False
          )
      else:
          overflow_checker = ds_utils.CheckOverflow([parameters])
          overflow = overflow_checker.check()
      assert overflow
  
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 TestCheckOverflow obj's underlying member method code with example usage: obj.test(check_using_norm) without return types