File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed
torchtitan/experiments/simple_fsdp Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff 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 (
You can’t perform that action at this time.
0 commit comments