docstr: Update the model parameters.
  
  .. note::
      This method will be called internally by ZeRO-Offload. DeepSpeed
      users should still use ``engine.step()`` as shown in the
      `Getting Started
      <https://www.deepspeed.ai/getting-started/#training>`_ guide.
  
  Args:
      closure (callable, optional): closure to compute the loss.
          Defaults to ``None``.
      fp16_param_groups: FP16 GPU parameters to update. Performing the
          copy here reduces communication time. Defaults to ``None``.
  
  Returns:
      loss: if ``closure`` is provided. Otherwise ``None``.
docstr_tok: ["Update","the","model","parameters",".","..","note",":",":","This","method","will","be","called","internally","by","ZeRO-Offload",".","DeepSpeed","users","should","still","use","``","engine.step","(",")","``","as","shown","in","the","`","Getting","Started","<","https",":","\/\/www.deepspeed.ai\/getting-started\/","#","training",">","`","_","guide",".","Args",":","closure","(","callable",",","optional",")",":","closure","to","compute","the","loss",".","Defaults","to","``","None","``",".","fp16_param_groups",":","FP16","GPU","parameters","to","update",".","Performing","the","copy","here","reduces","communication","time",".","Defaults","to","``","None","``",".","Returns",":","loss",":","if","``","closure","``","is","provided",".","Otherwise","``","None","``","."]
code: def step(self, closure=None, fp16_param_groups=None):
      """Update the model parameters.
  
      .. note::
          This method will be called internally by ZeRO-Offload. DeepSpeed
          users should still use ``engine.step()`` as shown in the
          `Getting Started
          <https://www.deepspeed.ai/getting-started/#training>`_ guide.
  
      Args:
          closure (callable, optional): closure to compute the loss.
              Defaults to ``None``.
          fp16_param_groups: FP16 GPU parameters to update. Performing the
              copy here reduces communication time. Defaults to ``None``.
  
      Returns:
          loss: if ``closure`` is provided. Otherwise ``None``.
      """
  
      loss = None
      if closure is not None:
          with torch.enable_grad():
              loss = closure()
  
      # intended device for step
      device = torch.device("cpu")
  
      # converting the fp16 params to a group of parameter
      if type(fp16_param_groups) is list:
          if type(fp16_param_groups[0]) is not list:
              fp16_param_groups = [fp16_param_groups]
      elif fp16_param_groups is not None:
          fp16_param_groups = [[fp16_param_groups]]
  
      for group_id, group in enumerate(self.param_groups):
          for param_id, p in enumerate(group["params"]):
              if p.grad is None:
                  continue
  
              assert p.device == device, (
                  f"CPUAdam param is on {p.device} and must be 'cpu', make "
                  "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config."
              )
  
              state = self.state[p]
              # State initialization
              if len(state) == 0:
                  # print(f'group {group_id} param {param_id} = {p.numel()}')
                  state["step"] = 0
  
                  # use full precision by default unless self.fp32_optimizer_states is off
                  state_dtype = (
                      torch.float
                      if self.fp32_optimizer_states
                      else p.dtype
                  )
  
                  # gradient momentums
                  state["exp_avg"] = torch.zeros_like(
                      p.data, dtype=state_dtype, device=device
                  )
                  # memory_format=torch.preserve_format)
                  # gradient variances
                  state["exp_avg_sq"] = torch.zeros_like(
                      p.data, dtype=state_dtype, device=device
                  )
                  # memory_format=torch.preserve_format)
  
              state["step"] += 1
              beta1, beta2 = group["betas"]
  
              if fp16_param_groups is not None:
                  self.ds_opt_adam.adam_update_copy(
                      self.opt_id,
                      state["step"],
                      group["lr"],
                      beta1,
                      beta2,
                      group["eps"],
                      group["weight_decay"],
                      group["bias_correction"],
                      p.data,
                      p.grad.data,
                      state["exp_avg"],
                      state["exp_avg_sq"],
                      fp16_param_groups[group_id][param_id].data,
                  )
              else:
                  self.ds_opt_adam.adam_update(
                      self.opt_id,
                      state["step"],
                      group["lr"],
                      beta1,
                      beta2,
                      group["eps"],
                      group["weight_decay"],
                      group["bias_correction"],
                      p.data,
                      p.grad.data,
                      state["exp_avg"],
                      state["exp_avg_sq"],
                  )
      return loss
  
file_path: DeepSpeed/deepspeed/ops/adam/cpu_adam.py
imports: import torch
  from cpuinfo import get_cpu_info
  from deepspeed.utils import logger
  from deepspeed.utils.logging import should_log_le
  from deepspeed.ops.op_builder import CPUAdamBuilder
label_desc: Use image node_id 4 for calling the DeepSpeedCPUAdam obj's underlying member method code with example usage: obj.step(closure, fp16_param_groups) and returns: loss