class_docstr: FP16 Optimizer for training fp16 models. Handles loss scaling.
  
  For usage example please see, TODO:  DeepSpeed V2 Tutorial
   
class_docstr_tok: ["FP16","Optimizer","for","training","fp16","models",".","Handles","loss","scaling",".","For","usage","example","please","see",",","TODO",":","DeepSpeed","V2","Tutorial"]
docstr: Loads a state_dict created by an earlier call to state_dict().
  If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``,
  whose parameters in turn came from ``model``, it is expected that the user
  will call ``model.load_state_dict()`` before
  ``fp16_optimizer_instance.load_state_dict()`` is called.
  Example::
      model = torch.nn.Linear(D_in, D_out).cuda().half()
      optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
      optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
      ...
      checkpoint = torch.load("saved.pth")
      model.load_state_dict(checkpoint['model'])
      optimizer.load_state_dict(checkpoint['optimizer'])
docstr_tok: ["Loads","a","state_dict","created","by","an","earlier","call","to","state_dict","(",")",".","If","``","fp16_optimizer_instance","``","was","constructed","from","some","``","init_optimizer","``",",","whose","parameters","in","turn","came","from","``","model","``",",","it","is","expected","that","the","user","will","call","``","model.load_state_dict","(",")","``","before","``","fp16_optimizer_instance.load_state_dict","(",")","``","is","called",".","Example",":",":","model","=","torch.nn.Linear","(","D_in",",","D_out",")",".cuda","(",")",".half","(",")","optimizer","=","torch.optim.SGD","(","model.parameters","(",")",",","lr=1e-3",")","optimizer","=","FP16_Optimizer","(","optimizer",",","static_loss_scale","=","128.0",")","...","checkpoint","=","torch.load","(","``","saved.pth","''",")","model.load_state_dict","(","checkpoint","[","'model","'","]",")","optimizer.load_state_dict","(","checkpoint","[","'optimizer","'","]",")"]
code: def load_state_dict(self, state_dict, load_optimizer_states=True):
      """
      Loads a state_dict created by an earlier call to state_dict().
      If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``,
      whose parameters in turn came from ``model``, it is expected that the user
      will call ``model.load_state_dict()`` before
      ``fp16_optimizer_instance.load_state_dict()`` is called.
      Example::
          model = torch.nn.Linear(D_in, D_out).cuda().half()
          optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
          optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
          ...
          checkpoint = torch.load("saved.pth")
          model.load_state_dict(checkpoint['model'])
          optimizer.load_state_dict(checkpoint['optimizer'])
      """
      # I think it should actually be ok to reload the optimizer before the model.
      self.dynamic_loss_scale = state_dict["dynamic_loss_scale"]
      self.cur_scale = state_dict["cur_scale"]
      self.cur_iter = state_dict["cur_iter"]
      if state_dict["dynamic_loss_scale"]:
          self.last_overflow_iter = state_dict["last_overflow_iter"]
          self.scale_factor = state_dict["scale_factor"]
          self.scale_window = state_dict["scale_window"]
      if load_optimizer_states:
          self.optimizer.load_state_dict(
              state_dict[OPTIMIZER_STATE_DICT]
          )
      self.clip_grad = state_dict[CLIP_GRAD]
      # At this point, the optimizer's references to the model's fp32 parameters are up to date.
      # The optimizer's hyperparameters and internal buffers are also up to date.
      # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
      # out of date.  There are two options.
      # 1:  Refresh the master params from the model's fp16 params.
      # This requires less storage but incurs precision loss.
      # 2:  Save and restore the fp32 master copies separately.
      # We choose option 2.
      #
      # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device
      # of their associated parameters, because it's possible those buffers might not exist yet in
      # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been
      # constructed in the same way as the one whose state_dict we are loading, the same master params
      # are guaranteed to exist, so we can just copy_() from the saved master params.
      for current, saved in zip(
          self.fp32_groups_flat, state_dict["fp32_groups_flat"]
      ):
          current.data.copy_(saved.data)
  
imports: import torch
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  from deepspeed.runtime import DeepSpeedOptimizer
  from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
  from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
  from deepspeed.utils import groups, logger, log_dist
  from deepspeed import comm
  from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD
label_desc: Use image node_id 22 for calling the FP16_Optimizer obj's underlying member method code with example usage: obj.load_state_dict(state_dict, load_optimizer_states) without return types