code: def warmup_jit_fusion(
      batch_size: int,
      hidden_size: int,
      seq_length: int = 512,
      vocab_size: int = 32768,
      dtype: torch.dtype = torch.float32,
  ):
      """Compile JIT functions before the main training steps"""
  
      embed = Embedding(vocab_size, hidden_size).to(
          get_current_device()
      )
      linear_1 = Linear(
          hidden_size, hidden_size * 4, skip_bias_add=True
      ).to(get_current_device())
      linear_2 = Linear(
          hidden_size * 4, hidden_size, skip_bias_add=True
      ).to(get_current_device())
  
      x = torch.randint(
          vocab_size,
          (batch_size, seq_length),
          dtype=torch.long,
          device=get_current_device(),
      )
      x = embed(x)
      y, y_bias = linear_1(x)
      z, z_bias = linear_2(y)
      # 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]):
          for _ in range(10):
              bias = torch.rand_like(
                  y_bias, dtype=dtype, device=get_current_device()
              )
              input_ = torch.rand_like(
                  y, dtype=dtype, device=get_current_device()
              )
              bias.requires_grad, input_.requires_grad = (
                  bias_grad,
                  input_grad,
              )
              bias_gelu_impl(input_, bias)
  
      # Warmup fused bias+dropout+add
      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]
      ):
          for _ in range(10):
              input_ = torch.rand_like(
                  z, dtype=dtype, device=get_current_device()
              )
              residual = torch.rand_like(
                  x, dtype=dtype, device=get_current_device()
              )
              bias = torch.rand_like(
                  z_bias, dtype=dtype, device=get_current_device()
              )
              input_.requires_grad = input_grad
              bias.requires_grad = bias_grad
              residual.requires_grad = residual_grad
              bias_dropout_add_fused_train(
                  input_, bias, residual, dropout_rate
              )
  
      torch.cuda.empty_cache()
  
imports: import torch
  from colossalai.legacy.nn.layer.colossalai_layer import Embedding, Linear
  from colossalai.utils import get_current_device
  from .bias_dropout_add import bias_dropout_add_fused_train
  from .bias_gelu import bias_gelu_impl
label_desc: Use image node_id 2 for calling a global function with example usage: warmup_jit_fusion(batch_size, hidden_size, seq_length, vocab_size, dtype) without return types