From 3d7ffac68a79144b30f396a8d123f38fd20ceddb Mon Sep 17 00:00:00 2001 From: Kevin Stephano Date: Thu, 3 Oct 2024 07:20:08 -0700 Subject: [PATCH] Cleaning up nvFuser's repro script api. (#1231) --- thunder/examine/__init__.py | 6 +++++- thunder/tests/test_nvfuser.py | 10 +++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/thunder/examine/__init__.py b/thunder/examine/__init__.py index f2f960a16f..ad284e7526 100644 --- a/thunder/examine/__init__.py +++ b/thunder/examine/__init__.py @@ -263,7 +263,11 @@ def get_nvfuser_repro(trace: TraceCtx, fusion_name: str, /) -> str: ) fd = fusion.last_used - get_repro = getattr(fd, "getReproString", None) + # The API for nvFuser version >=2.14 + get_repro = getattr(fd, "repro_script_for", None) + # The legacy nvFuser API + if get_repro is None: + get_repro = getattr(fd, "getReproString", None) if get_repro is None: raise RuntimeError("The installed version of nvFuser does not support repro generation unless on crash.") diff --git a/thunder/tests/test_nvfuser.py b/thunder/tests/test_nvfuser.py index 5963e66209..2224bfc9ba 100644 --- a/thunder/tests/test_nvfuser.py +++ b/thunder/tests/test_nvfuser.py @@ -1050,8 +1050,16 @@ def sdpa_fn(q, k, v, dropout_p, is_causal, scale): # Check nv_sdpfa_fwd is not in bwd_fusion -> that would indicate rematerialization assert "nv_sdpfa_bwd" in bwd_fusion[-1][-1].name and "nv_sdpfa_fwd" not in bwd_fusion[-1][-1].name + + nvf_fd = bwd_fusion[-1][-1].last_used + repro_script = None + # Legacy repro script API + if nvfuser_version() < LooseVersion("0.2.14"): + repro_script = nvf_fd.getReproString() + else: + repro_script = nvf_fd.repro_script_for() assert ( - bwd_fusion[-1][-1].last_used.getReproString().count("is_cpu=True") == 2 + repro_script.count("is_cpu=True") == 2 ), "Expected philox_seed and philox_offset inputs to be CPU scalar tensors." # Torch reference computation