class_docstr: A shard aware Grad Scaler which enables loss scaling with/without cpu_offload. This is a
  slight modification of the pytorch grad scaler.
  https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler
class_docstr_tok: ["A","shard","aware","Grad","Scaler","which","enables","loss","scaling","with\/without","cpu_offload",".","This","is","a","slight","modification","of","the","pytorch","grad","scaler",".","https",":","\/\/pytorch.org\/docs\/stable\/amp.html","#","torch.cuda.amp.GradScaler"]
docstr: :meth:`step` carries out the following two operations:
  
  1.  Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``
      earlier in the iteration).  As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.
  2.  If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled
      gradients.  Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.
  
  ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.
  
  Returns the return value of ``optimizer.step(*args, **kwargs)``.
  
  Args:
      optimizer (torch.optim.Optimizer):  Optimizer that applies the gradients.
      args:  Any arguments.
      kwargs:  Any keyword arguments.
  
  .. warning::
      Closure use is not currently supported.
  
  Note: This is an exact copy of the step function in grad_scaler.py. If this copy is deleted then the
  unittest test_cpu_offload_and_cpu_grads fails. This is because the parent class step function calls
  the parent class unscale_ function which does not handle torch.distributed.all_reduce on cpu.
docstr_tok: [":","meth",":","`","step","`","carries","out","the","following","two","operations",":","1",".","Internally","invokes","``","unscale_","(","optimizer",")","``","(","unless",":","meth",":","`","unscale_","`","was","explicitly","called","for","``","optimizer","``","earlier","in","the","iteration",")",".","As","part","of","the",":","meth",":","`","unscale_","`",",","gradients","are","checked","for","infs\/NaNs",".","2",".","If","no","inf\/NaN","gradients","are","found",",","invokes","``","optimizer.step","(",")","``","using","the","unscaled","gradients",".","Otherwise",",","``","optimizer.step","(",")","``","is","skipped","to","avoid","corrupting","the","params",".","``","*","args","``","and","``","*","*","kwargs","``","are","forwarded","to","``","optimizer.step","(",")","``",".","Returns","the","return","value","of","``","optimizer.step","(","*","args",",","*","*","kwargs",")","``",".","Args",":","optimizer","(","torch.optim.Optimizer",")",":","Optimizer","that","applies","the","gradients",".","args",":","Any","arguments",".","kwargs",":","Any","keyword","arguments",".","..","warning",":",":","Closure","use","is","not","currently","supported",".","Note",":","This","is","an","exact","copy","of","the","step","function","in","grad_scaler.py",".","If","this","copy","is","deleted","then","the","unittest","test_cpu_offload_and_cpu_grads","fails",".","This","is","because","the","parent","class","step","function","calls","the","parent","class","unscale_","function","which","does","not","handle","torch.distributed.all_reduce","on","cpu","."]
code: def step(self, optimizer: SGD, *args, **kwargs) -> Optional[float]:  # type: ignore
      """
      :meth:`step` carries out the following two operations:
  
      1.  Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``
          earlier in the iteration).  As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.
      2.  If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled
          gradients.  Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.
  
      ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.
  
      Returns the return value of ``optimizer.step(*args, **kwargs)``.
  
      Args:
          optimizer (torch.optim.Optimizer):  Optimizer that applies the gradients.
          args:  Any arguments.
          kwargs:  Any keyword arguments.
  
      .. warning::
          Closure use is not currently supported.
  
      Note: This is an exact copy of the step function in grad_scaler.py. If this copy is deleted then the
      unittest test_cpu_offload_and_cpu_grads fails. This is because the parent class step function calls
      the parent class unscale_ function which does not handle torch.distributed.all_reduce on cpu.
      """
      if not self._enabled:
          return optimizer.step(*args, **kwargs)
  
      if "closure" in kwargs:
          raise RuntimeError(
              "Closure use is not currently supported if GradScaler is enabled."
          )
  
      self._check_scale_growth_tracker("step")  # type: ignore
  
      optimizer_state = self._per_optimizer_states[id(optimizer)]
  
      if optimizer_state["stage"] is OptState.STEPPED:
          raise RuntimeError(
              "step() has already been called since the last update()."
          )
  
      retval = None
  
      if (
          hasattr(optimizer, "_step_supports_amp_scaling")
          and optimizer._step_supports_amp_scaling
      ):
          # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.
          # The contract with custom optimizers is that their step() should accept an additional,
          # optional grad_scaler kwarg.  We append self to the kwargs so the custom optimizer has full information:
          # it can query its own state, invoke unscale_ on itself, etc
          retval = optimizer.step(
              *args, **dict(kwargs, grad_scaler=self)
          )
          optimizer_state["stage"] = OptState.STEPPED
          return retval
  
      if optimizer_state["stage"] is OptState.READY:
          self.unscale_(optimizer)
  
      assert (
          len(optimizer_state["found_inf_per_device"]) > 0
      ), "No inf checks were recorded for this optimizer."
      retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)  # type: ignore
      optimizer_state["stage"] = OptState.STEPPED
      return retval
  
imports: from collections import defaultdict
  import collections.abc
  from enum import Enum
  import logging
  from typing import Any, Dict, List, Optional, Union
  import warnings
  import torch
  from torch.cuda import FloatTensor
  from torch.cuda.amp.common import amp_definitely_not_available
  from torch.cuda.amp.grad_scaler import GradScaler
  import torch.distributed
  from torch.optim import Optimizer
  from torch.optim.sgd import SGD
  from fairscale.internal import torch_version
label_desc: Use image node_id 6 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj.step(optimizer) and returns: retval, optimizer, retval