code: def check_vocab_parallel_loss():
      device = get_current_device()
      dtype = torch.float32
  
      i = gpc.get_local_rank(ParallelMode.PARALLEL_1D)
  
      criterion = VocabParallelCrossEntropyLoss1D()
      criterion_master = torch.nn.CrossEntropyLoss()
  
      out_shape = (BATCH_SIZE, SEQ_LENGTH, NUM_CLASSES)
      out_master = torch.randn(out_shape, dtype=dtype, device=device)
      target_master = torch.randint(
          NUM_CLASSES,
          (BATCH_SIZE, SEQ_LENGTH),
          dtype=torch.long,
          device=device,
      )
      torch.distributed.broadcast(out_master, src=0)
      torch.distributed.broadcast(target_master, src=0)
      out = torch.chunk(out_master, DEPTH, dim=-1)[i]
      out = out.clone()
      out.requires_grad = True
  
      loss = criterion(out, target_master)
  
      out_master = out_master.clone()
      out_master.requires_grad = True
      loss_master = criterion_master(out_master, target_master)
      check_equal(loss, loss_master)
      print_rank_0("vocab parallel loss forward: pass")
  
      loss.backward()
      loss_master.backward()
  
      out_grad = out_master.grad
      out_grad = torch.chunk(out_grad, DEPTH, dim=-1)[i]
      check_equal(out_grad, out.grad)
      print_rank_0("vocab parallel loss backward: pass")
  
file_path: colossalai/tests/test_legacy/test_layers/test_1d/checks_1d/check_layer_1d.py
imports: import torch
  import torch.distributed
  from torch.nn import Parameter
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.global_variables import tensor_parallel_env
  from colossalai.legacy.nn import Classifier1D, Embedding1D, Linear1D_Col, Linear1D_Row, VanillaClassifier, VocabParallelClassifier1D, VocabParallelCrossEntropyLoss1D, VocabParallelEmbedding1D
  from colossalai.legacy.utils import print_rank_0
  from colossalai.utils import get_current_device
  from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, NUM_CLASSES, SEQ_LENGTH, VOCAB_SIZE, check_equal
label_desc: Use image node_id 9 for calling a global function with example usage: check_vocab_parallel_loss() without return types