class_docstr: Float16 optimizer wrapper for half precision (fp16 and bf16) data types.
  This optimizer wrapper holds main parameters and gradients in fp32 to support
  stable convergence.
  
  Arguments:
      optimizer: base optimizer such as Adam or SGD.
      fp32_grad_accum: to enable the use of fp32 in gradient accumulation and allreduce.
      contiguous_grad_bucket: to enable allocating the master gradients in the 
          contiguous memory space to reduce memory fragmentation.
      async_grad_allreduce: enable asynchronous gradient allreduce that is executed
          along with the training step backprop.
docstr: Create the grad accumulation and all-reduce hook for backprop.
code: def _make_param_hook(self, param, main_param, i, grad_chunk_info):
      """Create the grad accumulation and all-reduce hook for backprop."""
  
      # Hook used for back-prop.
      def param_hook(*unused):
          # Accumulates gradients on main gradients
          if param.grad is not None:
              if main_param.grad is None:
                  main_param.grad = param.grad.float()
              else:
                  main_param.grad.add_(param.grad.data)
              # Deallocate grad memory.
              param.grad = None
  
          # Asynchronous gradients allreduce accross data_parallel ranks
          if self._require_backward_grad_sync:
              if self._grad_allreduce_chunk_size_mb > 0:
                  self._main_grad_buffers[i].update_chunk_info(
                      grad_chunk_info
                  )
                  while True:
                      allreduce_tensor = self._main_grad_buffers[
                          i
                      ].get_allreduce_tensor()
                      if allreduce_tensor is None:
                          break
                      if self._grad_div_ar_fusion:
                          torch.distributed.all_reduce(
                              allreduce_tensor,
                              group=get_data_parallel_group(
                                  with_context_parallel=True
                              ),
                              async_op=True,
                              op=torch.distributed._make_nccl_premul_sum(
                                  self._grad_divisor
                              ),
                          )
                      else:
                          allreduce_tensor.div_(
                              get_data_parallel_world_size(
                                  with_context_parallel=True
                              )
                          )
                          torch.distributed.all_reduce(
                              allreduce_tensor,
                              group=get_data_parallel_group(
                                  with_context_parallel=True
                              ),
                              async_op=True,
                          )
              else:
                  if self._grad_div_ar_fusion:
                      torch.distributed.all_reduce(
                          main_param.grad,
                          group=get_data_parallel_group(
                              with_context_parallel=True
                          ),
                          async_op=True,
                          op=torch.distributed._make_nccl_premul_sum(
                              self._grad_divisor
                          ),
                      )
                  else:
                      main_param.grad.div_(
                          get_data_parallel_world_size(
                              with_context_parallel=True
                          )
                      )
                      torch.distributed.all_reduce(
                          main_param.grad,
                          group=get_data_parallel_group(
                              with_context_parallel=True
                          ),
                          async_op=True,
                      )
  
      return param_hook
  
imports: from contextlib import contextmanager
  import torch
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling the MainParamsOptimizerWrapper obj's underlying member method code with example usage: obj._make_param_hook(param, main_param, i, grad_chunk_info) and returns: param_hook