Skip to content

Commit 4ebce58

Browse files
author
Han Wang
committed
fix: add missing inductor options for torch.compile
Add max_autotune, epilogue_fusion, triton.cudagraphs, max_fusion_size options to match the reference implementation.
1 parent a183f95 commit 4ebce58

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

deepmd/pt_expt/train/training.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,12 @@ def fn(
222222
compile_opts.pop("backend", None)
223223
if "options" not in compile_opts:
224224
compile_opts["options"] = {}
225-
compile_opts["options"].setdefault("shape_padding", True)
225+
opts = compile_opts["options"]
226+
opts.setdefault("max_autotune", False)
227+
opts.setdefault("epilogue_fusion", False)
228+
opts.setdefault("triton.cudagraphs", False)
229+
opts.setdefault("shape_padding", True)
230+
opts.setdefault("max_fusion_size", 8)
226231

227232
compiled_lower = torch.compile(
228233
traced_lower,

0 commit comments

Comments
 (0)