code: def forward(
      ctx,
      input_: Tensor,
      weight: Tensor,
      bias: Tensor,
      weight_id: int,
      bias_id: int,
      normalized_shape: int,
      eps: float,
      output_parallel_mode: ParallelMode,
      input_x_weight_parallel_mode: ParallelMode,
  ) -> Tensor:
      ctx.weight_id = weight_id
      ctx.bias_id = bias_id
  
      sum_ = torch.sum(input_, dim=-1, keepdim=True)
      sqr_sum = torch.sum(input_**2, dim=-1, keepdim=True)
      mean, sqr_mean = (
          all_reduce(torch.stack((sum_, sqr_sum)), output_parallel_mode)
          / normalized_shape
      )
  
      output, mu, sigma = norm_forward(
          input_, mean, sqr_mean, weight, bias, eps
      )
  
      ctx.save_for_backward(mu, sigma, weight)
  
      ctx.normalized_shape = normalized_shape
      ctx.output_parallel_mode = output_parallel_mode
      ctx.input_x_weight_parallel_mode = input_x_weight_parallel_mode
  
      return output
  
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 1 for calling the _Layernorm3D obj's underlying member method code with example usage: obj.forward(ctx, input_, weight, bias, weight_id, bias_id, normalized_shape, eps, output_parallel_mode, input_x_weight_parallel_mode) and returns: output