Commit 8fed349
committed
Update base for Update on "[graph_trainer] Copy forward metadata to backward 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]1 parent a44c97f commit 8fed349
0 file changed
0 commit comments