[graph_trainer][aot_fx_trace][graph_pp] Add FSDP collective splitting (#3088)

Add GraphPP extraction for explicit FSDP UNSHARD and REDUCE_GRAD graphs. The splitters use shared FSDP trace-pattern helpers to identify validated all-gather/wait reconstruction chains and reduce-scatter epilogues, keep the no-FSDP forward/backward graphs on the GraphPP flat calling convention, and leave fused edge graphs out of the runtime contract.

Teach the default SAC policy to force-save the same helper-selected unshard output when FSDP runs with reshard_after_forward=False, matching the AutoParallel force-save behavior instead of saving every all-gather target.

Add FSDP split tests for collective extraction, view and split-cat reconstruction, malformed unshard chains, no-op paths, input-source metadata, reduce-grad input naming, and pre-reduce cast placement.

Author: sanketpurandare

torchtitan/experiments/graph_trainer/fsdp_passes.py

@@ -48,12 +48,76 @@
     _is_backward_node,
     _MODULE_FQN,
 )
+from torchtitan.experiments.graph_trainer.fsdp_patterns import find_fsdp_unshard_outputs
 from torchtitan.tools.logging import logger
 
 
 _FSDP_BUCKET_META = "fsdp_bucket"
 
 
+def _chain_nodes_to_placeholder(
+    output: fx.Node,
+    placeholder: fx.Node,
+) -> set[fx.Node]:
+    """Return graph nodes in one FSDP unshard chain, excluding the placeholder."""
+    chain_nodes: set[fx.Node] = set()
+    queue = [output]
+    while queue:
+        node = queue.pop()
+        if node is placeholder or node in chain_nodes:
+            continue
+        chain_nodes.add(node)
+        queue.extend(node.all_input_nodes)
+    return chain_nodes
+
+
+def deduplicate_fsdp_unshard_chains_pass(
+    gm: torch.fx.GraphModule,
+    example_inputs: tuple | None = None,
+) -> torch.fx.GraphModule:
+    """Canonicalize duplicate SimpleFSDP unshard chains per flat parameter.
+
+    A traced parametrized module can read the same FSDP parameter more than
+    once. Each read materializes an equivalent all-gather/wait reconstruction
+    chain from the same flat parameter placeholder. Downstream FSDP passes
+    assume one unsharded value per flat parameter, so this pass rewrites all
+    duplicate chains to the first chain and removes the now-dead duplicate
+    collective infrastructure.
+    """
+    del example_inputs
+
+    removable_nodes: set[fx.Node] = set()
+    num_duplicate_chains = 0
+    for placeholder in gm.graph.find_nodes(op="placeholder"):
+        unshard_outputs = find_fsdp_unshard_outputs(placeholder)
+        if len(unshard_outputs) <= 1:
+            continue
+        canonical_output = unshard_outputs[0]
+        for duplicate_output in unshard_outputs[1:]:
+            removable_nodes.update(
+                _chain_nodes_to_placeholder(duplicate_output, placeholder)
+            )
+            duplicate_output.replace_all_uses_with(canonical_output)
+            num_duplicate_chains += 1
+
+    if num_duplicate_chains == 0:
+        return gm
+
+    def _is_impure_for_fsdp_dedup(node: fx.Node) -> bool:
+        if node in removable_nodes:
+            return False
+        return node.is_impure()
+
+    gm.graph.eliminate_dead_code(is_impure_node=_is_impure_for_fsdp_dedup)
+    gm.graph.lint()
+    gm.recompile()
+    logger.info(
+        "Canonicalized %d duplicate FSDP unshard chain(s)",
+        num_duplicate_chains,
+    )
+    return gm
+
+
 def is_wait_tensor_from_fsdp(node: torch.fx.Node) -> bool:
     """
     Returns True if the node is a wait_tensor node that is the result of an all_gather

torchtitan/experiments/graph_trainer/fsdp_patterns.py

@@ -0,0 +1,236 @@
+# 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.
+"""Pattern helpers for GraphPP's current simple-FSDP collective traces.
+
+The matchers intentionally follow c10d functional traces produced by FSDP2:
+
+    param_shard -> all_gather -> wait -> view*/split-cat -> compute
+    local_grad -> cast/view* -> reduce_scatter -> wait -> param_grad
+    local_grad -> cast/view* -> all_reduce -> wait -> param_grad
+
+They are structural helpers for today's trace shape. A future upstream FSDP or
+torch.pipelining annotation should replace this with explicit collective-region
+metadata instead of broader pattern matching.
+"""
+
+import operator
+from typing import Any
+
+import torch
+import torch.fx as fx
+
+
+def is_wait_tensor(node: fx.Node) -> bool:
+    return (
+        node.op == "call_function"
+        and node.target == torch.ops._c10d_functional.wait_tensor.default
+    )
+
+
+def is_all_gather_into_tensor(node: fx.Node) -> bool:
+    return (
+        node.op == "call_function"
+        and node.target == torch.ops._c10d_functional.all_gather_into_tensor.default
+    )
+
+
+def is_reduce_scatter_tensor(node: fx.Node) -> bool:
+    return (
+        node.op == "call_function"
+        and node.target is torch.ops._c10d_functional.reduce_scatter_tensor.default
+    )
+
+
+def is_all_reduce(node: fx.Node) -> bool:
+    return (
+        node.op == "call_function"
+        and node.target is torch.ops._c10d_functional.all_reduce.default
+    )
+
+
+def is_reduce_grad_collective(node: fx.Node) -> bool:
+    return is_reduce_scatter_tensor(node) or is_all_reduce(node)
+
+
+def _find_last_all_gather_in_chain(start_node: fx.Node) -> fx.Node | None:
+    """Find the final all-gather in a linear FSDP unshard launch chain."""
+    node = start_node
+    last_all_gather = None
+    while True:
+        if is_all_gather_into_tensor(node):
+            last_all_gather = node
+        if len(node.users) != 1:
+            break
+        user = next(iter(node.users))
+        if len(user.all_input_nodes) > 1:
+            break
+        node = user
+    return last_all_gather
+
+
+def _find_last_user_in_wait_chain(wait_node: fx.Node) -> fx.Node:
+    """Find the last FSDP unshard node before the value enters real compute.
+
+    The traced FSDP unshard has a mostly linear shape:
+
+        flat_param -> ... -> all_gather -> wait -> view* -> compute
+
+    Some models reshape the gathered flat buffer through a split/cat fanout:
+
+        wait -> split -> getitem_0 --+
+                      -> getitem_1 --+-> cat -> view* -> compute
+
+    In both cases the FSDP region ends before the first consumer with multiple
+    FX inputs. The split/getitem/cat fanout is still part of reconstructing the
+    unsharded parameter value, so it is included in the chain.
+    """
+    node = wait_node
+    while True:
+        if len(node.users) != 1:
+            if (
+                node.op == "call_function"
+                and node.target == torch.ops.aten.split.Tensor
+                and all(
+                    user.op == "call_function"
+                    and user.target == operator.getitem
+                    and len(user.users) == 1
+                    for user in node.users
+                )
+            ):
+                getitem_users = [next(iter(user.users)) for user in node.users]
+                potential_cat = getitem_users[0]
+                if all(user == potential_cat for user in getitem_users) and (
+                    potential_cat.op == "call_function"
+                    and potential_cat.target == torch.ops.aten.cat.default
+                ):
+                    node = potential_cat
+                    continue
+            break
+
+        user = next(iter(node.users))
+        if len(user.all_input_nodes) > 1:
+            break
+        node = user
+    return node
+
+
+def _find_last_non_view_node_in_chain(node: fx.Node) -> fx.Node:
+    """Return the value GraphPP should pass across the unshard boundary."""
+    result = node
+    while hasattr(result.target, "is_view") and result.target.is_view:
+        if len(result.all_input_nodes) != 1:
+            raise ValueError(f"View node {result.name} should have exactly one input")
+        result = result.all_input_nodes[0]
+    return result
+
+
+def _unshard_output_from_all_gather(last_all_gather: fx.Node) -> fx.Node:
+    if len(last_all_gather.users) != 1:
+        raise ValueError(
+            f"Expected one wait_tensor user for all_gather node {last_all_gather.name}, "
+            f"got {len(last_all_gather.users)}"
+        )
+    wait_node = next(iter(last_all_gather.users))
+    if not is_wait_tensor(wait_node):
+        raise ValueError(
+            f"Expected wait_tensor after all_gather node {last_all_gather.name}, "
+            f"got {wait_node.name}"
+        )
+
+    wait_chain_user = _find_last_user_in_wait_chain(wait_node)
+    return _find_last_non_view_node_in_chain(wait_chain_user)
+
+
+def find_fsdp_unshard_outputs(param_placeholder: fx.Node) -> tuple[fx.Node, ...]:
+    """Return all FSDP unshard outputs launched from one flat parameter input.
+
+    Most parameters have one linear all-gather chain. Some real traces read the
+    same parametrized value more than once, which produces multiple equivalent
+    all-gather/wait chains from the same placeholder.
+    ``deduplicate_fsdp_unshard_chains_pass`` canonicalizes those duplicate
+    chains before downstream FSDP passes rely on a single unsharded value.
+    """
+    last_all_gather = _find_last_all_gather_in_chain(param_placeholder)
+    if last_all_gather is not None:
+        return (_unshard_output_from_all_gather(last_all_gather),)
+
+    outputs: list[fx.Node] = []
+    seen: set[fx.Node] = set()
+    for user in param_placeholder.users:
+        if len(user.all_input_nodes) > 1:
+            continue
+        last_all_gather = _find_last_all_gather_in_chain(user)
+        if last_all_gather is None:
+            continue
+        output = _unshard_output_from_all_gather(last_all_gather)
+        if output not in seen:
+            seen.add(output)
+            outputs.append(output)
+    return tuple(outputs)
+
+
+def find_fsdp_unshard_output(param_placeholder: fx.Node) -> fx.Node | None:
+    """Return the extracted unshard output for one flat parameter input.
+
+    GraphPP and the SAC force-save policy must agree on this node. It is the
+    same value AutoParallel saves for ``reshard_after_forward=False``: the last
+    non-view node in the FSDP all-gather/wait reconstruction chain. Parameters
+    without an all-gather are replicated or otherwise already local, so callers
+    should keep the original placeholder as that parameter's unsharded value.
+
+    TODO(sanketpurandare): requires upstream change: FSDP trace/passes should
+    annotate unshard collective regions for downstream graph extraction.
+    """
+    outputs = find_fsdp_unshard_outputs(param_placeholder)
+    if not outputs:
+        return None
+    return outputs[0]
+
+
+def find_fsdp_unshard_save_node(param_placeholder: fx.Node) -> fx.Node | None:
+    return find_fsdp_unshard_output(param_placeholder)
+
+
+def find_fsdp_unshard_save_nodes(param_placeholder: fx.Node) -> tuple[fx.Node, ...]:
+    """Return all FSDP unshard values that SAC must save for one parameter."""
+    return find_fsdp_unshard_outputs(param_placeholder)
+
+
+def find_fsdp_reduce_grad_input(param_grad_output: Any) -> fx.Node | None:
+    """Return the split point before an FSDP reduce-grad epilogue.
+
+    The backward FSDP/DDP/HSDP tail is traced as a unary chain ending in the
+    synced grad output:
+
+        local_grad -> cast/view* -> reduce_scatter -> wait -> sharded_grad
+        local_grad -> cast/view* -> all_reduce -> wait -> replicated_grad
+        local_grad -> cast/view* -> all_reduce -> wait -> reduce_scatter
+          -> wait -> grad
+
+    GraphPP splits at the input to the earliest grad-sync collective in that
+    suffix. The cast remains in ``bw_no_fsdp`` so microbatch accumulation
+    happens in FSDP's reduce dtype, and ``reduce_grad`` contains only the
+    scheduled collective epilogue. Values that are not FX nodes, such as
+    ``None`` parameter-grad slots, are not collective outputs and are preserved
+    by the caller.
+
+    TODO(sanketpurandare): requires upstream change: FSDP trace/passes should
+    annotate reduce-grad collective regions for downstream graph extraction.
+    """
+    if not isinstance(param_grad_output, fx.Node):
+        return None
+
+    node = param_grad_output
+    reduce_grad_input = None
+    while isinstance(node, fx.Node) and len(node.all_input_nodes) == 1:
+        input_node = node.all_input_nodes[0]
+        if len(input_node.users) > 1:
+            break
+        previous_node = node
+        node = input_node
+        if is_reduce_grad_collective(previous_node):
+            reduce_grad_input = node
+    return reduce_grad_input

torchtitan/experiments/graph_trainer/graph_pp/__init__.py

@@ -12,10 +12,20 @@
     GraphPPDiDwSplit,
     split_di_dw_graph,
 )
+from torchtitan.experiments.graph_trainer.graph_pp.split_fsdp_collectives import (
+    GraphPPFSDPBackwardSplit,
+    GraphPPFSDPForwardSplit,
+    split_backward_fsdp_collectives,
+    split_forward_fsdp_collectives,
+)
 
 __all__ = [
     "GraphPPDiDwSplit",
+    "GraphPPFSDPBackwardSplit",
+    "GraphPPFSDPForwardSplit",
     "GraphMeta",
     "partition_joint_graph",
+    "split_backward_fsdp_collectives",
     "split_di_dw_graph",
+    "split_forward_fsdp_collectives",
 ]