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: Initialize optimizer state for parameters
  
  Ignores parameters that have already been initialized.
  
  Arguments:
      params (iterable, optional): parameters to initialize
          (default: all parameters)
docstr_tok: ["Initialize","optimizer","state","for","parameters","Ignores","parameters","that","have","already","been","initialized",".","Arguments",":","params","(","iterable",",","optional",")",":","parameters","to","initialize","(","default",":","all","parameters",")"]
code: def init_params(
      self,
      params: Optional[Iterable[torch.nn.Parameter]] = None,
      dtype: Optional[torch.dtype] = None,
      grad_sync_dtype: Optional[torch.dtype] = None,
      param_sync_dtype: Optional[torch.dtype] = None,
  ) -> None:
      """Initialize optimizer state for parameters
  
      Ignores parameters that have already been initialized.
  
      Arguments:
          params (iterable, optional): parameters to initialize
              (default: all parameters)
  
      """
  
      # Default cases
      if params is None:
          params = self.parameters()
      elif isinstance(params, torch.Tensor):
          params = [params]
  
      # Ignore parameters that have already been initialized
      params = [
          param
          for param in params
          if "fragments" not in self.state[param]
      ]
      if not params:
          return
  
      # Get indices corresponding to parameters
      id_map = dict()
      for param_group_id, group in enumerate(self.param_groups):
          for param_id, param in enumerate(group["params"]):
              id_map[param] = (param_group_id, param_id)
  
      # Initialize parameters
      for param in params:
          if param in id_map:
              param_group_id, param_id = id_map[param]
              self._init_param_state(
                  param,
                  param_group_id,
                  param_id,
                  dtype=dtype,
                  grad_sync_dtype=grad_sync_dtype,
                  param_sync_dtype=param_sync_dtype,
              )
  
      # Check if buckets are underutilized
      num_params = sum(1 for param in self.parameters())
      num_initialized_params = sum(
          1
          for param in self.parameters()
          if "fragments" in self.state[param]
      )
      if num_initialized_params == num_params:
          bucket_size = sum(
              bucket.bucket_size for bucket in self.state["buckets"]
          )
          filled_size = sum(
              bucket.filled_size for bucket in self.state["buckets"]
          )
          buckets_utilization = filled_size / bucket_size
          if buckets_utilization < 0.7:
              warnings.warn(
                  f"Only {buckets_utilization:.1%} of buckets are used. "
                  "Consider decreasing the bucket_cap_mb argument."
              )
  
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 11 for calling the DistributedFusedAdam obj's underlying member method code with example usage: obj.init_params(params, dtype, grad_sync_dtype, param_sync_dtype) without return types