code: def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None):
      """Clips gradient norm of an iterable of parameters.
  
      This has been adapted from Nvidia megatron. We add norm averaging
      to consider MoE params when calculating norm as they will result
      in different norms across different ranks.
  
      This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
      added functionality to handle model parallel parameters. Note that
      the gradients are modified in place.
  
      Arguments:
          parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
              single Tensor that will have gradients normalized
          max_norm (float or int): max norm of the gradients
          norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
              infinity norm.
  
      Returns:
          Total norm of the parameters (viewed as a single vector).
      """
      if isinstance(parameters, torch.Tensor):
          parameters = [parameters]
      parameters = list(
          filter(lambda p: p.grad is not None, parameters)
      )
      max_norm = float(max_norm)
      norm_type = float(norm_type)
      if norm_type == inf:
          total_norm = max(p.grad.data.abs().max() for p in parameters)
          total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
          # Take max across all GPUs.
          if mpu is not None:
              dist.all_reduce(
                  total_norm_cuda,
                  op=dist.ReduceOp.MAX,
                  group=mpu.get_model_parallel_group(),
              )
          total_norm = total_norm_cuda[0].item()
      else:
          total_norm = 0
          for p in parameters:
              if mpu is not None:
                  if (
                      mpu.get_model_parallel_rank() == 0
                  ) or is_model_parallel_parameter(p):
                      param_norm = p.grad.data.norm(norm_type)
                      total_norm += param_norm.item() ** norm_type
              else:
                  param_norm = p.grad.data.float().norm(norm_type)
                  total_norm += param_norm.item() ** norm_type
  
          # Sum across all model parallel GPUs.
          total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
          if mpu is not None:
              dist.all_reduce(
                  total_norm_cuda,
                  op=dist.ReduceOp.SUM,
                  group=mpu.get_model_parallel_group(),
              )
          total_norm = total_norm_cuda[0].item() ** (1.0 / norm_type)
  
      # Need to average total_norm across different GPUs due to the presence of moe params
      pg = groups._get_data_parallel_group()
      scaled_norm = (
          total_norm * 1.0 / float(dist.get_world_size(group=pg))
      )
  
      scaled_norm_tensor = torch.cuda.FloatTensor([float(scaled_norm)])
      dist.all_reduce(scaled_norm_tensor, group=pg)
      total_norm = scaled_norm_tensor.item()
  
      clip_coef = max_norm / (total_norm + 1e-6)
      if clip_coef < 1:
          for p in parameters:
              p.grad.data.mul_(clip_coef)
      return total_norm
  
imports: from collections.abc import Iterable
  from deepspeed.moe.utils import is_moe_param
  import os
  import psutil
  import gc
  from math import sqrt
  from math import floor
  from bisect import bisect_left
  import torch
  from torch import inf
  from deepspeed import comm
  from deepspeed.utils import groups, logger
  from deepspeed.runtime.constants import PIPE_REPLICATED
  from numpy import prod
label_desc: Use image node_id 10 for calling a global function with example usage: clip_grad_norm_(parameters, max_norm, norm_type, mpu) and returns: total_norm