code: def count_zeros_fp32(parameters, model_parallel_group):
      if isinstance(parameters, torch.Tensor):
          parameters = [parameters]
  
      # Filter parameters based on:
      #   - grad should not be none
      #   - parameter should not be shared
      #   - should not be a replica due to tensor model parallelism
      total_num_zeros = torch.tensor(
          [0.0], dtype=torch.float, device="cuda"
      )
      for param in parameters:
          grad_not_none = param.grad is not None
          is_not_shared = param_is_not_shared(param)
          is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(
              param
          )
          if grad_not_none and is_not_shared and is_not_tp_duplicate:
              grad = param.grad.detach()
              num_zeros = grad.numel() - torch.count_nonzero(grad)
              total_num_zeros = num_zeros + total_num_zeros
  
      # Sum across all model-parallel GPUs.
      torch.distributed.all_reduce(
          total_num_zeros,
          op=torch.distributed.ReduceOp.SUM,
          group=model_parallel_group,
      )
  
      total_num_zeros = total_num_zeros.item()
  
      return total_num_zeros
  
file_path: megatron-lm/megatron/optimizer/clip_grads.py
imports: import os
  import torch
  from torch import inf
  from apex.multi_tensor_apply import multi_tensor_applier
  import amp_C
  from megatron.model.module import param_is_not_shared
  from megatron.core.tensor_parallel import param_is_not_tensor_parallel_duplicate
label_desc: Use image node_id 2 for calling a global function with example usage: count_zeros_fp32(parameters, model_parallel_group) and returns: total_num_zeros