code: def parallel_lm_logits(
      input_, word_embeddings_weight, parallel_output, bias=None
  ):
      """LM logits using word embedding weights."""
      args = get_args()
      # Parallel logits.
      if (
          args.async_tensor_model_parallel_allreduce
          or args.sequence_parallel
      ):
          input_parallel = input_
          model_parallel = (
              mpu.get_tensor_model_parallel_world_size() > 1
          )
          async_grad_allreduce = (
              args.async_tensor_model_parallel_allreduce
              and model_parallel
              and not args.sequence_parallel
          )
      else:
          input_parallel = (
              tensor_parallel.copy_to_tensor_model_parallel_region(
                  input_
              )
          )
          async_grad_allreduce = False
  
      # Matrix multiply.
      logits_parallel = tensor_parallel.linear_with_grad_accumulation_and_async_allreduce(
          input=input_parallel,
          weight=word_embeddings_weight,
          bias=bias,
          gradient_accumulation_fusion=args.gradient_accumulation_fusion,
          async_grad_allreduce=async_grad_allreduce,
          sequence_parallel=args.sequence_parallel,
      )
      # Gather if needed.
  
      if parallel_output:
          return logits_parallel
  
      return tensor_parallel.gather_from_tensor_model_parallel_region(
          logits_parallel
      )
  
imports: import torch
  import torch.nn.functional
  from megatron import get_args
  from megatron.core import mpu, tensor_parallel
  from megatron.core.enums import ModelType
  from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
  from .enums import AttnMaskType, LayerType
  from .module import MegatronModule
  from .transformer import ParallelTransformer
  from .utils import get_linear_layer
  from .utils import init_method_normal, scaled_init_method_normal
label_desc: Use image node_id 1 for calling a global function with example usage: parallel_lm_logits(input_, word_embeddings_weight, parallel_output, bias) and returns: tensor_parallel, logits_parallel