code: def norm_backward(
      grad: Tensor, mu: Tensor, sigma: Tensor, weight: Tensor
  ):
      # dbias, dweight = grad, grad * mu / sigma
      dz = grad * weight
      dmu = dz / sigma
      dvar = dz * mu * (-0.5) * sigma ** (-3)
      dmean = -dmu
      dvar = torch.sum(dvar, -1, keepdim=True)
      dmean = torch.sum(dmean, -1, keepdim=True)
  
      return dmu, dmean, dvar
  
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication import all_gather, all_reduce, broadcast, reduce, reduce_scatter
  from colossalai.legacy.constants import INPUT_GROUP_3D, WEIGHT_GROUP_3D
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from ._utils import get_parallel_mode_from_env, push_async_grad
label_desc: Use image node_id 5 for calling a global function with example usage: norm_backward(grad, mu, sigma, weight) and returns: dmu,  dmean,  dvar