From 24cf9f6a85d018099f88da34c16653606e975616 Mon Sep 17 00:00:00 2001 From: Omar Elayan Date: Tue, 25 Jun 2024 11:34:17 +0300 Subject: [PATCH] Disable nvtx decorator to avoid graph break as in #5697 --- deepspeed/inference/engine.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index 2f8e545f3390..cc93ec147c52 100755 --- 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