code: def backward(ctx, grad_output, *args):
      grad_output = grad_output.contiguous()
      if ctx.return_residual:
          (grad_input,) = args
          grad_input = grad_input.contiguous()
      process_group = ctx.process_group
      sequence_parallel = ctx.sequence_parallel
      if ctx.compute_weight_gradient:
          x, weight = ctx.saved_tensors
          if process_group is not None and sequence_parallel:
              total_x, handle_x = all_gather_raw(
                  x, process_group, async_op=True
              )
          else:
              total_x = x
      else:
          (weight,) = ctx.saved_tensors
          total_x = None
      batch_shape = grad_output.shape[:-1]
      batch_dim = batch_shape.numel()
      grad_output = grad_output.reshape(
          batch_dim, grad_output.shape[-1]
      )
      if ctx.needs_input_grad[0]:
          if not ctx.return_residual:
              grad_input = F.linear(grad_output, weight.t())
          else:
              grad_input = torch.addmm(
                  grad_input.reshape(batch_dim, grad_input.shape[-1]),
                  grad_output,
                  weight,
              )
          grad_input = grad_input.reshape(
              *batch_shape, grad_input.shape[-1]
          )
          if process_group is not None:
              reduce_fn = (
                  reduce_scatter_raw
                  if sequence_parallel
                  else all_reduce_raw
              )
              grad_input, handle_grad_input = reduce_fn(
                  grad_input, process_group, async_op=True
              )
      else:
          grad_input = None
      if ctx.needs_input_grad[1]:
          assert ctx.compute_weight_gradient
          if process_group is not None and sequence_parallel:
              handle_x.wait()
          grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad(
              total_x.reshape(batch_dim, total_x.shape[-1]),
              grad_output,
              ctx.needs_input_grad[2],
          )
      else:
          grad_weight = None
          grad_bias = grad_output if ctx.needs_input_grad[2] else None
      if process_group is not None and ctx.needs_input_grad[0]:
          handle_grad_input.wait()
      return grad_input, grad_weight, grad_bias, None, None, None
  
imports: from functools import partial
  from typing import Optional
  import fused_dense_lib
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.distributed import ProcessGroup
  from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd
  from flash_attn.utils.distributed import all_gather_raw, all_reduce, all_reduce_raw, reduce_scatter, reduce_scatter_raw
label_desc: Use image node_id 2 for calling the FusedDenseFunc obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  grad_weight,  grad_bias,  None,  None,  None