[RFC] adding Tensor and Pipeline Parallelism to transformers
Author: stas00
Pull request: no
Labels: name: Pipeline Parallel

name: Performance

name: Tensor Parallel

name: WIP

Following up on this proposal https://github.com/huggingface/transformers/issues/12772 I just had a discussion with @hyunwoongko (with great help from @JakeTae who patiently translated for us), and we tried to discuss a strategy of how to best integrate Tensor Parallelism (TP) and Pipeline Parallelism (PP) into `transformers`, making it easy for reviewers and the contributors. Note that 
 [parallelformers](https://github.com/tunib-ai/parallelformers) currently implements only TP.
  
So here is a great example of how the TP can be added, as @hyunwoongko  already implemented it in his fork for `GPTNeo`
https://github.com/tunib-ai/transformers/commit/5bf8655be624b3aeda799b80fddd220213491b04 (he didn't use `GPT2` since it already has the naive PP implemented). So you can see exactly what we want to merge. It's a very thin layer to the model and most of the functionality is in the helper parallel utils. The end of the change is multiple tests/examples that need to be converted to our test framework.

Now, while adding TP is relatively easy, adding PP is very complex in the current state of HF models because they include many features that interfere with implementing PP - due to the requirements:
1. for the model to be `nn.Sequential `and 
2. inputs/outputs to be simple tensors with the first dimension of batch size.

So to implement PP we will most likely have to fork each model, strip the unnecessary for scalability features and only then be able to implement PP.

So my thinking is that perhaps we do it from the get the going? Instead of integrating TP into the normal model - say `GPTNeo`, we fork it to say `GTPNeo3D` from the get going and do all the work including TP and PP on that new model. Once everybody is happy we can rinse and repeat for other models. 

I added 3D to `GPTNeo` to make `GTPNeo3D` - 3D = DP/TP/PP - not exactly sure about this particular name or attached to it, this is just something to start with.

Also once TP is implemented in say `GTPNeo3D` we can start replicating it to other models. Because [parallelformers](https://github.com/tunib-ai/parallelformers) has them all covered already. PP will be much harder and we can do this in parallel.

I wanted to check in with the team to see if this approach resonates better, rather than modifying the existing models.

Thank you!

Also see [this blog post explaining parallelforms](https://tunib.notion.site/TECH-2021-07-26-Parallelformers-Journey-to-deploying-big-models_TUNiB-32b19a599c38497abaad2a98727f6dc8).

-----------------------

Additionally see the main pytorch Parallelism discussion at https://github.com/pytorch/rfcs/pull/32

@LysandreJik, @sgugger, @patrickvonplaten 


Comments:
@stas00 - I like this a lot! And as we're been dragging our feet with implementing some of the Megatron 3D parallelism into `mistral` - I think it might be a great way for us to collaborate; we can just start with the base GPT-2 model perhaps?

I think my (and the Mistral team's) addition the next few weeks will be trying to do some benchmarking of Megatron and existing gains with various subsets of parallelism (at a very fundamental level - profiling which kernels are being called, etc.) and maybe creating a set of unit tests to verify correctness?

Separately - might be worth keeping logs of how to "3D-ify" new models, and ways we might make that procedure even easier moving forward.

Let me know if this makes sense!

@stas00 @siddk If we are creating a new class, we do not need to modify the existing `parallelize()` method, so we do not need to work with GPTNeo. I think GPT2 would be better.

Thanks for the feedback, Sidd.

The reason @hyunwoongko thought of starting with GPTNeo was because GPT2 already has the naive PP `parallelize()`. But the problem is that it's not just in the model, it's also in the Trainer. So we probably need to choose some other action name for that function altogether. At least for a time being so that we could move forward.

Note that the intention is to do simple things first and not do too many things at once. So I think starting with GPTNeo on a clean slate is a better idea. Once it's happy it'd be trivial to replicate that to GPT2. And it's already done as you can see from the link in OP.

Here is my vision of _3Difying_ `transformers`:

step 1. implement TP in one model
step 2a. start replicating TP to other models
step 2b. start working on PP in one model
step 3a. start replicating PP to other models.

note how step 2 can be done in parallel by different people.

So I can see that Mistral's team efforts would be parallel work and not sequential. So for example:

step 3b. implement Mistral's GPT2 improvements to GPT2
step 4a. start replicating it to other models.

If were were to start with GPT2 we would interfere with your work, Sidd, so I think it's actually best if we pick 2 different starting models.

But let's stay focused in this discussion on TP+PP, otherwise it'd be too easy to get side-tracked. We already spent too much time talking - let's see some code going into `transformers`! :)

wrt trainers, it'll be a natural part of the work - I'm not worried too much about it. I don't know much about accelerate yet, but HF Trainer should be relatively easy.


This makes a lot of sense to me - thanks @stas00 and @hyunwoongko for the clarifications! The steps above form a pretty good concrete plan - but if you both are already planning on tackling it, maybe it makes sense for us to tackle some of the other Megatron-LM improvements first, like the custom loss scaling/kernels/etc. (in mistral, so we can break things 😅)? And as y'all build the "main API" for 3D parallelism, we can just drop that in, and train larger models! 

The PR with the mistral's first set GPT-2 improvements is waiting on approval right now - once that's in we can move a bit faster as well.

That sounds like a perfect plan to me, Sidd.

@stas00 I think the following method is not good for megatron-friendly method.

```
step 1. implement megatron-friendly TP in one model
step 2a. start replicating megatron-friendly TP to other models
step 2b. start working on megatron-friendly PP in one model
step 3a. start replicating megatron-friendly PP to other models.
```

Ultimately, implementing PP requires rewriting all modeling code. (including `GPT2Attention`, `GPT2MLP`, `GPT2Model`, ...) I wasn't familiar with PP until not long ago. but recently, I became very familiar with PP and found out that we had to rewrite all the code. (`generation_utils.py` used for inference should also be changed.) Therefore, it is recommended that megatron-friendly TP and PP be implemented together. (I think it's inefficient to implement megatron-friendly TP alone.)
 


The transformers-friendly method (=parallelformers) has the advantage of being able to extend the model quickly because it does not need to rewrite the modeling code (it uses the existing transformers code), but it is not compatible with PP. So we have to remove all the transformers-friendly TP when implementing PP. Which strategy we take is a matter of choice. We can quickly expand them in a transformers friendly way, and then change them one by one to be megatron friendly like 

```
step 1. implement transformers-friendly TP in one model
step 2a. start replicating transformers-friendly TP to other models
step 2b. start working on megatron-friendly TP + PP in one model
step 3a. start replicating megatron-friendly TP + PP to other models.
```

Or there is a way to not implement transformers-friendly methods because they will be removed anyway. But, since there are thousands of lines of code to write for megatron-friendly and tens of lines of code for transformers-friendly, the megatron-friendly approach will scale very slowly. 

```
step 1. start working on megatron-friendly TP + PP in one model
step 2. start replicating megatron-friendly TP + PP to other models.
```

One thing to note is that the transformers-friendly TP implementation is completely eliminated when implementing the megatron-friendly TP. A megatron-friendly TP is implemented differently from a transformers-friendly TP. 

Adding a GPTNeo3D to experiment seems like a good idea to me. At the end of the day, that modeling file can leave in the same folder as `modeling_gptneo.py`.

Note that while you experiment, you can leverage #13467 to share models on the Hub that have no implementation in Transformers and still work with the auto-model API.

> Adding a GPTNeo3D to experiment seems like a good idea to me. At the end of the day, that modeling file can leave in the same folder as `modeling_gptneo.py`.

Great!

> Note that while you experiment, you can leverage #13467 to share models on the Hub that have no implementation in Transformers and still work with the auto-model API.

The 3D GPTNeo model's weights are the same as a normal GPTNeo model's - i.e. it can be used w/ or w/ PP/TP, so I'm not sure why we need a special API?

And I guess we won't be able to use `AutoModel`, because the `config.model_type` will say 'gpt_neo', but we will want to load it with `GPTNeo3D*` classes.


@hyunwoongko, you're bringing up excellent points.

I suppose the main question is how much of a benefit we can give to users by having just TP. My thinking is that if it's easy to add TP to all models and since you have already done this, let's do it.

I'm concerned that adding PP will be a very slow process because as you said it requires massive rewrites to the model's code, and meanwhile those models that are waiting their turn won't be very scalable (except with Deepspeed ZeRO).

Besides we can delegate the TP adding to the rest of the models to others (other developers and even community) since it's mostly just replaying the code you have already written. But it still requires work, at least in adding tests and documentation, and then PRs.

The only concern with adding the transformers-friendly way is that the external API remains the same when we add PP.

How does that sound?

@stas00 But anyway, I don't prefer PP. As you know, PP is memory inefficient because it is not compatible with ZeRO 2, 3. In fact, we also decided not to use PP when developing language models. So adding just TP would be helpful for many people. So let's go with the following strategy. but, as you said, the API for both methods should remain the same.

```
step 1. implement transformers-friendly TP in one model
step 2a. start replicating transformers-friendly TP to other models
step 2b. start working on megatron-friendly TP + PP in one model
step 3a. start replicating megatron-friendly TP + PP to other models.
```

But transformers-friendly TPs have no reason to rewrite their modeling code. What should we do?


That's great, @hyunwoongko!

And once we complete `GPTNeo3D` with TP we can decide whether to fold it back to the normal `GPTNeo` model or keep it separate. I'm saying that if at the end we will do PP only for a few select models (which is too a real possibiilty), then there is absolutely no need to fork 60 models and create a lot more maintenance work for `transformers`, if they will have just TP+DP.

@stas00 

In my opinion, transformers-friendly TP have no reason to write their own modeling code like GPTNeo3D.

1. So the transformers-friendly TP will just use the existing model
2. And let's make a new modeling class such as GPT2For3D when we develop the megatron-friendly TP + PP (GPT2, Bert, T5, etc, It will probably be some models, not all.) 

I'm thinking of an API like this.
```python
from transformers import GPTNeoModel

model = GPTNeoModel.from_pretrained("elutherai/gpt-neo-1.3B", tensor_model_parallel_size=4)

or

model = GPTNeoModel.from_pretrained("elutherai/gpt-neo-1.3B", tp=4)
```

I implemented megatron friendly model internally like 

```python

@classmethod
def from_yaml(
    cls,
    cfg_path: str,
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
    tp: int = None,
    pp: int = None,
):
    """
    Create model from yaml config file

    Args:
        cfg_path: path of configurations
        tensor_model_parallel_size: tensor model parallel world size
        pipeline_model_parallel_size: pipeline model parallel world size
        tp (int): equivalent with `tensor_model_parallel_size`
        pp (int): equivalent with `pipeline_model_parallel_size`
    """

    if tp is not None:
        assert tensor_model_parallel_size == 1, (
            "you can't use param `tensor_model_parallel_size` and `tp` at the same time. "
            "they are equivalent. so please use one of them."
        )
        tensor_model_parallel_size = tp

    if pp is not None:
        assert pipeline_model_parallel_size == 1, (
            "you can't use param `pipeline_model_parallel_size` and `pp` at the same time. "
            "they are equivalent. so please use one of them."
        )
        pipeline_model_parallel_size = pp
```

I totally agree, that this is a much better way to proceed.

@sgugger, is it ok if we change the initial proposal and add TP to the normal model classes? As we continued discussing this and based on my experience with trying to add PP to transformers it'll be a huge amount of work to do it for all models, and so it's very likely many models will never get it. And since TP requires no changes to the models then there is no reason to make it difficult on users and maintainers to fork the model for that feature to work.

And we believe just having TP+DP will already be a great boon to the scalability of the models (if Deepspeed ZeRO doesn't already address this for whatever reason).

For PP new classes will be needed 100%.

Thank you.

As long as the changes are minimal, no objection from my side. I agree it makes much more sense to get that out if it's faster and deliver the PP later on.

the problem is the 'parallelize()' method, the API for layerwise naive parallelism in GPT2 and T5. Do you agree to remove this method? The megatron-friendly TP + PP cannot handle it that way. This is because in the case of PP, parallelization occurs at the time of model creation. That's why I let `from_pretrained` takes the tp and pp sizes as input.

> I'm thinking of an API like this.
> 
> ```python
> from transformers import GPTNeoModel
> 
> model = GPTNeoModel.from_pretrained("elutherai/gpt-neo-1.3B", tensor_model_parallel_size=4)
> 
> or
> 
> model = GPTNeoModel.from_pretrained("elutherai/gpt-neo-1.3B", tp=4)
> ```

I think `transformers` tends to go with more spelled out args, but not too too long, so perhaps `tensor_parallel_size=4`

> the problem is the 'parallelize()' method, the naive parallelism (layer-wise) implementation. Do you agree to remove this method? The megatron-friendly TP + PP cannot handle it that way. This is because in the case of PP, parallelization occurs at the time of model creation. That's why I let from_pretrained take the tp and pp sizes as input.

The naive PP is experimental:

https://github.com/huggingface/transformers/blob/50c746eeb71f7b8f95a264b09249c9555cdd2e17/src/transformers/models/gpt2/modeling_gpt2.py#L527-L529

but we shouldn't remove it until we replace it with real PP, because users actively use the naive PP at the moment.

That's why we proposed to work on NeoGPT first so that it's easier to take time and not need to have that older code interfere.


@stas00 

> I think `transformers` tends to go with more spelled out args, but not too too long, so perhaps `tensor_parallel_size=4`

So I made it support both variables (long name and short name). not good?

> but we shouldn't remove it until we replace it with real PP, because users actively use the naive PP at the moment. That's why we proposed to work on NeoGPT first so that it's easier to take time and not need to have that older code interfere.

I totally agree with you. Let's start from GPTNeo.

---

The second thing to discuss is the embedding layer. When I implemented parallelformers, I didn't actually parallelize the embedding layer. In this case, the embedding layer is copied to all GPUs. Therefore, it is memory inefficient. But in fact we can apply `VocabParallelEmbedding` and `VocabParallelCrossEntropy`. (However, we should not use the original CrossEntropy in this case) we also need to decide whether or not to add `VocabParallelEmbedding` to the transforemrs-friendly TP. 

I didn't tell you guys, but I actually experimented little by little. I already figured out that I can do `VocabParallelEmbedding` internally with transformers-friendly TPs.

> @stas00
> 
> > I think `transformers` tends to go with more spelled out args, but not too too long, so perhaps `tensor_parallel_size=4`
> 
> So I made it support both variables (long name and short name). not good?

At the moment I don't recall `transformers` using shortcut aliases for arg names, so probably just having `tensor_parallel_size` is fine. (no need to repeat "`model_`" as the shorter name I proposed is not ambiguous)

> The second thing to discuss is the embedding layer. When I implemented parallelformers, I didn't actually parallelize the embedding layer. In this case, the embedding layer is copied to all GPUs. Therefore, it is memory inefficient. But in fact we can apply `VocabParallelEmbedding` and `VocabParallelCrossEntropy`. (However, we should not use the original CrossEntropy in this case) we also need to decide whether or not to add `VocabParallelEmbedding` to the transformers-friendly TP.

Was CrossEntropy the reason for not doing it in the first place in parallelformers? I guess the integration will allow to overcome this then if I understood your comment correctly.

But otherwise by all means let's make TP as efficient as possible.


1. I like the name `tensor_parallel_size` more, but I named it `tensor_model_parallel_size` because I wanted to follow the Megatron-LM nomenclature. In fact, if we input the mpu to DeepSpeed, methods such as `mpu.XXX_model_parallel_rank()` are called inside it. Therefore, it is better to unify the names.

2. Since `parallelformers` is inference only toolkit, there was no reason to worry about CrossEntropy. The reason I didn't do it at the time was because it was a bit complicated. (But it's not difficult.)

How about implementing it with options first? 

`from_pretrained(tensor_model_parallel_size=4, embedding_parallelism=True)`

> 1. I like the name `tensor_parallel_size` more, but I named it `tensor_model_parallel_size` because I wanted to follow the Megatron-LM nomenclature. In fact, if we input the mpu to DeepSpeed, methods such as `mpu.XXX_model_parallel_rank()` are called inside it.

Ah, ok, we can use `tensor_model_parallel_size` then to make things easier to cross-code. May be then add a note at why this particular name has been chosen.

> 2. Since `parallelformers` are inference only in the first place, there was no reason to worry about CrossEntropy. The reason I didn't do it at the time was because it was a bit complicated. (But it's not difficult.)

Ah, right, I forgot that `parallelformers` was intended for inference only in the first place. Yes, so what you proposed is a good idea.

> How about implementing it with options first?
> 
> `from_pretrained(tensor_model_parallel_size=4, embedding_parallelism=True)`

Is there a technical reason for not always doing the latter?

Because of `VocabParallelCrossEntropy`. the user should be able to use a loss function other than CrossEntropy by using the Transformers model. (RMS, Center Loss, Large-margin softmax, ...) With VocabParallelEmbedding, the Loss function should handle this appropriately. You can check this https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/cross_entropy.py 

So I thought the default value of embedding_parallelism as false and turning it on when the user wants to.



Thank you for the explanation, Hyunwooongko.

Then yes we need that arg. Should the default be `False` then, so the priority is for the user code to work out of the box and we document `embedding_parallelism=True` as an optimization?

Further, `embedding` is ambiguous since we have different types, should we say explicitly `word_embed_parallelism`?

Oh I was wrong. Only tying the output embeddings have problems with the loss function. I checked and it doesn't matter since neither gpt2 nor gpt neo are tying output embeddings.

In most cases, we don't need to worry about the loss function. Therefore, I will implement embedding parallelism works everytime so this option is unnecessary. and users do not need to worry about it. If I find a model that tying input and output embeddings without an lm head later, I will think about it then. 

But maybe Meg-DS and GPT NeoX use embedsing tying. So this option will be needed in the future.

If I'm not mistaken many models have input and output embeddings tied.

Hi all, I helped implement pipeline parallelism in Megatron (and was also one of the lead authors on the PipeDream project). Happy to answer any questions.

I had a question too: what is the current plan for the new PP-friendly model classes? What is going into these, and how will they be different from the vanilla model classes?

Thanks!

> Hi all, I helped implement pipeline parallelism in Megatron (and was also one of the lead authors on the PipeDream project). Happy to answer any questions.

Thank you for joining in and offering to support this endeavour, Deepak!

Have you looked at the recent version of the PP in the core pytorch? I tried to suggest to make the API much more flexible over the last spring - at least wrt inputs and outputs, which is very limited in most current PP implementations. so the new API is much more encouraging. You can even pass non-tensor inputs/outputs.

They have some other interesting tech in there. For example stashes to which you can push / pop at different stages, which could for example help pass around complex structures.

> I had a question too: what is the current plan for the new PP-friendly model classes? What is going into these, and how will they be different from the vanilla model classes?

The main obstacles to making HF models PP-friendly are:
1. variety of complex inputs: like a tuple of a tuple of tensors, inputs that aren't tensors, inputs that are tensors but aren't of first dimension of batch size, etc. Most of the variables are very optional and are there to support research.
2. models weren't written with `nn.Sequential` in mind and are very difficult to make into such. for example some models have conditionals on running the encoder or decoder stage or not, which is tricky for `nn.Sequential`

So based on my earlier attempts to implement PP in `transformers` we have to fork the existing models, strip down all the unnecessary features, convert to `nn.Sequential` while adjusting the inputs/outputs to work with it.

As I suggested above pytorch core's new PP API should make this work easier as it's more flexible. But of course, there are other options, more on that later.

The other approach is to start from scratch and build the model with PP in mind from the ground up using Megatron and Deepspeed as a reference, and then to try to adjust the outcome to re-use as much as possible from the current `transformers` model arsenal. The goal is to avoid maintaining 2 separate code bases. I think building from scratch would be preferable to @hyunwoongko - but then we have only a few models w/ PP to borrow from (gpt2, bert - not even t5). Remember in `transformers` we have some 50+ models. Some are slightly different, others are quite different from each other.

So these are the 2 ways I have in my mind.

I'm very open to hearing any other propositions. The key need is the ability to replicate the solution to several dozen of models.

And of course, the other essential question is which PP API to use:
1. write our own / borrow from Megatron
2. Deepspeed
3. Pytorch core (will probably require pt-1.9 or even pt-1.10 for some of the recent features, but it should be no problem)
4. I think FairScale has an API as well, but they have been upstreaming it into the pytorch core, so it's probably the best to rely on the latter.

Beside the ease of use, we also want to make sure that the API allows for the most efficient incarnation of the PP tech - since there are quite a few of them. The goal is to minimize the idling bubble.

BTW, for those who perhaps are new to the topic I wrote this doc: https://huggingface.co/transformers/parallelism.html So that you can quickly understand what we are talking about.

@sgugger @stas00 

We need to name the new class. Do you prefer `GPT2Parallel` or `ParallelGPT2` or `GPT23D`? Or any other good names? I think `GPT23D` is very weird. It looks like 23-dimensional parallelism and names like `GPT2Parallel` are also weird. `GPT2ParallelForSequenceClassification` or `GPT2ParallelWithLMHead` looks like Parallel is for SequenceClassification or Parallel is with LMHead. But putting `Parallel` in the prefix makes everything good. (e.g. `ParallelGPT2ForSequenceClassification` or `ParallelGPT2WithLMHead`) If later extended to TF or Flax, naming such as `ParallelTFGPT2` is also possible.

The initially proposed 3D addition is awkward when it becomes `GPT23DForSequenceClassification`

`ParallelGPT2ForSequenceClassification` rings nicely to me.

Another alternative is to use a postfix: `GPT2ForSequenceClassificationParallel`

> The initially proposed 3D addition is awkward when it becomes GPT23DForSequenceClassification

Well... It looks like GPT twenty three lol

Well, the name doesn't really matter. But we have to decide. would you like to vote? What's the best way?

Let's wait for @sgugger to follow up.

We can vote for the appendix variations, but the main structure (should it be prefix/postfix/infix) is up to Sylvain as he is overseeing the big structure. 

I like the `Parallel` prefix for those new models.

> Have you looked at the recent version of the PP in the core pytorch? I tried to suggest to make the API much more flexible over the last spring - at least wrt inputs and outputs, which is very limited in most current PP implementations. so the new API is much more encouraging. You can even pass non-tensor inputs/outputs.

I have not, but I will try to soon. I agree that flexibility in terms of the number of input and output tensors (as well as types) is good.

> The main obstacles to making HF models PP-friendly are:

This makes sense to me. I agree that having a sanitized model with far fewer optional arguments is nice and easier to support.

I think one way to go is to have the `Parallel*` classes just inherit (or borrow implementations) from the corresponding existing classes (but with the `forward()` method not supporting all the optional arguments). This will ensure that the model implementation really only lives in a single place, but the necessary if guards and other code is elsewhere in the `Parallel*` class implementation. The nice thing about these transformer models is that they are pretty repetitive, so the amount of PP-specific code in the guts of the model implementation is not a lot (e.g., don't run through the embedding layer if the current stage is not the first one, etc.).

For models like T5 with an encoder and decoder, it is important to think through how tensors should be passed through the different stages (e.g., the `encoder_hidden_state` is an input to every decoder stage).

Down the road, you might want to also support interleaved schedules (that trade off a smaller pipeline bubble size for more communication). This will require different execution schedules, so it is also perhaps worthwhile thinking through how to specify these in an easy way.

HF Transformers doesn't use inheritance in models to help the readers of the model to understand it better, so we will either copy-and-strip or we will build ground-up and then copy what we can to match the original class.

> For models like T5 with an encoder and decoder, it is important to think through how tensors should be passed through the different stages (e.g., the encoder_hidden_state is an input to every decoder stage).

How do you propose to deal with conditional runs of encoder (T5)? I see Megatron-LM added T5 but last I checked it didn't support PP.

> Down the road, you might want to also support interleaved schedules (that trade off a smaller pipeline bubble size for more communication). This will require different execution schedules, so it is also perhaps worthwhile thinking through how to specify these in an easy way.

I thought that was exactly the question of choosing the right PP framework/API, since some of those support interleaved schedule and others don't. It's best to delegate such things to the external API I'd think.

I realize that Megatron-LM's approach is different since it doesn't use any API but builds its own.

I made first draft PR for this: https://github.com/huggingface/transformers/pull/13726

> HF Transformers doesn't use inheritance in models to help the readers of the model to understand it better.

Fair enough. I guess there will be some code duplication then.

> How do you propose to deal with conditional runs of encoder (T5)? I see Megatron-LM added T5 but last I checked it didn't support PP.

What do you mean by "conditional run of the encoder"? I have looked at pipeline parallelism with T5, but we always ran inputs through the encoder (and then as I mentioned above, passed the `encoder_hidden_state` through the pipeline stages with decoder layers).

> It's best to delegate such things to the external API I'd think.

I agree that this would be ideal. But unfortunately, the pipeline-parallelism schedule used is pretty problem- and hardware-dependent, so it might make sense to expose a couple of different options to users (especially since `transformers` supports so many different models with different computation characteristics). Even better would be a way to allow a user to specify their own new schedule if they wanted, but this can almost definitely be tabled to later.

Additionally, I don't think any pipeline parallelism API out there is robust and has enough features to be truly useful. For example, Torch's PP support seems pretty nice, but only currently supports an all-forward, all-backward schedule that has high memory footprint; 1F1B is strictly better than this (Section 2.2.1 in https://arxiv.org/pdf/2104.04473.pdf). The all-forward, all-backward schedule won't work super well for really large models.

> > How do you propose to deal with conditional runs of encoder (T5)? I see Megatron-LM added T5 but last I checked it didn't support PP.
> 
> What do you mean by "conditional run of the encoder"? I have looked at pipeline parallelism with T5, but we always ran inputs through the encoder (and then as I mentioned above, passed the `encoder_hidden_state` through the pipeline stages with decoder layers).

It runs the encoder only on the first pass, and not afterwards. You can see the conditional here:
https://github.com/huggingface/transformers/blob/469b80d4e7f9d0ca9411d77845600839e5edf113/src/transformers/models/t5/modeling_t5.py#L1367-L1376

So my question is how to build `nn.Sequential` when half of it is conditional.

> > It's best to delegate such things to the external API I'd think.
> 
> I agree that this would be ideal. But unfortunately, the pipeline-parallelism schedule used is pretty problem- and hardware-dependent, so it might make sense to expose a couple of different options to users (especially since `transformers` supports so many different models with different computation characteristics). Even better would be a way to allow a user to specify their own new schedule if they wanted, but this can almost definitely be tabled to later.

Do you mean that once the models are converted to a straightforward `nn.Sequential` then it could be fed to a variety of PP APIs w/o altering the model itself? 

How would that work? e.g. Deepspeed's PP uses all kinds of special APIs for Tied layers (e.g. `TiedSpec` if I remember the name correctly) and other features that they model has to call explicitly. So it's far from being a generic plug-n-play.

Perhaps pytorch's PP API is slightly more so.

> Additionally, I don't think any pipeline parallelism API out there is robust and has enough features to be truly useful. For example, Torch's PP support seems pretty nice, but only currently supports an all-forward, all-backward schedule that has high memory footprint; 1F1B is strictly better than this (Section 2.2.1 in https://arxiv.org/pdf/2104.04473.pdf). The all-forward, all-backward schedule won't work super well for really large models.

Doesn't Deepspeed PP do 1F1B as well?

You're making an excellent point about torch's PP not supporting progressive PP protocols. Definitely need to inquire about them supporting interleaved PP. I wonder if faiscale has been working on that. 

@hyunwoongko, what's your take - use an API, and which one you resonate the most with, or develop our own, or borrow an internal implementation from Megatron-LM? 

Perhaps we should prepare a table of pros and cons for the different approaches. But somehow I feel you already have something you feel is the best in mind.


> You're making an excellent point about torch's PP not supporting progressive PP protocols. Definitely need to inquire about them supporting interleaved PP. I wonder if faiscale has been working on that.

@pritamdamania, if I may ask - do you have plans to support the interleaved PP protocol in pytorch?

We are having a discussion on which PP framework we should use in `transformers`. As you know I favour pytorch core because you made it much more user friendly than most other PP frameworks I have seen, but as Deepak says above it may have trouble with huge models because it uses all-forward, all-backward schedule. And we happen to work a lot with huge models as of recent (currently using Megatron-Deepspeed for that). Or perhaps my notion is outdated and the interleaved PP is in works already in pytorch?

Thank you!


> @hyunwoongko, what's your take - use an API, and which one you resonate the most with, or develop our own, or borrow an internal implementation from Megatron-LM?

my plan is DeepSpeed pipeline module. because we should consider ZeRO. Note that my implementation is a variant of Megatron-DeepSpeed.

plus) DeepSpeed PP is based on 1F1B https://github.com/microsoft/DeepSpeed/issues/1110

So Deepspeed PP with ZeRO-1, correct?

From what I understand while ZeRO-2/3 could technically work, they won't give any performance improvements over ZeRO-1.

But the key feature is that we could easily turn off PP and enable Z2/3 + offload for inference, which is why we use Megatron-Deepspeed for the BigScience and not just Megatron-LM. That is we use DS PP + Z1.


Currently, DeepSpeed PP and ZeRO 2 and 3 are incompatible. If users don't want to use PP, TP + ZeRO DP is enough. However, the reason we want to provide `ParallelGPT2` with DeepSpeed PP is because there are many other things. (PP, Kernel fusion, Sparse attention, Activation checkpoint offloading ...). I think DeepSpeed is the best API to provide all these features and I think it is recommended to unify it with one toolkit as much as possible.

> It runs the encoder only on the first pass, and not afterwards. You can see the conditional here.

I believe this is only on the first pass _during inference_. You would run it every time for training.


> Do you mean that once the models are converted to a straightforward nn.Sequential then it could be fed to a variety of PP APIs w/o altering the model itself? How would that work?

The only real things we need to know are: a) What computation should run on each "virtual stage" (and how these stages are mapped to actual GPUs)? b) How should tensors be routed between stages? Using this, you should in theory be able to implement any pipeline schedule that takes care of forward and backward passes.

Things like weight tying can be thought of as postprocessing steps at the end of these forward and backward steps (and before the optimizer step runs). For example, in Megatron, we perform an all-reduce of the gradients of all copies of the embedding layers after completing the forward and backward passes in a given batch (and the logic here is the same regardless of the schedule used for forward and backward passes).

I think names such as GPT2Model3D is better. ParallelGPT2 is also nice, but it seems like an ambiguous name in that the existing GPT2 is capable of tensor model parallelism.

Additionally, on slack we have now started discussing a totally new direction for PP and that's where we don't touch the model and get it parallelized automatically.

For example see  SageMaker's PP overview https://aws.amazon.com/blogs/aws/amazon-sagemaker-simplifies-training-deep-learning-models-with-billions-of-parameters/ Note that it now doesn't even mention `nn.Sequential ` (it used to back in the spring) - it now totally automates the process. Some months back it preferred nn.Sequential but didn't require it - now it doesn't even mention it. Unfortunately, there is no disclosure on how they do it.  I think we should be able to do something similar.

The closest approach to doing it that is publicly disclosed is FlexFlow https://huggingface.co/transformers/parallelism.html#flexflow 
https://github.com/flexflow/FlexFlow


> @pritamdamania, if I may ask - do you have plans to support the interleaved PP protocol in pytorch?

@stas00 I'm assuming you are referring to Interleaved Schedule mentioned here: https://developer.nvidia.com/blog/scaling-language-model-training-to-a-trillion-parameters-using-megatron/. As you maybe aware there is a lot of research out there in terms of improving pipeline parallel performance on various axes (memory, speed) and each of them have their tradeoffs. It is not feasible to include all those algorithms in Pytorch. Although there are a few options here:

1) If there is an algorithm that is clearly superior than the rest and if there is enough demand for it from the community, we can implement that in Pytorch.
2) We could evaluate making pipeline parallelism in PyTorch extensible so that the community can quickly try out different algorithms on top of a core pipelining framework without having to reimplement the whole algorithm from scratch themselves.

