class_docstr: Linear layer execution with asynchronous communication and gradient accumulation fusion in backprop.
code: def backward(ctx, grad_output):
      if ctx.compute_weight_gradient:
          input, weight = ctx.saved_tensors
      else:
          weight = ctx.saved_tensors[0]
          input = None
  
      use_bias = ctx.use_bias
  
      # only get sequence parallel inputs if need to calculate weight grad
      handle = None
      if ctx.compute_weight_gradient:
          if ctx.sequence_parallel_enabled:
              world_size = get_tensor_model_parallel_world_size()
              shape = list(input.shape)
              shape[0] *= world_size
  
              all_gather_buffer = torch.empty(
                  shape,
                  dtype=input.dtype,
                  device=torch.cuda.current_device(),
                  requires_grad=False,
              )
              handle = torch.distributed.all_gather_into_tensor(
                  all_gather_buffer,
                  input,
                  group=get_tensor_model_parallel_group(),
                  async_op=True,
              )
              total_input = all_gather_buffer
          else:
              total_input = input
  
      grad_input = grad_output.matmul(weight)
  
      if handle is not None:
          handle.wait()
  
      if ctx.async_grad_allreduce:
          # Asynchronous all-reduce
          handle = torch.distributed.all_reduce(
              grad_input,
              group=get_tensor_model_parallel_group(),
              async_op=True,
          )
  
      # if no weight gradient, immediately return
      if not ctx.compute_weight_gradient:
          if ctx.sequence_parallel_enabled:
              assert not ctx.async_grad_allreduce
              world_size = get_tensor_model_parallel_world_size()
              shape = list(grad_input.shape)
              shape[0] //= world_size
  
              sub_grad_input = torch.empty(
                  torch.Size(shape),
                  dtype=grad_input.dtype,
                  device=torch.cuda.current_device(),
                  requires_grad=False,
              )
              handle = torch.distributed.reduce_scatter_tensor(
                  sub_grad_input,
                  grad_input,
                  group=get_tensor_model_parallel_group(),
                  async_op=True,
              )
              handle.wait()
              return sub_grad_input, None, None, None, None, None, None
          if ctx.async_grad_allreduce:
              handle.wait()
          return grad_input, None, None, None, None, None, None
  
      # Convert the tensor shapes to 2D for execution compatibility
      grad_output = grad_output.contiguous()
      grad_output = grad_output.view(
          grad_output.shape[0] * grad_output.shape[1],
          grad_output.shape[2],
      )
      total_input = total_input.view(
          total_input.shape[0] * total_input.shape[1],
          total_input.shape[2],
      )
  
      if ctx.sequence_parallel_enabled:
          assert not ctx.async_grad_allreduce
          sub_grad_input = torch.empty(
              input.shape,
              dtype=input.dtype,
              device=torch.cuda.current_device(),
              requires_grad=False,
          )
          handle = torch.distributed.reduce_scatter_tensor(
              sub_grad_input,
              grad_input,
              group=get_tensor_model_parallel_group(),
              async_op=True,
          )
  
      if ctx.gradient_accumulation_fusion:
          if not hasattr(weight, "main_grad"):
              raise RuntimeError(
                  "attempted to perform gradient accumulation fusion on param without setting main_grad"
              )
          if weight.main_grad.dtype == torch.float32:
              fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
                  total_input, grad_output, weight.main_grad
              )
          elif weight.main_grad.dtype in (
              torch.float16,
              torch.bfloat16,
          ):
              fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
                  total_input, grad_output, weight.main_grad
              )
          else:
              raise RuntimeError(
                  f"unsupported dtype for main_grad ({weight.main_grad.dtype})"
              )
          grad_weight = None
      else:
          grad_weight = grad_output.t().matmul(total_input)
      grad_bias = grad_output.sum(dim=0) if use_bias else None
      if ctx.sequence_parallel_enabled:
          handle.wait()
          return (
              sub_grad_input,
              grad_weight,
              grad_bias,
              None,
              None,
              None,
              None,
          )
      if ctx.async_grad_allreduce:
          handle.wait()
      return grad_input, grad_weight, grad_bias, None, None, None, None
  
imports: from typing import Optional, Dict, Tuple, List
  import warnings
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.nn.parameter import Parameter
  from apex._autocast_utils import _cast_if_autocast_enabled
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
  from apex.transformer.utils import divide
  from apex.transformer.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import scatter_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_scatter_to_sequence_parallel_region
  from apex.transformer.tensor_parallel.random import get_cuda_rng_tracker
  from apex.transformer.tensor_parallel.utils import VocabUtility
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 2 for calling the LinearWithGradAccumulationAndAsyncCommunication 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,  None, grad_input,  None,  None,  None,  None,  None,  None, sub_grad_input,  grad_weight,  grad_bias,  None,  None,  None,  None, sub_grad_input,  None,  None,  None,  None,  None,  None