class_docstr: Optimizer used for ZeRO-1 and ZeRO-2.
docstr: Compute and return the gradient norm for gradient clipping.
  
  Args:
      gradients (List[Tensor]): The gradients to compute norm
      norm_type (int, optional): type of the used p-norm, Can be ``'inf'`` for infinity norm. Defaults to 2.
  
  Returns:
      float: The total norm of given gradients
docstr_tok: ["Compute","and","return","the","gradient","norm","for","gradient","clipping",".","Args",":","gradients","(","List","[","Tensor","]",")",":","The","gradients","to","compute","norm","norm_type","(","int",",","optional",")",":","type","of","the","used","p-norm",",","Can","be","``","'inf","'","``","for","infinity","norm",".","Defaults","to","2",".","Returns",":","float",":","The","total","norm","of","given","gradients"]
code: def _compute_grad_norm(
      self, gradients: List[Tensor], norm_type: int = 2
  ) -> float:
      r"""
      Compute and return the gradient norm for gradient clipping.
  
      Args:
          gradients (List[Tensor]): The gradients to compute norm
          norm_type (int, optional): type of the used p-norm, Can be ``'inf'`` for infinity norm. Defaults to 2.
  
      Returns:
          float: The total norm of given gradients
      """
  
      if len(gradients) == 0:
          return 0.0
  
      norm_type = float(norm_type)
      if norm_type == inf:
          total_norm = max(grad.data.abs().max() for grad in gradients)
          total_norm_cuda = torch.tensor(
              [float(total_norm)],
              device=get_current_device(),
              dtype=torch.float,
          )
          dist.all_reduce(
              total_norm_cuda,
              op=torch.distributed.ReduceOp.MAX,
              group=self.dp_pg,
          )
          total_norm = total_norm_cuda.item()
  
      else:
          total_norm_exponentiated = 0.0
          for grad in gradients:
              grad_norm_exponentiated = (
                  grad.data.double().norm(norm_type) ** norm_type
              )
              total_norm_exponentiated += grad_norm_exponentiated
  
          # Sum across all model parallel GPUs.
          total_norm_exponentiated_cuda = torch.tensor(
              [float(total_norm_exponentiated)],
              device=get_current_device(),
              dtype=torch.float,
          )
          torch.distributed.all_reduce(
              total_norm_exponentiated_cuda,
              op=torch.distributed.ReduceOp.SUM,
              group=self.dp_pg,
          )
          total_norm = total_norm_exponentiated_cuda.item() ** (
              1.0 / norm_type
          )
  
      return total_norm
  
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 17 for calling the LowLevelZeroOptimizer obj's underlying member method code with example usage: obj._compute_grad_norm(gradients, norm_type) and returns: total_norm, int