class_docstr: Gradient scaler for model-parallel inf check. The inf in gradients are checked across tensor-parallel
  ranks in (1) executing optimizer step and (2) gradient scaler update.
code: def _maybe_opt_step(
      self, optimizer, optimizer_state, *args, **kwargs
  ):
      retval = None
      found_inf = torch.cuda.FloatTensor(
          [
              sum(
                  v.item()
                  for v in optimizer_state[
                      "found_inf_per_device"
                  ].values()
              )
          ]
      )
  
      # Update across all model parallel instances.
      torch.distributed.all_reduce(
          found_inf,
          op=torch.distributed.ReduceOp.MAX,
          group=parallel_state.get_model_parallel_group(),
      )
  
      if found_inf.item() == 0:
          retval = optimizer.step(*args, **kwargs)
      return retval
  
file_path: apex/apex/transformer/amp/grad_scaler.py
imports: from collections import defaultdict
  import torch
  from apex.transformer import parallel_state
label_desc: Use image node_id 3 for calling the GradScaler obj's underlying member method code with example usage: obj._maybe_opt_step(optimizer, optimizer_state) and returns: retval