From 05c2550b718d4592ce01b7ed83bcdfe2289faa2c Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Fri, 21 Aug 2026 16:31:39 -0700 Subject: [PATCH] [graph_trainer] Add reusable dI/dW backward-node classification Generalize the GraphPP dI/dW dependency analysis into a standalone helper the EP-overlap scheduler can consume: partition_backward_nodes classifies every node in the joint graph's gradient closures into di_nodes / dw_only_nodes / shared_nodes by pure ancestor algebra (closure of the input-gradient outputs vs closure of the parameter-gradient outputs), never inspecting op names or dtypes, and exposes movable_nodes, the subset of dw-only work a pass may DEFER (move later, never earlier). Safety rules keep deferral provable: - pinned per node: collectives, mutations, CPU synchronizations (_local_scalar_dense or device-to-host reads), side-effectful ops, and producers with unclassified LIVE users; a pure dead subtree (e.g. the unused getitem outputs of a multi-output quantize op) does not pin its producer. - mutation hazards: a reader of a mutated buffer's alias family that sits BEFORE the write without a data edge through it is pinned (deferring it could push the read past the write); readers downstream of the write or after it keep their order against the fixed write. A write whose target cannot be resolved to graph values empties movable_nodes entirely. split_di_dw.py stays untouched: its extraction-based dI membership deliberately includes dW-only placeholders as pass-through live-ins, and rebasing it on the closure helper would change which values it saves. Validated on a real tiny-MXFP8-GroupedExperts joint trace (EP mesh wired, world size 1): all wgrad GEMMs and their quantize/rearrange chains classify movable, dgrad chains stay on the dI side, and collectives are never movable. 13 CPU unit tests. Co-Authored-By: Claude Fable 5 --- .../graph_trainer/backward_partition.py | 264 ++++++++++++++++ .../tests/test_backward_partition.py | 293 ++++++++++++++++++ 2 files changed, 557 insertions(+) create mode 100644 torchtitan/experiments/graph_trainer/backward_partition.py create mode 100644 torchtitan/experiments/graph_trainer/tests/test_backward_partition.py diff --git a/torchtitan/experiments/graph_trainer/backward_partition.py b/torchtitan/experiments/graph_trainer/backward_partition.py new file mode 100644 index 0000000000..ee1391fda7 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/backward_partition.py @@ -0,0 +1,264 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Dependency-based dI/dW classification of backward graph nodes. + +Contract: + Given a backward (or joint) FX graph and the flat output value nodes for + input gradients and parameter gradients, classify every node in either + ancestor closure: + + di_ancestors = ancestors(input_grad_outputs) + dw_ancestors = ancestors(param_grad_outputs) + di_nodes = di_ancestors - dw_ancestors + dw_only_nodes = dw_ancestors - di_ancestors + shared_nodes = di_ancestors & dw_ancestors + + Classification is pure dataflow: it never inspects op names or dtypes, so a + BF16 backward and a quantized backward with the same dependency structure + classify identically. Only ``dw_only_nodes`` may be deferred past the dI + outputs; ``shared_nodes`` stay on the dgrad path. ``movable_nodes`` is the + subset of ``dw_only_nodes`` a scheduler may DEFER — move later, never + earlier — which is what the safety rules assume. Symbolic-shape nodes + classify like any other node; a scheduler that emits deferred nodes in + topological order keeps sym producers ahead of their consumers. +""" + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch.fx as fx + +from torchtitan.experiments.graph_trainer.ep_pass_utils import is_c10d_functional_node +from torchtitan.experiments.graph_trainer.graph_pp.utils import ( + base_tensor_for_mutation_target, + is_mutation_node, + node_closure, +) + +_COLLECTIVE_NAMESPACES = ("_c10d_functional_autograd", "_dtensor") + + +@dataclass(frozen=True, slots=True) +class BackwardNodePartition: + """dI/dW classification of backward nodes plus the deferrable subset. + + ``di_nodes``, ``dw_only_nodes``, and ``shared_nodes`` are disjoint and + together cover the union of both ancestor closures. ``movable_nodes`` is + the subset of ``dw_only_nodes`` that a deferral pass may move past the dI + outputs; it is empty when the graph cannot be proven safe to reorder. + """ + + di_nodes: set[fx.Node] + dw_only_nodes: set[fx.Node] + shared_nodes: set[fx.Node] + movable_nodes: set[fx.Node] + + +def _is_collective_node(node: fx.Node) -> bool: + if is_c10d_functional_node(node): + return True + return ( + node.op == "call_function" + and isinstance(node.target, torch._ops.OpOverload) + and node.target.namespace in _COLLECTIVE_NAMESPACES + ) + + +def _device_type(value: object) -> str | None: + device = getattr(value, "device", None) + return None if device is None else device.type + + +def _is_cpu_sync_node(node: fx.Node) -> bool: + """Return whether ``node`` reads device values back to the host.""" + if node.op != "call_function": + return False + if node.target is torch.ops.aten._local_scalar_dense.default: + return True + device = node.kwargs.get("device") + out_type = ( + torch.device(device).type + if device is not None + else _device_type(node.meta.get("val")) + ) + if out_type != "cpu": + return False + return any( + _device_type(inp.meta.get("val")) not in (None, "cpu") + for inp in node.all_input_nodes + ) + + +def _is_inert(node: fx.Node, memo: dict[fx.Node, bool]) -> bool: + """Return whether ``node`` is a pure value whose uses are all inert. + + An inert user cannot pin its producer: def-before-use survives any + topological reorder, and no effect or classified consumer depends on the + node's position. The common case is a dead ``getitem`` of a multi-output + op whose other outputs are the ones actually consumed. + """ + cached = memo.get(node) + if cached is not None: + return cached + memo[node] = False + result = ( + node.op == "call_function" + and not _is_collective_node(node) + and not is_mutation_node(node) + and not _is_cpu_sync_node(node) + and not node.is_impure() + and all(_is_inert(user, memo) for user in node.users) + ) + memo[node] = result + return result + + +def _written_value_bases(node: fx.Node) -> list[fx.Node] | None: + """Return view bases of the graph values ``node`` writes to. + + Returns ``None`` when a written argument cannot be resolved to graph + nodes; callers must treat that as unsafe. + """ + schema = getattr(node.target, "_schema", None) + if schema is None: + return None + bases: list[fx.Node] = [] + for index, arg_spec in enumerate(schema.arguments): + if arg_spec.alias_info is None or not arg_spec.alias_info.is_write: + continue + if arg_spec.name in node.kwargs: + value = node.kwargs[arg_spec.name] + elif index < len(node.args): + value = node.args[index] + else: + # Defaulted mutable argument: no graph value is written. + continue + values = value if isinstance(value, (list, tuple)) else (value,) + for item in values: + if item is None: + continue + base = base_tensor_for_mutation_target(item) + if base is None: + return None + bases.append(base) + return bases + + +def _value_descendants(node: fx.Node) -> set[fx.Node]: + """Return all nodes deriving a value from ``node`` (transitive users).""" + seen: set[fx.Node] = set() + stack = [node] + while stack: + for user in stack.pop().users: + if user not in seen: + seen.add(user) + stack.append(user) + return seen + + +def _mutation_deferral_pins(graph: fx.Graph) -> tuple[bool, set[fx.Node]]: + """Return ``(unresolvable, pinned_readers)`` for graph mutations. + + Deferral only moves nodes LATER, and mutations are never movable, so a + write is a hazard only for a node that reads the written buffer's alias + family WITHOUT a data edge through the write and sits BEFORE the write: + deferring that reader could push its read past the write. Readers + downstream of the write (data edge keeps them after it) or originally + after the fixed write keep their relative order. A write whose target + cannot be resolved to graph nodes makes every deferral unprovable. + """ + order = {node: index for index, node in enumerate(graph.nodes)} + pinned: set[fx.Node] = set() + for node in graph.nodes: + if not is_mutation_node(node): + continue + bases = _written_value_bases(node) + if bases is None: + return True, set() + post_write = _value_descendants(node) + for base in bases: + for reader in _value_descendants(base): + if ( + reader is not node + and reader not in post_write + and order[reader] < order[node] + ): + pinned.add(reader) + return False, pinned + + +def partition_backward_nodes( + graph_or_gm: fx.Graph | fx.GraphModule, + *, + input_grad_outputs: Sequence[fx.Node], + param_grad_outputs: Sequence[fx.Node], +) -> BackwardNodePartition: + """Classify backward nodes into dI-only, dW-only, and shared sets. + + Slicing the flat backward outputs into input-gradient and + parameter-gradient value nodes is the caller's job; ``None`` grad slots + are ignored. + + ``movable_nodes`` keeps only ``dw_only_nodes`` that are provably safe to + defer. A dw-only node is pinned (classified but not movable) when it is: + + * not a ``call_function`` node (placeholders and get_attr constants are + live-ins, not deferrable computation); + * a collective (``_c10d_functional``/``_dtensor`` namespaces); + * a mutation (writes through an aliased argument); + * a CPU synchronization (``aten._local_scalar_dense`` or a device-to-host + read); + * side-effectful per ``fx.Node.is_impure``; + * used by an unclassified LIVE node other than the graph output, since + deferring it would move a definition past that untracked use. An inert + user (a pure dead subtree, e.g. the unused ``getitem`` outputs of a + multi-output quantize op) does not pin its producer; + * a pre-write reader of a mutated buffer's alias family (see + ``_mutation_deferral_pins``) — deferring it could push the read past + the write. + + ``movable_nodes`` is empty when a mutation's written target cannot be + resolved to graph nodes. + """ + graph = ( + graph_or_gm.graph if isinstance(graph_or_gm, fx.GraphModule) else graph_or_gm + ) + di_ancestors = node_closure(input_grad_outputs) + dw_ancestors = node_closure(param_grad_outputs) + dw_only = dw_ancestors - di_ancestors + inert_memo: dict[fx.Node, bool] = {} + + def is_movable(node: fx.Node) -> bool: + if node.op != "call_function": + return False + if _is_collective_node(node): + return False + if is_mutation_node(node): + return False + if _is_cpu_sync_node(node): + return False + if node.is_impure(): + return False + return all( + user in dw_ancestors or user.op == "output" or _is_inert(user, inert_memo) + for user in node.users + ) + + unresolvable, hazard_pins = _mutation_deferral_pins(graph) + if unresolvable: + movable: set[fx.Node] = set() + else: + movable = { + node for node in dw_only if node not in hazard_pins and is_movable(node) + } + + return BackwardNodePartition( + di_nodes=di_ancestors - dw_ancestors, + dw_only_nodes=dw_only, + shared_nodes=di_ancestors & dw_ancestors, + movable_nodes=movable, + ) diff --git a/torchtitan/experiments/graph_trainer/tests/test_backward_partition.py b/torchtitan/experiments/graph_trainer/tests/test_backward_partition.py new file mode 100644 index 0000000000..a6ab7c0979 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/test_backward_partition.py @@ -0,0 +1,293 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import operator + +import torch +import torch.fx as fx +from torch.testing._internal.common_utils import TestCase + +from torchtitan.experiments.graph_trainer.backward_partition import ( + BackwardNodePartition, + partition_backward_nodes, +) +from torchtitan.experiments.graph_trainer.graph_pp.utils import node_closure + +aten = torch.ops.aten + + +def _custom_quantize(x): + return x + + +def _custom_grouped_mm(a, b): + return a + + +class TestBackwardNodePartition(TestCase): + """Unit tests for dependency-based dI/dW backward node classification.""" + + def _build_backward_graph( + self, + *, + trunk_target=aten.mul.Tensor, + di_target=aten.mm.default, + dw_quant_target=aten.relu.default, + dw_target=aten.mm.default, + ): + """Build a synthetic backward graph: shared trunk -> dI + dW branches. + + Returns the graph and a role -> node mapping. ``d_input`` is the input + gradient output and ``d_weight`` is the parameter gradient output. + """ + graph = fx.Graph() + grad_out = graph.placeholder("grad_out") + saved_act = graph.placeholder("saved_act") + weight = graph.placeholder("weight") + trunk = graph.call_function(trunk_target, args=(grad_out, saved_act)) + d_input = graph.call_function(di_target, args=(trunk, weight)) + dw_quant = graph.call_function(dw_quant_target, args=(trunk,)) + d_weight = graph.call_function(dw_target, args=(dw_quant, saved_act)) + graph.output((d_input, d_weight)) + nodes = { + "grad_out": grad_out, + "saved_act": saved_act, + "weight": weight, + "trunk": trunk, + "di_out": d_input, + "dw_quant": dw_quant, + "dw_out": d_weight, + } + return graph, nodes + + def _partition(self, graph, nodes) -> BackwardNodePartition: + return partition_backward_nodes( + graph, + input_grad_outputs=[nodes["di_out"]], + param_grad_outputs=[nodes["dw_out"]], + ) + + def _assert_disjoint_partition(self, partition, nodes): + classified = node_closure([nodes["di_out"]]) | node_closure([nodes["dw_out"]]) + self.assertEqual( + partition.di_nodes | partition.dw_only_nodes | partition.shared_nodes, + classified, + ) + self.assertEqual(partition.di_nodes & partition.dw_only_nodes, set()) + self.assertEqual(partition.di_nodes & partition.shared_nodes, set()) + self.assertEqual(partition.dw_only_nodes & partition.shared_nodes, set()) + self.assertTrue(partition.movable_nodes <= partition.dw_only_nodes) + + def test_partition_is_disjoint_union_of_closures(self): + graph, nodes = self._build_backward_graph() + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + + def test_trunk_and_branch_membership(self): + graph, nodes = self._build_backward_graph() + partition = self._partition(graph, nodes) + self.assertEqual(partition.di_nodes, {nodes["weight"], nodes["di_out"]}) + self.assertEqual(partition.dw_only_nodes, {nodes["dw_quant"], nodes["dw_out"]}) + self.assertEqual( + partition.shared_nodes, + {nodes["grad_out"], nodes["saved_act"], nodes["trunk"]}, + ) + self.assertEqual(partition.movable_nodes, {nodes["dw_quant"], nodes["dw_out"]}) + + def test_collective_in_dw_branch_is_pinned(self): + graph, nodes = self._build_backward_graph() + with graph.inserting_before(nodes["dw_out"]): + coll = graph.call_function( + torch.ops._c10d_functional.all_to_all_single.default, + args=(nodes["dw_quant"], [1], [1], "0"), + ) + wait = graph.call_function( + torch.ops._c10d_functional.wait_tensor.default, args=(coll,) + ) + nodes["dw_out"].replace_input_with(nodes["dw_quant"], wait) + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + self.assertTrue({coll, wait} <= partition.dw_only_nodes) + self.assertEqual(partition.movable_nodes, {nodes["dw_quant"], nodes["dw_out"]}) + + def test_pre_write_reader_of_mutated_buffer_is_pinned(self): + graph, nodes = self._build_backward_graph() + # In-place write into saved_act, which both branches read. dw_quant + # reads the buffer BEFORE the write without a data edge to it, so + # deferring it could push the read past the write: pinned. dw_out + # reads THROUGH the write (data edge keeps it after): movable. + with graph.inserting_before(nodes["dw_out"]): + mut = graph.call_function(aten.add_.Tensor, args=(nodes["saved_act"], 1.0)) + nodes["dw_out"].replace_input_with(nodes["saved_act"], mut) + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + self.assertIn(nodes["dw_quant"], partition.dw_only_nodes) + self.assertNotIn(nodes["dw_quant"], partition.movable_nodes) + self.assertIn(mut, partition.dw_only_nodes) + self.assertNotIn(mut, partition.movable_nodes) + self.assertEqual(partition.movable_nodes, {nodes["dw_out"]}) + + def test_unresolvable_mutation_target_empties_movable(self): + graph, nodes = self._build_backward_graph() + # A write whose target is not a graph value cannot be reasoned about. + with graph.inserting_before(nodes["dw_out"]): + graph.call_function(aten.add_.Tensor, args=(3.0, 1.0)) + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + self.assertEqual(partition.movable_nodes, set()) + + def test_dw_local_mutation_pins_only_itself(self): + graph, nodes = self._build_backward_graph() + with graph.inserting_after(nodes["weight"]): + dw_buf = graph.placeholder("dw_buf") + with graph.inserting_before(nodes["dw_out"]): + mut = graph.call_function(aten.add_.Tensor, args=(dw_buf, 1.0)) + nodes["dw_out"].replace_input_with(nodes["saved_act"], mut) + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + self.assertTrue({dw_buf, mut} <= partition.dw_only_nodes) + self.assertEqual(partition.movable_nodes, {nodes["dw_quant"], nodes["dw_out"]}) + + def test_cpu_scalar_read_is_pinned(self): + graph, nodes = self._build_backward_graph() + with graph.inserting_before(nodes["dw_out"]): + sync = graph.call_function( + aten._local_scalar_dense.default, args=(nodes["dw_quant"],) + ) + scaled = graph.call_function( + aten.mul.Tensor, args=(nodes["dw_quant"], sync) + ) + nodes["dw_out"].replace_input_with(nodes["dw_quant"], scaled) + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + self.assertIn(sync, partition.dw_only_nodes) + self.assertNotIn(sync, partition.movable_nodes) + self.assertEqual( + partition.movable_nodes, + {nodes["dw_quant"], scaled, nodes["dw_out"]}, + ) + + def test_device_to_host_copy_is_pinned(self): + graph, nodes = self._build_backward_graph() + fake_mode = torch._subclasses.FakeTensorMode(allow_non_fake_inputs=True) + with fake_mode: + nodes["dw_quant"].meta["val"] = torch.empty(4, device="cuda") + with graph.inserting_before(nodes["dw_out"]): + host_copy = graph.call_function( + aten._to_copy.default, + args=(nodes["dw_quant"],), + kwargs={"device": torch.device("cpu")}, + ) + nodes["dw_out"].replace_input_with(nodes["dw_quant"], host_copy) + partition = self._partition(graph, nodes) + self.assertIn(host_copy, partition.dw_only_nodes) + self.assertNotIn(host_copy, partition.movable_nodes) + + def test_sym_placeholder_feeding_both_sides_is_shared(self): + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + graph, nodes = self._build_backward_graph() + shape_env = ShapeEnv() + fake_mode = torch._subclasses.FakeTensorMode( + allow_non_fake_inputs=True, shape_env=shape_env + ) + with fake_mode: + sym_batch = shape_env.create_unbacked_symint() + with graph.inserting_after(nodes["weight"]): + sym = graph.placeholder("sym_batch") + sym.meta["val"] = sym_batch + with graph.inserting_before(nodes["di_out"]): + di_view = graph.call_function( + aten.view.default, args=(nodes["di_out"].args[0], [sym, -1]) + ) + nodes["di_out"].replace_input_with(nodes["trunk"], di_view) + with graph.inserting_before(nodes["dw_out"]): + dw_view = graph.call_function( + aten.view.default, args=(nodes["dw_quant"], [sym, -1]) + ) + nodes["dw_out"].replace_input_with(nodes["dw_quant"], dw_view) + partition = self._partition(graph, nodes) + self._assert_disjoint_partition(partition, nodes) + self.assertIn(sym, partition.shared_nodes) + self.assertNotIn(sym, partition.movable_nodes) + self.assertIn(dw_view, partition.movable_nodes) + + def test_untracked_side_effect_user_pins_producer(self): + graph, nodes = self._build_backward_graph() + with graph.inserting_before(nodes["dw_out"]): + graph.call_function( + aten._assert_scalar.default, args=(nodes["dw_quant"], "msg") + ) + partition = self._partition(graph, nodes) + self.assertIn(nodes["dw_quant"], partition.dw_only_nodes) + self.assertNotIn(nodes["dw_quant"], partition.movable_nodes) + self.assertIn(nodes["dw_out"], partition.movable_nodes) + + def test_dead_getitem_users_do_not_pin_producer(self): + # Mirrors the real MXFP8 wgrad chain: a multi-output quantize op + # whose rowwise outputs are never consumed. The dead getitems are + # inert and must not pin the producer. + graph, nodes = self._build_backward_graph() + with graph.inserting_after(nodes["dw_quant"]): + multi = graph.call_function(_custom_quantize, args=(nodes["dw_quant"],)) + with graph.inserting_after(multi): + dead_b = graph.call_function(operator.getitem, args=(multi, 1)) + used = graph.call_function(operator.getitem, args=(multi, 2)) + dead_a = graph.call_function(operator.getitem, args=(multi, 0)) + nodes["dw_out"].update_arg(0, used) + partition = self._partition(graph, nodes) + self.assertIn(multi, partition.movable_nodes) + self.assertIn(used, partition.movable_nodes) + self.assertIn(nodes["dw_quant"], partition.movable_nodes) + classified = ( + partition.di_nodes | partition.dw_only_nodes | partition.shared_nodes + ) + self.assertNotIn(dead_a, classified) + self.assertNotIn(dead_b, classified) + + def test_classification_ignores_call_targets(self): + aten_graph, aten_nodes = self._build_backward_graph() + quant_graph, quant_nodes = self._build_backward_graph( + dw_quant_target=_custom_quantize, dw_target=_custom_grouped_mm + ) + aten_partition = self._partition(aten_graph, aten_nodes) + quant_partition = self._partition(quant_graph, quant_nodes) + + def roles(partition, nodes, node_set): + members = getattr(partition, node_set) + return {role for role, node in nodes.items() if node in members} + + for node_set in ( + "di_nodes", + "dw_only_nodes", + "shared_nodes", + "movable_nodes", + ): + self.assertEqual( + roles(aten_partition, aten_nodes, node_set), + roles(quant_partition, quant_nodes, node_set), + ) + + def test_no_input_grads_puts_everything_on_dw_side(self): + graph, nodes = self._build_backward_graph() + partition = partition_backward_nodes( + graph, + input_grad_outputs=[], + param_grad_outputs=[nodes["di_out"], nodes["dw_out"]], + ) + self.assertEqual(partition.di_nodes, set()) + self.assertEqual(partition.shared_nodes, set()) + self.assertEqual(partition.dw_only_nodes, set(nodes.values())) + self.assertEqual( + partition.movable_nodes, + {nodes["trunk"], nodes["di_out"], nodes["dw_quant"], nodes["dw_out"]}, + ) + + +if __name__ == "__main__": + from torch.testing._internal.common_utils import run_tests + + run_tests()