Disable nvtx decorator to avoid graph break as in #5697

microsoft · Jun 25, 2024 · 24cf9f6 · 24cf9f6
1 parent 2ce397c
commit 24cf9f6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
@@ -6,6 +6,7 @@
 import torch
 import time
 import os
+import deepspeed.utils.nvtx
 from deepspeed import comm as dist
 from deepspeed.utils.logging import log_dist
 
@@ -646,6 +647,9 @@ def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwarg
 
         if self._is_compiled:
             return
+
+        # Avoid graph breaks
+        nvtx.enable_nvtx = False
         self.module.compile(backend=backend, **compile_kwargs)
         self._is_compiled = True