[GraphTrainer] Support transformer_block_bucketing in aot_fx_trace mode by yiming0416 · Pull Request #2934 · pytorch/torchtitan

yiming0416 · 2026-04-10T18:43:18Z

Summary

Enable transformer_block_bucketing pass in aot_fx_trace compile mode. Previously this only worked in aot and jit modes because aot_fx_trace traces a joint fwd+bwd graph where nn_module_stack metadata alone isn't sufficient to correctly identify and separate forward vs backward nodes for bucketing.

Root cause: In the joint graph produced by aot_fx_trace, forward and backward collectives coexist. The existing manual_overlap_bucketing pass needs to bucket each direction independently — forward collectives in forward execution order, backward collectives in reverse order — but had no way to distinguish them.

Solution: Leverage the existing annotate_fn mechanism (already used for AC region tagging) to:

Tag each bucket-level submodule's forward with its FQN (annotate_module_fqns)
Mark backward nodes with is_bwd=True during metadata propagation
Provide a direction-aware module_stack_fn to manual_overlap_bucketing that filters by fwd/bwd

A new joint_transformer_block_bucketing_reordering_pass applies bucketing in two passes — once for forward nodes, once for backward — so each direction's collectives are bucketed and reordered independently.

SherlockNoMad · 2026-04-15T23:05:52Z

    traced_result: "TracedResult",
+    model: torch.nn.Module,


Keep the interface clean if possible.

trace_result has some internal GraphModule, can we use that?

if not, alternative is to cache an instance of model inside TracedResult

I attached an instance of model to TracedResult.

We can use the internal GraphModule from TracedResult as they has been flattened and lost the module hierarchy info

…ackward subgraphs" Graph-trainer was only copying seq_nr metadata within the outer FX graph. That misses the way old AOT behaves for flex attention: the matching forward source for a backward flex node can live inside a nested GraphModule, including placeholder nodes from the captured flex subgraphs. In practice that meant backward flex nodes in the graph-trainer path could miss the forward custom / nn_module_stack metadata that old AOT used to recover. The old outer-graph heuristic was also unreliable once remat-tagged nodes showed up in the same seq_nr bucket, because the first outer node with that seq_nr was no longer a reliable proxy for the true forward source. Fix _copy_fwd_metadata_to_bw_nodes() to mirror the old AOT behavior more closely: - walk all nested GraphModules - build the seq_nr -> forward-node map across those subgraphs - only copy onto nodes identified as backward via partitioner_tag or remat metadata 1) This should address failures in #2934 [ghstack-poisoned]

Graph-trainer was only copying seq_nr metadata within the outer FX graph. That misses the way old AOT behaves for flex attention: the matching forward source for a backward flex node can live inside a nested GraphModule, including placeholder nodes from the captured flex subgraphs. In practice that meant backward flex nodes in the graph-trainer path could miss the forward custom / nn_module_stack metadata that old AOT used to recover. The old outer-graph heuristic was also unreliable once remat-tagged nodes showed up in the same seq_nr bucket, because the first outer node with that seq_nr was no longer a reliable proxy for the true forward source. Fix _copy_fwd_metadata_to_bw_nodes() to mirror the old AOT behavior more closely: - walk all nested GraphModules - build the seq_nr -> forward-node map across those subgraphs - only copy onto nodes identified as backward via partitioner_tag or remat metadata 1) This should address failures in #2934 [ghstack-poisoned]

