[wip] [pipeline parallel] t5 - experiment #2
Author: stas00
Pull request: yes
Labels: name: Pipeline Parallel

name: WIP

The first attempt at t5/pp using pytorch-nightly Pipe https://github.com/huggingface/transformers/pull/9765 was successful to a degree, but at the moment can't be combined with any other Parallel solutions.

All the examples of Pipeline conversion use trivial examples or models that lend easily to being converted to `Sequential`. `transformers` models or at least `t5` doesn't easily lend to this transformation due to complex intertwined logic and a huge number of variables passed around.

The main challenge: In order to build a Pipeline one needs to convert the Module stack into a `Sequential` list.

So in the case of t5, we need to convert this logic:
```
T5ForConditionalGeneration->
  logic
  T5Stack-> 
     logic 
     loop(T5Block, T5Block, T5Block, ...) ->
     logic
  logic
  T5Stack-> 
     logic 
     loop(T5Block, T5Block, T5Block, ...) ->
     logic
  logic
```
into 
```
Pipe(
  Sequential(
    T5ForConditionalGeneration,
    T5ForConditionalGeneration_p1,
    T5Stack,
    T5Stack_p1,
    T5Block,
    T5Block,
    T5Block,
    ...
    T5Stack_p2,
    T5ForConditionalGeneration_p2,
    T5Stack,
    T5Stack_p1,
    T5Block,
    T5Block,
    T5Block,
    ...
    T5Stack_p2,
    T5ForConditionalGeneration_p3,
  )
)
```
I think we don't need to Sequentialize any further beyond T5Block, but we will have to see down the road.

Problems:
1. Can't change the structure of the model because of the pre-trained weights.
2. The inputs/outputs are very complicated because the entry into the Pipeline (first and last stages) can only be a tuple of pure Tensors.
3. The inputs/outputs besides required to be Tensors have to expose first dimension to be batch-dimension since it slices all inputs and restores all outputs on that dimension on the way to/from forward (but only on the very first and last stages of the sequence)

I did successfully implement a t5-pipeline version https://github.com/huggingface/transformers/pull/9765  that uses 2 shorter pipes, as it was natural to convert a loop over `T5Block`s to `Sequential` and it now looks like this 
```
T5ForConditionalGeneration->
  logic
  T5Stack-> Pipe(Sequential(T5Block, T5Block, T5Block))
  logic
  T5Stack-> Pipe(Sequential(T5Block, T5Block, T5Block))
  logic
```
using pytorch pipe in a very painful way overcoming problem n2. But it's doubtful this approach will work with any other 1D Parallel side (e.g. combining with Sharded DDP) - definitely doesn't work with DeepSpeed Zero-DP.

But that implementation won't work with DeepSpeed pipeline - it has to be Sequential from the top-level. Not sure about fairscale yet.

So I'm trying again, this time starting by just trying to Sequentialize the layers while overcoming problem n1. 

If you do look at the code, please ignore everything in the diff but `modeling_t5.py` (and I removed a lot of the model parallel code as it is getting in the way and it won't be needed if we figure out the pipe - since `pipe(chunks=1) == naive vertical MP`, so we get all the complex things that MP currently does for free. But we have to do even more complicated things instead. Naive vertical MP appears to be trivial compared to the changes required to make pipe work.

You can see the process of conversion in this PR, I Sequentialized:
1. the `T5Block`-loop
2. the 2nd half of `T5Stack`, 

now I need to continue breaking up the structure upstream. At this stage there is no Pipe in the code, the 1st main difficulty is to Sequentilize the layers.

If you want to see just how I converted the `T5Block`-loop into Sequential, it is this commit - might be easier to see: https://github.com/huggingface/transformers/pull/9940/commits/4c0ea522157f693bccce80c4cbecc24019186676 The input/output have to be the same because Sequential sends the output of one stage to the input of another. 

If you have some brilliant ideas that I'm perhaps missing at how to easily Sequentialize t5 layers I'm all ears.

@patrickvonplaten, @sgugger, @LysandreJik 

Comments:
This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.

go away bad bot

too long. closing.