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.
docstr: Initialize optimizer state for parameters in one effective bucket
  
  If any FP8 params are detected, all non-FP8 params are removed
  from the bucket and their overlapped grad syncs are disabled.
  This assumes that weight matrices are FP8 params and that
  non-FP8 params are small (e.g. biases and layer norm params).
docstr_tok: ["Initialize","optimizer","state","for","parameters","in","one","effective","bucket","If","any","FP8","params","are","detected",",","all","non-FP8","params","are","removed","from","the","bucket","and","their","overlapped","grad","syncs","are","disabled",".","This","assumes","that","weight","matrices","are","FP8","params","and","that","non-FP8","params","are","small","(","e.g",".","biases","and","layer","norm","params",")","."]
code: def init_params_bucket(
      self,
      params: Iterable[torch.nn.Parameter],
      param_sync_dtype: Optional[torch.dtype] = None,
      **kwargs,
  ) -> None:
      """Initialize optimizer state for parameters in one effective bucket
  
      If any FP8 params are detected, all non-FP8 params are removed
      from the bucket and their overlapped grad syncs are disabled.
      This assumes that weight matrices are FP8 params and that
      non-FP8 params are small (e.g. biases and layer norm params).
  
      """
  
      # Ignore parameters that have already been initialized
      if isinstance(params, torch.Tensor):
          params = [params]
      params = [
          param
          for param in params
          if "fragments" not in self.state[param]
      ]
      if not params:
          return
  
      # Ignore non-FP8 params if there are any FP8 params
      if any(_is_fp8_tensor(param) for param in params):
          for param in params:
              if not _is_fp8_tensor(param):
                  param._disable_overlap_grad_sync = True
          params = filter(_is_fp8_tensor, params)
          param_sync_dtype = torch.uint8
  
      # Initialize parameter buckets
      super().init_params_bucket(
          params,
          param_sync_dtype=param_sync_dtype,
          **kwargs,
      )
  
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 5 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj.init_params_bucket(params, param_sync_dtype) without return types