remat and recomp are exclusive

Lightning-AI · Dec 18, 2024 · 1159797 · 1159797
1 parent af68c02
commit 1159797
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/thunder/tests/test_nvfuser_remat.py b/thunder/tests/test_nvfuser_remat.py
@@ -308,7 +308,7 @@ def func(t0):
         return t3, t4
 
     t0 = make_tensor(2, 2, dtype=torch.float32, device=device)
-    compiled_func = executor.make_callable(func)
+    compiled_func = executor.make_callable(func, enable_saved_for_backward_recomputation=False)
     _ = compiled_func(t0)
     traces = thunder.last_traces(compiled_func)
     trace = traces[0]
@@ -369,7 +369,10 @@ def func(
 
     try:
         compiled_func = thunder.jit(
-            func, executors=(torch_compile_cat_ex, thunder.nvfuser_executor), use_rematerialization=True
+            func,
+            executors=(torch_compile_cat_ex, thunder.nvfuser_executor),
+            use_rematerialization=True,
+            enable_saved_for_backward_recomputation=False,
         )
         _ = compiled_func(
             t0,