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: Allocate contiguous buffers for param buckets
  
  For FP8 params, the FP8 data buffer is made a view into a
  contiguous buffer.
docstr_tok: ["Allocate","contiguous","buffers","for","param","buckets","For","FP8","params",",","the","FP8","data","buffer","is","made","a","view","into","a","contiguous","buffer","."]
code: def init_param_buffer(self) -> None:
      """Allocate contiguous buffers for param buckets
  
      For FP8 params, the FP8 data buffer is made a view into a
      contiguous buffer.
  
      """
  
      # Make sure all params are initialized
      self.contiguous_param_buffer = True
      self.init_params()
  
      # Construct param buffers
      buffer_sizes = collections.defaultdict(lambda: 0)
      for bucket in self.state["buckets"]:
          dtypes = bucket.dtypes()
          buffer_sizes[dtypes] = max(
              bucket.contiguous_buffer_offset + bucket.bucket_size,
              buffer_sizes[dtypes],
          )
      for dtypes, buffer_size in buffer_sizes.items():
          _, _, param_sync_dtype = dtypes
          self._param_buffers[dtypes] = torch.zeros(
              [buffer_size],
              dtype=param_sync_dtype,
              device=self.device,
          )
  
      # Figure out corresponding positions in params and param buffer
      params = list(self.parameters())
      param_flat_views = []
      param_buffer_views = []
      for i, param in enumerate(params):
          fragment = self.state[param]["fragments"][0]
          bucket_id = fragment.bucket_id
          bucket = self.state["buckets"][bucket_id]
          param_size = param.numel()
          bucket_start, _ = fragment.bucket_range
          buffer_offset = bucket.contiguous_buffer_offset
          buffer_start = buffer_offset + bucket_start
          buffer_end = buffer_start + param_size
          param_buffer = self._param_buffers[bucket.dtypes()]
          param_buffer_view = param_buffer[
              buffer_start:buffer_end
          ].detach()
          if param_buffer_view.device != param.device:
              raise RuntimeError(
                  "Attempted to change a parameter with device={param.device} "
                  f"into a buffer view with device={param_buffer_view.device}"
              )
          if _is_fp8_tensor(param):
              param_flat_views.append(param._data.detach().view(-1))
          else:
              if param_buffer_view.dtype != param.dtype:
                  raise RuntimeError(
                      f"Attempted to change a parameter with dtype={param.dtype} "
                      f"into a buffer view with dtype={param_buffer_view.dtype}"
                  )
              param_flat_views.append(param.detach().view(-1))
          param_buffer_views.append(param_buffer_view)
  
      # Copy values into param buffer
      _multi_tensor_copy(
          param_flat_views,
          param_buffer_views,
          dummy_overflow_buf=self._dummy_overflow_buf,
      )
  
      # Make all params a view into the param buffer
      for param, buffer_view in zip(params, param_buffer_views):
          if _is_fp8_tensor(param):
              param._data.data = buffer_view.view(param.size())
          else:
              param.data = buffer_view.view(param.size())
  
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 7 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj.init_param_buffer() without return types