[GraphTrainer] Remove inductor_region from RMSNorm annotation pass

Address review feedback:

1. Default rmsnorm_compile_config to None (was required kwarg).
2. Remove per-norm inductor_region assignment — the CapabilityBasedPartitioner
   in regional_inductor already separates disconnected norm subgraphs into
   distinct partitions. Explicit region IDs caused each norm to compile as a
   separate subgraph, inflating memory by ~17 GiB.

Benchmark (Llama3 8B, 8×H100, FSDP4+TP2, c4_test, 20 steps):

| Branch                         | tps   | MFU    | Memory    |
|--------------------------------|-------|--------|-----------|
| main                           | 6,538 | 38.29% | 30.92 GiB |
| branch (with inductor_region)  | 6,559 | 38.41% | 47.74 GiB |
| branch (without inductor_region)| 6,600 | 38.65% | 30.92 GiB |

Author: SherlockNoMad

torchtitan/experiments/graph_trainer/passes.py

@@ -539,7 +539,7 @@ def annotate_rmsnorm_for_regional_inductor_pass(
     gm: torch.fx.GraphModule,
     example_inputs: tuple | None = None,
     *,
-    rmsnorm_compile_config: dict | None,
+    rmsnorm_compile_config: dict | None = None,
 ) -> torch.fx.GraphModule:
     """Tag RMSNorm ops with compile_with_inductor for regional_inductor.
 
@@ -548,10 +548,6 @@ def annotate_rmsnorm_for_regional_inductor_pass(
     users) so that ``regional_inductor_pass`` compiles each norm as a
     fused Inductor kernel.
 
-    Each ``_fused_rms_norm`` / ``_fused_rms_norm_backward`` call is
-    assigned a separate ``inductor_region`` so that norms at different
-    points in the graph are compiled independently.
-
     Args:
         gm: The graph module to annotate.
         example_inputs: Example inputs (unused, required by pass interface).
@@ -570,31 +566,26 @@ def annotate_rmsnorm_for_regional_inductor_pass(
         torch.ops.aten._fused_rms_norm_backward.default,
     }
 
-    next_region_id = 0
     num_tagged = 0
 
     for node in gm.graph.nodes:
         if node.op != "call_function" or node.target not in _RMSNORM_TARGETS:
             continue
 
-        annotation = dict(compile_annotation)
-        annotation["inductor_region"] = next_region_id
-        next_region_id += 1
-
-        # Tag the fused norm node itself.
-        node.meta.setdefault("custom", {})["compile_with_inductor"] = annotation
+        node.meta.setdefault("custom", {})["compile_with_inductor"] = compile_annotation
         num_tagged += 1
 
         # Tag getitem users that extract outputs from the fused op.
         for user in node.users:
             if user.op == "call_function" and user.target is operator.getitem:
-                user.meta.setdefault("custom", {})["compile_with_inductor"] = annotation
+                user.meta.setdefault("custom", {})[
+                    "compile_with_inductor"
+                ] = compile_annotation
                 num_tagged += 1
 
     if num_tagged > 0:
         logger.info(
-            f"Tagged {num_tagged} RMSNorm nodes across "
-            f"{next_region_id} regions for regional Inductor compilation"
+            f"Tagged {num_tagged} RMSNorm nodes for regional Inductor compilation"
         )
 
     return gm

torchtitan/experiments/graph_trainer/tests/test_passes.py

@@ -1356,17 +1356,6 @@ def _count_tagged_nodes(self, gm):
                 count += 1
         return count
 
-    def _get_tagged_regions(self, gm):
-        """Return a dict mapping inductor_region -> list of node names."""
-        regions = {}
-        for node in gm.graph.nodes:
-            custom = node.meta.get("custom", {})
-            annotation = custom.get("compile_with_inductor")
-            if annotation is not None:
-                region = annotation.get("inductor_region")
-                regions.setdefault(region, []).append(node.name)
-        return regions
-
     def test_tags_fused_rmsnorm_and_getitems(self):
         """_fused_rms_norm nodes and their getitem users are tagged."""
         gm = self._build_rmsnorm_gm(
@@ -1401,8 +1390,8 @@ def test_does_not_tag_non_rmsnorm_nodes(self):
 
         self.assertEqual(self._count_tagged_nodes(gm), 0)
 
-    def test_separate_regions_per_fused_norm(self):
-        """Each _fused_rms_norm call gets its own inductor_region ID."""
+    def test_multiple_fused_norms_all_tagged(self):
+        """Multiple _fused_rms_norm calls and their getitems are all tagged."""
         gm = self._build_rmsnorm_gm(
             [
                 torch.ops.aten._fused_rms_norm.default,
@@ -1416,12 +1405,8 @@ def test_separate_regions_per_fused_norm(self):
             rmsnorm_compile_config={"max_autotune": True},
         )
 
-        regions = self._get_tagged_regions(gm)
-        # Should have exactly 2 distinct regions
-        self.assertEqual(len(regions), 2)
-        # Each region: 1 fused node + 2 getitem users = 3 nodes
-        for region_id, node_names in regions.items():
-            self.assertEqual(len(node_names), 3)
+        # 2 fused nodes + 2*2 getitem users = 6 tagged nodes
+        self.assertEqual(self._count_tagged_nodes(gm), 6)
 
     def test_backward_fused_norm_tagged(self):
         """_fused_rms_norm_backward nodes and their getitem users are tagged."""
@@ -1438,8 +1423,6 @@ def test_backward_fused_norm_tagged(self):
 
         # 1 fused backward node + 2 getitem users = 3 tagged nodes
         self.assertEqual(self._count_tagged_nodes(gm), 3)
-        regions = self._get_tagged_regions(gm)
-        self.assertEqual(len(regions), 1)
 
     def test_no_fused_norm_is_noop(self):
         """When no fused RMSNorm ops exist, no nodes are tagged."""
@@ -1475,7 +1458,6 @@ def test_none_compile_config_tags_with_empty_annotation(self):
             annotation = node.meta.get("custom", {}).get("compile_with_inductor")
             if annotation is not None:
                 self.assertNotIn("inductor_configs", annotation)
-                self.assertIn("inductor_region", annotation)
 
     def test_compile_config_propagated_to_annotation(self):
         """The compile config dict is wrapped under inductor_configs in the annotation."""
@@ -1496,8 +1478,8 @@ def test_compile_config_propagated_to_annotation(self):
             if annotation is not None:
                 self.assertEqual(annotation["inductor_configs"], config)
 
-    def test_fwd_and_bwd_get_distinct_regions(self):
-        """Forward and backward fused norms get distinct region IDs."""
+    def test_fwd_and_bwd_both_tagged(self):
+        """Forward and backward fused norms are both tagged."""
         gm = self._build_rmsnorm_gm(
             [
                 torch.ops.aten._fused_rms_norm.default,
@@ -1511,26 +1493,8 @@ def test_fwd_and_bwd_get_distinct_regions(self):
             rmsnorm_compile_config={"max_autotune": True},
         )
 
-        regions = self._get_tagged_regions(gm)
-        self.assertEqual(len(regions), 2)
-
-    def test_multiple_fused_norms_get_distinct_regions(self):
-        """Multiple fused norm calls each get their own region ID."""
-        gm = self._build_rmsnorm_gm(
-            [
-                torch.ops.aten._fused_rms_norm.default,
-                torch.ops.aten._fused_rms_norm.default,
-                torch.ops.aten._fused_rms_norm.default,
-            ]
-        )
-
-        annotate_rmsnorm_for_regional_inductor_pass(
-            gm,
-            rmsnorm_compile_config={"max_autotune": True},
-        )
-
-        regions = self._get_tagged_regions(gm)
-        self.assertEqual(len(regions), 3)
+        # 2 fused nodes + 2*2 getitem users = 6 tagged nodes
+        self.assertEqual(self._count_tagged_nodes(gm), 6)
 
 
 if __name__ == "__main__":