class_docstr: Checks for overflow in gradient across parallel process
class_docstr_tok: ["Checks","for","overflow","in","gradient","across","parallel","process"]
code: def _has_inf_or_nan(x, i):
      try:
          # if x is half, the .float() incurs an additional deep copy, but it's necessary if
          # Pytorch's .sum() creates a one-element tensor of the same type as x
          # (which is true for some recent version of pytorch).
          cpu_sum = float(x.float().sum())
          # More efficient version that can be used if .sum() returns a Python scalar
          # cpu_sum = float(x.sum())
      except RuntimeError as instance:
          # We want to check if inst is actually an overflow exception.
          # RuntimeError could come from a different error.
          # If so, we still want the exception to propagate.
          if "value cannot be converted" not in instance.args[0]:
              raise
          return True
      else:
          if (
              cpu_sum == float("inf")
              or cpu_sum == -float("inf")
              or cpu_sum != cpu_sum
          ):
              return True
          return False
  
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 6 for calling the CheckOverflow obj's underlying member method code with example usage: obj._has_inf_or_nan(x, i) and returns: False, True, True