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"]
docstr: Get dictionary containing optimizer state (deprecated v1 format)
  
  Default behavior is to perform communication so that the
  entire optimizer state is returned on the root rank in the
  process group. In this case, all ranks in the process group
  must enter this function and no value is returned on non-root
  ranks.
  
  Arguments:
      gather_on_root (bool, optional): Gather state from all
          ranks on the root rank (default: True)
docstr_tok: ["Get","dictionary","containing","optimizer","state","(","deprecated","v1","format",")","Default","behavior","is","to","perform","communication","so","that","the","entire","optimizer","state","is","returned","on","the","root","rank","in","the","process","group",".","In","this","case",",","all","ranks","in","the","process","group","must","enter","this","function","and","no","value","is","returned","on","non-root","ranks",".","Arguments",":","gather_on_root","(","bool",",","optional",")",":","Gather","state","from","all","ranks","on","the","root","rank","(","default",":","True",")"]
code: def _state_dict_v1(
      self, gather_on_root: bool = True
  ) -> Optional[dict]:
      """Get dictionary containing optimizer state (deprecated v1 format)
  
      Default behavior is to perform communication so that the
      entire optimizer state is returned on the root rank in the
      process group. In this case, all ranks in the process group
      must enter this function and no value is returned on non-root
      ranks.
  
      Arguments:
          gather_on_root (bool, optional): Gather state from all
              ranks on the root rank (default: True)
  
      """
      warnings.warn(
          "Making optimizer state dictionary in deprecated v1 format. "
          "Future support is not guaranteed."
      )
  
      state_dict = super().state_dict()
      if not gather_on_root:
          return state_dict
  
      # Finish any asynchronous communication
      self.grad_sync()
      self.param_sync()
  
      # Export local state to byte string
      state_bytes = io.BytesIO()
      torch.save(state_dict, state_bytes)
      state_bytes.seek(0)
      state_bytes_view = state_bytes.getbuffer()
  
      # Get data sizes on all ranks
      local_state_size = len(state_bytes_view)
      state_sizes = [None] * self.distributed_size
      torch.distributed.all_gather_object(
          state_sizes,
          local_state_size,
          group=self.process_group,
      )
      max_state_size = max(state_sizes)
  
      # Construct workspace buffers
      chunk_size = (
          self.default_shard_size
          * torch.finfo(self.grad_sync_dtype).bits
          // 8
      )
      if self.distributed_rank == 0:
          gathered_state_bytes = [
              torch.empty([size], dtype=torch.uint8, device="cpu")
              for size in state_sizes
          ]
          gathered_state_bytes[0].copy_(
              torch.frombuffer(state_bytes_view, dtype=torch.uint8)
          )
          gathered_chunks_buffers = [
              torch.empty(
                  [chunk_size * self.distributed_size],
                  dtype=torch.uint8,
                  device=self.device,
              )
              for _ in range(self.pipeline_size)
          ]
      else:
          chunk_buffers = [
              torch.empty(
                  [chunk_size],
                  dtype=torch.uint8,
                  device=self.device,
              )
              for _ in range(self.pipeline_size)
          ]
  
      # Split data into chunks and gather on root rank
      # Note: Assuming we are using the NCCL backend, communication
      # must happen on the GPU. We split the data into fixed-size
      # chunks to limit GPU memory usage.
      main_stream = torch.cuda.current_stream()
      for stream in self._pipeline_streams:
          stream.wait_stream(main_stream)
      for stream_id, offset in enumerate(
          range(0, max_state_size, chunk_size)
      ):
          stream_id %= self.pipeline_size
          stream = self._pipeline_streams[stream_id]
          with torch.cuda.stream(stream):
              # Buffers for chunk
              if self.distributed_rank == 0:
                  gathered_chunks = [
                      gathered_chunks_buffers[stream_id][
                          i * chunk_size : (i + 1) * chunk_size
                      ]
                      for i in range(self.distributed_size)
                  ]
              else:
                  chunk = chunk_buffers[stream_id]
  
              # Copy to GPU
              if (
                  self.distributed_rank != 0
                  and offset < local_state_size
              ):
                  local_chunk_size = min(
                      chunk_size, local_state_size - offset
                  )
                  chunk[:local_chunk_size].copy_(
                      torch.frombuffer(
                          state_bytes_view,
                          dtype=torch.uint8,
                          count=local_chunk_size,
                          offset=offset,
                      ),
                      non_blocking=True,
                  )
  
              # Gather on root
              # Note: Call in main stream to avoid memory pool
              # overheads from internal memory allocations in
              # gather.
              main_stream.wait_stream(stream)
              with torch.cuda.stream(main_stream):
                  if self.distributed_rank == 0:
                      if self._gather_no_copy:
                          no_copy_kwarg = {"no_copy": True}
                      else:
                          no_copy_kwarg = {}
                      torch.distributed.gather(
                          gathered_chunks[0],
                          gathered_chunks,
                          dst=self.process_group_root,
                          group=self.process_group,
                          **no_copy_kwarg,
                      )
                  else:
                      torch.distributed.gather(
                          chunk,
                          dst=self.process_group_root,
                          group=self.process_group,
                      )
              stream.wait_stream(main_stream)
  
              # Copy back to CPU
              if self.distributed_rank == 0:
                  for rank in range(1, self.distributed_size):
                      rank_chunk_start = offset
                      rank_chunk_end = min(
                          offset + chunk_size, state_sizes[rank]
                      )
                      rank_chunk_size = (
                          rank_chunk_end - rank_chunk_start
                      )
                      if rank_chunk_size > 0:
                          src = gathered_chunks[rank][:rank_chunk_size]
                          dst = gathered_state_bytes[rank][
                              rank_chunk_start:rank_chunk_end
                          ]
                          dst.copy_(src, non_blocking=True)
  
      # Synchronize GPU
      for stream in self._pipeline_streams:
          main_stream.wait_stream(stream)
      main_stream.synchronize()
  
      # Return gathered state data on root rank
      if self.distributed_rank == 0:
          return {"gathered_states": gathered_state_bytes}
      else:
          return None
  
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 38 for calling the DistributedFusedAdam obj's underlying member method code with example usage: obj._state_dict_v1(gather_on_root) and returns: state_dict, dict, None