class_docstr: Implements LAMB algorithm.
  
  Currently GPU-only.  Requires Apex to be installed via
  ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
  
  This version of fused LAMB implements 2 fusions.
  
    * Fusion of the LAMB update's elementwise operations
    * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
  
  :class:`apex.optimizers.FusedLAMB`'s usage is identical to any ordinary Pytorch optimizer::
  
      opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
      ...
      opt.step()
  
  :class:`apex.optimizers.FusedLAMB` may be used with or without Amp.  If you wish to use :class:`FusedLAMB` with Amp,
  you may choose any ``opt_level``::
  
      opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
      model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
      ...
      opt.step()
  
  In general, ``opt_level="O1"`` is recommended.
  
  LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups.
      lr (float, optional): learning rate. (default: 1e-3)
      betas (Tuple[float, float], optional): coefficients used for computing
          running averages of gradient and its norm. (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve
          numerical stability. (default: 1e-8)
      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`_
          NOT SUPPORTED now! (default: False)
      adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
          True for decoupled weight decay(also known as AdamW) (default: True)
      grad_averaging (bool, optional): whether apply (1-beta2) to grad when
          calculating running averages of gradient. (default: True)
      set_grad_none (bool, optional): whether set grad to None when zero_grad()
          method is called. (default: True)
      max_grad_norm (float, optional): value used to clip global grad norm
          (default: 1.0)
      use_nvlamb (boolean, optional): Apply adaptive learning rate to 0.0
          weight decay parameter (default: False)
      step_supports_amp_scaling(boolean, optional): whether to use customized
          gradient unscaling logic (default: True)
  
  .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
      https://arxiv.org/abs/1904.00962
  .. _On the Convergence of Adam and Beyond:
      https://openreview.net/forum?id=ryQu7f-RZ
class_docstr_tok: ["Implements","LAMB","algorithm",".","Currently","GPU-only",".","Requires","Apex","to","be","installed","via","``","pip","install","-v","--","no-cache-dir","--","global-option=","''","--","cpp_ext","''","--","global-option=","''","--","cuda_ext","''",".\/","``",".","This","version","of","fused","LAMB","implements","2","fusions",".","*","Fusion","of","the","LAMB","update","'s","elementwise","operations","*","A","multi-tensor","apply","launch","that","batches","the","elementwise","updates","applied","to","all","the","model","'s","parameters","into","one","or","a","few","kernel","launches",".",":","class",":","`","apex.optimizers.FusedLAMB","`","'s","usage","is","identical","to","any","ordinary","Pytorch","optimizer",":",":","opt","=","apex.optimizers.FusedLAMB","(","model.parameters","(",")",",","lr","=","....",")","...","opt.step","(",")",":","class",":","`","apex.optimizers.FusedLAMB","`","may","be","used","with","or","without","Amp",".","If","you","wish","to","use",":","class",":","`","FusedLAMB","`","with","Amp",",","you","may","choose","any","``","opt_level","``",":",":","opt","=","apex.optimizers.FusedLAMB","(","model.parameters","(",")",",","lr","=","....",")","model",",","opt","=","amp.initialize","(","model",",","opt",",","opt_level=","''","O0","''","or","``","O1","or","``","O2","''",")","...","opt.step","(",")","In","general",",","``","opt_level=","''","O1","''","``","is","recommended",".","LAMB","was","proposed","in","`","Large","Batch","Optimization","for","Deep","Learning",":","Training","BERT","in","76","minutes","`","_",".","Arguments",":","params","(","iterable",")",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups",".","lr","(","float",",","optional",")",":","learning","rate",".","(","default",":","1e-3",")","betas","(","Tuple","[","float",",","float","]",",","optional",")",":","coefficients","used","for","computing","running","averages","of","gradient","and","its","norm",".","(","default",":","(","0.9",",","0.999",")",")","eps","(","float",",","optional",")",":","term","added","to","the","denominator","to","improve","numerical","stability",".","(","default",":","1e-8",")","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","`","_","NOT","SUPPORTED","now","!","(","default",":","False",")","adam_w_mode","(","boolean",",","optional",")",":","Apply","L2","regularization","or","weight","decay","True","for","decoupled","weight","decay","(","also","known","as","AdamW",")","(","default",":","True",")","grad_averaging","(","bool",",","optional",")",":","whether","apply","(","1-beta2",")","to","grad","when","calculating","running","averages","of","gradient",".","(","default",":","True",")","set_grad_none","(","bool",",","optional",")",":","whether","set","grad","to","None","when","zero_grad","(",")","method","is","called",".","(","default",":","True",")","max_grad_norm","(","float",",","optional",")",":","value","used","to","clip","global","grad","norm","(","default",":","1.0",")","use_nvlamb","(","boolean",",","optional",")",":","Apply","adaptive","learning","rate","to","0.0","weight","decay","parameter","(","default",":","False",")","step_supports_amp_scaling","(","boolean",",","optional",")",":","whether","to","use","customized","gradient","unscaling","logic","(","default",":","True",")","..","_Large","Batch","Optimization","for","Deep","Learning","-","Training","BERT","in","76","minutes",":","https",":","\/\/arxiv.org\/abs\/1904.00962","..","_On","the","Convergence","of","Adam","and","Beyond",":","https",":","\/\/openreview.net\/forum","?","id=ryQu7f-RZ"]
code: def __init__(
      self,
      params,
      lr=1e-3,
      bias_correction=True,
      grad_averaging=True,
      betas=(0.9, 0.999),
      eps=1e-8,
      weight_decay=0.0,
      max_grad_norm=0.0,
      adam_w_mode=True,
      use_nvlamb=False,
      step_supports_amp_scaling=True,
      overlap_reductions=True,
      dwu_group_size=0,
      dwu_num_blocks=4,
      dwu_num_chunks=4,
      dwu_num_rs_pg=1,
      dwu_num_ar_pg=4,
      dwu_num_ag_pg=0,
      fused_norm=False,
      e5m2_allgather=False,
      verbose=False,
      clip_after_ar=True,
      full_ar=False,
      set_param_views_to_flat_buffer=False,
      skip_allgather=False,
      fuse_scale=False,
      param_order=None,
      nccl_allgather_channels=0,
  ):
      defaults = dict(
          lr=lr,
          bias_correction=bias_correction,
          betas=betas,
          eps=eps,
          weight_decay=weight_decay,
          grad_averaging=grad_averaging,
          max_grad_norm=max_grad_norm,
      )
  
      super(DistributedFusedLAMB, self).__init__(params, defaults)
  
      global fused_adam_cuda, distributed_lamb_cuda
      fused_adam_cuda = importlib.import_module("fused_adam_cuda")
      distributed_lamb_cuda = importlib.import_module(
          "distributed_lamb_cuda"
      )
  
      self._overflow_buf = torch.cuda.IntTensor([0])
      self._has_overflow = False
      self.multi_tensor_lamb_compute_update_term = (
          distributed_lamb_cuda.multi_tensor_lamb_compute_update_term
      )
      self.multi_tensor_lamb_update_weights = (
          distributed_lamb_cuda.multi_tensor_lamb_update_weights
      )
      import amp_C
  
      self.multi_tensor_l2norm = amp_C.multi_tensor_l2norm
  
      self._grad_averaging = grad_averaging
      self._adam_w_mode = 1 if adam_w_mode else 0
      self._use_nvlamb = use_nvlamb
      self._step_supports_amp_scaling = step_supports_amp_scaling
      self._is_accumulation_step = False
      self._last_step = False
      self._overlap_reductions = overlap_reductions
      self._global_scale = None
      self._num_blocks = dwu_num_blocks
      self._num_chunks = dwu_num_chunks
      self._e5m2_allgather = e5m2_allgather
      self._verbose = verbose
      self._clip_after_ar = clip_after_ar
      self._full_ar = full_ar
      self._fuse_scale = fuse_scale
      self._L2_grad_norm = None
      self._set_flat_param_view = set_param_views_to_flat_buffer
      self._skip_ag = skip_allgather
      self._fused_norm = fused_norm if not clip_after_ar else False
      self._current_process_group = c10d._get_default_group()
      self._available_ranks = get_process_group_ranks(
          self._current_process_group
      )
      self._group_size = (
          torch.cuda.device_count()
          if dwu_group_size <= 0
          else dwu_group_size
      )
      self._world_size = torch.distributed.get_world_size()
      self._num_groups = self._world_size // self._group_size
      self._rank_in_group = (
          torch.distributed.get_rank() % self._group_size
      )
  
      self._lr = torch.tensor(0.0, dtype=torch.float32, device="cuda")
  
      self._resume_from_checkpoint = False
      self._step = torch.cuda.IntTensor([0])
  
      # Master weight, moment, gradient buffers
      (
          self._fp32_p,
          self._fp32_m,
          self._fp32_v,
          self._fp16_p,
          self._fp16_g,
      ) = (None, None, None, None, None)
  
      # Check if collectives have no_copy option
      self._reduce_scatter_no_copy = (
          "no_copy"
          in inspect.getfullargspec(
              torch.distributed.reduce_scatter
          ).args
      )
      self._all_gather_no_copy = (
          "no_copy"
          in inspect.getfullargspec(torch.distributed.all_gather).args
      )
  
      if "reduce_scatter_tensor" not in dir(torch.distributed):
          torch.distributed.reduce_scatter_tensor = (
              torch.distributed._reduce_scatter_base
          )
      if "all_gather_into_tensor" not in dir(torch.distributed):
          torch.distributed.all_gather_into_tensor = (
              torch.distributed._all_gather_base
          )
  
      self._num_rs_pg = dwu_num_rs_pg
      self._num_ar_pg = dwu_num_ar_pg
      self._num_ag_pg = dwu_num_ag_pg
  
      if self._full_ar:  # full all reduce, only need AR and AG groups
          # l2_grad_norm may be reduced within a node to limit from memory reads
          for group_i in range(self._num_groups):
              ranks = [
                  group_i * self._group_size + j
                  for j in range(self._group_size)
              ]
              l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks)
              if torch.distributed.get_rank() in ranks:
                  self._l2_grad_norm_pg = l2_grad_norm_pg
  
          self._ar_pg = []
          # consider all the ranks
          ranks = list(range(0, self._world_size))
          for i in range(self._num_ar_pg):
              if self._verbose:
                  print(f"creating new AR group {i}: {ranks}")
              grp = torch.distributed.new_group(ranks=ranks)
              if grp != torch.distributed.GroupMember.NON_GROUP_MEMBER:
                  if self._verbose:
                      print(
                          f"group {i}: init barrier (device: {torch.cuda.current_device()})"
                      )
                  torch.distributed.barrier(
                      group=grp,
                      device_ids=[torch.cuda.current_device()],
                  )
              if self._verbose:
                  print(f"created new AR group {i}: {ranks}")
  
              if torch.distributed.get_rank() in ranks:
                  self._ar_pg.append(grp)
          self._ar_st = [
              torch.cuda.Stream() for _ in range(self._num_ar_pg)
          ]
          if nccl_allgather_channels > 0:
              os.putenv(
                  "NCCL_MAX_NCHANNELS", str(nccl_allgather_channels)
              )
          if self._num_ag_pg == 0:
              self._ag_pg = self._ar_pg
              self._ag_st = self._ar_st
              self._num_ag_pg = self._num_ar_pg
          else:
              self._ag_pg = []
              ranks = []
              stride = torch.cuda.device_count()
              for i in range(self._num_groups):
                  rs = list(range(i * stride, (i + 1) * stride))
                  ranks.append(rs)
              for rs in ranks:
                  for i in range(self._num_ag_pg):
                      grp = torch.distributed.new_group(ranks=rs)
                      if torch.distributed.get_rank() in rs:
                          if self._verbose:
                              print(f"creating AG group {i}: {rs}")
                          self._ag_pg.append(grp)
  
              self._ag_st = [
                  torch.cuda.Stream() for _ in range(self._num_ag_pg)
              ]
      else:  # reduce-scatter + all-reduce, need RS, AR, AG groups
          if self._num_groups > 1:
              self._ar_pg = []
              for dev_i in range(self._group_size):
                  ranks = [
                      dev_i + j * self._group_size
                      for j in range(self._num_groups)
                  ]
                  for i in range(self._num_ar_pg):
                      if self._verbose:
                          print(f"creating new AR group {i}: {ranks}")
                      grp = torch.distributed.new_group(ranks=ranks)
                      if (
                          grp
                          != torch.distributed.GroupMember.NON_GROUP_MEMBER
                      ):
                          if self._verbose:
                              print(
                                  f"group {i}: init barrier (device: {torch.cuda.current_device()})"
                              )
                          torch.distributed.barrier(
                              group=grp,
                              device_ids=[torch.cuda.current_device()],
                          )
                      if self._verbose:
                          print(f"created new AR group {i}: {ranks}")
  
                      if torch.distributed.get_rank() in ranks:
                          self._ar_pg.append(grp)
              self._ar_st = [
                  torch.cuda.Stream() for _ in range(self._num_ar_pg)
              ]
          rs_ranks = []
          for group_i in range(self._num_groups):
              rs_ranks.append(
                  [
                      group_i * self._group_size + j
                      for j in range(self._group_size)
                  ]
              )
          self._rs_pg = []
          for group_i in range(self._num_groups):
              ranks = rs_ranks[group_i]
              for i in range(self._num_rs_pg):
                  grp = torch.distributed.new_group(ranks=ranks)
                  if torch.distributed.get_rank() in ranks:
                      self._rs_pg.append(grp)
                      if self._verbose:
                          print(f"creating RS group : {ranks}")
              l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks)
              if torch.distributed.get_rank() in ranks:
                  self._l2_grad_norm_pg = l2_grad_norm_pg
          self._rs_st = [
              torch.cuda.Stream() for _ in range(self._num_rs_pg)
          ]
          if self._num_ag_pg == 0:
              self._ag_pg = self._rs_pg
              self._ag_st = self._rs_st
              self._num_ag_pg = self._num_rs_pg
          else:
              self._ag_pg = []
              for group_i in range(self._num_groups):
                  ranks = rs_ranks[group_i]
                  for i in range(self._num_ag_pg):
                      grp = torch.distributed.new_group(ranks=ranks)
                      if torch.distributed.get_rank() in ranks:
                          self._ag_pg.append(grp)
                          if self._verbose:
                              print(f"creating AG group : {ranks}")
              self._ag_st = [
                  torch.cuda.Stream() for _ in range(self._num_ag_pg)
              ]
      for ag_pg in self._ag_pg:
          torch.distributed.barrier(group=ag_pg)
  
      self._l2_grad_norm_st = torch.cuda.Stream()
      self._completion_st = torch.cuda.Stream()
      self._step.record_stream(self._completion_st)
  
      self._reductions_works = [None] * self._num_blocks
      self._allgather_works = [None] * self._num_blocks
  
      self._one = torch.cuda.IntTensor([1])
  
      self._first_step = True
      self._lazy_init_stage1_done, self._lazy_init_stage2_done = (
          False,
          False,
      )
      self._param_order = self.AtomicCounter()
  
      p_offset = 0
      p_i = 0
      self._model_params = []
      self._grad_accs = []
      self._group_properties = []
      for group in self.param_groups:
          prev = None
          beta1, beta2 = group["betas"]
          beta3 = 1.0 - beta1 if self._grad_averaging else 1.0
          bias_correction = 1 if group["bias_correction"] else 0
          eps = group["eps"]
          weight_decay = group["weight_decay"]
          for p in group["params"]:
              if not p.requires_grad:
                  continue
              self._model_params.append(p)
              self._group_properties.append(
                  (
                      weight_decay,
                      bias_correction,
                      beta1,
                      beta2,
                      beta3,
                      eps,
                  )
              )
              p_grads_size = p.numel()
              if self._set_flat_param_view:
                  if param_order:
                      # this is executed when param_order is specified by the user
                      self._param_order.add(param_order[p])
                  else:
                      self._param_order.add(p_i)
              p_offset += p_grads_size
              # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
              # RNN is one example of consecutive parameters:
              # (weight_ih, weight_hh, bias_ih, bias_hh)
              if prev is not None and (
                  prev.data_ptr() + prev.numel() * prev.element_size()
                  != p.data_ptr()
              ):
                  p_offset = ((p_offset + 63) // 64) * 64
              prev = p
              p_i += 1
      if param_order:
          self._param_order.order = torch.argsort(
              torch.tensor(self._param_order.order)
          ).tolist()
      self._grads_generated = [False] * len(self._model_params)
      self._grads_fp16, self._grads_fp32 = [], []
      if self._overlap_reductions:
          self._current_block = self._num_blocks
  
      self._net_total_param_size = p_offset
      self._total_param_size = p_offset
      dwu_min_page_size = (
          256 * self._num_blocks * self._num_chunks * self._group_size
      )
      self._total_param_size = (
          (self._total_param_size + dwu_min_page_size - 1)
          // dwu_min_page_size
      ) * dwu_min_page_size
      self._new_params = torch.zeros(
          [self._total_param_size],
          dtype=torch.uint8 if self._e5m2_allgather else torch.float16,
          device="cuda",
      )
  
imports: import os
  import math
  import inspect
  import torch
  import importlib
  import amp_C
  from apex.multi_tensor_apply import multi_tensor_applier
  import torch.distributed.distributed_c10d
label_desc: Use image node_id 1 to create a new DistributedFusedLAMB object from inherited base classes: torch.optim with example: obj = DistributedFusedLAMB(params, lr, bias_correction, grad_averaging, betas, eps, weight_decay, max_grad_norm, adam_w_mode, use_nvlamb, step_supports_amp_scaling, overlap_reductions, dwu_group_size, dwu_num_blocks, dwu_num_chunks, dwu_num_rs_pg, dwu_num_ar_pg, dwu_num_ag_pg, fused_norm, e5m2_allgather, verbose, clip_after_ar, full_ar, set_param_views_to_flat_buffer, skip_allgather, fuse_scale, param_order, nccl_allgather_channels)