Update

[ghstack-poisoned]

Author: ruisizhang123

autoparallel/auto_bucketing.py

@@ -76,16 +76,14 @@ def simple_fsdp_autobucketing_reordering_pass(
         print("Reorder scheduler nodes with autobucketing algroithm")
         node_length = len(snodes)
         snodes = reorder.reorder_all_gather(
-            snodes,
-            bucketable_nodes,
-            all_gather_before_last_wait=True
-        )
-        assert node_length == len(snodes), (
-            f"Missed nodes in reordering all gather: expected {node_length}, but got {len(snodes)}"
+            snodes, bucketable_nodes, all_gather_before_last_wait=False
         )
+        assert node_length == len(
+            snodes
+        ), f"Missed nodes in reordering all gather: expected {node_length}, but got {len(snodes)}"
         snodes = reorder.reorder_reduce_scatter(snodes, bucketable_nodes)
-        assert node_length == len(snodes), (
-            f"Missed nodes in reordering reduce scatter: expected {node_length}, but got {len(snodes)}"
-        )
+        assert node_length == len(
+            snodes
+        ), f"Missed nodes in reordering reduce scatter: expected {node_length}, but got {len(snodes)}"
 
     return snodes

autoparallel/autobucketing_util/bucket_func.py

@@ -9,7 +9,7 @@
 
 import torch
 from torch._inductor import ir, scheduler
-from torch._inductor.comms import bucket_all_gathers, bucket_reduce_scatters, get_op_idx
+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
@@ -23,6 +23,8 @@
     _replace_scheduler_buffer,
     _schedule_fallback_operation,
     _schedule_snode,
+    bucket_all_gathers,
+    bucket_reduce_scatters,
     check_ir_node_bucketable,
 )
 

autoparallel/autobucketing_util/bucket_plan.py

@@ -108,7 +108,9 @@ def get_simplefsdp_auto_plan(
         comp_time_dict,
         memory_dict,
         peak_memory_per_step_dict,
-    ) = benchmark_and_sync_runtime(sched, snodes, bucketable_nodes)
+    ) = benchmark_and_sync_runtime(
+        sched, snodes, name_to_buf, name_to_fused_node, bucketable_nodes, configs
+    )
     future_comp_time = sum(comp_time_dict.values())
     peak_memory = max(peak_memory_per_step_dict.values()) + configs.peak_memory_offset
 
@@ -139,7 +141,7 @@ def get_simplefsdp_auto_plan(
                 current_ag_bucket,
                 schedule_fallback_operation,
                 name_to_buf,
-                torch.ops._c10d_functional.all_gather_into_tensor.default,
+                "torch.ops._c10d_functional.all_gather_into_tensor.default",
                 comm_cache,
             )
 
@@ -243,7 +245,7 @@ def get_simplefsdp_auto_plan(
                         current_rs_bucket,
                         schedule_fallback_operation,
                         name_to_buf,
-                        torch.ops._c10d_functional.reduce_scatter_tensor.default,
+                        "torch.ops._c10d_functional.reduce_scatter_tensor.default",
                         comm_cache,
                         ReduceOp.AVG,
                     )
@@ -290,7 +292,7 @@ def get_simplefsdp_auto_plan(
                 current_rs_bucket,
                 schedule_fallback_operation,
                 name_to_buf,
-                torch.ops._c10d_functional.reduce_scatter_tensor.default,
+                "torch.ops._c10d_functional.reduce_scatter_tensor.default",
                 comm_cache,
                 ReduceOp.AVG,
             )
@@ -333,10 +335,10 @@ def get_simplefsdp_auto_plan(
                     ]
                     seen_new_bucketable_ag = False
 
-    if len(current_ag_bucket) > 0 or len(all_gather_plan) == 0:
+    if len(current_ag_bucket) > 0:
         all_gather_plan.append(current_ag_bucket)
 
