From c8fd1cfad12a9378c7454446fd60e334d14b97a2 Mon Sep 17 00:00:00 2001 From: Xiaodong Wang Date: Fri, 19 May 2023 21:50:33 +0000 Subject: [PATCH] [pt2] Turn off lazy reinit when cuda graph is on (#101848) Summary: cuda graph doesn't work with cuda 11's cupti lazy reinit. So we'll turn it off if any modules turn on cudagraph Test Plan: test with cuda graph on Reviewed By: aaronenyeshi Differential Revision: D45967197 Pull Request resolved: https://github.com/pytorch/pytorch/pull/101848 Approved by: https://github.com/aaronenyeshi --- torch/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torch/__init__.py b/torch/__init__.py index 531930296ef..dfd027e94f5 100644 --- a/torch/__init__.py +++ b/torch/__init__.py @@ -1496,6 +1496,10 @@ class _TorchCompileInductorWrapper: options or () ), "triton.cudagraphs does not support dynamic shapes. Please set dynamic=False or triton.cudagraphs=False" + # FIXME: CUPTI Lazy Re-init and CUDA Graph crashes with CUDA 11. + if self.config.get("triton.cudagraphs", False): + os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1" + def __eq__(self, other): return (isinstance(other, _TorchCompileInductorWrapper) and self.config == other.config and