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
code: def _unscale_grads_(
      self,
      optimizer: torch.optim.Optimizer,
      inv_scale: torch.Tensor,
      found_inf: torch.Tensor,
      allow_fp16: bool = True,
  ) -> Dict[torch.device, torch.Tensor]:
      per_device_inv_scale = _GeneralMultiDeviceReplicator(inv_scale)
      per_device_found_inf = _GeneralMultiDeviceReplicator(found_inf)
  
      # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.
      # There could be thousands of grads, so we'd like to iterate through them just once.
      # However, we don't know their devices or dtypes in advance.
  
      # https://stackoverflow.com/questions/5029934/defaultdict-of-defaultdict
      # Google says mypy struggles with defaultdicts type annotations.
      per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))  # type: ignore[var-annotated]
      with torch.no_grad():
          for group in optimizer.param_groups:
              for param in group["params"]:
                  if param.grad is None:
                      continue
                  if (
                      not allow_fp16
                  ) and param.grad.dtype == torch.float16:
                      raise ValueError(
                          "Attempting to unscale FP16 gradients."
                      )
                  if param.grad.is_sparse:
                      # is_coalesced() == False means the sparse grad has values with duplicate indices.
                      # coalesce() deduplicates indices and adds all values that have the same index.
                      # For scaled fp16 values, there's a good chance coalescing will cause overflow,
                      # so we should check the coalesced _values().
                      if param.grad.dtype is torch.float16:
                          # coalesce is not supported in torch.float16
                          param_grad_fp32 = param.grad.type(
                              torch.float32
                          ).coalesce()
                          param.grad = param_grad_fp32.type(
                              torch.float16
                          )
                      to_unscale = param.grad._values()
                  else:
                      to_unscale = param.grad
  
                  per_device_and_dtype_grads[to_unscale.device][
                      to_unscale.dtype
                  ].append(to_unscale)
  
          for (
              device,
              per_dtype_grads,
          ) in per_device_and_dtype_grads.items():
              for grads in per_dtype_grads.values():
                  if grads[0].device.type == "cpu":
                      self._foreach_non_finite_check_and_unscale_cpu_(
                          grads,
                          per_device_found_inf.get(device),
                          per_device_inv_scale.get(device),
                      )
                  else:
                      torch._amp_foreach_non_finite_check_and_unscale_(
                          grads,
                          per_device_found_inf.get(device),
                          per_device_inv_scale.get(device),
                      )
      # There exist contexts (e.g. w/ `use_orig_params=True`) wherein some
      # ranks may have no (non-zero sized) parameter shards, necessitating the
      # initialization of `per_device_found_inf._per_device_tensors` here
      if not per_device_found_inf._per_device_tensors:
          assert self._scale is not None
          per_device_found_inf.get(self._scale.device)
      return per_device_found_inf._per_device_tensors
  
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 8 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj._unscale_grads_(optimizer, inv_scale, found_inf, allow_fp16) and returns: per_device_found_inf