…nstance (#180579) Problem: When `manual_overlap_bucketing` is called multiple times on the same GraphModule (e.g. separate forward and backward passes), the old implementation stored bucketed node types in node.meta, which persists on the graph nodes across calls. The second call's _manual_reorder_graph iterates all graph nodes and misclassifies nodes bucketed by the first call, building incorrect overlap dependency edges. Fix: Moved bucketed node type tracking from `node.meta["manual_bucket_node_type"]` to `ManualOverlapPreservingBucketer.bucketed_node_types` dict, preventing metadata leaking across multiple `manual_overlap_bucketing` calls on the same graph. Test: In pytorch/torchtitan#2934 ``` MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_debugmodel ./run_train.sh --compile.mode aot_fx_trace ``` Pull Request resolved: #180579 Approved by: https://github.com/SherlockNoMad

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom): * __->__ #2875 Graph-trainer was only copying seq_nr metadata within the outer FX graph. That misses the way old AOT behaves for flex attention: the matching forward source for a backward flex node can live inside a nested GraphModule, including placeholder nodes from the captured flex subgraphs. In practice that meant backward flex nodes in the graph-trainer path could miss the forward custom / nn_module_stack metadata that old AOT used to recover. The old outer-graph heuristic was also unreliable once remat-tagged nodes showed up in the same seq_nr bucket, because the first outer node with that seq_nr was no longer a reliable proxy for the true forward source. Fix _copy_fwd_metadata_to_bw_nodes() to mirror the old AOT behavior more closely: - walk all nested GraphModules - build the seq_nr -> forward-node map across those subgraphs - only copy onto nodes identified as backward via partitioner_tag or remat metadata 1) This should address failures in #2934

SherlockNoMad · 2026-04-20T17:25:52Z

        output_spec: Original output pytree spec used during reconstruction.
+        model: The original nn.Module, cached for downstream passes that need
+            module structure (e.g. transformer block bucketing). Only set when
+            traced via :func:`trace_train_step`.


Only set when traced via :func:`trace_train_step`.

I am not a fan of conditional behavior.

Also TracedResult is now becoming a contract between Trainer and pre-compile, so adding model is not feasible now.

So this is pointing towards something more fundamental, we probably need a different way to come up with bucketing than get_transformer_block_buckets(traced_result.model).

Let's land for now, but add to TODO here to remove the model field from Traced Result.

aot_fx_trace mode traces with record_module_stack=False, so nodes lack nn_module_stack metadata that manual_overlap_bucketing uses to match nodes to modules. This change uses the existing annotate_fn mechanism (already used for AC region tagging) to tag module FQNs during tracing, then provides a custom module_stack_fn to manual_overlap_bucketing. Changes: - Add annotate_module_fqns() and get_module_stack_from_annotation() to common_utils.py - Call annotate_module_fqns() in both llama3 and deepseek_v3 annotate functions - Add module_stack_fn parameter to transformer_block_bucketing_reordering_pass - Update construct_default_graph_passes to accept compile_config, model, and parallel_dims, and wire up the bucketing pass with reassign_to_pg - Pass config/model/parallel_dims from GraphTrainer to construct_default_graph_passes - Add numerics tests for both llama3 and deepseek_v3 with aot_fx_trace + transformer_block_bucketing

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.14.0) (oldest at bottom): * #2916 * __->__ #3043 Precompile integration tests (run_precompile_tests.py) were never run in CI, which allowed the bucketing + CooR incompatibility from #2934 (enabled in #3036) to go undetected. Add them to both CI workflows: - A10 workflow: Llama3 aot_fx_trace precompile test (FSDP+TP) - H100 workflow: DSv3 aot_fx_trace precompile test (FSDP+TP+EP)

Precompile integration tests (run_precompile_tests.py) were never run in CI, which allowed the bucketing + CooR incompatibility from #2934 to go undetected. Add them to both CI workflows: - A10 workflow: Llama3 aot_fx_trace precompile test (FSDP+TP) - H100 workflow: DSv3 aot_fx_trace precompile test (FSDP+TP+EP) Also fix the DSv3 precompile test config from debugmodel to debugmodel_ep — EP tests require AllToAllTokenDispatcher which is only created when moe_comm_backend="standard". ghstack-source-id: 35a8421 Pull-Request: #3043

