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 _lazy_init_stage2(self):
      if self._lazy_init_stage2_done:
          return
      if not self._set_flat_param_view:
          # reversing is needed for overlapping allreduce and backprop, but currently not supported for flat param view
          self._param_order.order.reverse()
  
          # re-order model_params, grad_accs, group_properties lists
      self._model_params = [
          self._model_params[i] for i in self._param_order.order
      ]
      self._grad_accs = [
          self._grad_accs[i] for i in self._param_order.order
      ]
      self._group_properties = [
          self._group_properties[i] for i in self._param_order.order
      ]
  
      def _get_flat_view(param):
          if param.is_contiguous(memory_format=torch.channels_last):
              K, C, H, W = param.shape
              pv = param.as_strided(
                  size=(K, H, W, C), stride=(H * W * C, W * C, C, 1)
              )
          elif param.is_contiguous(
              memory_format=torch.channels_last_3d
          ):
              K, C, D, H, W = param.shape
              pv = param.as_strided(
                  size=(K, D, H, W, C),
                  stride=(D * H * W * C, H * W * C, W * C, C, 1),
              )
          else:
              pv = param
          return pv.view(-1)
  
      # re-collect grads info (size, offset) after ordering
      prev = None
      p_offset = 0
      self._grads_info = []
      self._individual_flat_grads = []
      for i, p in enumerate(self._model_params):
          p_grads_size = p.numel()
          self._grads_info.append(
              {
                  "param_grads_size": p_grads_size,
                  "param_offset": p_offset,
              }
          )
          self._individual_flat_grads.append(
              self._flat_grads[
                  p_offset : p_offset + p_grads_size
              ].view_as(p)
          )
          # for the first iteration
          self._do_overlapped_reduction(i, p)
          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
  
      self._low_param_i = [0] * self._num_blocks
      for block_id in range(self._num_blocks - 1, -1, -1):
          p_i = len(self._grads_info) - 1
          while (
              p_i > 0
              and self._grads_info[p_i]["param_offset"]
              > block_id * self._block_size
          ):
              p_i -= 1
          self._low_param_i[block_id] = p_i
      # print("self._low_param_i", self._low_param_i)
  
      # This paragraph does two things:
      # 1) Copy model parameters into master buffer
      # 2) Create tensor lists for unpacking new parameter tensor after all-gather
      self._packed_flat_to_model_params_fp16 = []
      self._packed_flat_to_model_params_fp32 = []
      self._model_params_num = len(self._model_params)
      self._contrib_tensor_list = []
      self._contrib_min_param_i, self._contrib_max_param_i = -1, -1
      self._contrib_update_frag_for_norm = []
      self._contrib_model_param_for_norm_fp16 = []
      self._contrib_model_param_for_norm_fp32 = []
      self._contrib_model_param_for_norm_is_fp16 = []
      self._model_param_is_contrib = []
      self._contrib_group_properties = []
      for shard_id in range(self._group_size):
          for block_id in range(self._num_blocks):
              for chunk_id in range(self._num_chunks):
                  flat_shard_start = (
                      (
                          (block_id * self._num_chunks + chunk_id)
                          * self._group_size
                      )
                      + shard_id
                  ) * self._shard_size
                  flat_shard_end = flat_shard_start + self._shard_size
                  for param_i, (
                      p,
                      grads_info,
                      group_props,
                  ) in enumerate(
                      zip(
                          self._model_params,
                          self._grads_info,
                          self._group_properties,
                      )
                  ):
                      flat_grad_start = grads_info["param_offset"]
                      flat_grad_end = (
                          flat_grad_start
                          + grads_info["param_grads_size"]
                      )
                      clipped_start = (lambda a, b: a if a > b else b)(
                          flat_grad_start, flat_shard_start
                      )
                      clipped_end = (lambda a, b: a if a < b else b)(
                          flat_grad_end, flat_shard_end
                      )
                      if clipped_start < clipped_end:
                          grad_offset = clipped_start - flat_grad_start
                          grad_length = clipped_end - clipped_start
                          shard_offset = (
                              clipped_start - flat_shard_start
                          )
                          pf = _get_flat_view(p)
                          model_param_fragment = pf[
                              grad_offset : grad_offset + grad_length
                          ]
                          new_param_packed_fragment = (
                              self._new_params_mega_chunks[shard_id][
                                  block_id
                              ][chunk_id][
                                  shard_offset : shard_offset
                                  + grad_length
                              ]
                          )
                          if (
                              model_param_fragment.dtype
                              == torch.float16
                          ):
                              self._packed_flat_to_model_params_fp16.append(
                                  (
                                      new_param_packed_fragment,
                                      model_param_fragment,
                                  )
                              )
                          else:
                              self._packed_flat_to_model_params_fp32.append(
                                  (
                                      new_param_packed_fragment,
                                      model_param_fragment,
                                  )
                              )
                          if shard_id == self._rank_in_group:
                              self._model_param_is_contrib.append(
                                  param_i
                              )
                              # copy model parameters into master buffer
                              master_param_fragment = (
                                  self._fp32_p_chunks[block_id][
                                      chunk_id
                                  ][
                                      shard_offset : shard_offset
                                      + grad_length
                                  ]
                              )
                              opti_state_m_fragment = (
                                  self._fp32_m_chunks[block_id][
                                      chunk_id
                                  ][
                                      shard_offset : shard_offset
                                      + grad_length
                                  ]
                              )
                              opti_state_v_fragment = (
                                  self._fp32_v_chunks[block_id][
                                      chunk_id
                                  ][
                                      shard_offset : shard_offset
                                      + grad_length
                                  ]
                              )
                              opti_state_u_fragment = (
                                  self._fp32_u_chunks[block_id][
                                      chunk_id
                                  ][
                                      shard_offset : shard_offset
                                      + grad_length
                                  ]
                              )
                              opti_state_g_fragment = (
                                  self._fp16_g_chunks[block_id][
                                      chunk_id
                                  ][
                                      shard_offset : shard_offset
                                      + grad_length
                                  ]
                              )
                              opti_state_p_fragment = (
                                  self._fp16_p_chunks[block_id][
                                      chunk_id
                                  ][
                                      shard_offset : shard_offset
                                      + grad_length
                                  ]
                              )
                              # print("model_param_fragment.size()=%s, new_param_packed_fragment.size()=%s, master_param_fragment.size()=%s" % (str(model_param_fragment.size()), str(new_param_packed_fragment.size()), str(master_param_fragment.size())))
                              if not self._resume_from_checkpoint:
                                  master_param_fragment.copy_(
                                      model_param_fragment
                                  )
                              self._contrib_group_properties.append(
                                  group_props
                              )
                              self._contrib_tensor_list.append(
                                  (
                                      master_param_fragment,
                                      opti_state_m_fragment,
                                      opti_state_v_fragment,
                                      opti_state_u_fragment,
                                      opti_state_g_fragment,
                                      opti_state_p_fragment,
                                  )
                              )  # p, m, v, u, g, p_copy
                              self._contrib_update_frag_for_norm.append(
                                  opti_state_u_fragment
                              )
                              if p.dtype == torch.float16:
                                  self._contrib_model_param_for_norm_fp16.append(
                                      p
                                  )
                              else:
                                  self._contrib_model_param_for_norm_fp32.append(
                                      p
                                  )
                              self._contrib_model_param_for_norm_is_fp16.append(
                                  True
                                  if p.dtype == torch.float16
                                  else False
                              )
                              if self._contrib_min_param_i < 0:
                                  self._contrib_min_param_i = param_i
                              self._contrib_max_param_i = param_i
      self._contrib_model_param_for_norm_num = len(
          self._contrib_model_param_for_norm_is_fp16
      )
      if len(self._contrib_model_param_for_norm_fp16) == 0:
          self._contrib_model_param_for_norm_fp16 = None
      if len(self._contrib_model_param_for_norm_fp32) == 0:
          self._contrib_model_param_for_norm_fp32 = None
      self._contrib_model_param_for_norm_is_fp32 = torch.tensor(
          [
              not is_fp16
              for is_fp16 in self._contrib_model_param_for_norm_is_fp16
          ],
          dtype=torch.bool,
          device="cuda",
      )
      self._contrib_model_param_for_norm_is_fp16 = torch.tensor(
          [
              is_fp16
              for is_fp16 in self._contrib_model_param_for_norm_is_fp16
          ],
          dtype=torch.bool,
          device="cuda",
      )
      self._offsets = torch.tensor(
          self._model_param_is_contrib, dtype=torch.int64, device="cuda"
      )
  
      p, m, v, u, g, p_copy = list(zip(*self._contrib_tensor_list))
      self._contrib_compute_update_term_tensor_list = [g, p, m, v, u]
      self._contrib_update_weights_tensor_list = [u, p, p_copy]
  
      math_type = self._fp32_u.dtype
      decay, bias_correction, beta1, beta2, beta3, epsilon = list(
          zip(*self._contrib_group_properties)
      )
      self._contrib_beta1 = torch.tensor(
          beta1, dtype=math_type, device="cuda"
      )
      self._contrib_beta2 = torch.tensor(
          beta2, dtype=math_type, device="cuda"
      )
      self._contrib_beta3 = torch.tensor(
          beta3, dtype=math_type, device="cuda"
      )
      self._contrib_bias_correction = torch.tensor(
          bias_correction, dtype=torch.int, device="cuda"
      )
      self._contrib_epsilon = torch.tensor(
          epsilon, dtype=math_type, device="cuda"
      )
      self._contrib_weight_decay = torch.tensor(
          decay, dtype=math_type, device="cuda"
      )
  
      self._packed_flat_to_model_params_fp16 = (
          list(zip(*self._packed_flat_to_model_params_fp16))
          if len(self._packed_flat_to_model_params_fp16) > 0
          else None
      )
      self._packed_flat_to_model_params_fp32 = (
          list(zip(*self._packed_flat_to_model_params_fp32))
          if len(self._packed_flat_to_model_params_fp32) > 0
          else None
      )
  
      self._lazy_init_stage2_done = True
  
      self.complete_reductions()
      self._first_step = False
  
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 3 for calling the DistributedFusedLAMB obj's underlying member method code with example usage: obj._lazy_init_stage2() and returns: pv