class_docstr: Implements the 1-bit Adam algorithm. Currently GPU-only.
  For usage example please see https://www.deepspeed.ai/tutorials/onebit-adam/
  For technical details please read https://arxiv.org/abs/2102.02888
  
  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)
      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 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","1-bit","Adam","algorithm",".","Currently","GPU-only",".","For","usage","example","please","see","https",":","\/\/www.deepspeed.ai\/tutorials\/onebit-adam\/","For","technical","details","please","read","https",":","\/\/arxiv.org\/abs\/2102.02888","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",")","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","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.freeze_step
      ):
          if dist.get_rank() == 0:
              print(
                  "Checkpoint loaded and OnebitAdam warmup stage starts/continues."
              )
          if self.adam_freeze_key is True:
              self.adam_freeze_key = False
              if self.using_pipeline:
                  self.deepspeed.pipeline_enable_backward_allreduce = (
                      True
                  )
              else:
                  self.deepspeed.enable_backward_allreduce = True
      else:
          if dist.get_rank() == 0:
              print(
                  "Checkpoint loaded and OnebitAdam compression stage starts/continues."
              )
          if self.adam_freeze_key is False:
              self.adam_freeze_key = True
              if self.using_pipeline:
                  self.deepspeed.pipeline_enable_backward_allreduce = (
                      False
                  )
              else:
                  self.deepspeed.enable_backward_allreduce = False
      # We reset the compression errors when loading checkpoints for 3 reasons:
      # 1) The worker and server error at each GPU are distinct, so in current implementation
      # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors.
      # If we want to save them correctly we need O(num_gpu*model_size) memory in order to
      # gather all the error, which is a very large memory requirement. It's possible to save
      # them in a distributed way, but it will make the checkpoint saving/loading much more complicated.
      # 2) Even if we are able to save the compression errors correctly, you need to have the
      # exact same number of GPUs in order to load them correctly.
      # 3) We verified on BERT pre-training that occasionally resetting the compression error
      # at checkpoint loading does not affect the convergence.
      # However, please avoid frequent checkpoint loading which could break the error
      # compensation mechanism thus affect the convergence.
      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")
  
file_path: deepspeed/deepspeed/runtime/fp16/onebit/adam.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 OnebitAdam obj's underlying member method code with example usage: obj.load_state_dict(state_dict) without return types