class_docstr: Adam optimizer with ZeRO algorithm.
  
  Currently GPU-only. Requires Apex to be installed via
  ``python setup.py install --cuda_ext --cpp_ext --distributed_adam --deprecated_fused_adam``.
  
  This implements the ZeRO-2 algorithm, which distributes the
  optimizer state and gradients between parallel processes. In
  particular, the parameters are flattened, grouped into fixed-size
  buckets, and the optimizer state for each bucket is sharded over
  the parallel processes. Options are provided to overlap the
  gradient synchronization with the backward pass compute.
  
  Adam was proposed in `Adam: A Method for Stochastic
  Optimization`_, AdamW in `Decoupled Weight Decay Regularization`_,
  and ZeRO in `ZeRO: Memory Optimizations Toward Training Trillion
  Parameter Models`_.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts
          defining parameter groups.
      lr (float, optional): learning rate. (default: 1e-3)
      bias_correction (bool, optional): apply correction factor to
          moment estimates. (default: True)
      betas (Tuple[float, float], optional): coefficients used for
          computing running averages of gradient and its square.
          (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve
          numerical stability. (default: 1e-8)
      adam_w_mode (boolean, optional): Decouple weight decay
          regularization (also known as AdamW algorithm) (default:
          True)
      weight_decay (float, optional): weight decay (L2 penalty)
          (default: 0)
      amsgrad (boolean, optional): whether to use the AMSGrad
          variant of this algorithm from the paper
          `On the Convergence of Adam and Beyond`_ (default: False).
          This is not yet supported.
      dtype (torch.dtype, optional): datatype for optimizer state
          (default: torch.float32)
      grad_sync_dtype (torch.dtype, optional): datatype for gradient
          synchronization (default: same as dtype)
      param_sync_dtype (torch.dtype, optional): datatype for
          parameter synchronization (default: same as dtype)
      device (torch.device, optional): device for optimizer state
          (default: cuda). Currently only supports GPU with one GPU
          per process.
      process_group (torch.distributed.ProcessGroup, optional):
          parallel processes participating in optimizer (default:
          default group in torch.distributed). This group is
          interpreted as a 2D grid with dimensions
          distributed_size x redundant_size.
      distributed_process_group (torch.distributed.ProcessGroup,
          optional): parallel processes to distribute optimizer
          state over (default: same as process_group)
      redundant_process_group (torch.distributed.ProcessGroup,
          optional): parallel processes to replicate optimizer state
          over (default: group only containing calling process)
      average_grad_sync (bool, optional): whether to use average
          reduction for gradient synchronization rather than sum
          (default: True)
      overlap_grad_sync (boolean, optional): whether to overlap
          gradient synchronization with backward pass compute
          (default: True)
      overlap_param_sync (boolean, optional): whether to overlap
          parameter synchronization with forward pass compute
          (default: False). This is an experimental feature.
      bucket_cap_mb (float, optional): bucket size in megabytes
          (default: 100)
      pipeline_size (int, optional): number of buckets to process
          simultaneously in optimizer step (default: 2)
      contiguous_param_buffer (bool, optional): convert parameters
          into views into large persistent buffers (default: False).
          This enables some performance optimizations (e.g. avoiding
          some memory copies), but may add memory overhead (e.g. if
          the memory allocator can't reuse the original parameter
          buffers).
      contiguous_grad_buffer (bool, optional): allocate gradient
          buckets out of a large persistent buffers (default:
          False). This allows individual parameter gradients to be
          accessed externally (see grad_buffer_view function). It
          enables some performance optimizations (e.g. avoiding some
          memory copies), but prevents some memory optimizations
          (e.g. the memory allocator can't reuse buffers for
          gradient buckets).
      store_params (bool, optional): store a distributed copy of the
          parameters as optimizer state (default: True). This may be
          desirable if the optimizer dtype has higher precision than
          the parameter dtype.
      store_param_remainders (bool, optional): if model is BF16 and
          optimizer is FP32, store bits required to reconstruct FP32
          params (default: False). This is an experimental feature.
  
  .. _Adam\: A Method for Stochastic Optimization:
      https://arxiv.org/abs/1412.6980
  .. _On the Convergence of Adam and Beyond:
      https://openreview.net/forum?id=ryQu7f-RZ
  .. _Decoupled Weight Decay Regularization: https://arxiv.org/abs/1711.05101
  .. _ZeRO\: Memory Optimizations Toward Training Trillion Parameter Models:
      https://arxiv.org/abs/1910.02054
class_docstr_tok: ["Adam","optimizer","with","ZeRO","algorithm",".","Currently","GPU-only",".","Requires","Apex","to","be","installed","via","``","python","setup.py","install","--","cuda_ext","--","cpp_ext","--","distributed_adam","--","deprecated_fused_adam","``",".","This","implements","the","ZeRO-2","algorithm",",","which","distributes","the","optimizer","state","and","gradients","between","parallel","processes",".","In","particular",",","the","parameters","are","flattened",",","grouped","into","fixed-size","buckets",",","and","the","optimizer","state","for","each","bucket","is","sharded","over","the","parallel","processes",".","Options","are","provided","to","overlap","the","gradient","synchronization","with","the","backward","pass","compute",".","Adam","was","proposed","in","`","Adam",":","A","Method","for","Stochastic","Optimization","`","_",",","AdamW","in","`","Decoupled","Weight","Decay","Regularization","`","_",",","and","ZeRO","in","`","ZeRO",":","Memory","Optimizations","Toward","Training","Trillion","Parameter","Models","`","_",".","Arguments",":","params","(","iterable",")",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups",".","lr","(","float",",","optional",")",":","learning","rate",".","(","default",":","1e-3",")","bias_correction","(","bool",",","optional",")",":","apply","correction","factor","to","moment","estimates",".","(","default",":","True",")","betas","(","Tuple","[","float",",","float","]",",","optional",")",":","coefficients","used","for","computing","running","averages","of","gradient","and","its","square",".","(","default",":","(","0.9",",","0.999",")",")","eps","(","float",",","optional",")",":","term","added","to","the","denominator","to","improve","numerical","stability",".","(","default",":","1e-8",")","adam_w_mode","(","boolean",",","optional",")",":","Decouple","weight","decay","regularization","(","also","known","as","AdamW","algorithm",")","(","default",":","True",")","weight_decay","(","float",",","optional",")",":","weight","decay","(","L2","penalty",")","(","default",":","0",")","amsgrad","(","boolean",",","optional",")",":","whether","to","use","the","AMSGrad","variant","of","this","algorithm","from","the","paper","`","On","the","Convergence","of","Adam","and","Beyond","`","_","(","default",":","False",")",".","This","is","not","yet","supported",".","dtype","(","torch.dtype",",","optional",")",":","datatype","for","optimizer","state","(","default",":","torch.float32",")","grad_sync_dtype","(","torch.dtype",",","optional",")",":","datatype","for","gradient","synchronization","(","default",":","same","as","dtype",")","param_sync_dtype","(","torch.dtype",",","optional",")",":","datatype","for","parameter","synchronization","(","default",":","same","as","dtype",")","device","(","torch.device",",","optional",")",":","device","for","optimizer","state","(","default",":","cuda",")",".","Currently","only","supports","GPU","with","one","GPU","per","process",".","process_group","(","torch.distributed.ProcessGroup",",","optional",")",":","parallel","processes","participating","in","optimizer","(","default",":","default","group","in","torch.distributed",")",".","This","group","is","interpreted","as","a","2D","grid","with","dimensions","distributed_size","x","redundant_size",".","distributed_process_group","(","torch.distributed.ProcessGroup",",","optional",")",":","parallel","processes","to","distribute","optimizer","state","over","(","default",":","same","as","process_group",")","redundant_process_group","(","torch.distributed.ProcessGroup",",","optional",")",":","parallel","processes","to","replicate","optimizer","state","over","(","default",":","group","only","containing","calling","process",")","average_grad_sync","(","bool",",","optional",")",":","whether","to","use","average","reduction","for","gradient","synchronization","rather","than","sum","(","default",":","True",")","overlap_grad_sync","(","boolean",",","optional",")",":","whether","to","overlap","gradient","synchronization","with","backward","pass","compute","(","default",":","True",")","overlap_param_sync","(","boolean",",","optional",")",":","whether","to","overlap","parameter","synchronization","with","forward","pass","compute","(","default",":","False",")",".","This","is","an","experimental","feature",".","bucket_cap_mb","(","float",",","optional",")",":","bucket","size","in","megabytes","(","default",":","100",")","pipeline_size","(","int",",","optional",")",":","number","of","buckets","to","process","simultaneously","in","optimizer","step","(","default",":","2",")","contiguous_param_buffer","(","bool",",","optional",")",":","convert","parameters","into","views","into","large","persistent","buffers","(","default",":","False",")",".","This","enables","some","performance","optimizations","(","e.g",".","avoiding","some","memory","copies",")",",","but","may","add","memory","overhead","(","e.g",".","if","the","memory","allocator","ca","n't","reuse","the","original","parameter","buffers",")",".","contiguous_grad_buffer","(","bool",",","optional",")",":","allocate","gradient","buckets","out","of","a","large","persistent","buffers","(","default",":","False",")",".","This","allows","individual","parameter","gradients","to","be","accessed","externally","(","see","grad_buffer_view","function",")",".","It","enables","some","performance","optimizations","(","e.g",".","avoiding","some","memory","copies",")",",","but","prevents","some","memory","optimizations","(","e.g",".","the","memory","allocator","ca","n't","reuse","buffers","for","gradient","buckets",")",".","store_params","(","bool",",","optional",")",":","store","a","distributed","copy","of","the","parameters","as","optimizer","state","(","default",":","True",")",".","This","may","be","desirable","if","the","optimizer","dtype","has","higher","precision","than","the","parameter","dtype",".","store_param_remainders","(","bool",",","optional",")",":","if","model","is","BF16","and","optimizer","is","FP32",",","store","bits","required","to","reconstruct","FP32","params","(","default",":","False",")",".","This","is","an","experimental","feature",".","..","_Adam\\",":","A","Method","for","Stochastic","Optimization",":","https",":","\/\/arxiv.org\/abs\/1412.6980","..","_On","the","Convergence","of","Adam","and","Beyond",":","https",":","\/\/openreview.net\/forum","?","id=ryQu7f-RZ","..","_Decoupled","Weight","Decay","Regularization",":","https",":","\/\/arxiv.org\/abs\/1711.05101","..","_ZeRO\\",":","Memory","Optimizations","Toward","Training","Trillion","Parameter","Models",":","https",":","\/\/arxiv.org\/abs\/1910.02054"]
code: def __init__(
      self,
      params: Union[Iterable[torch.nn.Parameter], Iterable[dict]],
      lr: float = 1e-3,
      bias_correction: bool = True,
      betas: Tuple[float, float] = (0.9, 0.999),
      eps: float = 1e-8,
      adam_w_mode: bool = True,
      weight_decay: float = 0.0,
      amsgrad: bool = False,
      dtype: torch.dtype = torch.float32,
      grad_sync_dtype: Optional[torch.dtype] = None,
      param_sync_dtype: Optional[torch.dtype] = None,
      device: Optional[torch.device] = "cuda",
      process_group: Optional[torch.distributed.ProcessGroup] = None,
      distributed_process_group: Optional[
          torch.distributed.ProcessGroup
      ] = None,
      redundant_process_group: Optional[
          torch.distributed.ProcessGroup
      ] = None,
      average_grad_sync: bool = True,
      overlap_grad_sync: bool = True,
      overlap_param_sync: bool = False,
      bucket_cap_mb: float = 100.0,
      pipeline_size: int = 2,
      contiguous_param_buffer: bool = False,
      contiguous_grad_buffer: bool = False,
      store_params: bool = True,
      store_param_remainders: bool = False,
  ):
      defaults = dict(
          lr=lr,
          bias_correction=bias_correction,
          betas=betas,
          eps=eps,
          weight_decay=weight_decay,
      )
      super().__init__(params, defaults)
  
      # Adam options
      self.adam_w_mode: bool = adam_w_mode
      if amsgrad:
          raise RuntimeError(
              "DistributedFusedAdam does not support the AMSGrad variant."
          )
  
      # Datatype options
      if grad_sync_dtype is None:
          grad_sync_dtype = dtype
      if param_sync_dtype is None:
          param_sync_dtype = dtype
      supported_dtypes = (torch.float32, torch.float16, torch.bfloat16)
      if (
          dtype not in supported_dtypes
          or grad_sync_dtype not in supported_dtypes
      ):
          raise ValueError(
              "Unsupported dtypes for DistributedFusedAdam "
              f"(dtype={dtype}, "
              f"grad_sync_dtype={grad_sync_dtype}, "
              f"param_sync_dtype={param_sync_dtype}))"
          )
      self.dtype: torch.dtype = dtype
      self.grad_sync_dtype: torch.dtype = grad_sync_dtype
      self.param_sync_dtype: torch.dtype = param_sync_dtype
  
      # Device options
      if not _devices_match(device, "cuda"):
          raise RuntimeError(
              "Invalid device for DistributedFusedAdam "
              f"(device={device})"
          )
      self.device: torch.device = torch.device(
          "cuda", torch.cuda.current_device()
      )
  
      # Process groups
      self.process_group: torch.distributed.ProcessGroup = (
          _get_default_group()
          if process_group is None
          else process_group
      )
      self.distributed_process_group: torch.distributed.ProcessGroup = (
          self.process_group
          if distributed_process_group is None
          else distributed_process_group
      )
      self.redundant_process_group: Optional[
          torch.distributed.ProcessGroup
      ] = redundant_process_group
      self.process_group_size: int = torch.distributed.get_world_size(
          self.process_group
      )
      self.distributed_rank: int = torch.distributed.get_rank(
          self.distributed_process_group
      )
      self.distributed_size: int = torch.distributed.get_world_size(
          self.distributed_process_group
      )
      self.redundant_size: int = (
          1
          if self.redundant_process_group is None
          else torch.distributed.get_world_size(
              self.redundant_process_group
          )
      )
      if (
          self.process_group_size
          != self.distributed_size * self.redundant_size
      ):
          raise RuntimeError(
              "Invalid process group configuration "
              f"(process group size = {self.process_group_size}, "
              f"distributed process group size = {self.distributed_size}, "
              f"redundant process group size = {self.redundant_size})"
          )
      self.process_group_root: int = get_global_rank(
          self.process_group, 0
      )
  
      # Use average reduction for grad sync
      self.average_grad_sync: bool = average_grad_sync
      # Copy param grads to bucket as soon as available
      self.greedy_grad_copy: bool = True
      # Synchronize grad buckets as soon as their grads are available
      self.overlap_grad_sync: bool = overlap_grad_sync
      # Try synchronizing param buckets just before param is needed
      self.overlap_param_sync: bool = overlap_param_sync
      # Number of buckets to synchronize at a time
      self.pipeline_size: int = pipeline_size
  
      # Store params or param remainders
      if store_param_remainders:
          if store_params:
              raise RuntimeError(
                  "Attempted to construct DistributedFusedAdam "
                  "with store_params=True and store_param_remainders=True"
              )
          if (
              self.dtype != torch.float32
              or self.param_sync_dtype != torch.bfloat16
          ):
              raise RuntimeError(
                  "DistributedFusedAdam requires "
                  "BF16 params and FP32 optimizer state "
                  "when storing parameter remainders "
                  f"(dtype={self.dtype}, "
                  f"param_sync_dtype={self.param_sync_dtype}))"
              )
      self.store_params: bool = store_params
      self.store_param_remainders: bool = store_param_remainders
  
      # Determine bucket sizes
      dtype_size = torch.finfo(self.grad_sync_dtype).bits // 8
      self.alignment: int = 128 // dtype_size
      bucket_size = 1024 * 1024 * bucket_cap_mb / dtype_size
      shard_size = int(bucket_size / self.distributed_size)
      shard_size = _round_to_multiple(
          shard_size, self.alignment, round_up=False
      )
      shard_size = max(shard_size, self.alignment)
      self.default_shard_size: int = shard_size
  
      # Optimizer state
      self.state["buckets"]: List[StateBucket] = []
      self.state["step"]: int = 0
  
      # Gradient state
      self._grads_buckets: Dict[
          int, GradientBucket
      ] = collections.defaultdict(self.GradientBucket)
      # Param state
      self._params_buckets: Dict[
          int, ParameterBucket
      ] = collections.OrderedDict()
  
      # Whether to allocate contiguous buffers for parameters
      self.contiguous_param_buffer: bool = contiguous_param_buffer
      # Whether to allocate contiguous buffers for gradients
      self.contiguous_grad_buffer: bool = contiguous_grad_buffer
      # Contiguous buffers for parameters
      self._param_buffers: Dict[
          Tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor
      ] = {}
      # Contiguous buffers for gradients
      self._grad_buffers: Dict[
          Tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor
      ] = {}
  
      # Side streams for optimizer step and communication
      self._pipeline_streams: List[torch.cuda.Stream] = [
          torch.cuda.Stream() for _ in range(self.pipeline_size + 1)
      ]
  
      # Scale by factor before optimizer step. Used for grad
      # clipping and gradient scaler.
      self._grad_scale: torch.Tensor = torch.full(
          [], 1.0, dtype=torch.float32, device=self.device
      )
      # Norm of parameter gradients. Used for gradient clipping and
      # gradient scaler.
      self._grad_norm: Optional[torch.Tensor] = None
  
      # Dummy flag for multi-tensor kernels
      # Note: Apex multi-tensor kernels have a noop_flag argument
      # that is intended to detect non-finite values. It shouldn't
      # have any effect with the kernels used in the optimizer, but
      # we still set it to zero out of an abundance of caution.
      self._dummy_overflow_buf: torch.Tensor = torch.zeros(
          [1], dtype=torch.int32, device=self.device
      )
  
      # Check if collectives have no_copy option
      self._gather_no_copy: bool = (
          "no_copy"
          in inspect.getfullargspec(torch.distributed.gather).args
      )
  
      # Make sure parameter values are same across processes
      self._broadcast_params()
  
      # Lock for callbacks
      self._lock: threading.Lock = threading.Lock()
      # Attach hooks for gradient synchronization
      self._register_post_backward_hooks()
      # Attach hooks for param synchronization
      if self.overlap_param_sync:
          self._register_pre_forward_hooks()
  
imports: import collections
  import contextlib
  from dataclasses import dataclass
  import enum
  import inspect
  import io
  import itertools
  import threading
  from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
  import warnings
  import torch
  from torch.distributed.distributed_c10d import _get_default_group
  from apex.multi_tensor_apply import multi_tensor_applier
  import amp_C
  import distributed_adam_cuda
  from torch.distributed.distributed_c10d import _coalescing_manager
label_desc: Use image node_id 1 to create a new DistributedFusedAdam object from inherited base classes: torch.optim with example: obj = DistributedFusedAdam(params, lr, bias_correction, betas, eps, adam_w_mode, weight_decay, amsgrad, dtype, grad_sync_dtype, param_sync_dtype, device, process_group, distributed_process_group, redundant_process_group, average_grad_sync, overlap_grad_sync, overlap_param_sync, bucket_cap_mb, pipeline_size, contiguous_param_buffer, contiguous_grad_buffer, store_params, store_param_remainders)