[kernel] added jit warmup (#1792)

colossalai/kernel/jit/option.py

@@ -1,5 +1,11 @@
 import torch
 
+from colossalai.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
+
 JIT_OPTIONS_SET = False
 
 
@@ -30,3 +36,44 @@ def set_jit_fusion_options():
             torch._C._jit_override_can_fuse_on_gpu(True)
 
         JIT_OPTIONS_SET = True
+
+
+def warmup_jit_fusion(batch_size: int,
+                      hidden_size: int,
+                      seq_length: int = 512,
+                      vocab_size: int = 32768,
+                      dtype: torch.dtype = torch.float32):
+    """ Compilie 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()