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_stage1(self):
      if self._lazy_init_stage1_done:
          return
  
      p_i = 0
      # self._model_params = []
      # self._grad_accs = []
      # self._group_properties = []
      for group in self.param_groups:
          for p in group["params"]:
              torch.distributed.broadcast(p, 0)
              if not p.requires_grad:
                  continue
  
              def wrapper(param, param_i):
                  param_tmp = param.expand_as(param)
                  grad_acc = param_tmp.grad_fn.next_functions[0][0]
  
                  def allreduce_hook(*unused):
                      if not self._set_flat_param_view:
                          if self._first_step:
                              # first time
                              self._param_order.add(param_i)
                          else:
                              idx = self._param_order.order.index(
                                  param_i
                              )
                              self._do_overlapped_reduction(idx, param)
                      else:
                          if not self._first_step:
                              idx = self._param_order.order.index(
                                  param_i
                              )
                              self._do_overlapped_reduction(idx, param)
  
                  grad_acc.register_hook(allreduce_hook)
                  self._grad_accs.append(grad_acc)
  
              wrapper(p, p_i)
              p_i += 1
  
      self._block_size = self._total_param_size // self._num_blocks
      self._chunk_size = self._block_size // self._num_chunks
      self._shard_size = self._chunk_size // self._group_size
  
      self._flat_grads = torch.zeros(
          [self._total_param_size], dtype=torch.float16, device="cuda"
      )
      self._mega_shard_size = (
          self._num_blocks * self._num_chunks * self._shard_size
      )
      # initialize master weights, moments buffers if not loaded from checkpoint
      if self._fp32_p is None:
          self._fp32_p = torch.zeros(
              [self._mega_shard_size],
              dtype=torch.float32,
              device="cuda",
          )
          self._fp32_m = torch.zeros(
              [self._mega_shard_size],
              dtype=torch.float32,
              device="cuda",
          )
          self._fp32_v = torch.zeros(
              [self._mega_shard_size],
              dtype=torch.float32,
              device="cuda",
          )
          self._fp32_u = torch.zeros(
              [self._mega_shard_size],
              dtype=torch.float32,
              device="cuda",
          )
      # FIXME: Rethink fp16 label since it's either uint8 or fp16
      self._fp16_p = torch.zeros(
          [self._mega_shard_size],
          dtype=torch.uint8 if self._e5m2_allgather else torch.float16,
          device="cuda",
      )
      self._fp16_g = torch.zeros(
          [self._mega_shard_size], dtype=torch.float16, device="cuda"
      )
  
      def _flat_split(p):
          def __blockify(p):
              return [
                  p[
                      block_id
                      * self._block_size : (block_id + 1)
                      * self._block_size
                  ]
                  for block_id in range(self._num_blocks)
              ]
  
          def __chunkify(p):
              return [
                  p[
                      chunk_id
                      * self._chunk_size : (chunk_id + 1)
                      * self._chunk_size
                  ]
                  for chunk_id in range(self._num_chunks)
              ]
  
          def __shardify(p):
              return [
                  p[
                      shard_id
                      * self._shard_size : (shard_id + 1)
                      * self._shard_size
                  ]
                  for shard_id in range(self._group_size)
              ]
  
          list_of_blocks = __blockify(p)
          list_of_list_of_chunks = [
              __chunkify(block) for block in list_of_blocks
          ]
          list_of_list_of_list_of_shards = [
              [__shardify(chunk) for chunk in chunks]
              for chunks in list_of_list_of_chunks
          ]
          return (
              list_of_blocks,
              list_of_list_of_chunks,
              list_of_list_of_list_of_shards,
          )
  
      # note(crcrpar): the function below doesn't seem to be used at all.
      # def _flat_split_no_shards(p):
      #     def __blockify(p):
      #         return [p[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
      #     def __chunkify(p):
      #         return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._num_chunks)]
      #     list_of_blocks = __blockify(self._flat_grads)
      #     list_of_list_of_chunks = [__chunkify(block) for block in list_of_blocks]
      #     return list_of_blocks, list_of_list_of_chunks
  
      def _full_packed_split(p):
          def __shardify(p):
              return [
                  p[
                      mega_shard
                      * self._mega_shard_size : (mega_shard + 1)
                      * self._mega_shard_size
                  ]
                  for mega_shard in range(self._group_size)
              ]
  
          def __blockify(p):
              return [
                  p[
                      block_id
                      * self._num_chunks
                      * self._shard_size : (block_id + 1)
                      * self._num_chunks
                      * self._shard_size
                  ]
                  for block_id in range(self._num_blocks)
              ]
  
          def __chunkify(p):
              return [
                  p[
                      chunk_id
                      * self._shard_size : (chunk_id + 1)
                      * self._shard_size
                  ]
                  for chunk_id in range(self._num_chunks)
              ]
  
          list_of_mega_shards = __shardify(p)
          list_of_list_of_mega_blocks = [
              __blockify(mega_shard)
              for mega_shard in list_of_mega_shards
          ]
          list_of_list_of_list_of_mega_chunks = [
              [__chunkify(mega_block) for mega_block in mega_blocks]
              for mega_blocks in list_of_list_of_mega_blocks
          ]
          return (
              list_of_mega_shards,
              list_of_list_of_mega_blocks,
              list_of_list_of_list_of_mega_chunks,
          )
  
      def _packed_split(p):
          def __packed_blockify(p):
              packed_block_size = self._num_chunks * self._shard_size
              return [
                  p[
                      block_id
                      * packed_block_size : (block_id + 1)
                      * packed_block_size
                  ]
                  for block_id in range(self._num_blocks)
              ]
  
          def __packed_chunkify(p):
              # in the packed format, each chunk contains one shard, so packed_chunk_size == self._shard_size
              return [
                  p[
                      chunk_id
                      * self._shard_size : (chunk_id + 1)
                      * self._shard_size
                  ]
                  for chunk_id in range(self._num_chunks)
              ]
  
          list_of_blocks = __packed_blockify(p)
          list_of_list_of_chunks = [
              __packed_chunkify(block) for block in list_of_blocks
          ]
          return list_of_blocks, list_of_list_of_chunks
  
      def _split_assign(shards):
          packed_block_size = self._num_chunks * self._shard_size
          list_of_list_of_chunks = []
          for block_id in range(self._num_blocks):
              list_of_chunks = []
              for chunk_id in range(self._num_chunks):
                  # self._fp16_g[block_id*packed_block_size+chunk_id*self._shard_size:block_id*packed_block_size+(chunk_id+1)*self._shard_size] = shards[block_id][chunk_id][self._rank_in_group]
                  list_of_chunks.append(
                      shards[block_id][chunk_id][self._rank_in_group]
                  )
              list_of_list_of_chunks.append(list_of_chunks)
          return list_of_list_of_chunks
  
      (
          self._new_params_mega_shards,
          self._new_params_mega_blocks,
          self._new_params_mega_chunks,
      ) = _full_packed_split(self._new_params)
      # this splitting scheme is needed when allgather needs to be split into multiple chunks in a contiguous way
      (
          self._new_params2_blocks,
          self._new_params2_chunks,
          self._new_params2_shards,
      ) = _flat_split(self._new_params)
  
      self._fp32_p_blocks, self._fp32_p_chunks = _packed_split(
          self._fp32_p
      )
      self._fp32_m_blocks, self._fp32_m_chunks = _packed_split(
          self._fp32_m
      )
      self._fp32_v_blocks, self._fp32_v_chunks = _packed_split(
          self._fp32_v
      )
      self._fp32_u_blocks, self._fp32_u_chunks = _packed_split(
          self._fp32_u
      )
      self._fp16_p_blocks, self._fp16_p_chunks = _packed_split(
          self._fp16_p
      )
  
      if self._full_ar:
          # for gradient all-reduce
          (
              self._flat_grads_blocks,
              self._flat_grads_chunks,
              self._flat_grads_shards,
          ) = _flat_split(self._flat_grads)
          # for weight update
          self._fp16_g_chunks = _split_assign(self._flat_grads_shards)
      else:
          (
              self._flat_grads_blocks,
              self._flat_grads_chunks,
              self._flat_grads_shards,
          ) = _flat_split(self._flat_grads)
          self._fp16_g_blocks, self._fp16_g_chunks = _packed_split(
              self._fp16_g
          )
  
      self._lazy_init_stage1_done = True
  
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 2 for calling the DistributedFusedLAMB obj's underlying member method code with example usage: obj._lazy_init_stage1() and returns: list_of_blocks,  list_of_list_of_chunks,  list_of_list_of_list_of_shards, list_of_mega_shards,  list_of_list_of_mega_blocks,  list_of_list_of_list_of_mega_chunks, list_of_blocks,  list_of_list_of_chunks, list_of_list_of_chunks, unknown, unknown, unknown, unknown, unknown, unknown, unknown, unknown