Skip to content

Commit 8fed349

Browse files
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

File tree

    0 commit comments

    Comments
     (0)