Add pipeline parallel plan to PretrainedConfig and PreTrainedModel

[hmellor]

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[NouamaneTazi]

Finally PP is becoming native in transformers!! 🔥🔥🔥

For the attr to add i was thinking something more like:

base_model_pp_plan = OrderedDict([
    ("embed_tokens", {"input_tensors": [...], "output_tensors": [...]}),
    ("some_other_op_that_comes_before_the_first_layer", {"input_tensors": [...], "output_tensors": [...]}),
    ("layers.*", {"input_tensors": [...], "output_tensors": [...]}), # maybe leave ".*" to highlight it's a modulelist
    ("norm", {"input_tensors": [...], "output_tensors": [...]}),
    ("cast_to_fp32", {"input_tensors": [...], "output_tensors": [...]}),
    ("some_other_op_that_comes_after_the_last_layer", {"input_tensors": [...], "output_tensors": [...]})
])

And in LlamaForCausalLM

_pp_plan = base_model_pp_plan.insert(-1, ("lm_head", {"input_tensors": [...], "output_tensors": [...]})

Or for example Deepseek can add multiple heads etc

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Super nice for an initial design!

[hmellor]

I've simplified the schema as suggested. Perhaps we could also add an Enum to so that users can index these schemas without knowing that [0] means inputs and [1] means outputs?

[ArthurZucker]

Super nice!
maybe the only thing missing is a small test for models that do support PP!
this could for example be using this:

from torch.distributed.pipelining import pipeline, SplitPoint

# An example micro-batch input
x = torch.LongTensor([1, 2, 4, 5])

pipe = pipeline(
    module=mod,
    mb_args=(x,),
    split_spec={
        "layers.1": SplitPoint.BEGINNING,
    }
)

with something similar as what we are doing for TP! Tho this could be in another PR!
(or from here

[NouamaneTazi]

Very clean! Left minor comment to use tuple instead of list
Thanks!

[NouamaneTazi]

Very nice!

[hmellor]

Thanks both for the reviews!

Since there is currently no pipeline logic added by this PR let's save testing for a follow up PR.

[ArthurZucker]

Nice! I'll try to work on the follow up PR asap to futur proof the design but LGTM