> Or perhaps my notion is outdated and the interleaved PP is in works already in pytorch?

Currently we are not working on interleaved PP, but as I mentioned above we have a couple of options in terms of how we can enable this.

@deepakn94 Would love to get your thoughts on this regarding the two options I mentioned above since a lot of your research work is in this area :) My initial feeling is that there probably isn't one algorithm that would be the best for all use cases and even if that is true today, new research a few months later might make that algorithm obsolete. Do you feel it is valuable to have an extensible pipelining framework in PyTorch where researchers like yourself can quickly try out different algorithms/schedules? :)

Thank you for your follow up, @pritamdamania87!

Yes, Megatron, Deepspeed and Sagemaker all support the interleaved schedule.

> * If there is an algorithm that is clearly superior than the rest and if there is enough demand for it from the community, we can implement that in Pytorch.

I'd defer to @deepakn94 as he has much more experience with the various schedules. 

Perhaps @ShadenSmith has some insights to share as well, as he has built the PP framework in Deepspeed.

(Perhaps we need a PP-creators/users thread where we can share what works the best and how to make the different implementations interchangeable - i.e. creating a standard API).
 
> * We could evaluate making pipeline parallelism in PyTorch extensible so that the community can quickly try out different algorithms on top of a core pipelining framework without having to reimplement the whole algorithm from scratch themselves.

