class_docstr: ShardedGradScaler helps perform gradient scaling in a shard aware manner. It extends
  functionality from GradScaler:
  * Supports Pytorch DDP and FSDP implementations
  * Support CPU offloaded tensors (as used in fully sharded data parallel[FSDP])
  * Supports the custom Mixed Precision loss dtype (fp16, bf16) that FSDP returns
  * Sync inf/nan for scaled gradient tensors on any torch.device (where tensors are placed) across
  nodes
  
  Example::
  
      # Creates a ShardedGradScaler once at the beginning of training.
      scaler = ShardedGradScaler()
  
      for epoch in epochs:
          for input, target in data:
              optimizer.zero_grad()
              output = model(input)
              loss = loss_fn(output, target)
  
              # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
              scaler.scale(loss).backward()
  
              # scaler.step() first unscales gradients of the optimizer's params.
              # If gradients don't contain infs/NaNs, optimizer.step() is then called,
              # otherwise, optimizer.step() is skipped.
              scaler.step(optimizer)
  
              # Updates the scale for next iteration.
              scaler.update()
  
  See :class:`GradScaler` for explanation of scaling/unscaling and more use cases.
  
  Args:
      init_scale (float, optional, default=2.**16):  Initial scale factor.
      growth_factor (float, optional, default=2.0):  Factor by which the scale is multiplied during
          :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations.
      backoff_factor (float, optional, default=0.5):  Factor by which the scale is multiplied during
          :meth:`update` if inf/NaN gradients occur in an iteration.
      growth_interval (int, optional, default=2000):  Number of consecutive iterations without inf/NaN gradients
          that must occur for the scale to be multiplied by ``growth_factor``.
      enabled (bool, optional):  If ``False``, disables gradient scaling. :meth:`step` simply
          invokes the underlying ``optimizer.step()``, and other methods become no-ops.
          Default: ``True``
      process_group (ProcessGroup, optional, default=torch.distributed.group.WORLD):
          process group for sharding
docstr: Updates the scale factor.
  If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
  to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
  the scale is multiplied by ``growth_factor`` to increase it.
  Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
  used directly, it's used to fill GradScaler's internal scale tensor. So if
  ``new_scale`` was a tensor, later in-place changes to that tensor will not further
  affect the scale GradScaler uses internally.)
  Args:
      new_scale (float or :class:`torch.cuda.FloatTensor`, optional, default=None):  New scale factor.
  .. warning::
      :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
      been invoked for all optimizers used this iteration.
docstr_tok: ["Updates","the","scale","factor",".","If","any","optimizer","steps","were","skipped","the","scale","is","multiplied","by","``","backoff_factor","``","to","reduce","it",".","If","``","growth_interval","``","unskipped","iterations","occurred","consecutively",",","the","scale","is","multiplied","by","``","growth_factor","``","to","increase","it",".","Passing","``","new_scale","``","sets","the","new","scale","value","manually",".","(","``","new_scale","``","is","not","used","directly",",","it","'s","used","to","fill","GradScaler","'s","internal","scale","tensor",".","So","if","``","new_scale","``","was","a","tensor",",","later","in-place","changes","to","that","tensor","will","not","further","affect","the","scale","GradScaler","uses","internally",".",")","Args",":","new_scale","(","float","or",":","class",":","`","torch.cuda.FloatTensor","`",",","optional",",","default=None",")",":","New","scale","factor",".","..","warning",":",":",":","meth",":","`","update","`","should","only","be","called","at","the","end","of","the","iteration",",","after","``","scaler.step","(","optimizer",")","``","has","been","invoked","for","all","optimizers","used","this","iteration","."]
code: def update(
      self, new_scale: Optional[Union[float, torch.Tensor]] = None
  ) -> None:
      """
      Updates the scale factor.
      If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
      to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
      the scale is multiplied by ``growth_factor`` to increase it.
      Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
      used directly, it's used to fill GradScaler's internal scale tensor. So if
      ``new_scale`` was a tensor, later in-place changes to that tensor will not further
      affect the scale GradScaler uses internally.)
      Args:
          new_scale (float or :class:`torch.cuda.FloatTensor`, optional, default=None):  New scale factor.
      .. warning::
          :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
          been invoked for all optimizers used this iteration.
      """
  
      if not self._enabled:
          return
  
      _scale, _growth_tracker = self._check_scale_growth_tracker("update")  # type: ignore[var-annotated]
  
      if new_scale is not None:
          # Accept a new user-defined scale.
          if isinstance(new_scale, float):
              self._scale.fill_(new_scale)  # type: ignore[union-attr]
          else:
              reason = "new_scale should be a float or a 1-element torch.cuda.FloatTensor with requires_grad=False."
              assert isinstance(new_scale, torch.cuda.FloatTensor), reason  # type: ignore[attr-defined]
              assert new_scale.numel() == 1, reason
              assert new_scale.requires_grad is False, reason
              self._scale.copy_(new_scale)  # type: ignore[union-attr]
      else:
          # Consume shared inf/nan data collected from optimizers to update the scale.
          # If all found_inf tensors are on the same device as self._scale, this operation is asynchronous.
          found_infs = [
              found_inf.to(device=_scale.device, non_blocking=True)
              for state in self._per_optimizer_states.values()
              for found_inf in state["found_inf_per_device"].values()
          ]
  
          assert (
              len(found_infs) > 0
          ), "No inf checks were recorded prior to update."
  
          found_inf_combined = found_infs[0]
          if len(found_infs) > 1:
              for i in range(1, len(found_infs)):
                  found_inf_combined += found_infs[i]
  
          if _scale.device.type == "cpu":
              self._amp_update_scale_cpu_(found_inf_combined)
          else:
              torch._amp_update_scale_(
                  self._scale,  # type: ignore[arg-type]
                  self._growth_tracker,  # type: ignore[arg-type]
                  found_inf_combined,
                  self._growth_factor,  # type: ignore[arg-type]
                  self._backoff_factor,  # type: ignore[arg-type]
                  self._growth_interval,  # type: ignore[arg-type]
              )
  
      # To prepare for next iteration, clear the data collected from optimizers this iteration.
      self._per_optimizer_states = defaultdict(
          _refresh_per_optimizer_state
      )
  
imports: import logging
  from collections import abc, defaultdict
  from typing import Any, Dict, Iterable, List, Optional, overload, Sequence, Tuple, Union
  import torch
  import torch.distributed
  from torch.cuda.amp.grad_scaler import _MultiDeviceReplicator, GradScaler, OptState
  from torch.distributed.distributed_c10d import ProcessGroup
label_desc: Use image node_id 11 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj.update(new_scale) without return types