Repository: pytorch/PiPPy

Issue:
Exception when splitting model with "--autosplit"
Running `torchrun --nproc-per-node 4 examples/huggingface/pippy_gpt2.py --batches=50 --autosplit` results in the following assertion:

> [rank0]: Traceback (most recent call last):
> [rank0]:   File "PiPPy/examples/huggingface/pippy_gpt2.py", line 137, in <module>
> [rank0]:     run(args)
> [rank0]:   File "/PiPPy/examples/huggingface/pippy_gpt2.py", line 61, in run
> [rank0]:     gpt2_pipe = pipeline(
> [rank0]:   File "miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/torchpippy-0.2.0+2104a68-py3.10.egg/pippy/IR.py", line 1375, in pipeline
> [rank0]:     return Pipe.from_tracing(
> [rank0]:   File "miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/torchpippy-0.2.0+2104a68-py3.10.egg/pippy/IR.py", line 1192, in from_tracing
> [rank0]:     pipe = Pipe._from_traced(
> [rank0]:   File "miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/torchpippy-0.2.0+2104a68-py3.10.egg/pippy/IR.py", line 702, in _from_traced
> [rank0]:     traced = split_policy(traced)
> [rank0]:   File "miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/torchpippy-0.2.0+2104a68-py3.10.egg/pippy/ModelSplit.py", line 196, in _split_into_nstages_equal_size
> [rank0]:     gm, rv_nstages = _split_on_size_threshold_with_max_stages(
> [rank0]:   File "miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/torchpippy-0.2.0+2104a68-py3.10.egg/pippy/ModelSplit.py", line 75, in _split_on_size_threshold_with_max_stages
> [rank0]:     node_param_sizes = _analyze_node_size(gm)
> [rank0]:   File "miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/torchpippy-0.2.0+2104a68-py3.10.egg/pippy/ModelSplit.py", line 32, in _analyze_node_size
> [rank0]:     assert param_name in state_dict
> [rank0]: AssertionError

Verified on rev 2104a68619ec1162da4650e379545e275728c27a


Discussion:

Resolving patch:
diff --git a/examples/huggingface/pippy_gpt2.py b/examples/huggingface/pippy_gpt2.py
--- a/examples/huggingface/pippy_gpt2.py
+++ b/examples/huggingface/pippy_gpt2.py
@@ -9,8 +9,7 @@
 import torch
 import torch.distributed as dist
 
-from pippy import pipeline
-from pippy import SplitPoint, annotate_split_points
+from pippy import pipeline, SplitPoint
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 
@@ -19,19 +18,6 @@
 from hf_utils import generate_inputs_for_model, get_number_of_params
 
 
-def add_split_points(gpt2, nranks):
-    decoders_per_rank = (gpt2.config.n_layer + nranks - 1) // nranks
-    print(f"decoders_per_rank = {decoders_per_rank}")
-    nstages = 1
-    for i in range(1, gpt2.config.n_layer // decoders_per_rank):
-        annotate_split_points(
-            gpt2,
-            {f'transformer.h.{i * decoders_per_rank}': SplitPoint.BEGINNING},
-        )
-        nstages += 1
-    assert nstages == nranks, f"nstages = {nstages} nranks = {nranks}"
-
-
 def run(args):
     # Model configs
     config = GPT2Config()
@@ -55,25 +41,30 @@ def run(args):
     example_inputs = generate_inputs_for_model(
         model_class, gpt2, model_name, args.batch_size, args.device)
 
+    split_policy = None
+    split_spec = None
+
     if args.autosplit:
         # Automatic split
         from pippy import split_into_equal_size
-        gpt2_pipe = pipeline(
-            gpt2,
-            num_chunks=args.chunks,
-            example_args=(),
-            example_kwargs=example_inputs,
-            split_policy=split_into_equal_size(args.world_size),
-        )
+        split_policy = split_into_equal_size(args.world_size)
     else:
-        # Manually annotate split points
-        add_split_points(gpt2, args.world_size)
-        gpt2_pipe = pipeline(
-            gpt2,
-            num_chunks=args.chunks,
-            example_args=(),
-            example_kwargs=example_inputs,
-        )
+        # Use manual split spec
+        decoders_per_rank = (gpt2.config.n_layer + args.world_size - 1) // args.world_size
+        print(f"decoders_per_rank = {decoders_per_rank}")
+        split_spec = {
+            f'transformer.h.{i * decoders_per_rank}': SplitPoint.BEGINNING for i in range(1, args.world_size)
+        }
+
+    # Only one of `split_spec` and `split_policy` is used
+    gpt2_pipe = pipeline(
+        gpt2,
+        num_chunks=args.chunks,
+        example_args=(),
+        example_kwargs=example_inputs,
+        split_spec=split_spec,
+        split_policy=split_policy,
+    )
 
     assert gpt2_pipe.num_stages == args.world_size
     if args.rank == 0:
@@ -96,6 +87,10 @@ def run(args):
     else:
         out = schedule.step()
 
+    # Ideally we shouldn't need this barrier, but since this example only has
+    # one iteration, it's added for now to prevent the first few ranks from
+    # exiting before the last few ranks finish.
+    dist.barrier()
     print(f"Rank {args.rank} completes")
 
 
diff --git a/pippy/ModelSplit.py b/pippy/ModelSplit.py
--- a/pippy/ModelSplit.py
+++ b/pippy/ModelSplit.py
@@ -29,7 +29,9 @@ def _analyze_node_size(
     for node in gm.graph.nodes:
         if node.op == "get_attr":  # a parameter node
             param_name = node.target
-            assert param_name in state_dict
+            if param_name not in state_dict:
+                # In some cases, attr node is not a parameter or buffer, we just skip it
+                continue
             param = state_dict[param_name]
             # Find use site of this parameter
             for user in node.users:


Test patch: