Disable nvFuser setup with PyTorch 23.11 and later (#9837)

* Disable nvFuser setup with PyTorch 23.11 and later Signed-off-by: Abhishree <[email protected]> * Apply isort and black reformatting Signed-off-by: athitten <[email protected]> --------- Signed-off-by: Abhishree <[email protected]> Signed-off-by: athitten <[email protected]> Co-authored-by: athitten <[email protected]>
NVIDIA · Jul 24, 2024 · 94c7dc5 · 94c7dc5
1 parent cd0d2c2
commit 94c7dc5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -399,7 +399,9 @@ def _enable_nvidia_optimizations(self):
                 self.cfg.persist_layer_norm = False
 
             # NVFUSER available starting with 21.11
-            if NVIDIA_TORCH_MAJOR >= 21 or (NVIDIA_TORCH_MAJOR == 21 and NVIDIA_TORCH_MINOR >= 11):
+            if (NVIDIA_TORCH_MAJOR >= 21 or (NVIDIA_TORCH_MAJOR == 21 and NVIDIA_TORCH_MINOR >= 11)) and (
+                NVIDIA_TORCH_MAJOR < 23 or (NVIDIA_TORCH_MAJOR == 23 and NVIDIA_TORCH_MINOR < 11)
+            ):
 
                 # NVFUSER
                 torch._C._jit_set_profiling_executor(True)