class_docstr: Optimizer used for ZeRO-1 and ZeRO-2.
code: def step(self, closure=None):
      assert closure is None, "closure is not supported by step()"
      if not self.require_grad_sync:
          return
  
      if (
          self.mixed_precision_mixin is not None
          and self.mixed_precision_mixin.should_skip_step()
      ):
          self._grad_store.reset_all_gradients()
          if self._verbose:
              self._logger.info(f"Found overflow. Skip step")
          self.zero_grad()
          return
  
      # record all grads for unscale and clip
      grad_partition_groups = []
      norm_groups = []
  
      # sometimes not all params are 'really' working
      # for instance, when layer drop, the dropped layer has no grad
      # and should not be updated
      real_working_params = dict()
      real_master_params = dict()
      grad_index = 0 if self._partition_grads else self._local_rank
      for group_id in range(self.num_param_groups):
          master_params = self._master_param_groups_of_current_rank[
              group_id
          ]
          real_working_params[group_id] = []
          real_master_params[group_id] = []
          for splited_param in master_params:
              working_param = self._param_store.master_to_working_param[
                  id(splited_param)
              ]
              # if a working param requires grad and has no grad
              # it is not 'really' working, e.g. the droped layer
              # else the splited grad should be attached to the splited param
              grads = self._grad_store.get_partitioned_gradients_by_param_id(
                  group_id, id(working_param)
              )
              if len(grads) > 0:
                  # moe hybrid zero
                  if (
                      self.moe_extra_dp_pg is not None
                      and is_moe_tensor(working_param)
                  ):
                      real_working_params[group_id].append(
                          working_param
                      )
                      if self._partition_grads:
                          grad = grads
                      else:
                          param_slice = (
                              self._world_size
                              // self.moe_extra_dp_pg_size
                          )
                          grad = grads[
                              self.moe_extra_dp_pg_rank
                              * param_slice : (
                                  self.moe_extra_dp_pg_rank + 1
                              )
                              * param_slice
                          ]
                      grad = flatten(grad)
                  else:
                      real_working_params[group_id].append(
                          working_param
                      )
                      grad = grads[grad_index]
                  # no need to copy fp32 grad if master_weights is False
                  if self._master_weights:
                      grad = grad.to(splited_param.dtype).to(
                          splited_param.device
                      )
                  splited_param.grad = grad
                  grad_partition_groups.append(grad)
                  real_master_params[group_id].append(splited_param)
  
          # compute norm
          working_grads = (
              self._grad_store.get_working_grads_by_group_id(group_id)
          )
          norm_group = self._compute_grad_norm(gradients=working_grads)
          norm_groups.append(norm_group)
  
          self._grad_store.reset_grads_by_group_id(group_id)
  
          # update the params in the optimizer
          self.optim.param_groups[group_id][
              "params"
          ] = real_master_params[group_id]
  
      # unscale and clip grads
      global_norm = calculate_global_norm_from_list(
          norm_list=norm_groups
      )
      self._unscale_and_clip_grads(grad_partition_groups, global_norm)
  
      # TODO: we should store master param for ep
      if len(self.param_groups) > len(self._working_param_groups):
          for param in self.param_groups[-1]["params"]:
              param.data = param.data.to(torch.float32)
              param.grad = param.grad.to(torch.float32)
  
      # update the parameters
      self.optim.step()
  
      # release the moe gradm
      if len(self.param_groups) > len(self._working_param_groups):
          for param in self.param_groups[-1]["params"]:
              param.grad = None
              param.data = param.data.to(self._dtype)
  
      # release the grad
      grad_partition_groups = []
      for group_id in range(self.num_param_groups):
          release_param_grad(
              self._master_param_groups_of_current_rank[group_id]
          )
  
      # update working partition updated by the current rank
      device = get_current_device()
      for group_id in range(self.num_param_groups):
          master_working_param = self.optim.param_groups[group_id][
              "params"
          ]
          for idx, splited_param in enumerate(master_working_param):
              working_param = real_working_params[group_id][idx]
              if self.moe_extra_dp_pg is not None and is_moe_tensor(
                  working_param
              ):
                  all_splited_param = [
                      torch.zeros(
                          splited_param.shape,
                          device=device,
                          dtype=self._dtype,
                      )
                      for _ in range(self.moe_extra_dp_pg_size)
                  ]
                  dist.all_gather(
                      all_splited_param,
                      splited_param.to(device).to(self._dtype),
                      group=self.moe_extra_dp_pg,
                  )
              else:
                  all_splited_param = [
                      torch.zeros(
                          splited_param.shape,
                          device=device,
                          dtype=self._dtype,
                      )
                      for _ in range(self._world_size)
                  ]
                  dist.all_gather(
                      all_splited_param,
                      splited_param.to(device).to(self._dtype),
                      group=self.dp_pg,
                  )
              working_param.data.copy_(
                  flatten(all_splited_param)[
                      : working_param.numel()
                  ].reshape_as(working_param)
              )
          self.optim.param_groups[group_id][
              "params"
          ] = self._master_param_groups_of_current_rank[group_id]
  
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 16 for calling the LowLevelZeroOptimizer obj's underlying member method code with example usage: obj.step(closure) without return types