from __future__ import annotations import heapq from collections import Counter, defaultdict from typing import Any, Optional, TYPE_CHECKING import torch import torch.fx as fx from torch._dynamo.graph_deduplication import _stable_topological_sort from torch._inductor.fx_passes.bucketing import ( _schedulable_wait_node, is_all_gather_into_tensor as is_all_gather, is_fsdp_all_gather, is_fsdp_reduce_scatter, is_reduce_scatter_tensor as is_reduce_scatter, merge_all_gather_bucket, merge_reduce_scatter_bucket, ) from torch._inductor.fx_passes.overlap_preserving_bucketer import ( bucket_key, OverlapPreservingBucketer, ) from torch._inductor.fx_passes.overlap_scheduling import ( CollectiveInfo, is_compute_node, OverlapScheduler, ) from torch.utils._ordered_set import OrderedSet from .graph_view import get_subgraph_by_path, GraphView, make_graph_view if TYPE_CHECKING: from collections.abc import Callable import logging logger = logging.getLogger(__name__) class ManualOverlapPreservingBucketer(OverlapPreservingBucketer): """ Buckets collective operations based on user specifications. The actual bucket happens in bucket_collectives, where all-gathers/reduce-scatters in `nodes` will be buckted one single all-gather/reduce-scatter. """ def __init__( self, *args: Any, **kwargs: Any, ): super().__init__(*args, **kwargs) self.node_to_wait_map: dict[fx.Node, fx.Node] = defaultdict() def _bucket_group(self, coll_nodes: list[fx.Node]) -> None: assert len(coll_nodes) > 0, "bucketed coll_nodes should have nonzero node" waits = [self.collective_info[n].wait_node for n in coll_nodes] # Use earliest wait insertion point first_wait = min(waits, key=lambda w: self.node_idx[w]) # Find insertion location first = coll_nodes[0] next_node = first while next_node in coll_nodes: next_node = next_node.next if is_all_gather(first): new_nodes, replacements = merge_all_gather_bucket( self.graph, coll_nodes, wait_insertion_point=first_wait, insert_before=next_node, mode="custom_ops", ) elif is_reduce_scatter(first): new_nodes, replacements = merge_reduce_scatter_bucket( self.graph, coll_nodes, wait_insertion_point=first_wait, insert_before=next_node, mode="custom_ops", ) else: raise ValueError( "bucket non all_gather/reduce_scatter node is not supported" ) logger.debug(f"bucketing nodes: {coll_nodes} into {new_nodes}") # noqa: G004 # Identify the new wait and start new_waits = [n for n in new_nodes if _schedulable_wait_node(n)] assert len(new_waits) == 1, f"Expected exactly one new wait, got {new_waits}" new_wait = new_waits[0] new_start = new_wait.args[0] assert isinstance(new_start, fx.Node) # Set manual bucketing-specific metadata # Note: Generic metadata (nn_module_stack, fwd_nn_module_stack, custom, stack_trace) # is now preserved automatically by the bucketing functions in bucketing.py node_type = ( "bucketed_all_gather" if is_all_gather(first) else "bucketed_reduce_scatter" ) for n in new_nodes: if n == new_wait: node_type = node_type + "_wait" n.meta["manual_bucket_node_type"] = node_type if "wait" in node_type: self.node_to_wait_map[n] = new_wait def manual_bucket_collectives(self, nodes: list[fx.Node]) -> None: """ Bucket all all-gather/reduce-scatter nodes from nodes into one all-gather/reduce-scatter. """ # Filter out valid collectives collectives = [n for n in nodes if n in self.collective_info] if collectives == []: return grouped_collectives: dict[object, OrderedSet[fx.Node]] = defaultdict(OrderedSet) for node in collectives: if not ( is_fsdp_all_gather(node, self.node_ancestors) or is_fsdp_reduce_scatter(node) ): continue key = bucket_key(node) if key is not None: grouped_collectives[key].add(node) for key, nodes in grouped_collectives.items(): # type: ignore[arg-type] self._bucket_group(list(nodes)) class ManualOverlapScheduler(OverlapScheduler): """ Scheduler that manual buckets and reorders collective nodes based on module_bucket_plans """ def __init__( self, gm: fx.GraphModule, module_bucket_plans: list[list[str] | str], insert_overlap_deps: bool, module_stack_fn: Callable[[fx.Node], list[tuple[str, type[Any]]]] | None = None, ): super().__init__( gm, max_in_flight_gb=0.0, max_compute_pre_fetch=0, collective_bucketing=True, insert_overlap_deps=insert_overlap_deps, compute_overlap_multipler=0.0, max_coll_distance=0, custom_runtime_estimation=None, collective_estimator="analytical", max_memory_increase_gb=None, max_memory_increase_ratio=None, ) self.module_bucket_plans = module_bucket_plans self.nodes_in_subgraph: list[list[fx.Node]] = [] self.bucketer = ManualOverlapPreservingBucketer( graph=self.graph, collective_info=self.collective_info, scheduled=OrderedSet(self.graph.nodes), ) self.insert_overlap_deps = insert_overlap_deps self.module_stack_fn = module_stack_fn def _identify_collectives(self) -> None: """Identify all collective operations.""" for node in self.nodes: if _schedulable_wait_node(node): start = node.args[0] info = CollectiveInfo( start_node=start, wait_node=node, size_bytes=0, estimated_time_ms=0, exposed_time_ms=0, ) self.collective_info[start] = info self.wait_to_start[node] = start self.unscheduled_collectives.add(start) def _add_to_ready_queue(self, node: fx.Node) -> None: """Manual scheduling uses single queue ordered by original node index.""" heapq.heappush(self.on_path_ready, (self.node_idx[node], node)) def run(self) -> torch.fx.GraphModule: """Entry point to run the manual bucket algorithm""" # Bucket collectives in each bucket_module self._manual_bucket_collectives() # Reorder collectives with last/next bucket_module self._manual_reorder_graph() return self.gm def _manual_reorder_graph(self) -> None: """ Reorder nodes in the FX graph to enforce manual overlap dependencies. forward graph (all-gathers only): modules are processed in order: module 0, 1, 2, ... before reordering: ag_start_0 -> ag_wait_0 -> compute_0 -> ag_start_1 -> ag_wait_1 -> compute_1 -> ... Reordering prefetches module i+1's parameters while computing module i It adds dependencies: ag_wait_i should depend on ag_start_(i+1) This enforces ag_start_(i+1) to happen before ag_wait_i so it overlaps with module i's compute after reordering: ag_start_0 -> ag_start_1 -> ag_wait_0 -> compute_0 -> ag_wait_1 -> compute_1 -> ... backward graph (all-gathers and reduce-scatters): modules are processed in reverse order: module N, N-1, N-2, ... before reordering: ag_start_N -> ag_wait_N -> compute_N -> rs_start_N -> rs_wait_N -> ... For all-gathers, prefetch module i-1's parameters while computing module i Adds dependencies: ag_wait_i should depend on ag_start_(i-1) So ag_start_(i-1) overlaps with module i's compute For reduce-scatters, defer rs_wait_i to happen after rs_start_(i-1) Adds dependencies: rs_wait_i should depend on rs_start_(i-1) So rs_start_i overlaps with module i-1's compute """ delayed_rs_wait_nodes: list[fx.Node] = [] current_rs_start_nodes: list[fx.Node] = [] overlap_deps: dict[fx.Node, OrderedSet[fx.Node]] = defaultdict(OrderedSet) # Re-initialize after graph modification in _manual_bucket_collectives self.node_idx = {n: i for i, n in enumerate(self.nodes)} self.on_path_ready = [] self.scheduled = OrderedSet() for node in self.nodes: if self.in_degree[node] == 0: self._add_to_ready_queue(node) # schedule reduce scatter normally in self._schedule while self.on_path_ready: _, node = heapq.heappop(self.on_path_ready) node_type = node.meta.get("manual_bucket_node_type", "") if node in self.scheduled: continue if node_type == "bucketed_reduce_scatter": # Collect reduce scatter start nodes (pre_bucket_rs and rs) current_rs_start_nodes.append(node) elif node_type == "bucketed_reduce_scatter_wait": # When we see a wait node from a new RS, flush delayed waits # with dependencies on previously collected RS start nodes if current_rs_start_nodes: for delayed in delayed_rs_wait_nodes: for rs_start in current_rs_start_nodes: overlap_deps[delayed].add(rs_start) delayed_rs_wait_nodes.clear() current_rs_start_nodes.clear() delayed_rs_wait_nodes.append(node) self._schedule(node) self.scheduled = OrderedSet(reversed(list(self.scheduled))) picked_ag: list[fx.Node] = [] last_compute: Optional[fx.Node] = None for node in self.scheduled: node_type = node.meta.get("manual_bucket_node_type", "") if node_type == "bucketed_all_gather": picked_ag.append(node) continue if node_type == "bucketed_all_gather_wait": # Connect corresponding all_gather_wait -> all_gather edges if picked_ag: for ag in picked_ag: overlap_deps[self.bucketer.node_to_wait_map[node]].add(ag) picked_ag.clear() if is_compute_node(node): last_compute = node if last_compute is not None and not bool( OrderedSet(picked_ag) & OrderedSet(self.node_ancestors[last_compute]) ): for ag in picked_ag: overlap_deps[last_compute].add(ag) _stable_topological_sort(self.graph, overlap_deps) self.graph.lint() if self.insert_overlap_deps: from torch._inductor.fx_passes.control_dependencies import ( preserve_node_ordering, ) preserve_node_ordering(self.graph, overlap_deps) def _manual_bucket_collectives(self) -> None: """Bucket nodes in each module_bucket from module_bucket_plans.""" self._obtain_nodes_in_subgraph() for i, nodes in enumerate(self.nodes_in_subgraph): self.bucketer.manual_bucket_collectives(nodes=nodes) _stable_topological_sort(self.graph, {}) self.graph.lint() self.nodes = list(self.graph.nodes) self.in_degree = Counter(user for node in self.nodes for user in node.users) def _schedule(self, node: fx.Node) -> None: """Schedule a node.""" assert node not in self.scheduled assert all(n in self.scheduled for n in node.all_input_nodes) self.scheduled.add(node) for user in node.users: self.in_degree[user] -= 1 if self.in_degree[user] == 0: self._add_to_ready_queue(user) def _obtain_nodes_in_subgraph(self) -> None: """ Obtain nodes in each subgraph from module_bucket_plans """ graph_view: GraphView | None = make_graph_view(self.graph, self.module_stack_fn) if graph_view is None: return for module in self.module_bucket_plans: subgraph_view = get_subgraph_by_path(graph_view, module) self.nodes_in_subgraph.append(subgraph_view) all_subgraph_nodes = [ node for sublist in self.nodes_in_subgraph for node in sublist ] unique_subgraph_nodes = list(OrderedSet(all_subgraph_nodes)) assert len(all_subgraph_nodes) <= len(unique_subgraph_nodes), ( f"Overlapping FX nodes detected across subgraphs in `module_bucket_plans`. " f"Expected disjoint node sets but found " f"{len(all_subgraph_nodes) - len(unique_subgraph_nodes)} duplicated node(s)." ) def manual_overlap_bucketing( gm: torch.fx.GraphModule, module_bucket_plans: list[list[str] | str], insert_overlap_deps: bool = False, module_stack_fn: Callable[[fx.Node], list[tuple[str, type[Any]]]] | None = None, ) -> torch.fx.GraphModule: """Schedule nodes based on user specifications in module_bucket_plans The manual overlapping consists of two steps: Step 1: bucket all-gather/reduce-scatter in each module in module_bucket_plans Step 2: reorder all-gather to overlap with last module_bucket & reorder reduce-scatter to overlap with next module_bucket TODO(ruisizhang123): allow users to explicitly specify which module_bucket they want to overlap. Args: gm: input graph module to optimize. module_bucket_plans: user specified FQNs module_stack_fn: Optional callable for extracting module hierarchy from nodes. Used to construct a GraphView for identifying nodes in module_bucket_plans. The module_class component of the returned tuples is not used by this pass. See the `module_stack_fn` parameter in `make_graph_view` (graph_view.py) for detailed documentation on signature, return format, and usage examples. """ # decode abbreviated FQNs to actual FQNs overlapped_gm = ManualOverlapScheduler( gm, module_bucket_plans, insert_overlap_deps, module_stack_fn ).run() overlapped_gm.recompile() return overlapped_gm