code: def backward(ctx, grad_output):
      bsz = grad_output.shape[0]
      grad_output_shape = grad_output.size()
      if grad_output_shape[1] % 16 != 0:
          grad_output = torch.cat(
              (
                  grad_output,
                  torch.zeros(
                      (
                          bsz,
                          (16 - (grad_output_shape[1] % 16)),
                          grad_output_shape[2],
                      ),
                      device=grad_output.device,
                      dtype=grad_output.dtype,
                  ),
              ),
              1,
          )
  
      assert ctx.config.training
  
      if ctx.config.pre_layer_norm and ctx.config.normalize_invertible:
          (
              input_mask,
              attn_qkvw,
              attn_qkvb,
              attn_ow,
              attn_ob,
              attn_nw,
              attn_nb,
              inter_w,
              inter_b,
              output_w,
              output_b,
              norm_w,
              norm_b,
          ) = ctx.saved_tensors
      else:
          (
              output,
              input,
              input_mask,
              attn_qkvw,
              attn_qkvb,
              attn_ow,
              attn_ob,
              attn_nw,
              attn_nb,
              inter_w,
              inter_b,
              output_w,
              output_b,
              norm_w,
              norm_b,
          ) = ctx.saved_tensors
  
      cuda_module = (
          stochastic_transformer_cuda_module
          if ctx.config.stochastic_mode
          else transformer_cuda_module
      )
      backward_func = (
          cuda_module.backward_fp16
          if ctx.config.fp16
          else cuda_module.backward_fp32
      )
  
      (
          grad_input,
          grad_attn_qkvw,
          grad_attn_qkvb,
          grad_attn_ow,
          grad_attn_ob,
          grad_attn_nw,
          grad_attn_nb,
          grad_inter_w,
          grad_inter_b,
          grad_output_w,
          grad_output_b,
          grad_norm_w,
          grad_norm_b,
      ) = backward_func(
          ctx.config.layer_id,
          grad_output,
          (
              ctx.inp_norm
              if (
                  ctx.config.pre_layer_norm
                  and ctx.config.normalize_invertible
              )
              else output
          ),
          (
              ctx.inp_norm
              if (
                  ctx.config.pre_layer_norm
                  or not ctx.config.normalize_invertible
              )
              else input
          ),
          ctx.qkv_tf,
          ctx.soft_inp,
          (
              ctx.soft_inp
              if ctx.config.attn_dropout_checkpoint
              else ctx.ctx_bufB
          ),
          ctx.attn_o_inp,
          (
              ctx.ff1_inp
              if ctx.config.normalize_invertible
              else ctx.add_res
          ),
          ctx.ff1_inp,
          (ctx.ff2_inp if ctx.config.gelu_checkpoint else ctx.gelu_inp),
          ctx.ff2_inp,
          ctx.attn_prob_dropout_mask,
          ctx.attn_output_dropout_mask,
          ctx.layer_output_dropout_mask,
          ctx.attn_layer_norm_var,
          ctx.attn_layer_norm_mean,
          ctx.layer_norm_var,
          ctx.layer_norm_mean,
          (
              ctx.inp_norm
              if (
                  ctx.config.pre_layer_norm
                  and ctx.config.normalize_invertible
              )
              else input
          ),
          input_mask,
          attn_qkvw,
          attn_qkvb,
          attn_ow,
          attn_ob,
          attn_nw,
          attn_nb,
          inter_w,
          inter_b,
          output_w,
          output_b,
          norm_w,
          norm_b,
      )
  
      # This appears to be an effective way to release context memory
      ctx.qkv_tf = None
      ctx.soft_inp = None
      ctx.ctx_bufB = None
      ctx.gelu_inp = None
      ctx.ff2_inp = None
      ctx.attn_o_inp = None
      ctx.ff1_inp = None
      ctx.add_res = None
      ctx.inp_norm = None
      ctx.config = None
      ctx.attn_layer_norm_mean = None
      ctx.layer_norm_mean = None
      ctx.attn_prob_dropout_mask = None
      ctx.attn_output_dropout_mask = None
      ctx.layer_output_dropout_mask = None
      ctx.attn_layer_norm_var = None
      ctx.layer_norm_var = None
  
      if grad_output_shape[1] % 16 != 0:
          grad_input = torch.narrow(
              grad_input, 1, 0, grad_output_shape[1]
          )
  
      return (
          grad_input,
          None,
          None,
          None,
          None,
          grad_attn_qkvw,
          grad_attn_qkvb,
          grad_attn_ow,
          grad_attn_ob,
          grad_attn_nw,
          grad_attn_nb,
          grad_inter_w,
          grad_inter_b,
          grad_output_w,
          grad_output_b,
          grad_norm_w,
          grad_norm_b,
          None,
      )
  
imports: import json
  import math
  import torch
  from torch import nn
  from torch.autograd import Function
  from ..op_builder import TransformerBuilder, StochasticTransformerBuilder
label_desc: Use image node_id 2 for calling the DeepSpeedTransformerFunction obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  None,  None,  None,  None,  grad_attn_qkvw,  grad_attn_qkvb,  grad_attn_ow,  grad_attn_ob,  grad_attn_nw,  grad_attn_nb,  grad_inter_w,  grad_inter_b,  grad_output_w,  grad_output_b,  grad_norm_w,  grad_norm_b,  None