code: def get_grad_zeros(parameters, mpu=None):
      """Compute the number of grads with zero values.
  
      This is adapted from get_grad_norm
  
      Arguments:
          parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
              single Tensor that will have gradients normalized
  
      Returns:
          Total number of params with zero values (viewed as a single vector).
      """
      if isinstance(parameters, torch.Tensor):
          parameters = [parameters]
      parameters = list(
          filter(lambda p: p.grad is not None, parameters)
      )
  
      total_zeros = 0.0
      tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
      for p in parameters:
          # Pipeline parallelism may replicate parameters. Avoid multi-counting.
          if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
              continue
  
          # Filter to avoid over-counting replicated tensors from tensor
          # model parallelism
          if (tensor_mp_rank > 0) and not is_model_parallel_parameter(
              p
          ):
              continue
  
          count_zeros = p.grad.numel() - torch.count_nonzero(p.grad)
          total_zeros += count_zeros.item()
  
      # Sum across all model parallel GPUs.
      total_zeros_cuda = torch.cuda.FloatTensor([float(total_zeros)])
      if mpu is not None:
          dist.all_reduce(
              total_zeros_cuda,
              op=dist.ReduceOp.SUM,
              group=mpu.get_model_parallel_group(),
          )
      total_zeros = total_zeros_cuda[0].item()
  
      return total_zeros
  
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 12 for calling a global function with example usage: get_grad_zeros(parameters, mpu) and returns: total_zeros