code: def backward(ctx, output_grad: Tensor) -> Tuple[Tensor, ...]:
      input_, weight = ctx.saved_tensors
      output_grad = all_gather(output_grad, 0, ctx.output_parallel_mode)
  
      input_grad = torch.matmul(output_grad, weight.transpose(0, 1))
      input_grad, input_op = reduce_scatter(
          input_grad, 0, ctx.input_parallel_mode, async_op=True
      )
  
      weight_grad = torch.matmul(
          input_.reshape(-1, input_.shape[-1]).transpose(0, 1),
          output_grad.reshape(-1, output_grad.shape[-1]),
      )
      weight_grad, op = reduce_scatter(
          weight_grad, 0, ctx.weight_parallel_mode, async_op=True
      )
      weight_grad = push_async_grad(op, weight_grad, ctx.weight_id)
  
      input_op.wait()
  
      return input_grad, weight_grad, 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 _Linear3D obj's underlying member method code with example usage: obj.backward(ctx, output_grad) and returns: input_grad,  weight_grad,  None,  None,  None,  None