code: def backward(ctx, grad_output):
      # Retreive tensors from the forward path.
      softmax, target_mask, masked_target_1d = ctx.saved_tensors
      label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size
  
      # All the inputs have softmax as thier gradient.
      grad_input = softmax
      # For simplicity, work with the 2D gradient.
      partition_vocab_size = softmax.size()[-1]
      grad_2d = grad_input.view(-1, partition_vocab_size)
  
      # Add the gradient from matching classes.
      arange_1d = torch.arange(
          start=0, end=grad_2d.size()[0], device=grad_2d.device
      )
  
      softmax_update = 1.0 - target_mask.view(-1).float()
  
      if label_smoothing > 0:
          smoothing = label_smoothing * vocab_size / (vocab_size - 1)
          grad_2d[arange_1d, masked_target_1d] -= (
              1.0 - smoothing
          ) * softmax_update
          average_grad = 1 / vocab_size
          grad_2d[arange_1d, :] -= smoothing * average_grad
      else:
          grad_2d[arange_1d, masked_target_1d] -= softmax_update
  
      # Finally elementwise multiplication with the output gradients.
      grad_input.mul_(grad_output.unsqueeze(dim=-1))
  
      return grad_input, None, None
  
file_path: megatron-lm/megatron/core/tensor_parallel/cross_entropy.py
imports: import torch
  from megatron.core.parallel_state import get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from .utils import VocabUtility
label_desc: Use image node_id 2 for calling the _VocabParallelCrossEntropy obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  None,  None