code: def _warmup_jit_function():
      """Compilie JIT functions before the main training steps"""
      args = get_args()
      if args.bf16:
          dtype = torch.bfloat16
      elif args.fp16:
          dtype = torch.float16
      else:
          dtype = torch.float32
  
      # Warmup fused bias+gelu
      bias = torch.rand(
          args.ffn_hidden_size // args.tensor_model_parallel_size,
          dtype=dtype,
          device="cuda",
      )
      input = torch.rand(
          (
              args.seq_length,
              args.micro_batch_size,
              args.ffn_hidden_size // args.tensor_model_parallel_size,
          ),
          dtype=dtype,
          device="cuda",
      )
      # Warmup JIT fusions with the input grad_enable state of both forward
      # prop and recomputation
      for bias_grad, input_grad in zip([True, True], [False, True]):
          bias.requires_grad, input.requires_grad = (
              bias_grad,
              input_grad,
          )
          for _ in range(5):
              output = bias_gelu(bias, input)
      del bias, input, output
  
      # Warmup fused bias+dropout+add
      if args.sequence_parallel:
          seq_length = (
              args.seq_length
              // mpu.get_tensor_model_parallel_world_size()
          )
      else:
          seq_length = args.seq_length
      input = torch.rand(
          (seq_length, args.micro_batch_size, args.hidden_size),
          dtype=dtype,
          device="cuda",
      )
      residual = torch.rand(
          (seq_length, args.micro_batch_size, args.hidden_size),
          dtype=dtype,
          device="cuda",
      )
      bias = torch.rand(
          (args.hidden_size), dtype=dtype, device="cuda"
      ).expand_as(residual)
      dropout_rate = 0.1
      # Warmup JIT fusions with the input grad_enable state of both forward
      # prop and recomputation
      for input_grad, bias_grad, residual_grad in zip(
          [False, True], [True, True], [True, True]
      ):
          input.requires_grad = input_grad
          bias.requires_grad = bias_grad
          residual.requires_grad = residual_grad
          for _ in range(5):
              output = bias_dropout_add_fused_train(
                  input, bias, residual, dropout_rate
              )
      del bias, input, residual, output
      torch.cuda.empty_cache()
  
imports: import random
  import os
  import time
  import numpy
  import torch
  from datetime import timedelta
  from megatron import fused_kernels
  from megatron import get_adlr_autoresume
  from megatron import get_args
  from megatron import get_tensorboard_writer
  from megatron.core import mpu, tensor_parallel
  from megatron.arguments import parse_args, validate_args
  from megatron.checkpointing import load_args_from_checkpoint
  from megatron.global_vars import set_global_variables
  from megatron.model.transformer import bias_dropout_add_fused_train
  from megatron.model.fused_bias_gelu import bias_gelu
label_desc: Use image node_id 9 for calling a global function with example usage: _warmup_jit_function() without return types