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 a parameter
  
  Initializing the master weights requires slicing a flattened
  view of the param. FP8 tensors do not handle these operations
  gracefully, so we hack around it by explicitly casting to
  FP32.
docstr_tok: ["Initialize","optimizer","state","for","a","parameter","Initializing","the","master","weights","requires","slicing","a","flattened","view","of","the","param",".","FP8","tensors","do","not","handle","these","operations","gracefully",",","so","we","hack","around","it","by","explicitly","casting","to","FP32","."]
code: def _init_param_state(
      self,
      param: torch.nn.Parameter,
      param_group_id: int,
      param_id: int,
      param_sync_dtype: Optional[torch.dtype] = None,
      **kwargs,
  ) -> None:
      """Initialize optimizer state for a parameter
  
      Initializing the master weights requires slicing a flattened
      view of the param. FP8 tensors do not handle these operations
      gracefully, so we hack around it by explicitly casting to
      FP32.
  
      """
  
      # Initialize non-FP8 params as usual
      if not _is_fp8_tensor(param):
          super()._init_param_state(
              param,
              param_group_id,
              param_id,
              param_sync_dtype=param_sync_dtype,
              **kwargs,
          )
  
      # Return immediately if already initialized
      if "fragments" in self.state[param]:
          return
  
      # Initialize with FP32 copy of param
      fp32_param = param.float()
      super()._init_param_state(
          fp32_param,
          param_group_id,
          param_id,
          param_sync_dtype=torch.uint8,
          **kwargs,
      )
      self.state[param].update(self.state[fp32_param])
      del self.state[fp32_param]
  
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 6 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj._init_param_state(param, param_group_id, param_id, param_sync_dtype) without return types