Revert "[2/N] Autobucketing: add estimation utils for autobucketing (#128)"

This reverts commit 6206d9d.

Author: ruisizhang123

autoparallel/auto_bucketing.py

@@ -19,7 +19,6 @@ class simplefsdp_autobucketing_config:
     - load_cache: set to True to load cache from save_estimation_path
     - enable_bucket_ir: set to True to bucket all_gather/reduce_scatter
     - enable_reorder_ir: set to True to reorder all_gather/reduce_satter
-    - calibrate_number: number of samples to calibrate during comm estimation
     """
 
     relax_ratio = 0
@@ -28,7 +27,6 @@ class simplefsdp_autobucketing_config:
     save_estimation_path = "/mnt/mffuse/cache_ruisi/estimation_mast.pkl"
     enable_bucket_ir = True
     enable_reorder_ir = True
-    calibrate_number = 40
 
 
 def simple_fsdp_autobucketing_reordering_pass(

autoparallel/autobucketing_util/bucket_utils.py

@@ -4,23 +4,15 @@
 # LICENSE file in the root directory of this source tree.
 
 # mypy: ignore-errors
-from functools import reduce
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict
 
 import torch
-from torch._inductor import ir, scheduler
+from torch._inductor import scheduler
 from torch._inductor.dependencies import WeakDep
-from torch._inductor.ir import NoneLayout
-from torch._inductor.utils import buf_name_to_fused_snode, is_collective, is_wait
-from torch.distributed import ProcessGroup
-from torch.distributed.distributed_c10d import _resolve_process_group
+from torch._inductor.utils import buf_name_to_fused_snode, is_collective
 from torch.utils._ordered_set import OrderedSet
 
 
-def get_data_size(size):
-    return reduce(lambda x, y: x * y, size)
-
-
 def _find_recursive_deps_of_snode(
     snode: "scheduler.BaseSchedulerNode",
     collected_node_set: OrderedSet["scheduler.BaseSchedulerNode"],
@@ -125,119 +117,3 @@ def get_bucketable_ir_nodes(
                 bucketable_ir_nodes.add(snode.node.get_name())
 
     return bucketable_ir_nodes
-
-
-def check_ir_node_bucketable(
-    ir_node: "ir.IRNode", bucketable_ir_nodes: set[str]
-) -> bool:
-    """
-    Determine if the AG/RS & AG/RS wait node is from bucketable nodes or not
-    """
-    ir_node_origins = list(getattr(ir_node, "origins", None))
-    if len(ir_node_origins) == 0:
-        # bucketed AG and RS doesn't have origins
-        return True
-
-    if is_wait(ir_node):
-        ir_node = ir_node.inputs[0]
-
-    if is_collective(
-        ir_node, op=torch.ops._c10d_functional.all_gather_into_tensor.default
-    ):
-        ir_node_name = ir_node.get_name()
-    elif is_collective(
-        ir_node, op=torch.ops._c10d_functional.reduce_scatter_tensor.default
-    ):
-        ir_node_name = ir_node.get_name()
-    else:
-        return False
-
-    if ir_node_name in bucketable_ir_nodes:
-        return True
-
-    return False
-
-
-def _get_fx_node(
-    snode_or_ir_node: Union["scheduler.BaseSchedulerNode", "ir.IRNode"],
-    expected_op: Any,
-) -> torch.fx.Node:
-    origins = None
-    if isinstance(snode_or_ir_node, scheduler.BaseSchedulerNode):
-        origins = snode_or_ir_node.node.get_origins()
-    elif isinstance(snode_or_ir_node, ir.IRNode):
-        origins = snode_or_ir_node.origins
-    else:
-        raise ValueError(
-            f"Expected BaseSchedulerNode or IRNode, got {type(snode_or_ir_node)}. Offending value: {snode_or_ir_node}"
-        )
-    origins_with_expected_op = [o for o in origins if o.target == expected_op]
-    if len(origins_with_expected_op) != 1:
-        print(
-            "[Get FX exception] origins_with_expected_op",
-            origins_with_expected_op,
-            "expected_op",
-            expected_op,
-            "snode_or_ir_node",
-            snode_or_ir_node,
-        )
-        return None
-    return origins_with_expected_op[0]
-
-
-def get_snode_process_group_info(
-    snode: "scheduler.BaseSchedulerNode",
-    expected_op: Any,
-    resolve_pg: bool = False,
-) -> tuple[int, Union[str, ProcessGroup]]:
-    fx_node = _get_fx_node(snode, expected_op=expected_op)
-    # return None if the snode doesn't have a valid fx_node
-    if fx_node is None:
-        return None
-
-    if expected_op == torch.ops._c10d_functional.all_gather_into_tensor.default:
-        group_size, group_name = (
-            snode.node.constant_args[0],
-            snode.node.constant_args[1],
-        )
-    elif expected_op == torch.ops._c10d_functional.reduce_scatter_tensor.default:
-        group_size, group_name = (
-            snode.node.constant_args[1],
-            snode.node.constant_args[2],
-        )
-    elif expected_op == torch.ops._c10d_functional.all_reduce_.default:
-        group_size, group_name = fx_node.args[1], fx_node.args[2]
-    elif expected_op == torch.ops._c10d_functional.all_to_all_single.default:
-        group_size, group_name = 0, fx_node.args[3]
-    else:
-        raise ValueError(f"Unsupported op {expected_op}")
-
-    if resolve_pg:
-        group_name = _resolve_process_group(group_name)
-    return group_size, group_name
-
-
-def get_snode_tensor_info(
-    snode: "scheduler.BaseSchedulerNode", return_data_size: bool = False
-) -> tuple[Any, ...]:
-    input_dtype, input_device = (
-        snode.node.inputs[0].layout.dtype,
-        snode.node.inputs[0].layout.device,
-    )
-    input_size = get_data_size(snode.node.inputs[0].layout.size)
-
-    if not isinstance(snode.node.layout, NoneLayout):
-        output_dtype, output_device = (
-            snode.node.layout.dtype,
-            snode.node.layout.device,
-        )
-        output_size = get_data_size(snode.node.layout.size)
-    else:
-        # In all_reduce, layout is NoneLayout
-        # We set output info to be the same as input info as a special treatment
-        output_dtype, output_device, output_size = input_dtype, input_device, input_size
-
-    result = (input_dtype, input_device, output_dtype, output_device)
-    if return_data_size:
-        result += (input_size, output_size)
-    return result