code: def backward(ctx, grad_output, *args):
      grad_output = grad_output.contiguous()
      checkpoint_lvl = ctx.checkpoint_lvl
      activation = ctx.activation
      activation_fn = (
          partial(F.gelu, approximate="tanh")
          if activation == "gelu_approx"
          else (sqrelu_fwd if activation == "sqrelu" else F.relu)
      )
      if ctx.return_residual:
          (grad_input,) = args
          grad_input = grad_input.contiguous()
      process_group = ctx.process_group
      sequence_parallel = ctx.sequence_parallel
      x, weight1, weight2, *rest = ctx.saved_tensors
      if process_group is None or not sequence_parallel:
          total_x = x
      batch_shape = grad_output.shape[:-1]
      batch_dim = batch_shape.numel()
      if checkpoint_lvl in [0, 1]:
          if process_group is not None and sequence_parallel:
              total_x, handle_x = all_gather_raw(
                  x, process_group, async_op=True
              )
          if checkpoint_lvl == 0 or (
              checkpoint_lvl == 1 and activation == "relu"
          ):
              pre_act, output1 = rest
          elif checkpoint_lvl == 1:
              (pre_act,) = rest
              with torch.jit.fuser("fuser2"):
                  output1 = activation_fn(pre_act)
      elif checkpoint_lvl == 2:
          (bias1,) = rest
          if process_group is not None and sequence_parallel:
              total_x, _ = all_gather_raw(x, process_group)
          if ctx.heuristic == -1:
              pre_act = F.linear(total_x, weight1, bias1)
              with torch.jit.fuser("fuser2"):
                  output1 = activation_fn(pre_act)
          else:
              output1, pre_act = fused_dense_cuda.linear_act_forward(
                  total_x.reshape(batch_dim, total_x.shape[-1]),
                  weight1,
                  bias1,
                  activation == "gelu_approx",
                  True,
                  ctx.heuristic,
              )
  
      grad_output = grad_output.reshape(
          batch_dim, grad_output.shape[-1]
      )
      output1 = output1.reshape(batch_dim, output1.shape[-1])
      pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1])
      if ctx.needs_input_grad[3]:
          grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(
              output1, grad_output, ctx.needs_input_grad[4]
          )
      else:
          grad_weight2 = None
          grad_bias2 = grad_output if ctx.needs_input_grad[4] else None
      if ctx.heuristic == -1:
          # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act)
          grad_output1 = F.linear(grad_output, weight2.t())
          activation_grad_fn = (
              gelu_bwd
              if activation == "gelu_approx"
              else (sqrelu_bwd if activation == "sqrelu" else relu_bwd)
          )
          with torch.jit.fuser("fuser2"):
              grad_pre_act = activation_grad_fn(grad_output1, pre_act)
      else:
          # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't
          # just compute gelu/relu grad
          (
              grad_pre_act,
              grad_bias1,
          ) = fused_dense_cuda.bias_act_linear_dgrad_bgrad(
              weight2,
              grad_output,
              pre_act,
              activation == "gelu_approx",
              ctx.heuristic,
          )
          if not ctx.needs_input_grad[2]:
              grad_bias1 = None
      if ctx.needs_input_grad[0]:
          if not ctx.return_residual:
              grad_input = F.linear(grad_pre_act, weight1.t())
          else:
              grad_input = torch.addmm(
                  grad_input.reshape(batch_dim, grad_input.shape[-1]),
                  grad_pre_act,
                  weight1,
              )
          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.heuristic == -1:
          if ctx.needs_input_grad[1]:
              if (
                  process_group is not None
                  and sequence_parallel
                  and checkpoint_lvl != 2
              ):
                  handle_x.wait()
              (
                  grad_weight1,
                  grad_bias1,
              ) = fused_dense_cuda.linear_bias_wgrad(
                  total_x.reshape(batch_dim, total_x.shape[-1]),
                  grad_pre_act,
                  ctx.needs_input_grad[2],
              )
          else:
              grad_weight1 = None
              grad_bias1 = (
                  grad_pre_act if ctx.needs_input_grad[2] else None
              )
      else:
          if ctx.needs_input_grad[1]:
              if (
                  process_group is not None
                  and sequence_parallel
                  and checkpoint_lvl != 2
              ):
                  handle_x.wait()
              grad_weight1 = F.linear(
                  grad_pre_act.t(),
                  total_x.reshape(batch_dim, total_x.shape[-1]).t(),
              )
          else:
              grad_weight1 = None
      if process_group is not None and ctx.needs_input_grad[0]:
          handle_grad_input.wait()
      return (
          grad_input,
          grad_weight1,
          grad_bias1,
          grad_weight2,
          grad_bias2,
          None,
          None,
          None,
          None,
          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 FusedMLPFunc obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  grad_weight1,  grad_bias1,  grad_weight2,  grad_bias2,  None,  None,  None,  None,  None,  None,  None