class_docstr: Calculate AV in a ring-exchange style
code: def backward(ctx, grad_output):
      local_rank = gpc.get_local_rank(ParallelMode.SEQUENCE)
      local_world_size = gpc.get_world_size(ParallelMode.SEQUENCE)
      local_start_idx, local_end_idx = _calc_current_device_range(
          local_rank, ctx.sub_seq_length
      )
      attention_scores, sub_v = ctx.saved_tensors
  
      # calculate gradient of v
      grad_v = torch.matmul(
          attention_scores.transpose(2, 1), grad_output
      )
      dist.all_reduce(
          grad_v, group=gpc.get_group(ParallelMode.SEQUENCE)
      )
      grad_v = grad_v[:, local_start_idx:local_end_idx]
      grad_v /= local_world_size
  
      # calculate gradient for attention score
      grad_attention_score = torch.zeros_like(
          attention_scores,
          dtype=grad_output.dtype,
          device=get_current_device(),
      )
  
      # compute with local sub_k
      grad_attention_score[
          :, :, local_start_idx:local_end_idx
      ] += torch.matmul(grad_output, sub_v.transpose(2, 1))
  
      # compute QK^T in ring-all-reduce style
      for i in range(local_world_size - 1):
          sub_v = ring_forward(sub_v, ParallelMode.SEQUENCE)
          start_idx, end_idx = _calc_incoming_device_range(
              i, local_rank, local_world_size, ctx.sub_seq_length
          )
  
          # compute grad_q
          grad_attention_score[:, :, start_idx:end_idx] += torch.matmul(
              grad_output, sub_v.transpose(2, 1)
          )
  
      return grad_attention_score, grad_v, None, None, None, None
  
file_path: ColossalAI/colossalai/legacy/nn/layer/parallel_sequence/_operation.py
imports: import torch
  from torch import distributed
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication import ring_forward
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.parallel_sequence._utils import _calc_current_device_range, _calc_incoming_device_range
  from colossalai.utils import get_current_device
label_desc: Use image node_id 2 for calling the RingAV obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_attention_score,  grad_v,  None,  None,  None,  None