class_docstr: Implements the AdaScale_ algorithm for scaling the learning rate for
  distributed and large batch size training. Can be used in combination with
  ``torch.nn.parallel.DistributedDataParallel`` and ``torch.optim.SGD``.
  
  .. _AdaScale: https://proceedings.icml.cc/static/paper_files/icml/2020/4682-Supplemental.pdf
  
  This class subclasses `Optimizer` so that `torch.optim.lr_scheduler` can
  work with it. In other words, AdaScale is intended to be a complete wrapper of an
  torch Optimizer.
  
  Note that, AdaScale does *not* help increase per-GPU batch size.
  
  There are several ways to integrate AdaScale with your training loop.
  We show two examples below.
  
  Example 1: using PyTorch's `lr_scheduler` classes.
  
  .. code-block:: python
  
      optim = AdaScale(SGD(model.parameters(), lr=0.001))
      model = DistributedDataParallel(model)
      scheduler = LambdaLR(optim, lr_lambda=...)
  
      last_epoch = 0
      done = False
      step = 0
      while not done:
          for batch in dataset:
              optim.zero_grad()
              logits = model()
              loss = criterion(logits, ...)
              loss.backward()
              step += optim.gain()
              optim.step()
              epoch = step // len(dataset)
              if epoch > last_epoch:
                  scheduler.step()
                  last_epoch = epoch
              if epoch >= MAX_EPOCHS:
                  done = True
  
  Example 2: using a custom `update_lr()` function that update the learning
  rate based on the current step count per epoch.
  
  .. code-block:: python
  
      optim = AdaScale(SGD(model.parameters(), lr=0.001))
      model = DistributedDataParallel(model)
  
      step = 0
      while step < max_steps:
          for batch in ...:
              optim.zero_grad()
              logits = model()
              loss = criterion()
              loss.backward()
              step += optim.gain()
              optim.step()
              update_lr(step)
  
  Args:
      optimizer (torch.optim.Optimizer):
          Optimizer to apply AdaScale to.
      world_size (int):
          Number of world_size for distributed training.
          If None, defaults to ``dist.get_world_size()``.
      scale (float):
          Scaling factor of the batch size from scale equals 1, e.g. using a 10x
          larger batch size (summed across all ranks with gradient accumulation)
          means a scale of 10.
          If None, defaults to ``world_size * num_gradients_to_accumulate``.
      smoothing (float):
          Smoothing factor for moving average.
          If None, it defaults to ``max(1 - (world_size * num_gradients_to_accumulate)/1000, 0)``.
          Note, for very high scale training, higher smoothing value might be needed,
          esp at the begining of the training. Therefore, if your scale is close to or larger
          than 1000, try experimenting with smoothing value > 0 if the final accuracy is poor.
      num_gradients_to_accumulate (int):
          Number of passes that we accumulate gradients locally
          between each optimizer step. This can be changed during
          training as long as the train loop changes gradient accumulation
          accordingly.
          The loss in each pass can be either scaled or unscaled. See `is_scaled_loss` below.
          Default to 1, which does not accumulate gradients.
      is_scaled_loss (bool):
          If True, assume that the loss is scaled by `num_gradients_to_accumulate`.
          If False, the loss is not scaled.
          Default: True.
      debias_ewma (bool):
          (experimental) Use debias exponential moving average
          for smoothing and mu and sigma variables. False will
          use the method in the paper's Appendix B.3.
          Default: True, which is what have been validated so far.
code: def _final_callback(self) -> None:
      # This method should be invoked once for each backward pass, after
      # gradients have been synchronized between each worker, unless we
      # are in gradient accumulation mode, where grads are not all_reduced
      # between the GPUs.
      self._final_callback_queued = False
      assert isinstance(self._local_grad_sqr, torch.Tensor)
  
      # Keep track of number of backward calls for gradient accumulation.
      # TODO (min): this may not work with activation checkpointing when
      #             multiple backward calls happen in a big backward.
      self._num_backward_calls += 1
  
      # TODO (min, mike): We need to have a way to check that training loop & DDP
      #                   is doing the right thing where the gradient is reduced
      #                   in this backward pass.
      #                   Longer term, we may compute the gain and then inform
      #                   the training loop when it is a good time to step().
      assert (
          self._num_backward_calls - self._last_final_backward_call
      ) <= self._num_grads_to_accum, f"bug: {self._num_backward_calls} - {self._last_final_backward_call} should <= {self._num_grads_to_accum}"
      if (
          self._num_backward_calls - self._last_final_backward_call
      ) % self._num_grads_to_accum != 0:
          assert (
              self._local_grad_sqr is not None
          ), "We should still be in backward phase"
          return
  
      # Since self._local_grad_sqr is FP32, sum shouldn't overflow.
  
      # This vector has length of # of param_groups, so it is small, but we
      # use async to hide the all_reduce latency, esp when # of nodes is large.
      work = None
      if self._world_size > 1:
          work = dist.all_reduce(
              self._local_grad_sqr, async_op=True
          )  # SUM
  
      # Compute the sums of squares for reduced gradients.
      # Divide by _num_grads_to_accum since the gradients are accumulated.
      total_grad_sqr = np.array(
          [
              sum(
                  param.grad.pow(2).sum().item()
                  for param in group["params"]
              )
              for group in self._optimizer.param_groups
          ]
      )
  
      # Wait for all_reduce to be done and move it to cpu & np.
      if work:
          work.wait()
      local_grad_sqr = self._local_grad_sqr.cpu().numpy()
  
      if self._num_grads_to_accum > 1:
          # Handle scaling for for gradient accumulation
          if self._is_scaled_loss:
              # If loss is scaled down, we need to scale the local gradients back by a factor of _num_grads_to_accum squared;
              # total_grad_sqr is already scaled by _num_grads_to_accum squared.
              local_grad_sqr *= self._num_grads_to_accum**2
          else:
              # If loss is not scaled, local gradients are correct, but we need to scale the total_grad_sqr down to account for gradient accumulation.
              total_grad_sqr /= self._num_grads_to_accum**2
  
      # See appendix B.3 of the paper.
      # Modified to handle cases where scale != world_size
      #
      # local_grad_sqr is \sum_{i=1}^{c N} \norm{g_t_i}^2
      # where N is world size and c is num_grads_to_accum
      # total_grad_sqr is \norm{\bar{g}_t}^2
      S = self._scale
      cN = self._world_size * self._num_grads_to_accum
      grad_var = local_grad_sqr * (S / cN) / (
          cN - 1
      ) - total_grad_sqr * S / (cN - 1)
      grad_sqr = total_grad_sqr - grad_var / S
      grad_var = np.maximum(grad_var, 1e-6)
      grad_sqr = np.maximum(grad_sqr, 0.0)
      self._update_avg("grad_sqr_avg", grad_sqr, self.smoothing)
      self._update_avg("grad_var_avg", grad_var, self.smoothing)
      self._last_final_backward_call = self._num_backward_calls
      # Indicating backward is done.
      self._local_grad_sqr = None
  
imports: import functools
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
  import numpy
  import torch
  from torch.autograd import Variable
  import torch.distributed
  from torch.optim import SGD, Optimizer
label_desc: Use image node_id 17 for calling the AdaScale obj's underlying member method code with example usage: obj._final_callback() without return types