-    if len(current_rs_bucket) > 0 or len(reduce_scatter_plan) == 0:
+    if len(current_rs_bucket) > 0:
         reduce_scatter_plan.append(current_rs_bucket)
 
     return all_gather_plan, reduce_scatter_plan

autoparallel/autobucketing_util/bucket_utils.py

@@ -163,7 +163,7 @@ def check_ir_node_bucketable(
 
 def _get_fx_node(
     snode_or_ir_node: Union["scheduler.BaseSchedulerNode", "ir.IRNode"],
-    expected_op: Callable[[Any]],
+    expected_op: Any,
 ) -> torch.fx.Node:
     origins = None
     if isinstance(snode_or_ir_node, scheduler.BaseSchedulerNode):
@@ -190,7 +190,7 @@ def _get_fx_node(
 
 def get_snode_process_group_info(
     snode: "scheduler.BaseSchedulerNode",
-    expected_op: Callable[[Any]],
+    expected_op: Any,
     resolve_pg: bool = False,
 ) -> tuple[int, Union[str, ProcessGroup]]:
     fx_node = _get_fx_node(snode, expected_op=expected_op)
@@ -248,7 +248,7 @@ def get_snode_tensor_info(
 
 def _estimate_bucketed_node_list(
     current_node_list: list["scheduler.BaseSchedulerNode"],
-    schedule_fallback_operation: Callable[[Any]],
+    schedule_fallback_operation: Callable[[Any], Any],
     group_size: int,
     group_name: str,
     name_to_buf: Dict[str, "scheduler.SchedulerBuffer"],
@@ -272,7 +272,7 @@ def _estimate_bucketed_node_list(
         )
         return estimated_comm, comm_size_inp, comm_size_out
 
-    if comm_func == torch.ops._c10d_functional.all_gather_into_tensor.default:
+    if comm_func == "torch.ops._c10d_functional.all_gather_into_tensor.default":
         bucked_node = bucket_all_gathers(
             schedule_fallback_operation,
             group_size,
@@ -284,7 +284,7 @@ def _estimate_bucketed_node_list(
         )
         comm_size_inp = bucked_node[0].layout.size
         comm_size_out = bucked_node[1].layout.size
-    elif comm_func == torch.ops._c10d_functional.reduce_scatter_tensor.default:
+    elif comm_func == "torch.ops._c10d_functional.reduce_scatter_tensor.default":
         bucked_node = bucket_reduce_scatters(
             schedule_fallback_operation,
             group_size,
@@ -311,7 +311,7 @@ def _estimate_bucketed_node_list(
 
 def estimate_bucketed_snode_runtime(
     node_bucket_dict: Dict[tuple[Any, ...], list["scheduler.BaseSchedulerNode"]],
-    schedule_fallback_operation: Callable[[Any]],
+    schedule_fallback_operation: Callable[[Any], Any],
     name_to_buf: Dict[str, "scheduler.SchedulerBuffer"],
     comm_func: Callable[[Any], Any],
     comm_cache: Dict[Any, Any],
@@ -687,7 +687,7 @@ def _get_dim0_padded_size(tensor_size: torch.Size, dim0_factor: int) -> torch.Si
 
     reduce_scatter_tensor = schedule_fallback_operation(
         torch.ops._c10d_functional.reduce_scatter_tensor.default,
-        (reduce_scatter_input, reduce_op, group_size, group_name),
+        (reduce_scatter_input, str(reduce_op), group_size, group_name),
         {},
         dep_operations=chunk_cat,
     )

autoparallel/autobucketing_util/estimation.py

@@ -7,6 +7,7 @@
 import os
 import pickle
 from collections import defaultdict
+from typing import Any
 
 import torch
 import torch.distributed as c10d
@@ -15,7 +16,6 @@
 from torch._inductor.virtualized import V
 from torch.utils._ordered_set import OrderedSet
 
-from ..auto_bucketing import simplefsdp_autobucketing_config
 from .bucket_utils import (
     check_ir_node_bucketable,
     get_snode_process_group_info,
@@ -44,7 +44,7 @@ def benchmark_and_sync_runtime(
     name_to_buf: dict[str, "scheduler.SchedulerBuffer"],
     name_to_fused_node: dict[str, "scheduler.BaseSchedulerNode"],
     bucketable_nodes: set[str],
-    configs: "simplefsdp_autobucketing_config",
+    configs: Any,
 ):
     world_size = c10d.distributed_c10d.get_world_size()
 
@@ -220,6 +220,6 @@ def benchmark_and_sync_runtime(
     median_runtimes = sync_dict_across_ranks(comm_cache.cache, world_size)
     comm_cache.cache = median_runtimes
     comm_cache._update_max_size()
-    with open(configs.simplefsdp.save_estimation_path, "wb") as file:
+    with open(configs.save_estimation_path, "wb") as file:
         pickle.dump(comm_cache.cache, file)
     return comm_cache, comp_time_dict, memory_dict, peak_memory_per_step_dict

autoparallel/autobucketing_util/estimation_utils.py

@@ -200,6 +200,7 @@ def add_comm_time(self, tensor_input_size, tensor_output_size, comm_func, value)
     def get_comm_time(
         self, tensor_input_size, tensor_output_size, comm_func, calibrated=False
     ):
+        comm_func = str(comm_func)
         key = (tuple(tensor_input_size), tuple(tensor_output_size), comm_func)
         if key in self.cache:
             return self.cache[key]
@@ -368,7 +369,7 @@ def to_real_tensor(e: Any) -> Any:
                 return out
 
             def delete_tensor_in_list(tensor_list: list[Any]) -> None:
-                for i in range(len(tensor_list)):
+                for i in range(len(tensor_list) - 1, -1, -1):
                     if isinstance(tensor_list[i], torch.Tensor):
                         tensor_list[i].cpu()
                         del tensor_list[i]

autoparallel/autobucketing_util/reorder.py

@@ -9,9 +9,9 @@
 from typing import Dict, List, Optional, Tuple
 
 import torch
-from torch.utils._ordered_set import OrderedSet
 from torch._inductor import ir, scheduler
-from torch._inductor.utils import is_collective
+from torch._inductor.utils import contains_collective, contains_wait, is_collective
+from torch.utils._ordered_set import OrderedSet
 
 from .bucket_utils import check_ir_node_bucketable
 
@@ -143,7 +143,9 @@ def get_node_type(node: "scheduler.BaseSchedulerNode", bucketable_ir_nodes) -> N
 
     if isinstance(node, scheduler.GroupedSchedulerNode):
         # [Only for bucketing]: newly created AG and RS are grouped as GroupedSchedulerNode
-        child_nodes_type = [_get_ir_node_type(n.node, bucketable_ir_nodes) for n in node.snodes]
+        child_nodes_type = [
+            _get_ir_node_type(n.node, bucketable_ir_nodes) for n in node.snodes
+        ]
         if NodeType.AG_WAIT in child_nodes_type:
             return NodeType.AG_WAIT
         elif NodeType.RS_WAIT in child_nodes_type:
@@ -187,7 +189,11 @@ def reorder_all_gather(
             all_gather_list.append(node)
             inverse_user = list(inverse_users[node])
             inverse_user = [
-                n for n in inverse_user if node_to_type[n] == NodeType.COMPUTE
+                n
+                for n in inverse_user
+                if node_to_type[n] == NodeType.COMPUTE
+                and not contains_collective(n)
+                and not contains_wait(n)
             ]
             if len(inverse_user) > 0:
                 all_gather_list.extend(inverse_user)
@@ -244,7 +250,7 @@ def reorder_reduce_scatter(
             wait_list.append(node)
             node_user = node_users[node]
             node_user = [n for n in node_user if node_to_type[n] == NodeType.COMPUTE]
-            #wait_list.extend(node_user)
+            # wait_list.extend(node_user)
         elif node_type == NodeType.REDUCE_SCATTER:
             if len(wait_list) > 0:
                 # move the i-th wait node before (i+1)-th reduce scatter node