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]

Author: tugsbayasgalan

torchtitan/experiments/graph_trainer/tests/test_trace_module.py

@@ -680,10 +680,8 @@ def test_deepseek_v3(self):
         self._run_model_test(DeepSeekV3Model, config)
 
     def test_deepseek_v3_flex_attention(self):
-        """Multi-step bitwise test for DeepSeek MLA + flex attention + regional_inductor.
-
-        Uses a tiny model with small head dims to stay within triton shared
-        memory limits.  Annotates FlexAttention.forward via annotate_fn before
+        """Tests if we can propagate fwd node metadata reliably through backward.
+        Annotates FlexAttention.forward via annotate_fn before
         tracing so compile_with_inductor flows into the graph naturally.
         """
         from torch.fx.traceback import annotate_fn