Update (base update)

[ghstack-poisoned]

Author: ruisizhang123

autoparallel/autobucketing_util/bucket_func.py

@@ -9,8 +9,8 @@
 
 import torch
 from torch._inductor import ir, scheduler
-from torch._inductor.dependencies import StarDep, WeakDep
 from torch._inductor.comms import get_op_idx
+from torch._inductor.dependencies import StarDep, WeakDep
 from torch._inductor.utils import is_collective, is_wait
 from torch._inductor.virtualized import V
 from torch.utils._ordered_set import OrderedSet

autoparallel/autobucketing_util/bucket_plan.py

@@ -187,11 +187,13 @@ def get_simplefsdp_auto_plan(
             )
 
             # (3) Communication size criteria
-            break_comm_size_criteria = (
-                comm_cache.ag_max_inp_size <= comm_size_inp
-                or comm_cache.rs_max_out_size
-                <= heuristic_info["this_step_rs_comm_out_size"]
-            )
+            break_comm_size_criteria = comm_cache.ag_max_inp_size < comm_size_inp
+            if comm_cache.rs_max_out_size > 0:
+                break_comm_size_criteria = (
+                    break_comm_size_criteria
+                    or comm_cache.rs_max_out_size
+                    < heuristic_info["this_step_rs_comm_out_size"]
+                )
 
             if (
                 break_overlap_criteria