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 (default v2 format)
  
  All ranks in the process group must call this function since
  it performs communication. The same optimizer state is
  returned on all ranks.
docstr_tok: ["Get","dictionary","containing","optimizer","state","(","default","v2","format",")","All","ranks","in","the","process","group","must","call","this","function","since","it","performs","communication",".","The","same","optimizer","state","is","returned","on","all","ranks","."]
code: def _state_dict_v2(self) -> Optional[dict]:
      """Get dictionary containing optimizer state (default v2 format)
  
      All ranks in the process group must call this function since
      it performs communication. The same optimizer state is
      returned on all ranks.
  
      """
  
      # Make sure params are initialized
      self.init_params()
  
      # Finish any asynchronous communication
      self.grad_sync()
      self.param_sync()
  
      # Get state dict from base class
      state_dict = super().state_dict()
      state_dict["state"] = {"step": state_dict["state"]["step"]}
  
      # Initialize state dict with CPU buffers
      for param in self.parameters():
          # Get param index in state dict
          fragment = self.state[param]["fragments"][0]
          param_group_id = fragment.param_group_id
          param_id = fragment.param_id
          index = state_dict["param_groups"][param_group_id]["params"][
              param_id
          ]
  
          # Construct CPU buffers with optimizer state
          state_dict["state"][index] = dict(
              param=torch.zeros_like(
                  param, dtype=self.dtype, device="cpu"
              ),
              exp_avg=torch.zeros_like(
                  param, dtype=self.dtype, device="cpu"
              ),
              exp_avg_sq=torch.zeros_like(
                  param, dtype=self.dtype, device="cpu"
              ),
          )
  
      # Workspace buffers for gathering shards on root rank
      num_buckets = len(self.state["buckets"])
      max_bucket_size = max(
          bucket.bucket_size for bucket in self.state["buckets"]
      )
      bucket_buffers = [
          torch.empty(
              [max_bucket_size],
              dtype=self.dtype,
              device=self.device,
          )
          for _ in range(self.pipeline_size)
      ]
      if self.store_param_remainders:
          max_shard_size = max(
              bucket.shard_size for bucket in self.state["buckets"]
          )
          shard_bf16_buffers = [
              torch.empty(
                  [max_shard_size],
                  dtype=torch.bfloat16,
                  device=self.device,
              )
              for _ in range(self.pipeline_size)
          ]
  
      # Synchronize streams
      main_stream = torch.cuda.current_stream()
      for stream in self._pipeline_streams:
          stream.wait_stream(main_stream)
  
      def pack_param_shard(bucket_id: int) -> torch.Tensor:
          """Pack local shard of param values into contiguous buffer"""
  
          # Stream objects
          stream_id = bucket_id % self.pipeline_size
          stream = self._pipeline_streams[stream_id]
  
          # Bucket objects
          bucket = self.state["buckets"][bucket_id]
          shard_size = bucket.shard_size
  
          # Case 1: Param state is already packed
          if bucket.params_shard is not None:
              return bucket.params_shard
  
          # Case 2: Pack BF16 model params with 16-bit remainders
          if bucket.param_remainders_shard is not None:
              with torch.cuda.stream(stream):
                  # Pack bf16 param values
                  shard_bf16 = shard_bf16_buffers[stream_id][
                      :shard_size
                  ]
                  buffers_in = []
                  buffers_out = []
                  for fragment in bucket.fragments:
                      if not fragment.in_local_shard:
                          continue
                      param_range = slice(*fragment.shard_param_range)
                      shard_range = slice(*fragment.shard_range)
                      param = self.parameter(fragment)
                      buffers_in.append(param.view(-1)[param_range])
                      buffers_out.append(shard_bf16[shard_range])
                  _multi_tensor_copy(
                      buffers_in,
                      buffers_out,
                      dummy_overflow_buf=self._dummy_overflow_buf,
                  )
  
                  # Reconstruct fp32 from bf16 and remainders
                  shard_range = slice(
                      shard_size * self.distributed_rank,
                      shard_size * (self.distributed_rank + 1),
                  )
                  shard_fp32 = bucket_buffers[stream_id][shard_range]
                  _bf16_rem_to_fp32(
                      shard_bf16,
                      bucket.param_remainders_shard,
                      shard_fp32,
                  )
                  return shard_fp32
  
          # Case 3: Pack model params
          with torch.cuda.stream(stream):
              shard_range = slice(
                  shard_size * self.distributed_rank,
                  shard_size * (self.distributed_rank + 1),
              )
              shard = bucket_buffers[stream_id][shard_range]
              buffers_in = []
              buffers_out = []
              for fragment in bucket.fragments:
                  if not fragment.in_local_shard:
                      continue
                  param_range = slice(*fragment.shard_param_range)
                  shard_range = slice(*fragment.shard_range)
                  param = self.parameter(fragment)
                  buffers_in.append(param.view(-1)[param_range])
                  buffers_out.append(shard[shard_range])
              _multi_tensor_copy(
                  buffers_in,
                  buffers_out,
                  dummy_overflow_buf=self._dummy_overflow_buf,
              )
              return shard
  
      def start_all_gather(bucket_id: int, shard: torch.Tensor) -> None:
          """Launch all-gather on bucket shards
  
          Communication is done on main stream to ensure consistent
          ordering.
  
          """
  
          # Stream objects
          stream_id = bucket_id % self.pipeline_size
          stream = self._pipeline_streams[stream_id]
  
          # Workspace buffer
          bucket = self.state["buckets"][bucket_id]
          bucket_size = bucket.bucket_size
          bucket_buffer = bucket_buffers[stream_id][:bucket_size]
  
          # All-gather shards
          main_stream.wait_stream(stream)
          all_gather_into_tensor(
              bucket_buffer,
              shard,
              group=self.distributed_process_group,
          )
          stream.wait_stream(main_stream)
  
      def finish_all_gather(
          bucket_id: int, state_dict_key: str
      ) -> None:
          """Finish all-gather on bucket shards
  
          Data is copied into state dict CPU buffers.
  
          Splitting the NCCL all-gather and the CPU memcpys into
          separate stages helps achieve good overlap when kernel
          launches are serialized with
          CUDA_DEVICE_MAX_CONNECTIONS=1. In particular, the pipeline
          calls start_all_gather(bucket_id+1) before
          finish_all_gather(bucket_id).
  
          """
  
          # Stream objects
          stream_id = bucket_id % self.pipeline_size
          stream = self._pipeline_streams[stream_id]
  
          # Bucket objects
          bucket = self.state["buckets"][bucket_id]
          bucket_size = bucket.bucket_size
          bucket_buffer = bucket_buffers[stream_id][:bucket_size]
  
          # Update state dict
          with torch.cuda.stream(stream):
              for fragment in bucket.fragments:
                  param_range = slice(*fragment.param_range)
                  bucket_range = slice(*fragment.bucket_range)
                  param_group_id = fragment.param_group_id
                  param_id = fragment.param_id
                  index = state_dict["param_groups"][param_group_id][
                      "params"
                  ][param_id]
                  state_buffer = state_dict["state"][index][
                      state_dict_key
                  ]
                  state_fragment = state_buffer.view(-1)[param_range]
                  bucket_fragment = bucket_buffer[bucket_range]
                  state_fragment.copy_(
                      bucket_fragment, non_blocking=True
                  )
  
      # All-gather param state
      for bucket_id in range(num_buckets):
          shard = pack_param_shard(bucket_id)
          start_all_gather(bucket_id, shard)
          if bucket_id > 0:
              finish_all_gather(bucket_id - 1, "param")
          if bucket_id == num_buckets - 1:
              finish_all_gather(bucket_id, "param")
  
      # All-gather exp_avg state
      for bucket_id in range(num_buckets):
          shard = self.state["buckets"][bucket_id].exp_avg_shard
          start_all_gather(bucket_id, shard)
          if bucket_id > 0:
              finish_all_gather(bucket_id - 1, "exp_avg")
          if bucket_id == num_buckets - 1:
              finish_all_gather(bucket_id, "exp_avg")
  
      # All-gather exp_avg_sq state
      for bucket_id in range(num_buckets):
          shard = self.state["buckets"][bucket_id].exp_avg_sq_shard
          start_all_gather(bucket_id, shard)
          if bucket_id > 0:
              finish_all_gather(bucket_id - 1, "exp_avg_sq")
          if bucket_id == num_buckets - 1:
              finish_all_gather(bucket_id, "exp_avg_sq")
  
      # Synchronize GPU and return
      for stream in self._pipeline_streams:
          main_stream.wait_stream(stream)
      main_stream.synchronize()
      return state_dict
  
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 39 for calling the DistributedFusedAdam obj's underlying member method code with example usage: obj._state_dict_v2() and returns: state_dict, bucket, shard, shard_fp32