From 8196de1648dc73582ed9d43dab28cc002e89428b Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Wed, 21 Jun 2023 12:19:06 -0400
Subject: [PATCH] Skip unnecessary compilation

---
 megatron/initialize.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/megatron/initialize.py b/megatron/initialize.py
index 40b2370e8eb..e387c4ee786 100644
--- a/megatron/initialize.py
+++ b/megatron/initialize.py
@@ -158,6 +158,15 @@ def _compile_dependencies():
         print('>>> done with dataset index builder. Compilation time: {:.3f} '
               'seconds'.format(time.time() - start_time), flush=True)
 
+    try:
+        # Skip the rest if the kernels are unnecessary or already available (ex. from apex)
+        if args.use_flash_attn or args.masked_softmax_fusion:
+            import scaled_upper_triang_masked_softmax_cuda
+            import scaled_masked_softmax_cuda
+        return
+    except ImportError:
+        pass
+
     # ==================
     # Load fused kernels
     # ==================