class_docstr: Wrapper class that supports NeMo-Megatron optimizations
  
  Performs gradient clipping, unscaling, and optimizer step.
code: def __init__(self, *args, max_norm=0, norm_type=2, **kwargs):
      super().__init__(*args, **kwargs)
  
      assert (
          norm_type == 2
      ), "Currently only norm_type=2 is supported for MegatronFusedAdam"
  
      # Gradient clipping parameters
      self.max_norm = float(max_norm)
      self.norm_type = float(norm_type)
  
imports: import amp_C
  import torch
  from nemo.collections.nlp.modules.common.megatron.module import param_is_not_shared
label_desc: Use image node_id 1 to create a new MegatronFusedAdam object from inherited base classes: FusedAdam with example: obj = MegatronFusedAdam()