class_docstr: Wrap an optimizer to perform memory-efficient mixed precision training.
  
  This class wraps an optimizer to perform FP16 training.
  This implementation is heavily based on the Fairseq implementation
  of `MemoryEfficientFP16Optimizer`, which can be found here:
  <https://github.com/pytorch/fairseq/blob/main/fairseq/optim/fp16_optimizer.py#L382>
  
  This allows you to train bigger models on a single GPU, but can be unstable.
  Prefer the SafeFP16 implementation if you do not have concerns about memory.
  
  :param params:
      Model parameters
  :param optimizer:
      Any torch optimizer
  :param float loss_initial_scale:
      Initial loss scaling. Default chosen empirically, but models with very low
      or high loss values may need this adjusted. Stick with powers of 2
  :param float min_loss_scale:
      Throws an error if your loss scale goes below this threshold
class_docstr_tok: ["Wrap","an","optimizer","to","perform","memory-efficient","mixed","precision","training",".","This","class","wraps","an","optimizer","to","perform","FP16","training",".","This","implementation","is","heavily","based","on","the","Fairseq","implementation","of","`","MemoryEfficientFP16Optimizer","`",",","which","can","be","found","here",":","<","https",":","\/\/github.com\/pytorch\/fairseq\/blob\/main\/fairseq\/optim\/fp16_optimizer.py","#","L382",">","This","allows","you","to","train","bigger","models","on","a","single","GPU",",","but","can","be","unstable",".","Prefer","the","SafeFP16","implementation","if","you","do","not","have","concerns","about","memory",".",":","param","params",":","Model","parameters",":","param","optimizer",":","Any","torch","optimizer",":","param","float","loss_initial_scale",":","Initial","loss","scaling",".","Default","chosen","empirically",",","but","models","with","very","low","or","high","loss","values","may","need","this","adjusted",".","Stick","with","powers","of","2",":","param","float","min_loss_scale",":","Throws","an","error","if","your","loss","scale","goes","below","this","threshold"]
docstr: Load an optimizer state dict.
  
  Override from PyTorch implementation to avoid casting to FP32.
code: def load_state_dict(self, state_dict):
      """
      Load an optimizer state dict.
  
      Override from PyTorch implementation to avoid casting to FP32.
      """
      if "loss_scaler" in state_dict:
          # init from the state dict
          if isinstance(state_dict["loss_scaler"], float):
              # new method, restore the float
              self.scaler.loss_scale = state_dict["loss_scaler"]
          else:
              # old method, we stored the entire loss scaler
              self.scaler.loss_scale = state_dict[
                  "loss_scaler"
              ].loss_scale
  
      self.optimizer.load_state_dict(state_dict)
  
      # Hack: PyTorch automatically casts the optimizer state to match the
      # type of the current parameters. But with memory efficient fp16 the
      # params are FP16 while the optimizer state is FP32 and we don't want
      # to cast. A workaround is to manually copy back the original state
      # after the optimizer has been loaded.
      groups = self.optimizer.param_groups
      saved_groups = state_dict["param_groups"]
      id_map = {
          old_id: p
          for old_id, p in zip(
              chain(*(g["params"] for g in saved_groups)),
              chain(*(g["params"] for g in groups)),
          )
      }
      for k, v in state_dict["state"].items():
          if k in id_map:
              # make sure when we copy the original state back, we make sure
              # that original state is on the correct device
              param = id_map[k]
              like_device_v = {
                  j: w.to(param.device) if torch.is_tensor(w) else w
                  for j, w in v.items()
              }
              self.optimizer.state[param] = like_device_v
  
imports: import math
  from itertools import chain
  from typing import Optional
  import parlai.utils.logging
label_desc: Use image node_id 16 for calling the MemoryEfficientFP16Optimizer obj's underlying member method code with example usage: obj.load_state_dict(state_dict) without return types