Pipeline parallel naive inference
Author: 3outeille
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47289)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47289)
<!-- ci-dashboard-badge:end -->

# Summary

- Adding a very naive pipeline parallel inference engine as a baseline (no specific schedule such as gpipe, 1f1b yet)
- supports tied/untied weight embeddings
- To make it work seamless with `generate()`, we are broadcasting every logits to every ranks. In long term,  we want only the last stage to hold the logits (needs to change stuff in `generate()`)

# Left to do

- 2D-// with TP
- Think if `pp_plan` is useful or not  
- implement better scheduling method (gpipe, 1f1b)
- more optimize version of generate with PP
- handle pipeline parallel for backwards

# Repro

```python
#!/usr/bin/env python
"""Sanity check: run a model with transformers using pipeline parallelism.

Usage:
    torchrun --nproc-per-node=2 pp_inference.py
"""
import os
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed import DistributedConfig


def main(): 

    rank = int(os.environ.get("RANK", 0))
    pp_size = int(os.environ.get("WORLD_SIZE", 1))
    # model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
    model_id = "Qwen/Qwen2.5-0.5B-Instruct"
    is_tie_word_embeddings = model_id == "Qwen/Qwen2.5-0.5B-Instruct"

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        distributed_config=DistributedConfig(pp_size=pp_size),
        dtype=torch.bfloat16,
        tie_word_embeddings=is_tie_word_embeddings,
    ).eval()

    prompts = [
        "The capital of France is",
        "In one sentence, explain what pipeline parallelism is:",
    ]

    for prompt in prompts:
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
        output_ids = model.generate(input_ids, max_new_tokens=64, do_sample=False)
        if rank == 0:
            print(f"PROMPT : {prompt!r}")
            print(f"OUTPUT : {tokenizer.decode(output_ids[0], skip_special_tokens=True)!r}")
            print("-" * 80)


if __name__ == "__main__":
    main()
```

