class_docstr: Gather input from sequence parallel in forward and reduce-scatter gradient in backward
  
  Args:
      input_ (`torch.Tensor`): The input tensor from sequence parallel region.
      process_group (`torch.distributed.ProcessGroup`): The process group used for collective communication.
      overlap (`bool`): Whther to overlap the all_gather op and gradient calculate in backward.
class_docstr_tok: ["Gather","input","from","sequence","parallel","in","forward","and","reduce-scatter","gradient","in","backward","Args",":","input_","(","`","torch.Tensor","`",")",":","The","input","tensor","from","sequence","parallel","region",".","process_group","(","`","torch.distributed.ProcessGroup","`",")",":","The","process","group","used","for","collective","communication",".","overlap","(","`","bool","`",")",":","Whther","to","overlap","the","all_gather","op","and","gradient","calculate","in","backward","."]
code: def backward(ctx, grad_output):
      input_, weight, bias = ctx.saved_tensors
      use_bias = ctx.use_bias
      dim = ctx.dim
      process_group = ctx.process_group
      overlap = ctx.overlap
  
      # In order to be hooked into Gemini's '__torch_function__', adding a view operation to weight and bias. Used in FusedLayerNorm
      if use_bias:
          bias = bias.view(bias.shape)
  
      if not overlap:
          input_parallel = _gather(input_, dim, process_group)
  
          total_input = input_parallel
          grad_input = grad_output.matmul(weight)
          grad_output = grad_output.contiguous()
          # Convert the tensor shapes to 2D for execution compatibility
          if len(grad_output.shape) > 2:
              grad_output = grad_output.view(-1, grad_output.shape[-1])
              total_input = total_input.view(-1, total_input.shape[-1])
  
          if ctx.async_grad_reduce_scatter:
              # Asynchronous reduce-scatter
              input_list = [
                  item.contiguous()
                  for item in torch.chunk(
                      grad_input,
                      dist.get_world_size(process_group),
                      dim=dim,
                  )
              ]
              output = torch.empty(
                  input_.shape,
                  dtype=input_parallel.dtype,
                  device=input_parallel.device,
              ).contiguous()
              handle = dist.reduce_scatter(
                  output, input_list, group=process_group, async_op=True
              )
              # Delay the start of weight gradient computation shortly (3us) to have
              # reduce-scatter scheduled first and have GPU resources allocated
              _ = torch.empty(1, device=grad_output.device) + 1
  
          if _grad_accum_fusion_available and weight.grad is not None:
              grad = weight.grad
              if grad.dtype == torch.float32:
                  fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
                      total_input, grad_output, grad
                  )
                  grad_weight = None
              elif grad.dtype == torch.float16:
                  fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
                      total_input, grad_output, grad
                  )
                  grad_weight = None
              else:
                  grad_weight = grad_output.t().matmul(total_input)
          else:
              grad_weight = grad_output.t().matmul(total_input)
  
          grad_bias = grad_output.sum(dim=0) if use_bias else None
  
          if ctx.async_grad_reduce_scatter:
              handle.wait()
  
      else:
          input_ = input_.contiguous()
          world_size = dist.get_world_size(process_group)
          tensor_list = [
              torch.empty_like(input_) for _ in range(world_size)
          ]
  
          # do all gather in is async way
          gather_handle = dist.all_gather(
              tensor_list, input_, group=process_group, async_op=True
          )
          # calculate gradient and prepare data asynchronously with all-gather
          # calculate
          grad_input = grad_output.matmul(weight)
          grad_output = grad_output.contiguous()
          # Convert the tensor shapes to 2D for execution compatibility
          if len(grad_output.shape) > 2:
              grad_output = grad_output.view(-1, grad_output.shape[-1])
          grad_bias = grad_output.sum(dim=0) if use_bias else None
          # prepare data
          input_list = [
              item.contiguous()
              for item in torch.chunk(
                  grad_input,
                  dist.get_world_size(process_group),
                  dim=dim,
              )
          ]
          output = torch.empty(
              input_.shape, dtype=input_.dtype, device=input_.device
          ).contiguous()
          # wait until all-gather finished
          gather_handle.wait()
  
          # do reduce-scatter in async way
          reducescatter_handle = dist.reduce_scatter(
              output, input_list, group=process_group, async_op=True
          )
          input_parallel = torch.cat(tensor_list, dim=dim).contiguous()
          # calculate gradient
          if len(input_parallel.shape) > 2:
              input_parallel = input_parallel.view(
                  -1, input_parallel.shape[-1]
              )
  
          if _grad_accum_fusion_available and weight.grad is not None:
              grad = weight.grad
              if grad.dtype == torch.float32:
                  fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
                      input_parallel, grad_output, grad
                  )
                  grad_weight = None
              elif grad.dtype == torch.float16:
                  fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
                      input_parallel, grad_output, grad
                  )
                  grad_weight = None
              else:
                  grad_weight = grad_output.t().matmul(input_parallel)
          else:
              grad_weight = grad_output.t().matmul(input_parallel)
          # grad_weight = grad_output.t().matmul(input_parallel)
          # wait until reduce-scatter finished
          reducescatter_handle.wait()
  
      return output, grad_weight, grad_bias, None, None, None, None
  
imports: import torch
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the _LinearWithGatherForwardReduceScatterBackward obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: output,  grad_weight,  grad_bias,  None,  None,  None,  None