Add pipe.get_stage_module

pippy/IR.py

@@ -649,6 +649,11 @@ def forward(self, *args, **kwargs):
 
         return res
 
+    def get_stage_module(self, stage_idx: int) -> torch.nn.Module:
+        if stage_idx < 0 or stage_idx >= self.num_stages:
+            raise ValueError(f"Invalid stage index {stage_idx}!")
+        return getattr(self.split_gm, f"submod_{stage_idx}")
+
     @staticmethod
     def _number_and_count_forward_stages(gm: fx.GraphModule):
         num_stages = 0

test/test_transformer.py

@@ -1,5 +1,4 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
-import pippy
 import torch
 from pippy import annotate_split_points, Pipe, SplitPoint
 
@@ -60,7 +59,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     (x,),
 )
 assert pipe.num_stages == 2
-gm = pipe.split_gm
 
 
 def get_layers(module):
@@ -70,9 +68,10 @@ def get_layers(module):
 
 # Collect all layers in pipe
 layers = []
-for name, submod in gm.named_children():
-    print(f"\nStage {name}: \n", submod)
-    layers += get_layers(submod)
+for stage_idx in range(pipe.num_stages):
+    stage_mod = pipe.get_stage_module(stage_idx)
+    print(f"\nStage {stage_idx}: \n", stage_mod)
+    layers += get_layers(stage_mod)
 
 # Check layer completeness
 orig_layers = get_layers(transformer)

test/test_unflatten.py

@@ -1,5 +1,4 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
-import pippy
 import torch
 from pippy import Pipe, pipe_split
 
@@ -52,14 +51,14 @@ def forward(self, x: torch.Tensor, constant=None) -> torch.Tensor:
 )
 
 assert pipe.num_stages == 4
-gm = pipe.split_gm
 orig_state_dict = mod.state_dict()
 
 # Check qualnames
 print("\nParameters of each stage:")
-for name, submod in gm.named_children():
-    print(f"\nStage {name}:")
-    for param_name, param in submod.named_parameters():
+for stage_idx in range(pipe.num_stages):
+    print(f"\nStage {stage_idx}:")
+    stage_mod = pipe.get_stage_module(stage_idx)
+    for param_name, param in stage_mod.named_parameters():
         assert (
             param_name in orig_state_dict
         ), f"{param_name} not in original state dict"