code: def backward(ctx, output_grad):
      # Retrieve tensors from the forward path.
      softmax, target_mask, masked_target = ctx.saved_tensors
  
      # All the inputs have softmax as their gradient.
      input_grad = softmax
      # For simplicity, work with the 2D gradient.
      partition_vocab_size = softmax.size()[-1]
      grad_2d = input_grad.view(-1, partition_vocab_size)
  
      # Add the gradient from matching classes.
      arange_1d = torch.arange(
          start=0, end=grad_2d.size()[0], device=get_current_device()
      )
      grad_2d[arange_1d, masked_target] -= (
          1.0 - target_mask.view(-1).float()
      )
      input_grad.mul_(output_grad.unsqueeze(dim=-1))
  
      return input_grad, None, None, None
  
imports: import torch
  import torch.distributed
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.nn.functional import cross_entropy
  from torch.nn.modules.loss import _Loss
  from colossalai.legacy.constants import INPUT_GROUP_3D, OUTPUT_GROUP_3D, WEIGHT_GROUP_3D
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.parallel_3d import reduce_by_batch_3d, split_tensor_3d
  from colossalai.legacy.nn.layer.parallel_3d._utils import get_parallel_mode_from_env
  from colossalai.legacy.registry import LOSSES
  from colossalai.utils import get_current_device
label_desc: Use image node_id 2 for calling the _VocabParallelCrossEntropy3D obj's underlying member method code with example usage: obj.backward(ctx, output_grad) and returns: input_grad,  None,  None,  None