examples/inference/pippy/llama.py Assertion error about graphs

[685Degrees]

System Info

- `Accelerate` version: 1.1.1
- Platform: macOS-15.1.1-arm64-arm-64bit
- `accelerate` bash location: /Users/jason/Documents/GitHub/temp/.venv/bin/accelerate
- Python version: 3.11.9
- Numpy version: 2.1.3
- PyTorch version (GPU?): 2.5.1 (False)
- PyTorch XPU available: False
- PyTorch NPU available: False
- PyTorch MLU available: False
- PyTorch MUSA available: False
- System RAM: 36.00 GB
- `Accelerate` default config:
        Not found

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

Run python examples/inference/pippy/llama.py
The error message that results is:

Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:30<00:00, 15.31s/it]
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
Traceback (most recent call last):
  File "/Users/jason/Documents/GitHub/temp/test.py", line 38, in <module>
    model = prepare_pippy(model, split_points="auto", example_kwargs=inputs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/accelerate/inference.py", line 170, in prepare_pippy
    stage = build_pipeline(model, split_points, example_args, example_kwargs, num_chunks)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/accelerate/inference.py", line 87, in build_pipeline
    pipe = pipeline(
           ^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/distributed/pipelining/_IR.py", line 1231, in pipeline
    return Pipe.from_tracing(
           ^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/distributed/pipelining/_IR.py", line 1051, in from_tracing
    pipe = Pipe._from_traced(
           ^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/distributed/pipelining/_IR.py", line 750, in _from_traced
    new_submod = _outline_submodules(submodule.graph)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/distributed/pipelining/_unflatten.py", line 24, in _outline_submodules
    ).run_outer()
      ^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/export/unflatten.py", line 1014, in run_outer
    self.run_from(node_idx)
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/export/unflatten.py", line 1094, in run_from
    ).run_from(node_idx)
      ^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/export/unflatten.py", line 1094, in run_from
    ).run_from(node_idx)
      ^^^^^^^^^^^^^^^^^^
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/export/unflatten.py", line 1043, in run_from
    self.finalize_outputs()
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/export/unflatten.py", line 993, in finalize_outputs
    _verify_graph_equivalence(self.cached_graph_module, self.module)
  File "/Users/jason/Documents/GitHub/temp/.venv/lib/python3.11/site-packages/torch/export/unflatten.py", line 655, in _verify_graph_equivalence
    assert graph_dump(x.graph) == graph_dump(y.graph)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

cc: @JasonGross

Expected behavior

No assertion error