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: Make sure local shards of parameters are in expected datatypes
  
  For FP8 params, FP32 values are cast into FP8 using per-param
  scaling factors and per-param amaxes are computed and reduced.
docstr_tok: ["Make","sure","local","shards","of","parameters","are","in","expected","datatypes","For","FP8","params",",","FP32","values","are","cast","into","FP8","using","per-param","scaling","factors","and","per-param","amaxes","are","computed","and","reduced","."]
code: def _check_params_shard_dtypes(
      self,
      params_buckets: Dict[int, DistributedFusedAdam.ParameterBucket],
  ) -> None:
      """Make sure local shards of parameters are in expected datatypes
  
      For FP8 params, FP32 values are cast into FP8 using per-param
      scaling factors and per-param amaxes are computed and reduced.
  
      """
  
      # Just call base class function if there are no FP8 tensors
      num_fp8_params = sum(
          1 for param in self.parameters() if _is_fp8_tensor(param)
      )
      if num_fp8_params == 0:
          super()._check_params_shard_dtypes(params_buckets)
          return
  
      # FP8 scaling factors
      amaxes = []
      scales = []
      scale_invs = []
      i = -1
      for param in self.parameters():
          if not _is_fp8_tensor(param):
              continue
          i += 1
          fp8_meta = param._fp8_meta["scaling_fwd"]
          fp8_meta_index = param._fp8_meta_index
          amaxes.append(
              fp8_meta.amax_history[0][fp8_meta_index].view(1)
          )
          scales.append(fp8_meta.scale[fp8_meta_index].view(1))
          scale_invs.append(param._scale_inv.view(1))
  
      # Update cached scale-inverses
      packed_scales = torch.empty(
          num_fp8_params, dtype=torch.float32, device=self.device
      )
      packed_scale_views = [
          packed_scales[i].view(1) for i in range(num_fp8_params)
      ]
      _multi_tensor_copy(
          scales,
          packed_scale_views,
          dummy_overflow_buf=self._dummy_overflow_buf,
      )
      torch.reciprocal(packed_scales, out=packed_scales)
      _multi_tensor_copy(
          packed_scale_views,
          scale_invs,
          dummy_overflow_buf=self._dummy_overflow_buf,
      )
  
      # Cast local data to FP8
      fp8_params_shards = dict()
      for param in self.parameters():
          if not _is_fp8_tensor(param):
              continue
  
          # FP8 metadata
          fp8_meta = param._fp8_meta["scaling_fwd"]
          fp8_meta_index = param._fp8_meta_index
          fp8_dtype = param._fp8_dtype
  
          # Iterate through fragments with local data
          for fragment in self.state[param]["fragments"]:
              if not fragment.in_local_shard:
                  continue
              shard_start, shard_end = fragment.shard_range
              if shard_end <= shard_start:
                  continue
              shard_range = slice(shard_start, shard_end)
  
              # Get bucket containing fragment
              bucket_id = fragment.bucket_id
              if bucket_id not in params_buckets:
                  continue
              state_bucket = self.state["buckets"][bucket_id]
              param_bucket = params_buckets[bucket_id]
              if state_bucket.param_sync_dtype != torch.uint8:
                  continue
  
              # Allocate FP8 buffer if needed
              if bucket_id not in fp8_params_shards:
                  fp8_params_shards[bucket_id] = torch.empty_like(
                      param_bucket.params_shard, dtype=torch.uint8
                  )
  
              # FP8 cast and amax
              fp32_fragment = param_bucket.params_shard[
                  shard_range
              ].view(1, -1)
              fp8_fragment = fp8_params_shards[bucket_id][
                  shard_range
              ].view(1, -1)
              cast_to_fp8(
                  fp32_fragment,
                  fp8_meta,
                  fp8_meta_index,
                  fp8_dtype,
                  out=fp8_fragment,
              )
  
      # Update param shards with FP8 buffers
      for bucket_id, params_shard in fp8_params_shards.items():
          params_buckets[bucket_id].params_shard = params_shard
  
      # Reduce amaxes
      # Note: Assume each param has a separate amax
      packed_amaxes = torch.empty(
          num_fp8_params, dtype=torch.float32, device=self.device
      )
      packed_amax_views = [
          packed_amaxes[i].view(1) for i in range(num_fp8_params)
      ]
      _multi_tensor_copy(
          amaxes,
          packed_amax_views,
          dummy_overflow_buf=self._dummy_overflow_buf,
      )
      torch.distributed.all_reduce(
          packed_amaxes,
          op=torch.distributed.ReduceOp.MAX,
          group=self.distributed_process_group,
      )
      _multi_tensor_copy(
          packed_amax_views,
          amaxes,
          dummy_overflow_buf=self._dummy_overflow_buf,
      )
  
      # Handle any remaining dtype conversions
      super()._check_params_shard_dtypes(params_buckets)
  
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 12 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj._check_params_shard_dtypes(params_buckets) without return types