class_docstr: Overwrite the forward and backward function to calculate the cross entropy loss before gather
  
  Args:
      Function (:class:`torch.autograd.Function`): default
code: def backward(ctx, grad_output):
      # retrieve the saved tensors
      grad_output = grad_output * ctx.inv_num_non_zero
      exp_logits, mask, masked_target_1d = ctx.saved_tensors
  
      # use exp logits as the input grad
      grad_logits = exp_logits
      partion_vocab_size = grad_logits.shape[-1]
      grad_logits_2d = grad_logits.view(-1, partion_vocab_size)
  
      update = 1.0 - mask.view(-1).float()
      grad_logits_2d[
          torch.arange(0, grad_logits_2d.shape[0]), masked_target_1d
      ] -= update
  
      grad_logits.mul_(grad_output.unsqueeze(dim=-1))
      return grad_logits, None, None, None
  
file_path: ColossalAI/colossalai/shardformer/layer/loss.py
imports: import torch
  import torch.distributed
  from torch.autograd import Function
  from torch.distributed import ProcessGroup
label_desc: Use image node_id 2 for calling the DistCrossEntropy obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_logits,  None,  None,  None