Allow LN overlap for cc >=9 instead of ==9 (#443)

Signed-off-by: Guyue Huang <[email protected]> Co-authored-by: Guyue Huang <[email protected]> Co-authored-by: Dong Hyuk Chang <[email protected]>
NVIDIA · Dec 4, 2024 · ef3a798 · ef3a798
1 parent 0a4c22a
commit ef3a798
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/launcher_scripts/nemo_launcher/collections/conditional_cfgs.py b/launcher_scripts/nemo_launcher/collections/conditional_cfgs.py
@@ -30,7 +30,7 @@ def get_ln_sm_margin(cfg):
     Set SM margin to LayerNorm layer at H100. This is to overlap LN kernel with communication kernels.
     """
     global cuda_capability
-    if cuda_capability == 9:
+    if cuda_capability >= 9:
         print(8)
     else:
         print(0)