class_docstr: Implements the 0/1 Adam algorithm. Currently GPU-only.
  For usage example please see https://www.deepspeed.ai/tutorials/zero-one-adam/
  For technical details please read https://arxiv.org/abs/2202.06009
  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 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)
      var_freeze_step (int, optional): The latest step to update the variance,
          using the notation from https://arxiv.org/abs/2202.06009, it denotes the
          max{i|i in T_v}. Note that this is different from the freeze step from the
          1-bit Adam. The var_freeze_step is usually the end of the learning rate warmup
          and thus does not require tuning. (default: 100000)
      var_update_scaler (int, optional): The interval to update the variance. Note that
          the update policy for variance follows an exponential rule, where var_update_scaler
          denotes the kappa in the 0/1 Adam paper. (default: 16)
      local_step_scaler (int, optional): The interval to scale the local steps interval
          according to the learning rate policy. (default: 32678)
      local_step_clipper (int, optional): The largest interval for local steps with
          learning rate policy. This corresponds to the variable H in the 0/1 Adam paper.
          (default: 16)
      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 0/1 Adam!
      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')
  .. _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","0\/1","Adam","algorithm",".","Currently","GPU-only",".","For","usage","example","please","see","https",":","\/\/www.deepspeed.ai\/tutorials\/zero-one-adam\/","For","technical","details","please","read","https",":","\/\/arxiv.org\/abs\/2202.06009","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","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",")","var_freeze_step","(","int",",","optional",")",":","The","latest","step","to","update","the","variance",",","using","the","notation","from","https",":","\/\/arxiv.org\/abs\/2202.06009",",","it","denotes","the","max","{","i|i","in","T_v","}",".","Note","that","this","is","different","from","the","freeze","step","from","the","1-bit","Adam",".","The","var_freeze_step","is","usually","the","end","of","the","learning","rate","warmup","and","thus","does","not","require","tuning",".","(","default",":","100000",")","var_update_scaler","(","int",",","optional",")",":","The","interval","to","update","the","variance",".","Note","that","the","update","policy","for","variance","follows","an","exponential","rule",",","where","var_update_scaler","denotes","the","kappa","in","the","0\/1","Adam","paper",".","(","default",":","16",")","local_step_scaler","(","int",",","optional",")",":","The","interval","to","scale","the","local","steps","interval","according","to","the","learning","rate","policy",".","(","default",":","32678",")","local_step_clipper","(","int",",","optional",")",":","The","largest","interval","for","local","steps","with","learning","rate","policy",".","This","corresponds","to","the","variable","H","in","the","0\/1","Adam","paper",".","(","default",":","16",")","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","0\/1","Adam","!","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","'",")","..","_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"]
docstr: Overrides load_state_dict() to add special handling when loading checkpoints
docstr_tok: ["Overrides","load_state_dict","(",")","to","add","special","handling","when","loading","checkpoints"]
code: def load_state_dict(self, state_dict):
      """
      Overrides load_state_dict() to add special handling when loading checkpoints
      """
      # Because at different stage exp_avg_mask may change (e.g.,
      # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask
      # in checkpoints but always use the one user provided in training script.
      # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.)
      # Thus here we keep the exp_avg_mask unchanged when loading checkpoint
      for i, group in enumerate(self.param_groups):
          if "exp_avg_mask" in group:
              state_dict["param_groups"][i]["exp_avg_mask"] = group[
                  "exp_avg_mask"
              ]
          elif (
              "exp_avg_mask" not in group
              and "exp_avg_mask" in state_dict["param_groups"][i]
          ):
              state_dict["param_groups"][i].pop("exp_avg_mask")
      super().load_state_dict(state_dict)
      if (
          self.state[self.param_groups[0]["params"][0]]["step"]
          < self.var_freeze_step
      ):
          self.var_freeze_key = False
          if (
              self.state[self.param_groups[0]["params"][0]]["step"] + 1
          ) % self.state[self.param_groups[0]["params"][0]][
              "var_interval"
          ] == 0:
              if self.using_pipeline:
                  self.deepspeed.pipeline_enable_backward_allreduce = (
                      True
                  )
              else:
                  self.deepspeed.enable_backward_allreduce = True
          else:
              if self.using_pipeline:
                  self.deepspeed.pipeline_enable_backward_allreduce = (
                      False
                  )
              else:
                  self.deepspeed.enable_backward_allreduce = False
      else:
          self.var_freeze_key = True
          if self.using_pipeline:
              self.deepspeed.pipeline_enable_backward_allreduce = False
          else:
              self.deepspeed.enable_backward_allreduce = False
      self.reinitial_error_buffer = False
      for group in self.param_groups:
          for p in group["params"]:
              if "worker_error" in self.state[p]:
                  self.state[p].pop("worker_error")
              if "server_error" in self.state[p]:
                  self.state[p].pop("server_error")
              if "momentum_accumulator" in self.state[p]:
                  self.state[p].pop("momentum_accumulator")
  
file_path: DeepSpeed/deepspeed/runtime/fp16/onebit/zoadam.py
imports: import types
  import torch
  import numpy
  from deepspeed.accelerator import get_accelerator
  from deepspeed.runtime.utils import required_torch_version
  from deepspeed import comm
label_desc: Use image node_id 3 for calling the ZeroOneAdam obj's underlying member method code with example usage: obj.load_state_dict(state_dict) without return types