…nstance (pytorch#180579) Problem: When `manual_overlap_bucketing` is called multiple times on the same GraphModule (e.g. separate forward and backward passes), the old implementation stored bucketed node types in node.meta, which persists on the graph nodes across calls. The second call's _manual_reorder_graph iterates all graph nodes and misclassifies nodes bucketed by the first call, building incorrect overlap dependency edges. Fix: Moved bucketed node type tracking from `node.meta["manual_bucket_node_type"]` to `ManualOverlapPreservingBucketer.bucketed_node_types` dict, preventing metadata leaking across multiple `manual_overlap_bucketing` calls on the same graph. Test: In pytorch/torchtitan#2934 ``` MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_debugmodel ./run_train.sh --compile.mode aot_fx_trace ``` Pull Request resolved: pytorch#180579 Approved by: https://github.com/SherlockNoMad

pytorch-bot Bot added the ciflow/8gpu label Apr 10, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Apr 10, 2026

SherlockNoMad mentioned this pull request Apr 10, 2026

[graph_trainer] Nightly scout tracking issue #2856

Open

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch 4 times, most recently from e7a8559 to 3d07026 Compare April 10, 2026 22:42

SherlockNoMad reviewed Apr 10, 2026

View reviewed changes

Comment thread torchtitan/experiments/graph_trainer/make_fx_tracer.py Outdated

SherlockNoMad reviewed Apr 10, 2026

View reviewed changes

Comment thread torchtitan/experiments/graph_trainer/common_utils.py

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch 2 times, most recently from 63e130f to 89b5762 Compare April 13, 2026 21:16

yiming0416 changed the title ~~[GraphTrainer][AutoDev] Support transformer_block_bucketing in aot_fx_trace mode~~ [GraphTrainer] Support transformer_block_bucketing in aot_fx_trace mode Apr 13, 2026

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch 7 times, most recently from a95cd87 to 83d06fc Compare April 14, 2026 22:43

yushangdi mentioned this pull request Apr 15, 2026

[graph_trainer] Add CUDA graph kernel annotation pass #2926

Merged

yiming0416 marked this pull request as ready for review April 15, 2026 22:52

yiming0416 requested review from aditvenk, tianyu-l and xmfan as code owners April 15, 2026 22:52

SherlockNoMad reviewed Apr 15, 2026

View reviewed changes

Comment thread torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py

SherlockNoMad reviewed Apr 15, 2026

View reviewed changes

Comment thread torchtitan/experiments/graph_trainer/passes.py Outdated

SherlockNoMad reviewed Apr 15, 2026

View reviewed changes

Comment thread torchtitan/experiments/graph_trainer/passes.py Outdated

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch 2 times, most recently from 991e351 to 1daa927 Compare April 16, 2026 00:38

yiming0416 mentioned this pull request Apr 16, 2026

[bucketing] Move manual_bucket_node_type from node.meta to bucketer instance pytorch/pytorch#180579

Closed

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch 2 times, most recently from 3362a5e to 68737e0 Compare April 16, 2026 17:34

tugsbayasgalan mentioned this pull request Apr 16, 2026

[graph_trainer] Copy forward metadata to backward subgraphs #2875

Merged

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch 3 times, most recently from 236184d to 5ad46c8 Compare April 20, 2026 17:06

SherlockNoMad reviewed Apr 20, 2026

View reviewed changes

SherlockNoMad approved these changes Apr 20, 2026

View reviewed changes

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch from 5ad46c8 to 6e0b03c Compare April 20, 2026 18:11

yiming0416 force-pushed the graph_trainer/transformer-block-bucketing-aot-fx-trace branch from 6e0b03c to b83b975 Compare April 20, 2026 18:36

yiming0416 merged commit 000c178 into main Apr 20, 2026
11 of 15 checks passed

bobrenjc93 mentioned this pull request Apr 21, 2026

[graph_trainer] Add precompile integration tests to CI #3043

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[GraphTrainer] Support transformer_block_bucketing in aot_fx_trace mode#2934

[GraphTrainer] Support transformer_block_bucketing in aot_fx_trace mode#2934
yiming0416 merged 1 commit into
mainfrom
graph_trainer/transformer-block-bucketing-aot-fx-trace

yiming0416 commented Apr 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SherlockNoMad Apr 15, 2026

Uh oh!

yiming0416 Apr 16, 2026

Uh oh!

Uh oh!

SherlockNoMad Apr 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

yiming0416 commented Apr 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SherlockNoMad Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

yiming0416 Apr 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

SherlockNoMad Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yiming0416 commented Apr 10, 2026 •

edited

Loading