class_docstr: Optimizer used for ZeRO-1 and ZeRO-2.
code: def __init__(
      self,
      optimizer: Optimizer,
      initial_scale: int = 2**16,  # grad scaler config
      min_scale: int = 1,
      growth_factor: float = 2.0,
      backoff_factor: float = 0.5,
      growth_interval: int = 2000,
      hysteresis: int = 2,
      max_scale: int = 2**24,
      clip_grad_norm: float = 0.0,  # grad clipping
      verbose: bool = False,
      reduce_bucket_size: int = 1024 * 1024,  # communication
      communication_dtype: Optional[torch.dtype] = None,
      overlap_communication: bool = False,
      partition_grad: bool = False,  # stage 2 flag
      cpu_offload: bool = False,  # cpu offload
      dp_process_group: Optional[
          ProcessGroup
      ] = None,  # the dp pg for comm
      forced_dtype: Optional[torch.dtype] = None,
      moe_extra_dp_process_group: Optional[ProcessGroup] = None,
      master_weights: bool = True,  # master weights
  ):
      super(LowLevelZeroOptimizer, self).__init__(optim=optimizer)
      self._dtype = self.optim.param_groups[0]["params"][0].dtype
      self._logger = get_dist_logger()
      self._verbose = verbose
  
      # stage 2
      self._partition_grads = partition_grad
  
      self._cpu_offload = cpu_offload
  
      # grad accumulation
      self.require_grad_sync = True
  
      # if process_group is none, will use the default one
      self.dp_pg = dp_process_group
      self._local_rank = dist.get_rank(group=self.dp_pg)
      self._world_size = dist.get_world_size(group=self.dp_pg)
  
      # extra dp
      # This group is used to sync moe param, dp_world_size = moe_duplicates * extra_dp_size.
      # Non moe param will be sync by global dp pg, moe param will be sync by extra dp pg.
      # Moe param grad is be split as non moe param by global dp pg, and grad will be merged in step.
      # And moe working and master param are split by extra dp pg.
      self.moe_extra_dp_pg = moe_extra_dp_process_group
      if self.moe_extra_dp_pg is not None:
          self.moe_extra_dp_pg_size = dist.get_world_size(
              group=self.moe_extra_dp_pg
          )
          self.moe_extra_dp_pg_rank = dist.get_rank(
              group=self.moe_extra_dp_pg
          )
  
      # working and master params for mixed precision training
      self._working_param_groups = dict()
      self._master_param_groups_of_current_rank = dict()
  
      # communication params
      self._overlap_communication = overlap_communication
      self._reduce_bucket_size = reduce_bucket_size
      self._communication_dtype = communication_dtype
  
      # gradient clipping
      self._clip_grad_norm = clip_grad_norm
  
      # master weights copy
      self._master_weights = master_weights
  
      if forced_dtype:
          for group in self.optim.param_groups:
              group_params = group["params"]
              for param in group_params:
                  param.data = param.data.to(forced_dtype)
          self._dtype = forced_dtype
  
      # check argument conflict
      self._sanity_checks()
  
      # ParameterStore will manage the tensor buffers used for zero
      # it will not manage the tensors used by mixed precision training
      self._param_store = ParameterStore(self.dp_pg)
      self._grad_store = GradientStore(
          self.dp_pg, partition_grad=partition_grad
      )
      self._bucket_store = BucketStore(self.dp_pg)
  
      # moe param should not be stored in working_groups
      # because they have different parallel strategy
      # so we need to store them separately in param_groups
      # instead of working_groups
      moe_params = list()
  
      # iterate over the param group in the optimizer
      # partition these param groups for data parallel training
      # and add buffers to parameter store for future access
      for group_id, param_group in enumerate(self.optim.param_groups):
          group_params = list()
          for param in param_group["params"]:
              if param.requires_grad:
                  if self.moe_extra_dp_pg is None:
                      # skip moe param
                      if is_moe_tensor(param):
                          moe_params.append(param)
                          continue
                  group_params.append(param)
  
          # add the working params to working_param_groups for bookkeeping
          self._working_param_groups[group_id] = group_params
  
          master_param_current_rank = (
              self._create_master_param_current_rank(group_params)
          )
          self._master_param_groups_of_current_rank[
              group_id
          ] = master_param_current_rank
  
          # need to replace the params in the `params` field in the optimizer
          # so that when the optimizer calls step(), it only updates the tensors
          # managed by this data parallel rank
          param_group["params"] = master_param_current_rank
  
      # if there are moe params, store in addtional group in optim
      if len(moe_params) > 0:
          param_group = dict()
          for key, value in self.optim.param_groups[0].items():
              if key != "params":
                  param_group[key] = value
          param_group["params"] = moe_params
          self.optim.param_groups.append(param_group)
  
      # intialize communication stream for
      # communication-compuation overlapping
      if self._overlap_communication:
          self._comm_stream = device_utils.Stream()
  
      # reduction hook is only used if overlapping communication
      # or stage 2 is used
      # if it is stage 1 without overlapping, no hook will be attached
      if self._overlap_communication or self._partition_grads:
          self._attach_reduction_hook()
  
      # initialize mixed precision mixin
      self.mixed_precision_mixin: Optional[MixedPrecisionMixin] = None
      if self._dtype is torch.float16:
          self.mixed_precision_mixin = (
              LowLevelZeroFP16MixedPrecisionMixin(
                  self.num_param_groups,
                  self._grad_store,
                  initial_scale=initial_scale,
                  min_scale=min_scale,
                  growth_factor=growth_factor,
                  backoff_factor=backoff_factor,
                  growth_interval=growth_interval,
                  hysteresis=hysteresis,
                  max_scale=max_scale,
              )
          )
      elif self._dtype is torch.bfloat16:
          self.mixed_precision_mixin = BF16MixedPrecisionMixin()
  
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 1 to create a new LowLevelZeroOptimizer object from inherited base classes: OptimizerWrapper with example: obj = LowLevelZeroOptimizer(optimizer, initial_scale, min_scale, growth_factor, backoff_factor, growth_interval, hysteresis, max_scale, clip_grad_norm, verbose, reduce_bucket_size, communication_dtype, overlap_communication, partition_grad, cpu_offload, dp_process_group, forced_dtype, moe_extra_dp_process_group, master_weights)