class_docstr: Adam optimizer with ZeRO algorithm
  
  Child class of Apex DistributedFusedAdam, with optimizations for
  NeMo-Megatron.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts
          defining parameter groups.
      disable_distributed_parameters (bool, optional): use standard
          data-parallel communication instead of ZeRO.
          (default: False)
      **kwargs: keyword arguments to pass to Apex
          DistributedFusedAdam.
code: def _make_post_backward_hook(
      self,
      param: torch.nn.Parameter,
      param_group_id: int,
      param_id: int,
  ) -> Callable:
      def hook(*unused):
          if getattr(param, "_pre_forward_hook_is_enabled", False):
              raise RuntimeError(
                  "A parameter called its post-backward hook "
                  "before its pre-forward hook. "
                  "Please manually interact with the parameter "
                  "before the forward pass (e.g. by calling data_ptr) "
                  "or run DistributedFusedAdam with overlap_param_sync=False."
              )
          with self._lock:
              need_to_initialize = "fragments" not in self.state[param]
              if need_to_initialize:
                  self._init_param_state(
                      param, param_group_id, param_id
                  )
              if self.greedy_grad_copy and not getattr(
                  param, "_disable_greedy_grad_copy", False
              ):
                  self._grad_copy(param)
                  if self.overlap_grad_sync and not getattr(
                      param, "_disable_overlap_grad_sync", False
                  ):
                      self._try_start_bucket_grad_sync(
                          params=[param],
                          ignore_last_bucket=need_to_initialize,
                      )
  
      return hook
  
imports: import collections
  from typing import Callable, Dict, Iterable, Optional, Union
  import torch
  from apex.contrib.optimizers.distributed_fused_adam import DistributedFusedAdam, _disable_pre_forward_hook, _multi_tensor_copy
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace
  from megatron.core.dist_checkpointing.mapping import ShardedTensor
  from megatron.core.dist_checkpointing.optimizer import get_param_id_to_sharded_param_map, optim_state_to_sharding_state
  from nemo.utils import str_to_dtype
label_desc: Use image node_id 3 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj._make_post_backward_hook(param, param_group_id, param_id) and returns: hook