code: def backward(ctx, output_grad: Tensor) -> Tuple[Tensor, ...]:
      input_, weight = ctx.saved_tensors
      weight_grad = torch.matmul(
          output_grad.reshape(-1, output_grad.shape[-1]).transpose(
              0, 1
          ),
          input_.reshape(-1, input_.shape[-1]),
      )
      weight_grad = reduce(
          weight_grad, ctx.src_rank, ctx.input_parallel_mode
      )
      if gpc.get_local_rank(
          ctx.input_parallel_mode
      ) == gpc.get_local_rank(ctx.output_parallel_mode):
          weight_grad, op = all_reduce(
              weight_grad, ctx.weight_parallel_mode, async_op=True
          )
          weight_grad = push_async_grad(op, weight_grad, ctx.weight_id)
      else:
          weight_grad = None
  
      if ctx.use_bias:
          bias_grad = torch.sum(
              output_grad, dim=tuple(range(len(output_grad.shape))[:-1])
          )
          bias_grad = all_reduce(bias_grad, ctx.input_parallel_mode)
          bias_grad, op = all_reduce(
              bias_grad, ctx.weight_parallel_mode, async_op=True
          )
          bias_grad = push_async_grad(op, bias_grad, ctx.bias_id)
      else:
          bias_grad = None
  
      input_grad = torch.matmul(output_grad, weight)
  
      return (
          input_grad,
          weight_grad,
          bias_grad,
          None,
          None,
          None,
          None,
          None,
      )
  
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 2 for calling the _Classifier3D obj's underlying member method code with example usage: obj.backward(ctx, output_grad) and returns: input_grad,  weight_grad,  bias_grad,  None,  None,  None,  None,  None