Update on "[inductor] Conditionally copy args to cpu to minimize memory overhead of autotuning"

masnesral · masnesral · commit e987d75fbb12 · 2024-10-04T15:15:44.000-07:00
[ghstack-poisoned]
diff --git a/test/inductor/test_cuda_repro.py b/test/inductor/test_cuda_repro.py
@@ -422,8 +422,7 @@ def decorator(fn):
                     configs=configs,
                     save_cache_hook=False,
                     mutated_arg_names=["in_out_ptr0"],
-                    is_inference=True,
-                    is_backward=False,
+                    optimize_mem=True,
                     heuristic_type=HeuristicType.POINTWISE,
                 )
 
diff --git a/test/inductor/test_triton_heuristics.py b/test/inductor/test_triton_heuristics.py
@@ -126,8 +126,7 @@ def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
             "configs": configs,
             "save_cache_hook": False,
             "mutated_arg_names": [],
-            "is_inference": True,
-            "is_backward": False,
+            "optimize_mem": True,
             "heuristic_type": HeuristicType.POINTWISE,
             "inductor_meta": inductor_meta,
         }
diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py
@@ -2750,12 +2750,16 @@ def codegen_kernel(self, name=None):
             "constants": {},
         }
 
+        # Skip memory optimization for forward of the training loop where we expect
+        # every new node will increase the peak memory and our greedy approach would
+        # introduce a lot of unnecessary cpu copies.
+        optimize_mem = V.graph.is_inference or V.graph.is_backward
+
         inductor_meta = {
             "autotune_hints": set(self.autotune_hints),
             "kernel_name": str(Placeholder.DESCRIPTIVE_NAME),
             "mutated_arg_names": mutated_args,
-            "is_inference": V.graph.is_inference,
-            "is_backward": V.graph.is_backward,
+            "optimize_mem": optimize_mem,
             "no_x_dim": self.no_x_dim,
             "num_load": self.num_load,
             "num_reduction": self.num_reduction,
diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py
@@ -187,8 +187,7 @@ def __init__(
         configs,
         save_cache_hook,
         mutated_arg_names: List[str],  # see [Note: clone mutated buffers]
-        is_inference,
-        is_backward,
+        optimize_mem,
         heuristic_type,
         size_hints=None,
         inductor_meta=None,  # metadata not relevant to triton
@@ -212,8 +211,7 @@ def __init__(
         self.inductor_meta = {} if inductor_meta is None else inductor_meta
         self.save_cache_hook = save_cache_hook
         self.mutated_arg_names = mutated_arg_names
-        self.is_inference = is_inference
-        self.is_backward = is_backward
+        self.optimize_mem = optimize_mem
         self.configs = configs
         self.heuristic_type = heuristic_type
         self.custom_kernel = custom_kernel
@@ -718,12 +716,8 @@ def copy_args_to_cpu_if_needed(self, *args, **kwargs):
         If those clones would increase the peak memory usage, however, we instead
         copy to cpu and restore them after each iteratrion. Figure out the args
         to be copied and do the copying.
-
-        Skip this optimization for forward of the training loop where we expect
-        every new node will increase the peak memory and our greedy approach
-        would introduce a lot of unnecessary cpu copies.
         """
-        if not self.is_inference and not self.is_backward:
+        if not self.optimize_mem:
             return {}
 
         copies = {}
@@ -1132,8 +1126,7 @@ def cached_autotune(
             log.debug("autotune caching is disabled by config.force_disable_caches")
 
     mutated_arg_names = inductor_meta.pop("mutated_arg_names", ())
-    is_inference = inductor_meta.pop("is_inference", False)
-    is_backward = inductor_meta.pop("is_backward", False)
+    optimize_mem = inductor_meta.pop("optimize_mem", True)
 
     def decorator(fn):
         # Remove XBLOCK from config if it's not a function argument.
@@ -1160,8 +1153,7 @@ def decorator(fn):
                 configs=configs,
                 save_cache_hook=autotune_cache and autotune_cache.save,
                 mutated_arg_names=mutated_arg_names,
-                is_inference=is_inference,
-                is_backward=is_backward,
+                optimize_mem=optimize_mem,
                 heuristic_type=heuristic_type,
                 size_hints=size_hints,
                 custom_kernel=custom_kernel,
@@ -1174,8 +1166,7 @@ def decorator(fn):
             configs=configs,
             save_cache_hook=autotune_cache and autotune_cache.save,
             mutated_arg_names=mutated_arg_names,
-            is_inference=is_inference,
-            is_backward=is_backward,
+            optimize_mem=optimize_mem,
             heuristic_type=heuristic_type,
             size_hints=size_hints,
             custom_kernel=custom_kernel,

Original file line number	Diff line number	Diff line change
`@@ -422,8 +422,7 @@ def decorator(fn):`
`422`	`422`	`configs=configs,`
`423`	`423`	`save_cache_hook=False,`
`424`	`424`	`mutated_arg_names=["in_out_ptr0"],`
`425`		`- is_inference=True,`
`426`		`- is_backward=False,`
	`425`	`+ optimize_mem=True,`
`427`	`426`	`heuristic_type=HeuristicType.POINTWISE,`
`428`	`427`	`)`
`429`	`428`