code: def check_vocab_parallel_classifier_given_embed_weight():
      device = get_current_device()
      dtype = torch.float32
  
      j = gpc.get_local_rank(ParallelMode.PARALLEL_2D_ROW)
      i = gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)
  
      embed = VocabParallelEmbedding2D(VOCAB_SIZE, HIDDEN_SIZE)
      embed = embed.to(dtype).to(device)
      embed_master = torch.nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE)
      embed_master = embed_master.to(dtype).to(device)
  
      weight_master = embed_master.weight.data
      torch.distributed.broadcast(weight_master, src=0)
      weight = torch.chunk(weight_master, DEPTH, dim=-1)[j]
      weight = torch.chunk(weight, DEPTH, dim=0)[i]
      embed.weight.data.copy_(weight)
  
      layer = VocabParallelClassifier2D(
          HIDDEN_SIZE, VOCAB_SIZE, weight=embed.weight, bias=False
      )
      layer = layer.to(dtype).to(device)
      layer_master = VanillaClassifier(
          HIDDEN_SIZE,
          VOCAB_SIZE,
          weight=embed_master.weight,
          bias=False,
      )
      layer_master = layer_master.to(dtype).to(device)
  
      A_shape = (BATCH_SIZE, SEQ_LENGTH)
      A_master = torch.randint(VOCAB_SIZE, A_shape, device=device)
      torch.distributed.broadcast(A_master, src=0)
      A = A_master.clone()
      out = layer(embed(A))
  
      A_master = A_master.clone()
      C_master = layer_master(embed_master(A_master))
      C = torch.chunk(C_master, DEPTH, dim=0)[i]
      C = torch.chunk(C, DEPTH, dim=-1)[j]
      check_equal(out, C)
      print_rank_0(
          "vocab parallel classifier (given embed weight) forward: pass"
      )
  
      grad_shape = C_master.shape
      grad_master = torch.randn(grad_shape, dtype=dtype, device=device)
      torch.distributed.broadcast(grad_master, src=0)
      grad = torch.chunk(grad_master, DEPTH, dim=0)[i]
      grad = torch.chunk(grad, DEPTH, dim=-1)[j]
      grad = grad.clone()
      out.backward(grad)
  
      grad_master = grad_master.clone()
      C_master.backward(grad_master)
  
      W_grad = embed_master.weight.grad
      W_grad = torch.chunk(W_grad, DEPTH, dim=-1)[j]
      W_grad = torch.chunk(W_grad, DEPTH, dim=0)[i]
      check_equal(W_grad, embed.weight.grad)
      print_rank_0(
          "vocab parallel classifier (given embed weight) backward: pass"
      )
  
file_path: ColossalAI/tests/test_legacy/test_layers/test_2d/checks_2d/check_layer_2d.py
imports: import torch
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn import Classifier2D, CrossEntropyLoss2D, Embedding2D, LayerNorm2D, Linear2D, PatchEmbedding2D, VanillaClassifier, VanillaPatchEmbedding, VocabParallelClassifier2D, VocabParallelCrossEntropyLoss2D, VocabParallelEmbedding2D
  from colossalai.legacy.utils import print_rank_0
  from colossalai.utils import get_current_device
  from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, IMG_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_classifier_given_embed_weight() without return types