class_docstr: Implements the 1-bit Lamb algorithm. Currently GPU-only.
  For usage example please see https://www.deepspeed.ai/tutorials/onebit-lamb/
  For technical details please see our paper https://arxiv.org/abs/2104.06069.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups.
      lr (float, optional): learning rate. (default: 1e-3)
      freeze_step (int, optional): Number of steps for warmup (uncompressed)
          stage before we start using compressed communication. (default 100000)
      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)
      weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
      max_coeff(float, optional): maximum value of the lamb coefficient (default: 10.0)
      min_coeff(float, optional): minimum value of the lamb coefficient (default: 0.01)
      amsgrad (boolean, optional): whether to use the AMSGrad variant of this
          algorithm from the paper `On the Convergence of Adam and Beyond`_
          (default: False) NOT SUPPORTED in 1-bit Lamb!
      eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
          adds eps to the bias-corrected second moment estimate before
          evaluating square root instead of adding it to the square root of
          second moment estimate as in the original paper. (default: False)
      cuda_aware (boolean, required): Set True if the underlying MPI implementation
          supports CUDA-Aware communication. (default: False)
      comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl')
      coeff_beta (float, optional): coefficient used for computing
          running averages of lamb coefficient (default: 0.9) note that you may want to
          increase or decrease this beta depending on the freeze_step you choose, as
          1/(1 - coeff_beta) should be smaller than or equal to freeze_step
      factor_max (float, optional): maximum value of scaling factor to the frozen lamb
          coefficient during compression stage (default: 4.0)
      factor_min (float, optional): minimum value of scaling factor to the frozen lamb
          coefficient during compression stage (default: 0.5)
      factor_threshold (float, optional): threshold of how much the scaling factor can
          fluctuate between steps (default: 0.1)
  .. _Large Batch Optimization for Deep Learning\: Training BERT in 76 minutes:
      https://arxiv.org/abs/1904.00962
  .. _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
class_docstr_tok: ["Implements","the","1-bit","Lamb","algorithm",".","Currently","GPU-only",".","For","usage","example","please","see","https",":","\/\/www.deepspeed.ai\/tutorials\/onebit-lamb\/","For","technical","details","please","see","our","paper","https",":","\/\/arxiv.org\/abs\/2104.06069",".","Arguments",":","params","(","iterable",")",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups",".","lr","(","float",",","optional",")",":","learning","rate",".","(","default",":","1e-3",")","freeze_step","(","int",",","optional",")",":","Number","of","steps","for","warmup","(","uncompressed",")","stage","before","we","start","using","compressed","communication",".","(","default","100000",")","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",")","weight_decay","(","float",",","optional",")",":","weight","decay","(","L2","penalty",")","(","default",":","0",")","max_coeff","(","float",",","optional",")",":","maximum","value","of","the","lamb","coefficient","(","default",":","10.0",")","min_coeff","(","float",",","optional",")",":","minimum","value","of","the","lamb","coefficient","(","default",":","0.01",")","amsgrad","(","boolean",",","optional",")",":","whether","to","use","the","AMSGrad","variant","of","this","algorithm","from","the","paper","`","On","the","Convergence","of","Adam","and","Beyond","`","_","(","default",":","False",")","NOT","SUPPORTED","in","1-bit","Lamb","!","eps_inside_sqrt","(","boolean",",","optional",")",":","in","the","'update","parameters","'","step",",","adds","eps","to","the","bias-corrected","second","moment","estimate","before","evaluating","square","root","instead","of","adding","it","to","the","square","root","of","second","moment","estimate","as","in","the","original","paper",".","(","default",":","False",")","cuda_aware","(","boolean",",","required",")",":","Set","True","if","the","underlying","MPI","implementation","supports","CUDA-Aware","communication",".","(","default",":","False",")","comm_backend_name","(","string",",","optional",")",":","Set","to","'mpi","'","if","needed",".","(","default",":","'nccl","'",")","coeff_beta","(","float",",","optional",")",":","coefficient","used","for","computing","running","averages","of","lamb","coefficient","(","default",":","0.9",")","note","that","you","may","want","to","increase","or","decrease","this","beta","depending","on","the","freeze_step","you","choose",",","as","1\/","(","1","-","coeff_beta",")","should","be","smaller","than","or","equal","to","freeze_step","factor_max","(","float",",","optional",")",":","maximum","value","of","scaling","factor","to","the","frozen","lamb","coefficient","during","compression","stage","(","default",":","4.0",")","factor_min","(","float",",","optional",")",":","minimum","value","of","scaling","factor","to","the","frozen","lamb","coefficient","during","compression","stage","(","default",":","0.5",")","factor_threshold","(","float",",","optional",")",":","threshold","of","how","much","the","scaling","factor","can","fluctuate","between","steps","(","default",":","0.1",")","..","_Large","Batch","Optimization","for","Deep","Learning\\",":","Training","BERT","in","76","minutes",":","https",":","\/\/arxiv.org\/abs\/1904.00962","..","_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"]
code: def get_lamb_coeffs(self):
      return self.lamb_coeffs
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/runtime/fp16/onebit/lamb.py
imports: import types
  import torch
  import numpy
  from deepspeed import comm
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
label_desc: Use image node_id 4 for calling the OnebitLamb obj's underlying member method code with example usage: obj.get_lamb_coeffs() and returns: self