From 745dbb23fcd15cbc685605de7d2be08000c3ef9f Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Tue, 5 Nov 2024 10:02:26 -0800 Subject: [PATCH] Allow LN overlap for cc >=9 instead of ==9 Signed-off-by: Guyue Huang --- launcher_scripts/nemo_launcher/collections/conditional_cfgs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/launcher_scripts/nemo_launcher/collections/conditional_cfgs.py b/launcher_scripts/nemo_launcher/collections/conditional_cfgs.py index 18fcd5a59f..1f64ec0da6 100644 --- 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)