That is definitely the best approach it seems. 

I think for the plethora of HF Transformers users and uses - being able to choose the best schedule would be a great boon to the whole community.

@pritamdamania87, have you by chance contemplated an approach where any model could be made to support PP w/ only minor mods or none using automatic splitting based on the graph? For context please see 3 comments up: https://github.com/huggingface/transformers/issues/13690#issuecomment-934756008

> @pritamdamania87, have you by chance contemplated an approach where any model could be made to support PP w/ only minor mods or none using automatic splitting based on the graph? For context please see 3 comments up: #13690 (comment)

Yes, we have thought about an automated partitioning solution where you don't need to rewrite your model as `nn.Sequential` and you just need pass in an `nn.Module` and the pipelining framework takes care of the rest. One potential idea was to extract the graph using `torch.fx`, inspect the graph and appropriately partition it across devices. cc @wanchaol Was wondering if the design we brainstormed is in a state that we can share it publicly in an RFC?

> One potential idea was to extract the graph using `torch.fx`, inspect the graph and appropriately partition it across devices. 

Yes, this is precisely what we hope could be done in pytorch! Yes, please!

Looking forward to hearing about your research, Pritam.

And I also would like to invite @jiazhihao of [FlexFlow](https://github.com/flexflow/FlexFlow) to comment on this subject matter, since FF has been actively pursuing this direction, among various other scalability dimensions his project has been exploring. I wonder if the efforts could be joined as there is more than one group that could be already having something working. FF's paper is here: https://arxiv.org/abs/1807.05358

On the HF side we have made the required preparations to make our models `torch.fx` ready wrt to symbolic tracing. And @thomasw21 is now actively trying to figure out how to get FF working with HF Transformers. 

@stas00 Thanks for all the feedback! I had a general question reading the discussion on this issue. It looks like we are talking about using TP, PP and also ZeRO. I was curious to know more about how do you plan to leverage all of these parallelism strategies. Are you planning to use one or a combination of them for a particular model and a different strategy for a different model? If so, how are we making this decision? Is it based on some sort of experimental results?

As an example, why not just use something like ZeRO-3 for all of your models? Are there certain use cases you have where ZeRO-3 doesn't fit well and something like TP and PP is preferred? 

> @deepakn94 Would love to get your thoughts on this regarding the two options I mentioned above since a lot of your research work is in this area :) My initial feeling is that there probably isn't one algorithm that would be the best for all use cases and even if that is true today, new research a few months later might make that algorithm obsolete. Do you feel it is valuable to have an extensible pipelining framework in PyTorch where researchers like yourself can quickly try out different algorithms/schedules? :)

