Skip to content

Commit 3e084f4

Browse files
authored
[Reland] Fix SimpleFSDP aot_eager_autobucket (#1965)
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0) (oldest at bottom): * #1960 * #1959 * __->__ #1965 **Squash and Merge button won't work for this PR. I'll merge by myself.** #1963 was accdientally merge with Squash and Merge button. This is a reland PR.
1 parent 89c631c commit 3e084f4

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

torchtitan/experiments/simple_fsdp/backend.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,14 @@ def get_compile_backend(backend_name: str) -> Union[str, callable]:
2121
# Perform auto optimization in aten fx-level and execute code in aot_eager backend
2222
# The autobucketing logic is here: https://github.com/pytorch/pytorch/pull/163960
2323
from torch._dynamo.backends.common import aot_autograd as aot_autograd_backend
24+
25+
from torch._inductor.config import aten_distributed_optimizations as dist_opts
2426
from torch._inductor.fx_passes.overlap_scheduling import (
2527
schedule_overlap_bucketing,
2628
)
2729

28-
torch._inductor.config.test_configs.aten_fx_overlap_preserving_bucketing = True
29-
torch._inductor.config.test_configs.aten_fx_overlap_insert_overlap_deps = False
30+
dist_opts.collective_bucketing = True
31+
dist_opts.insert_overlap_deps = False
3032
torch._inductor.config.allow_buffer_reuse = False
3133

3234
def aten_autobucketing_reordering_pass(

0 commit comments

Comments
 (0)