code: def forward(
      ctx,
      input_: Tensor,
      weight: Tensor,
      bias: Optional[Tensor],
      weight_id: int,
      bias_id: Optional[int],
      input_parallel_mode: ParallelMode,
      weight_parallel_mode: ParallelMode,
      output_parallel_mode: ParallelMode,
  ) -> Tensor:
      ctx.use_bias = bias is not None
      ctx.weight_id = weight_id
  
      src_rank = gpc.get_ranks_in_group(input_parallel_mode)[
          gpc.get_local_rank(output_parallel_mode)
      ]
      weight = broadcast(weight, src_rank, input_parallel_mode)
      ctx.save_for_backward(input_, weight)
  
      output = torch.matmul(input_, weight.transpose(0, 1))
      output = all_reduce(output, output_parallel_mode)
  
      if bias is not None:
          ctx.bias_id = bias_id
          output += bias
  
      ctx.src_rank = src_rank
      ctx.input_parallel_mode = input_parallel_mode
      ctx.weight_parallel_mode = weight_parallel_mode
      ctx.output_parallel_mode = output_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 _Classifier3D obj's underlying member method code with example usage: obj.forward(ctx, input_, weight, bias, weight_id, bias_id, input_parallel_mode, weight_parallel_mode, output_parallel_mode) and returns: output