```
(env_pp) ➜  transformers_pp git:(pp_inference) ✗ torchrun --nproc-per-node=4 pp_inference.py
W0715 06:13:12.806000 1079087 torch/distributed/run.py:851] 
W0715 06:13:12.806000 1079087 torch/distributed/run.py:851] *****************************************
W0715 06:13:12.806000 1079087 torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
W0715 06:13:12.806000 1079087 torch/distributed/run.py:851] *****************************************
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_cutlass_90a.abi3.so
Failed to load /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so: Could not load this library: /fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/transformers_pp/env_pp/lib/python3.12/site-packages/torchao/_C_mxfp8.cpython-310-x86_64-linux-gnu.so
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 72/72 [00:00<00:00, 2159.60it/s]
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 72/72 [00:00<00:00, 2265.58it/s]
[transformers] Qwen2ForCausalLM LOAD REPORT from: Qwen/Qwen2.5-0.5B-Instruct
Key                                                               | Status  |          
------------------------------------------------------------------+---------+----------
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.q_proj.weight         | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.q_proj.bias           | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.k_proj.weight         | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.k_proj.bias           | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.v_proj.weight         | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.v_proj.bias           | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.self_attn.o_proj.weight         | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.mlp.gate_proj.weight            | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.mlp.up_proj.weight              | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.mlp.down_proj.weight            | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.input_layernorm.weight          | OWNED   | PP rank 1
model.layers.{6, 7, 8, 9, 10, 11}.post_attention_layernorm.weight | OWNED   | PP rank 1
model.layers.{0...23}.mlp.down_proj.weight                        | SKIPPED | PP rank 0
model.layers.{0...23}.mlp.up_proj.weight                          | SKIPPED | PP rank 0
model.layers.{0...23}.input_layernorm.weight                      | SKIPPED | PP rank 2
model.layers.{0...23}.self_attn.q_proj.bias                       | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.k_proj.bias                       | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.v_proj.weight                     | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.o_proj.weight                     | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.k_proj.weight                     | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.q_proj.weight                     | SKIPPED | PP rank 2
model.layers.{0...23}.post_attention_layernorm.weight             | SKIPPED | PP rank 3
model.layers.{0...23}.mlp.gate_proj.weight                        | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.v_proj.bias                       | SKIPPED | PP rank 2
model.embed_tokens.weight                                         | SKIPPED | PP rank 0
model.norm.weight                                                 | SKIPPED | PP rank 3

Notes:
- OWNED:        checkpoint weights loaded on this pipeline stage.
- SKIPPED:      checkpoint weights owned by another pipeline stage.
[transformers] Qwen2ForCausalLM LOAD REPORT from: Qwen/Qwen2.5-0.5B-Instruct
Key                                                                   | Status  |          
----------------------------------------------------------------------+---------+----------
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.q_proj.weight         | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.q_proj.bias           | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.k_proj.weight         | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.k_proj.bias           | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.v_proj.weight         | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.v_proj.bias           | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.self_attn.o_proj.weight         | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.mlp.gate_proj.weight            | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.mlp.up_proj.weight              | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.mlp.down_proj.weight            | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.input_layernorm.weight          | OWNED   | PP rank 2
model.layers.{12, 13, 14, 15, 16, 17}.post_attention_layernorm.weight | OWNED   | PP rank 2
model.layers.{0...23}.self_attn.o_proj.weight                         | SKIPPED | PP rank 0
model.layers.{0...23}.mlp.up_proj.weight                              | SKIPPED | PP rank 3
model.layers.{0...23}.post_attention_layernorm.weight                 | SKIPPED | PP rank 3
model.layers.{0...23}.mlp.gate_proj.weight                            | SKIPPED | PP rank 3
model.layers.{0...23}.self_attn.v_proj.bias                           | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.k_proj.bias                           | SKIPPED | PP rank 1
model.layers.{0...23}.self_attn.v_proj.weight                         | SKIPPED | PP rank 0
model.layers.{0...23}.mlp.down_proj.weight                            | SKIPPED | PP rank 1
model.layers.{0...23}.self_attn.k_proj.weight                         | SKIPPED | PP rank 1
model.embed_tokens.weight                                             | SKIPPED | PP rank 0
model.layers.{0...23}.input_layernorm.weight                          | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.q_proj.weight                         | SKIPPED | PP rank 0
model.layers.{0...23}.self_attn.q_proj.bias                           | SKIPPED | PP rank 3
model.norm.weight                                                     | SKIPPED | PP rank 3

Notes:
- OWNED:        checkpoint weights loaded on this pipeline stage.
- SKIPPED:      checkpoint weights owned by another pipeline stage.
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 73/73 [00:00<00:00, 1304.03it/s]
[transformers] Qwen2ForCausalLM LOAD REPORT from: Qwen/Qwen2.5-0.5B-Instruct
Key                                                             | Status  |          
----------------------------------------------------------------+---------+----------
model.embed_tokens.weight                                       | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.q_proj.weight         | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.q_proj.bias           | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.k_proj.weight         | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.k_proj.bias           | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.v_proj.weight         | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.v_proj.bias           | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.self_attn.o_proj.weight         | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.mlp.gate_proj.weight            | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.mlp.up_proj.weight              | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.mlp.down_proj.weight            | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.input_layernorm.weight          | OWNED   | PP rank 0
model.layers.{0, 1, 2, 3, 4, 5}.post_attention_layernorm.weight | OWNED   | PP rank 0
model.layers.{6...23}.self_attn.k_proj.bias                     | SKIPPED | PP rank 2
model.layers.{6...23}.self_attn.o_proj.weight                   | SKIPPED | PP rank 3
model.layers.{6...23}.mlp.up_proj.weight                        | SKIPPED | PP rank 1
model.layers.{6...23}.input_layernorm.weight                    | SKIPPED | PP rank 3
model.layers.{6...23}.mlp.down_proj.weight                      | SKIPPED | PP rank 1
model.layers.{6...23}.post_attention_layernorm.weight           | SKIPPED | PP rank 3
model.layers.{6...23}.self_attn.k_proj.weight                   | SKIPPED | PP rank 2
model.norm.weight                                               | SKIPPED | PP rank 3
model.layers.{6...23}.mlp.gate_proj.weight                      | SKIPPED | PP rank 2
model.layers.{6...23}.self_attn.q_proj.weight                   | SKIPPED | PP rank 2
model.layers.{6...23}.self_attn.v_proj.bias                     | SKIPPED | PP rank 2
model.layers.{6...23}.self_attn.q_proj.bias                     | SKIPPED | PP rank 3
model.layers.{6...23}.self_attn.v_proj.weight                   | SKIPPED | PP rank 1

Notes:
- OWNED:        checkpoint weights loaded on this pipeline stage.
- SKIPPED:      checkpoint weights owned by another pipeline stage.
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 74/74 [00:00<00:00, 858.83it/s]
[transformers] Qwen2ForCausalLM LOAD REPORT from: Qwen/Qwen2.5-0.5B-Instruct
Key                                                                   | Status  |          
----------------------------------------------------------------------+---------+----------
model.embed_tokens.weight                                             | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.q_proj.weight         | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.q_proj.bias           | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.k_proj.weight         | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.k_proj.bias           | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.v_proj.weight         | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.v_proj.bias           | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.self_attn.o_proj.weight         | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.mlp.gate_proj.weight            | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.mlp.up_proj.weight              | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.mlp.down_proj.weight            | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.input_layernorm.weight          | OWNED   | PP rank 3
model.layers.{18, 19, 20, 21, 22, 23}.post_attention_layernorm.weight | OWNED   | PP rank 3
model.norm.weight                                                     | OWNED   | PP rank 3
lm_head.weight                                                        | OWNED   | PP rank 3
model.layers.{0...17}.self_attn.k_proj.weight                         | SKIPPED | PP rank 1
model.layers.{0...17}.self_attn.v_proj.weight                         | SKIPPED | PP rank 0
model.layers.{0...17}.mlp.up_proj.weight                              | SKIPPED | PP rank 2
model.layers.{0...17}.post_attention_layernorm.weight                 | SKIPPED | PP rank 0
model.layers.{0...17}.mlp.gate_proj.weight                            | SKIPPED | PP rank 1
model.layers.{0...17}.self_attn.o_proj.weight                         | SKIPPED | PP rank 1
model.layers.{0...17}.self_attn.v_proj.bias                           | SKIPPED | PP rank 2
model.layers.{0...17}.self_attn.k_proj.bias                           | SKIPPED | PP rank 1
model.layers.{0...17}.input_layernorm.weight                          | SKIPPED | PP rank 2
model.layers.{0...17}.self_attn.q_proj.bias                           | SKIPPED | PP rank 2
model.layers.{0...17}.self_attn.q_proj.weight                         | SKIPPED | PP rank 2
model.layers.{0...17}.mlp.down_proj.weight                            | SKIPPED | PP rank 0

Notes:
- OWNED:        checkpoint weights loaded on this pipeline stage.
- SKIPPED:      checkpoint weights owned by another pipeline stage.
PROMPT : 'The capital of France is'
OUTPUT : "The capital of France is Paris. It was founded in 789 AD by Charlemagne, the last king of the Carolingian dynasty. The city has a long and rich history dating back to ancient times. In fact, it's one of the oldest cities in Europe.\nParis is famous for its beautiful architecture, including the E"
--------------------------------------------------------------------------------
PROMPT : 'In one sentence, explain what pipeline parallelism is:'
OUTPUT : 'In one sentence, explain what pipeline parallelism is: Pipeline parallelism refers to the ability of a system to execute multiple tasks concurrently within a single pipeline. It involves using multiple processors or threads to process data simultaneously, allowing for faster processing times and improved efficiency in certain applications. This technique can be particularly useful in scenarios where data needs to be processed quickly, such as in real'
--------------------------------------------------------------------------------
[rank0]:[W715 06:13:54.966502038 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_47289). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47289)
**Latest run:** [29395629171:1](https://github.com/huggingface/transformers/actions/runs/29395629171)
**Result:** `success` | **Jobs:** 16 | **Tests:** 172,166 | **Failures:** 0 | **Duration:** 13h 40m

<!-- ci-dashboard-recap:end -->