code: def forward(ctx, logits, targets, output_parallel_mode):
      # logits: [b/q^2, c/q]
      # labels: [b/q^2]
      # loss: [b/q^2]
      logits_max = torch.max(logits, dim=-1)[0]
      dist.all_reduce(
          logits_max,
          op=torch.distributed.ReduceOp.MAX,
          group=gpc.get_group(output_parallel_mode),
      )
      # Subtract the maximum value.
      logits = logits - logits_max.unsqueeze(dim=-1)
  
      vocab_size_per_partition = logits.size()[-1]
      rank = gpc.get_local_rank(output_parallel_mode)
      vocab_start = rank * vocab_size_per_partition
      vocab_end = (rank + 1) * vocab_size_per_partition - 1
  
      # loss[i] = 0 if targets[i] < vocab_start or targets[i] > vocab_end
      target_mask = (targets < vocab_start) | (targets > vocab_end)
      masked_target = targets.clone() - vocab_start
      masked_target[target_mask] = 0
      arange_1d = torch.arange(
          start=0, end=logits.size()[0], device=get_current_device()
      )
      predicted_logits = logits[arange_1d, masked_target]
      predicted_logits = (
          predicted_logits.clone().contiguous().view_as(targets)
      )
      predicted_logits[target_mask] = 0.0
      dist.all_reduce(
          predicted_logits, group=gpc.get_group(output_parallel_mode)
      )
  
      # Loss = log(sum(exp(logits))) - predicted-logit.
      exp_logits = torch.exp(logits)
      sum_exp_logits = exp_logits.sum(dim=-1)
      dist.all_reduce(
          sum_exp_logits, group=gpc.get_group(output_parallel_mode)
      )
      loss = torch.log(sum_exp_logits) - predicted_logits
  
      exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
      ctx.save_for_backward(exp_logits, target_mask, masked_target)
  
      return loss
  
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 1 for calling the _VocabParallelCrossEntropy3D obj's underlying member method code with example usage: obj.forward(ctx, logits, targets, output_parallel_mode) and returns: loss