I think it's better to provide an easy way for users to specify their own schedule. Of course you could have some initial set of defaults (and perhaps the community can add to these if there's enough interest). I don't think PyTorch should be in the business of trying to keep up with state-of-the-art here.

I am happy to provide input and help brainstorm on what the best way of specifying such schedules is (and perhaps even help write code if useful)!

> As an example, why not just use something like ZeRO-3 for all of your models? Are there certain use cases you have where ZeRO-3 doesn't fit well and something like TP and PP is preferred?

I don't think you should think of ZeRO-3 as a panacea. In particular, I think it's pretty communication heavy, so might not make sense for certain models and hardware deployments (some of the communication time can be hidden with clever overlapping, but not all of it). In such cases, PP + TP might be a better option.

> @stas00 Thanks for all the feedback! I had a general question reading the discussion on this issue. It looks like we are talking about using TP, PP and also ZeRO. I was curious to know more about how do you plan to leverage all of these parallelism strategies. Are you planning to use one or a combination of them for a particular model and a different strategy for a different model? If so, how are we making this decision? Is it based on some sort of experimental results?
> 
> As an example, why not just use something like ZeRO-3 for all of your models? Are there certain use cases you have where ZeRO-3 doesn't fit well and something like TP and PP is preferred?

Yes, of course. So there is a big group of us involved with the BigScience project https://github.com/bigscience-workshop/bigscience

