class_docstr: Optimizer used for ZeRO-1 and ZeRO-2.
code: def _run_reduction(self):
      if self._bucket_store.num_elements_in_bucket() > 0:
          self._bucket_store.build_grad_in_bucket()
  
          if self.moe_extra_dp_pg is None:
              flat_grads = self._bucket_store.get_flatten_grad()
              flat_grads /= self._world_size
          else:
              # record moe and non moe param
              moe_list = []
              for param in self._bucket_store._param_list:
                  moe_list.append(is_moe_tensor(param))
  
              # divide them into different groups
              moe_grad_list = []
              non_moe_grad_list = []
              for (
                  grad_list
              ) in self._bucket_store._grad_in_bucket.values():
                  non_moe_cur_grad = []
                  moe_cur_grad = []
                  for i in range(len(grad_list)):
                      if moe_list[i] == True:
                          moe_cur_grad.append(grad_list[i])
                      else:
                          non_moe_cur_grad.append(grad_list[i])
                  if len(moe_cur_grad) > 0:
                      moe_grad_list.append(moe_cur_grad)
                  if len(non_moe_cur_grad) > 0:
                      non_moe_grad_list.append(non_moe_cur_grad)
  
              if len(non_moe_grad_list) > 0:
                  non_moe_flat_grads = []
                  for grad_list in non_moe_grad_list:
                      non_moe_flat_grads.append(
                          _flatten_dense_tensors(grad_list)
                      )
                  non_moe_flat_grads = _flatten_dense_tensors(
                      non_moe_flat_grads
                  )
                  non_moe_flat_grads /= self._world_size
  
              if len(moe_grad_list) > 0:
                  moe_flat_grads = []
                  for grad_list in moe_grad_list:
                      moe_flat_grads.append(
                          _flatten_dense_tensors(grad_list)
                      )
                  moe_flat_grads = _flatten_dense_tensors(
                      moe_flat_grads
                  )
  
          # ready to add other tensors to bucket
          self._bucket_store.reset_num_elements_in_bucket()
  
          if self._overlap_communication:
              stream = self._comm_stream
              # in case of the memory being reused in the default stream
              if self.moe_extra_dp_pg is None:
                  flat_grads.record_stream(stream)
              else:
                  if len(non_moe_grad_list) > 0:
                      non_moe_flat_grads.record_stream(stream)
                  if len(moe_grad_list) > 0:
                      moe_flat_grads.record_stream(stream)
              # waiting for ops in the default stream finishing
              stream.wait_stream(device_utils.current_stream())
          else:
              stream = device_utils.current_stream()
  
          with device_utils.stream(stream):
              group_id = self._bucket_store.current_group_id
  
              if self.moe_extra_dp_pg is None:
                  grad_dtype = flat_grads.dtype
                  if self._communication_dtype is not None:
                      flat_grads = flat_grads.to(
                          self._communication_dtype
                      )
  
              if not self._partition_grads:
                  if self.moe_extra_dp_pg is None:
                      dist.all_reduce(flat_grads, group=self.dp_pg)
                      if flat_grads.dtype != grad_dtype:
                          flat_grads = flat_grads.to(grad_dtype)
  
                      flat_grads_per_rank = flat_grads.split(
                          flat_grads.numel() // self._world_size
                      )
                      grad_in_bucket = self._bucket_store.get_grad()
                      self._update_unpartitoned_grad(
                          grad_in_bucket.values(),
                          flat_grads_per_rank,
                          group_id,
                      )
  
                  # sync extra zero group
                  else:
                      # sync non moe param in global dp group
                      if len(non_moe_grad_list) > 0:
                          dist.all_reduce(
                              non_moe_flat_grads, group=self.dp_pg
                          )
                          flat_grads_per_rank = (
                              non_moe_flat_grads.split(
                                  non_moe_flat_grads.numel()
                                  // self._world_size
                              )
                          )
                          self._update_unpartitoned_grad(
                              non_moe_grad_list,
                              flat_grads_per_rank,
                              group_id,
                          )
  
                      # sync moe param only in zero group
                      if len(moe_grad_list) > 0:
                          dist.all_reduce(
                              moe_flat_grads, group=self.moe_extra_dp_pg
                          )
                          flat_grads_per_rank = moe_flat_grads.split(
                              moe_flat_grads.numel() // self._world_size
                          )
                          self._update_unpartitoned_grad(
                              moe_grad_list,
                              flat_grads_per_rank,
                              group_id,
                          )
  
              else:
                  if self.moe_extra_dp_pg is None:
                      flat_grads_list = list(
                          flat_grads.split(
                              len(flat_grads) // self._world_size
                          )
                      )
                      recieved_grad = torch.zeros_like(
                          flat_grads_list[0]
                      )
                      dist.reduce_scatter(
                          recieved_grad,
                          flat_grads_list,
                          group=self.dp_pg,
                      )
  
                      if recieved_grad.dtype != grad_dtype:
                          recieved_grad = recieved_grad.to(grad_dtype)
  
                      grad_in_bucket_current_rank = (
                          self._bucket_store.get_grad()[
                              self._local_rank
                          ]
                      )
                      self._update_partitoned_grad(
                          grad_in_bucket_current_rank,
                          recieved_grad,
                          group_id,
                          1,
                      )
                  else:
                      # categorize moe and non moe param
                      grad_in_bucket_current_rank = (
                          self._bucket_store.get_grad()[
                              self._local_rank
                          ]
                      )
                      moe_grad_in_bucket_current_rank = []
                      non_moe_grad_in_bucket_current_rank = []
                      for idx, grad in enumerate(
                          grad_in_bucket_current_rank
                      ):
                          if moe_list[idx] == True:
                              moe_grad_in_bucket_current_rank.append(
                                  grad
                              )
                          else:
                              non_moe_grad_in_bucket_current_rank.append(
                                  grad
                              )
  
                      if len(non_moe_grad_list) > 0:
                          flat_grads_list = list(
                              non_moe_flat_grads.split(
                                  len(non_moe_flat_grads)
                                  // self._world_size
                              )
                          )
                          recieved_grad = torch.zeros_like(
                              flat_grads_list[0]
                          )
                          dist.reduce_scatter(
                              recieved_grad,
                              flat_grads_list,
                              group=self.dp_pg,
                          )
                          self._update_partitoned_grad(
                              non_moe_grad_in_bucket_current_rank,
                              recieved_grad,
                              group_id,
                              1,
                          )
  
                      if len(moe_grad_list) > 0:
                          flat_grads_list = list(
                              moe_flat_grads.split(
                                  len(moe_flat_grads)
                                  // self.moe_extra_dp_pg_size
                              )
                          )
                          recieved_grad = torch.zeros_like(
                              flat_grads_list[0]
                          )
                          dist.reduce_scatter(
                              recieved_grad,
                              flat_grads_list,
                              group=self.moe_extra_dp_pg,
                          )
                          param_slice = (
                              self._world_size
                              // self.moe_extra_dp_pg_size
                          )
                          recieved_grad = list(
                              recieved_grad.split(
                                  len(recieved_grad) // param_slice
                              )
                          )
                          for split_recieved_grad in recieved_grad:
                              split_recieved_grad = (
                                  _unflatten_dense_tensors(
                                      split_recieved_grad,
                                      moe_grad_in_bucket_current_rank,
                                  )
                              )
                              for real_grad, grad in zip(
                                  split_recieved_grad,
                                  moe_grad_in_bucket_current_rank,
                              ):
                                  param_id = self._bucket_store.get_param_id_of_grad(
                                      grad
                                  )
                                  self._add_grad(
                                      real_grad,
                                      param_slice,
                                      group_id,
                                      param_id,
                                  )
  
              self._bucket_store.reset()
  
imports: import copy
  from contextlib import contextmanager
  from functools import partial
  from typing import Dict, Iterator, List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.nn
  from torch import Tensor, inf
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  from torch.distributed import ProcessGroup
  from torch.optim import Optimizer
  import colossalai.utils.device
  from colossalai.amp.naive_amp.mixed_precision_mixin import BF16MixedPrecisionMixin, FP16MixedPrecisionMixin, MixedPrecisionMixin
  from colossalai.interface import OptimizerWrapper
  from colossalai.logging import get_dist_logger
  from colossalai.tensor.moe_tensor.api import is_moe_tensor
  from colossalai.utils.device import IS_NPU_AVAILABLE, get_current_device
  from ._utils import calculate_global_norm_from_list, flatten, has_inf_or_nan, release_param_grad, sync_tensor
  from .bookkeeping import BucketStore, GradientStore, ParameterStore
label_desc: Use image node_id 8 for calling the LowLevelZeroOptimizer obj's underlying member method code with example usage: obj._run_reduction() without return types