Having found a working ZeRO-3 I was relieved we didn't have to make HF Transformers' models work with PP.

Until we did benchmarks before starting the first BigScience training and discovered that the HPC we were planning to use (JeanZay) had slow internode interconnects, which were further impacted by a shared network. And while on a fast interconnect ZeRO-3 performs on par with TP+PP, it was lagging quite behind in TFLOPS - if you're curious I have the summaries of all experiments here: https://github.com/bigscience-workshop/bigscience/blob/master/experiments/gpt2.md, so of course we ended up choosing Megatron-LM as our framework for its TP+PP which scored much better, and later added ZeRO-1 to it and Deespeed's PP, which has become https://github.com/microsoft/Megatron-DeepSpeed. This allowed us more flexibility, e.g. use of CPU offload for inference/finetuning and most importantly added an additional group of very experienced minds to our mindpool.

The trouble is now we had to figure out a new framework and all the new development had to be done for both Megatron and HF Transformers, we had to figure out how to convert checkpoints, and we are still sorting out some differences between gpt2 implementations as they aren't the same and a model trained with Megatron doesn't currently give the same performance with HF Transformers (I'm in the process of figuring out those differences). In other words lots of wasted time trying to develop a lot more than could have been.

That's why we want to make sure HF Transformers has its own TP ([almost there](https://github.com/huggingface/transformers/pull/13726)) and PP (this discussion), so that when we do research and apply it we don't need to do it twice and so that the models that we train perform identically under HF Transformers for inference and finetuning.

So long story short, it'd be much easier for us to move forward if we had similar capacity to Megatron-LM.

(and just to clarify in case it wasn't clear, Megatron-LM's team has been amazingly supportive of the BigScience project and they have an awesome framework, and the same goes for Deepspeed and their amazing support. We just need to have in-house tools to do the same.)

Please let me know if I have addressed your question to why we still want TP+PP.





> Yes, we have thought about an automated partitioning solution where you don't need to rewrite your model as nn.Sequential and you just need pass in an nn.Module and the pipelining framework takes care of the rest. 

Another approach that AWS Sagemaker's model parallelism employs (https://aws.amazon.com/blogs/aws/amazon-sagemaker-simplifies-training-deep-learning-models-with-billions-of-parameters/) is to trace the model on the level of submodules so the structure represents the call tree (i.e. like `model.modules()` but ordered at each level and annotated with input sizes). Then it swaps out modules for the remotely executed versions (for the pipeline).

You can poke around the source code by downloading the `smdistributed_modelparallel` wheel here: https://github.com/aws/deep-learning-containers/blob/5ecc625683887edd54c28431ac3cf6b95f30e84d/pytorch/training/docker/1.9/py3/cu111/Dockerfile.gpu#L51 . The relevant code is in `modelparallel/torch/{patches/tracing.py,module_partition.py}`.

The upside of this approach is that it doesn't require any model code changes. The downside is that the granularity is fixed on how model code was written, but in practice people write models in pretty granular form anyway.

@stas00 @deepakn94 Thanks for all of your feedback, this is very valuable for the PyTorch team!


> > @pritamdamania87, have you by chance contemplated an approach where any model could be made to support PP w/ only minor mods or none using automatic splitting based on the graph? For context please see 3 comments up: #13690 (comment)
> 
> Yes, we have thought about an automated partitioning solution where you don't need to rewrite your model as `nn.Sequential` and you just need pass in an `nn.Module` and the pipelining framework takes care of the rest. One potential idea was to extract the graph using `torch.fx`, inspect the graph and appropriately partition it across devices. cc @wanchaol Was wondering if the design we brainstormed is in a state that we can share it publicly in an RFC?

Let me refresh the RFC to address the latest comments in the coming days, and will share it to the public soon. 

> > Yes, we have thought about an automated partitioning solution where you don't need to rewrite your model as nn.Sequential and you just need pass in an nn.Module and the pipelining framework takes care of the rest.
> 
> Another approach that AWS Sagemaker's model parallelism employs (https://aws.amazon.com/blogs/aws/amazon-sagemaker-simplifies-training-deep-learning-models-with-billions-of-parameters/) is to trace the model on the level of submodules so the structure represents the call tree (i.e. like `model.modules()` but ordered at each level and annotated with input sizes). Then it swaps out modules for the remotely executed versions (for the pipeline).

@dzhulgakov I think in theory we can run `torch.fx` in order to swap out the to swap "patterns" for new blocks. I'm slightly unclear on how one can define the pipeline flow using only the graph and input/output sizes? After reading the blogpost, it seems one still needs to run dummy tensors to profile the job. 

> Amazon SageMaker runs an initial profiling job on your behalf in order to analyze the compute and memory requirements of your model. This information is then fed to a partitioning algorithm which decides how to split the model and how to map model partitions to GPUs, while minimizing communication.

Hi @pritamdamania87 @wanchaol , just wanted to check in and see if the RFC for pipeline parallelism has been made public? Thanks!

@deepakn94 Thanks for checking in, @jamesr66a was looking into this recently. @jamesr66a Would be great to share the RFC for pipeline parallelism with Deepak once we have something that we can share publicly. His feedback would be very valuable to ensure what we build here is valuable to the research community.

Hi folks, I think there are actually two RFC's in flight right now:

* @wanchaol has an RFC for pipeline parallelism front-end to better automate partitioning of models using techniques like `torch.fx` graph extraction
* I have an RFC for adding support for cross-host support to the `torch.distributed` pipeline parallel APIs, fleshing out the execution engine (e.g. adding support for 1F1B or programmable schedules), and further trying to drive coherence across the different types of parallelism (e.g. PP, TP, DP)

I can work on polishing my RFC and publishing it by, say, end of next week. Does that sound reasonable?

Also note the release of a new paper: [Varuna: Scalable, Low-cost Training of Massive Deep Learning Models](https://arxiv.org/abs/2111.04007) and its code base https://github.com/microsoft/varuna.

The relevant keys to this discussion:

1. is that it supposedly turns any model into PP on the fly - no changes required by the user. edit: I may have misunderstood the paper, I see now it requires manual inserting of cutpoints in the code: https://github.com/microsoft/varuna#cutpoint-demarcation), so this is not on the fly. I thought it said it can do it automatically)
2. it sports a superior PP scheduling algorithm which allows it to thrive in slow network environment


@jamesr66a That sounds great to me!

@wanchaol Would love to see your RFC as well!

@deepakn94, shared elsewhere a new relevant to our discussion paper:
**Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training**
Can Karakus, Rahul Huilgol, Fei Wu, Anirudh Subramanian, Cade Daniel, Derya Cavdar, Teng Xu, Haohan Chen, Arash Rahnama, Luis Quintela
https://arxiv.org/abs/2111.05972

Hi folks, as promised I have just PRed the RFCs mentioned above: https://github.com/pytorch/rfcs/pull/32

* RFC-0020 Pipeline Parallelism Strategic Plan
* RFC-0021 Pipeline Parallelism Technical Approach Proposal
* RFC-0022 Model Partitioning in Pipeline Parallelism Proposal (@wanchaol's RFC)

Please note that these RFCs only represent our current thinking given the information we have, and we encourage sharing feedback on the contents. We are open to revising our plans, particularly with respect to technical approach. We are considering between:

* Quickly delivering an API with a more constrained interface (i.e. `nn.Sequential`) to start design and implementation iteration quickly
* Deferring delivery of an API to explore techniques that will generalize the front-end of the pipeline parallelism API for better UX and applicability.

Your feedback on this^ question and on the RFCs in general will be greatly appreciated!

@jasmesr66a Hi jasmesr66a.

We made totally new pipeline parallel engine. We will share the news soon.
Thanks. 

@hyunwoongko any update on that?

@jamesr66a We'll release code 12/23. It might be faster ;) 

https://github.com/tunib-ai/oslo

We opened our codebase (framework) ! 

@hyunwoongko 🔥. Let me know if you need contributions to help get other models supported for this.

@KMFODA any way that you can help us is welcome. would you like to contribute to extend the models?


yes! I have a personal preference for adding t5-11b (or 3b) to the library.

I think it's good idea. Let me know if you need my help. I'm happy to help.

Thank you! Yes any help would be much appreciated. It would be good to have a guideline for adding new models that way anyone from the community can add more models later as well. 

Would copying the code for one of the GPT models and tweaking it for T5 be the best way to start? If so are there any modules in the GPT models that are not transferrable to T5?

Hi there, we are developing an auto 3D parallelism distributed software for transformers models.

We plan to integrate this software into `transfomers.trainer` class  and make it use as easily as the original trainer.

We use `torch.fx` and `transformers.utils.fx` for graph extraction, automatically partition the extracted graph and  use pipeline runtime engine from Deepspeed for pipeline parallelism. 

For tensor parallelism we use a config mapping to support `megatron.mpu.layers` in transfomers models automatically.

A prototype version is finished now. Theoretically, any fx traceable model could be run in 3D parallelism. More models are under testing and we are making this software more functional intensively. We will open source it very soon. 

Any advice will be appreciate~

Hey @lucasleesw - this sounds awesome! A team of us here at HuggingFace (@stas00 @thomasw21) have been slowly starting to understand the best path forward to 3D Parallelism (right now, we're looking at the Oslo framework linked above).

Do you have code you can point us to for your implementation, and after we've had time to go through it, maybe we can set up a Slack channel/have a call to talk about the possibility of integration?

Hi @siddk,we are busily polishing our code and completing functions. 
The code will be available within a mouth.

Hello @lucasleesw. We are also very interested in parallelization based on torch.fx. 
I hope to conduct collaborations and seminars when the codebase is opened.

from https://github.com/microsoft/DeepSpeed/pull/1512#discussion_r791312559
@stas00 

Context for others: I recently made a new flexible tensor parallelization engine with torch.fx. This is a discussion about it.

> I'm missing the full context. Do you suggest to have a policy record for each model like in the example you have shown here:

The following is the abstract class

```python
# oslo/parallelism/mpu.py

class LayerInfo:
    @staticmethod
    def base():
        """
        Returns:
            the base transformer block of model

        Examples:
            >>> return BertLayer
        """
        raise NotImplementedError

    @staticmethod
    def attention():
        """
        Returns:
            the last element of attention modules

        Examples:
            >>> return BertAttention
        """
        raise NotImplementedError

    @staticmethod
    def mlp():
        """
        Returns:
            the last element of mlp modules

        Examples:
            >>> return BertOutput
        """
        raise NotImplementedError

    @staticmethod
    def reducing_required():
        """
        Returns:
            arguments that are required reducing

        Examples:
            >>> return ["all_head_size", "num_attention_heads"]
        """
        raise NotImplementedError

```
and you can use like this. (I'll predefine almost all layer info classes)

BERT:

```python
from transformers.models.bert.modeling_bert import BertAttention, BertOutput, BertLayer
from oslo.parallelism.mpu import LayerInfo


class BertLayerInfo(LayerInfo):
    @staticmethod
    def base():
        return BertLayer

    @staticmethod
    def attention():
        return BertAttention

    @staticmethod
    def mlp():
        return BertOutput

    @staticmethod
    def reducing_required():
        return ["all_head_size", "num_attention_heads"]

```

GPT2:

```python
from transformers.models.gpt2.modeling_gpt2 import GPT2Block, GPT2Attention, GPT2MLP

from oslo.parallelism.mpu import LayerInfo


class GPT2LayerInfo(LayerInfo):
    @staticmethod
    def base():
        return GPT2Block

    @staticmethod
    def attention():
        return GPT2Attention

    @staticmethod
    def mlp():
        return GPT2MLP

    @staticmethod
    def reducing_required():
        return ["embed_dim", "split_size", "num_heads"]
```
The reason I created these classes is so that I can't modify the code inside the transformers because oslo is an external library now. but It would be easier if we can define them like `config_class` in XXXPreTrainedModel class like:

```python
class BertPreTrainedModel(PreTrainedModel):
    config_class = BertConfig # already exist.
    last_attention_class = BertAttention
    last_mlp_class = BertOutput
    base_class = BertLayer
    reducing_required = ["all_head_size", "num_attention_heads"]
```
Then, we don't need info class anymore because I can access them like `model.last_attention_class`.
To parallelize without parameter names, I at least needed to know what the attention class and mlp class were.

@lucasleesw could you let me know If you have some better ideas?

And you can use engine like this (distributed launcher is needed)

```python
from oslo import MPU, TensorParallelEngine

model = BertForXXX.from_pretrained(...)
mpu = MPU(tensor_parallel_size=4, pipeline_parallel_size=1)
enigne = TensorParallelEngine(model, BertLayerInfo, mpu)
engine.parallelize()

# rest code is same. training, inference, generation, ...
```

and this engine also supports training.


To deparallelize the model, you can do this:

```python
engine = TensorDeparallelEngine(model, BertLayerInfo)
engine.deparallelize()

# all parameters are deparallelized and are moved to cpu.
# likewise, training, inference, generation can be performed using usual code
```

If you have multiple bass transformer block classes, like Bart, you can do this.

```python
class BartEncoderLayerInfo(LayerInfo):
    @staticmethod
    def base():
        return BartDecoderLayer
    
     ... omitted

class BartDecoderLayerInfo(LayerInfo):
    @staticmethod
    def base():
        return BartEncoderLayer
    
     ... omitted

model = BartForXXX.from_pretrained(...)
mpu = MPU(tensor_parallel_size=4, pipeline_parallel_size=1)

encoder_enigne = TensorParallelEngine(model, BartEncoderLayerInfo, mpu)
encoder_engine.parallelize()

decoder_enigne = TensorParallelEngine(model, BartDecoderLayerInfo, mpu)
decoder_engine.parallelize()
```

I'll also make a pipeline parallelization engine during the week. 3d parallel training/inference/generation/deployent will be possible only with this level of info class.

I'm going to unify the two parallelization engines and make it look like this:
```python
from oslo import ParallelEngine, DeparalleEngine

# 2d parallelization
p_engine = ParallelEngine(model, tensor_parallel_size=4, pipeline_parallel_size=2, infos=the_list_of_info_classes)
p_engine.parallelize()

# 2d deparallelization
d_engine = DeparalleEngine(model, tensor_parallel_size=4, pipeline_parallel_size=2, infos=the_list_of_info_classes)
d_engine.deparallelize()

# if you want to parallelize data dimension, you can use DDP module in torch or ``deepspeed.initialize`` in deepspeed

# the rest code is same. 
# training/generation/inference/deployment can be performed without training code change
```

Since the current version code has not been added in OSLO yet, I will write some example code after deployment so that you can easily apply it. To make the integration easier, I will write some simple monkey-examples for anyone to use and put it in the `tests` folder of OSLO. Does it make sense?

To add a little more, currently, [some elements like "all_head_size", "num_attention_heads" are being managed inside deepspeed engine now](https://github.com/microsoft/DeepSpeed/pull/1512/files#diff-cf74152f0933aa49e65fc12cb0017a40f5b6e2880bd7c01e0d2b26693b6c88e7R400). but I think it's not good structure. so I pushed them towards info class. because they are model specific parts. This is because unexpected bugs may occur if this is managed at the engine level. For example, in T5, `inner_dim` is a variable that must be reduced, but it may not be the case in other models If another model uses a variable named `inner_dim`.

@lucasleew Likewise, if you know of a way to automatically detect these variables, please let me know ;)

But as I said before, I want it to be mounted as a transformers internal solution instead of `import oslo`. what do you think? I also want to move engine class into the transformers, so ultimately I want the following api.

```python
model.from_pretrained(..., tensor_parallel_size=4, pipeline_parallel_size=2)
```
I think it is good to have these APIs internally from a long-term perspective, and if you consider tensorflow and jax, it is not good to depend too much on the outside. As soon as the torch version integration is finished, I will create a tensorflow version, deploy it to oslo, test it, and move it to transformers.

I think your proposal for the models to provide additional meta data is great, especially if multiple frameworks could use that information.

So in your example:
```
class BertPreTrainedModel(PreTrainedModel):
    config_class = BertConfig # already exist.
    last_attention_class = BertAttention
    last_mlp_class = BertOutput
    base_class = BertLayer
    reducing_required = ["all_head_size", "num_attention_heads"]
```

we just need to think about good naming for each of these and then propose a PR to start adding those to the existing classes.

I'm not saying your naming is not good, just saying that others may have a different opinion.

@RezaYazdaniAminabadi, the above map can be used in Deepspeed-Inference as well, right? Are they any other bits that would help your side?

I wonder if this subject-matter is best split off into a separate issue here and discussed there:

[RFC] adding additional meta information for each model

with your example above and defining the target "consumers" - so in this case it's primarily frameworks that want to tap into TP, correct?

Or is some of it going to be used for PP as well? 

The specifics matter for making a well thought out proposal to the maintainers of HF Transformers - to why each config entry is added and who is it going to serve.

And we can start with proposing this to insert into 3-5 popular models and then down the road expand to more if it works out well. 


> I think your proposal for the models to provide additional meta data is great, especially if multiple frameworks could use that information. 

Cool ! Me and Jake will make PR about this.

> [RFC] adding additional meta information for each model

Also good.

> with your example above and defining the target "consumers" - so in this case it's primarily frameworks that want to tap into TP, correct? Or is some of it going to be used for PP as well?

Only for TP. My new PP engine does not require any meta information and modeling code changes.

> And we can start with proposing this to insert into 3-5 popular models and then down the road expand to more if it works out well.

I understand what you are saying. I'll make a new RFC, and let's discuss on there.

> @RezaYazdaniAminabadi, the above map can be used in Deepspeed-Inference as well, right? Are they any other bits that would help your side?

I think it's no. DS-inference and my engine have difference mechanism. 

> > @RezaYazdaniAminabadi, the above map can be used in Deepspeed-Inference as well, right? Are they any other bits that would help your side?
> 
> I think it's no. DS-inference and my engine have difference mechanism.

Right, so then let's think of a few consumers and declare additional metadata for those (if the maintainers will agree of course). That's also another reason why I'm proposing to start small with just a few models in case it doesn't accepted.

So looking at https://github.com/huggingface/transformers/blob/d169b8dadf67ac179c286257349fba319132873a/src/transformers/deepspeed.py#L35-L50 it looks like Reza's version gets all the info automatically for Bert, but needs hints for other models.

And so you, Kevin, won't have any need for these fields in your engine?

https://github.com/huggingface/transformers/blob/d169b8dadf67ac179c286257349fba319132873a/src/transformers/deepspeed.py#L42-L49

I'm asking that since if we are going to add meta-data to the models, so we will then want to fold the Deepspeed-Inference metadata into each model as well.

the other approach is to of course maintain this table separately, as I linked to in the comment above, so for your example, it could just be:

```
# src/transformers/oslo.py
oslo_pp_map = dict(
    bert=dict(
        last_attention_class = BertAttention,
        last_mlp_class = BertOutput,
        base_class = BertLayer,
        reducing_required = ["all_head_size", "num_attention_heads"],
    t5==dict(
        last_attention_class = T5Attention,
        last_mlp_class = T5Output,
        base_class = T5Layer,
        reducing_required = ["foo", "bar"],
    ),
    ...
)
```

which would be much easier to maintain should tomorrow you'd want to add/remove/rename some of these.

the detriment is that they are not following the locality rule so should the model change, this will break. This is of course is very unlikely to happen with stable models, and would be more of the case for recently added ones.

**Perhaps I may recommend to follow this path first, as it'd make things much more flexible for your work**. If you get to the new metadata past the maintainers it'd be very difficult to change this down the road due to backward compatibility requirement.

If you decide to change oslo and make a matching PR to transformers and all that is modified is `src/transformers/oslo.py` that doesn't impact anything else it'll be much much easier for you and all involved. This is based on my experience of maintaining `src/transformers/deepspeed.py` which requires barely any review from maintainers as it's self-contained and doesn't impact the rest of the project.



Hello @stas00.

> it looks like Reza's version gets all the info automatically for Bert, but needs hints for other models. And so you, Kevin, won't have any need for these fields in your engine?

Yes, I was automating the parameter name searching via torch.fx. but for this, I need to know the class name of attention and mlp. I thought defining class is easier than defining parameter names for users and integration point of view. However, for manageability of transformers, I will change the code to use the same tp map with deepspeed. It's not very difficult.

> the other approach is to of course maintain this table separately, as I linked to in the comment above, so for your example, it could just be:

I will try that method (configuring with dict object, not the class) and I'm going to use a map in much the same format as ds-inference.

> If you decide to change oslo and make a matching PR to transformers and all that is modified is src/transformers/oslo.py that doesn't impact anything else it'll be much much easier for you and all involved. This is based on my experience of maintaining src/transformers/deepspeed.py which requires barely any review from maintainers as it's self-contained and doesn't impact the rest of the project.

Yes I don't think it affects anything else. If we use oslo in `from_pretrained` as I suggested, we can use it by simply calling `engine.parallelize()` there. that's all. and it will be used briefly in the trainer class like deepspeed or sagemaker.



@stas00 how about this?

- `src/transformers/utils/model_parallel_utils.py`
   - this file already exists, let's insert tp_map to this file because it's not used independently only by deepspeed now.
- `src/transformers/oslo.py`
   - independant oslo code.
- `src/transformers/deepspeed.py `
  - independant ds code.

Currently transformers has `integrations.py`. It is also desirable to manage them here rather than separate them. There are various integration code such as `ray`.

@siddk @jaketae Let's set the role like this. I will only upload the new engine part to OSLO and I'll open a PR for this to the transformers, so please add the part that makes this available in `trainer`, `from_pretrained` and `save_pretrained (it's for deparallelization)`. In other words, I am in charge of the parallel engine part of oslo, and you guys are in charge of the transformers integration part. I was going to complete the PP engine and write a PR, but it would be better to put it on the transformers side by step rather than that. How about this?

to @RezaYazdaniAminabadi 



1. I've rewritten the code to use the same format map as yours. How about unifying our map in this format? The code on your side probably doesn't need to change much. 
2. And let's keep this map on the transformers side and let me and my colleagues (Sidd and Jake, if they agree) quickly expand models. This is perhaps the best form of extending the ds-inference model we discussed in the beginning.
3. It would also be better for scalability of ds-inference to manage reducing_required variables here as well. Otherwise you will have to modify the inference-engine code whenever a support model is added. This would be very cumbersome.
4. Since `linear_all_reduce` is not the official name, so let's use the official name `row_parallel_linear` from the Megatron-LM paper. 

Do these make sense?

---

to @stas00 

`BertForSeqeuenceClassification`, `BertForMaskedLM` and all the `BertForXXX..` are subclasses of `BertPreTrainedModel`, so it will be useful to detect them using `isinstance(model, cls)`. it would be better than managing all the keys as strings.

And I decided to manage OSLO as an external library and let transformers call it by `import oslo`. That seems like a more desirable structure. If something goes wrong, I'll fix the OSLO part. If all OSLO's code is managed on the transformers side, there is no one to manage it.


```python
src/transformers/utils/model_parallel_utils.py

from ..models.gpt2.modeling_gpt2 import GPT2Block, GPT2PreTrainedModel
from ..models.bert.modeling_bert import BertLayer, BertPreTrainedModel
from ..models.bart.modeling_bart import BartEncoderLayer, BartDecoderLayer, BartPretrainedModel


TENSOR_PARALLEL_MAPPING = {
    BertPreTrainedModel: {
        BertLayer: {
            "row_parallel_linear": ["output.dense"],
            "reducing_required": ["all_head_size", "num_attention_heads"],
        },
    },
    GPT2PreTrainedModel: {
        GPT2Block: {
            "row_parallel_linear": ["c_proj"],
            "reducing_required": ["embed_dim", "num_heads", "split_size"],
        },
    },
    BartPretrainedModel: {
        BartEncoderLayer: {
            "row_parallel_linear": ["out_proj", "fc2"],
            "reducing_required": ["embed_dim", "num_heads"],
        },
        BartDecoderLayer: {
            "row_parallel_linear": ["out_proj", "fc2"],
            "reducing_required": ["embed_dim", "num_heads"],
        },
    },
}

```


---

to @jaketae @siddk 

In addition, the similar maps for almost all models are already defined in [parallelformers](https://github.com/tunib-ai/parallelformers/tree/main/parallelformers/policies). So we can be able to increase the number of models very quickly. It would be nice if Sidd, Jake and I were responsible for extending the model. How about you guys?


@hyunwoongko I think the additional meta data is great! Our code for tensor parallelism is rely on a similar feature dict like yours. Our dict is look like:
```
    'bert':{
        'col_para_list':['query', 'key', 'value', 'intermediate.dense'],
        'row_para_list':['output.dense'],
        'mp_attr_list':['num_attention_heads','all_head_size']
    },
```
Obviously the feature dict is not good for a long-term perspective. 
Additional metadata is a great idea and I am wandering your opinion about modules apart from attention and mlp, how about other module such as linear module and convolution module? 

Hello @lucasleesw. Thanks to answer!

I did something like this:

1. Trace the model using torch.fx to make graph
2. Find the last forwarded linear or conv layer in the Attention module and row-parallelize them.
3. Find the last forwarded linear or conv layer in the MLP module and row-parallelize them.
4. Column-parallelize all linear and conv layers in the all module lists.
5. Parallelize embedding layer and do special processing for tied head module.

The reason this works is there is the premise that the output of the column-parallel linear layer is input to row-parallel linear layer. It worked well for almost all transformers models. Also, almost all attention modules and MLP modules of the transformers do not perform linear again after a row-parallelization required linear or conv layer. Of course there may be exceptions in the future. 

Currently, there are no implementations that use the tracing method like me. So I'm going to return to the previous version for ease of management of the transformers.

In addition, my implementation is rather risky, so it's a good idea to safely create a dict map. Prior to this version of implementation, I and Reza succeeded in omitting column parallel dict map by the following method: 

1. get the name of row parallel linear from user
2. row-parallelize the corresponding layers 
3. finally column-parallelize all the remaining layers of the module list.


In fact, I think that the method of omitting column parallel linear is also somewhat risky. Therefore, @lucasleesw's method, which defines both column parallel linear and row parallel linear names, is the safest. @stas00 Be aware that the more omissions there are, the less reliable the engine is. I think @lucasleesw's dictionary is the safest and best.

Hi @hyunwoongko. Thank you for your reply. I got your idea about linear or conv1d layer in the Attention module or MLP module and I think it is great. 
And indeed this would be kind of risky. For example, how about conv2d layers in `src/transformers/models/vit/modeling_vit.py`?
In our implementation, we built some api to help users built the dict map to handle this exception but we found it is hard to use for now.
I think it would be helpful if the metadata class has some api to let user extend the tensor parallelism methods.


@lucasleesw I totally agree with you. Defining both columns and rows is probably the safest and most extensible way.

@stas00 I initially used the method of defining both column and row parallel parameters, but since the process of defining them is quite difficult, I experimented with many ways to create a simple tensor parallelization map. But the simplification got the more possibility that can makes exceptions. So, like @lucasleesw's method, it would be best to use all three pieces of information: column, row, and mp_param.

We all parallelize in a similar way, and so does sagemaker too. Therefore, it would be convenient if we unify and manage this inside the transformers.

@lucasleesw Omitting column parallel linear and tracining method won't cause any problems in vit. Because embedding is not in the module list. https://github.com/huggingface/transformers/blob/master/src/transformers/models/vit/modeling_vit.py#L146 

They parallelize only the layers inside the base layer module (like BertLayer), not all existing layers. Even so, these simplifications can always make exceptions.

@lucasleesw I'm also wondering about your pp implementation. could you let me know? I used deepspeed pp in the beginning, but now we are implementing the same method with sagemaker pp.

@hyunwoongko You are right, thanks again for your inspiration. 
Our implementation will be available very soon, we look forward for your advice.