text: 
  
  TITLE
  [wip] [pipeline parallel] t5 - experiment
  
  COMMENTS
  21
  
  REACTIONS
  +1: 2
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  This PR is not ready for reviews. 
  
  I'm putting it up primarily for those who want an early preview of a possible Pipeline solution. @PeterAJansen, you wanted to see if you could get it working with 4x 40GB rig and t5-11b. Please give it a try.
  
  -------------------
  
  ## Intention
  
  We want to replace the naive model parallel (MP) implementation with a more efficient pipeline parallel (PP) implementation, which takes advantage of all participating gpus, and not not having one gpu run and the rest idling which is the case with the naive MP. 
  
  To give you a visual from the [GPipe paper](https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html), 
  
  ![mp-pp](https://user-images.githubusercontent.com/10676103/105620296-477e0300-5db0-11eb-8de8-7208be9b4c17.png)
  
  You will find a new argument `chunks`, which is how many pipeline stages you want to add, in the 2nd diagram of the image oabove you can see that `chunks=4`.
  
  So with `chunks=1` you get the naive mp, but it'd be even slower than the naive MP because of the RPC overhead.
  
  ## Overview
  
  Porting t5 to Pipeline Parallelism proved to be a study in hacking, due to the very restrictive original pipeline interface which only allows tensors or tuples of tensors as `input`/`output` arguments in `forward`, and in `transformers` we have a ton of very complex variables to pass to `forward` and return from it.
  
  We are trying to change the Pipeline design to be much more user-friendly: https://github.com/pytorch/pytorch/pull/50693
  
  This implementation tries to take advantage of 2 natural stacks, so I implemented it as 2 pipes:
  ```
  T5ForConditionalGeneration->
     T5Stack(encoder)->Pipe(Sequential([T5StackPipeSegment * 6])
     T5Stack(decoder)->Pipe(Sequential([T5StackPipeSegment * 6])
  ```
  6 for `t5-small`.
  
  Please don't even bother looking at the code, it is one big hack which took many hours to come up with to make the pipeline work, so clearly it is not something very portable or readable.
  
  ## Setup
  
  **important: you need pytorch-nightly to be able to use this.**
  ```
  pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu110/torch_nightly.html -U
  ```
  
  Just create another conda env not to mess up with your normal env, but pt-nightly is a solid peace of software, I use it all the time. here is a quick copy-n-paste of what you will need - just edit the location of the transformers checkout dir.
  
  ```
  conda create -y -n py38-pt18 python=3.8
  conda activate py38-pt18
  pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu110/torch_nightly.html -U
  
  git clone https://github.com/huggingface/transformers
  cd transformers
  gh pr checkout 9765 # or use whatever other method to checkout this PR
  pip install -e .[dev]
  pip install -r examples/_tests_requirements.txt
  ```
  
  Down the road I will look at using also fairscale/deepspeed but for now pytorch is just more accessible and hopefully will be more flexible soon.
  
  ## Deployment: script
  
  You can deploy PP directly via your own trainer/script, e.g. this is what I have been using while developing it:
  
  ```
  from transformers import T5Tokenizer, T5ForConditionalGeneration
  import transformers.models.t5.modeling_t5
  import transformers.utils.logging
  
  transformers.models.t5.modeling_t5.logger.setLevel(transformers.logging.INFO)
  
  mname = "t5-large"
  tokenizer = T5Tokenizer.from_pretrained(mname)
  model = T5ForConditionalGeneration.from_pretrained(mname, return_dict=True)
  model.to("cuda:0")
  model.pipeline_enable(chunks=2, device_map=None)
  
  texts = ["This is good", "This is bad", "This is really bad", "This is fantastic",]
  texts = ["translate English to French: "+x for x in texts]
  batch = tokenizer.prepare_seq2seq_batch(texts, return_tensors="pt")
  batch.to("cuda:0")
  outputs = model.generate(**batch)
  for x in outputs:
      decoded = tokenizer.decode(x, skip_special_tokens=True)
      print(decoded)
  
  model.pipeline_finalize()
  ```
  
  ## Deployment: HF Trainer
  
  But you can also use HF trainer. I tweaked the trainer to activate PP with:
  
  ```
  --pipeline "chunks=4"
  ```
  
  This will let the program do the partitioning for you. But you can control the partitioning manually by passing:
  
  ```
  --pipeline "chunks=4 device_map=0:0-3,1:3-12"
  ```
  
  Here we basically pass the equivalent of a dict  `{0: [0, 1, 2], 1: [3, 4, 5, 6, 7, 8, 9, 10, 11]}` which btw, you can pass in your script as:
  ```
  device_map = {0: [0, 1, 2], 1: [3, 4, 5, 6, 7, 8, 9, 10, 11]}
  model.pipeline_enable(chunks=30, device_map=device_map)
  ```
  The syntax is what you'd pass to `range`, so `device_map=0:0-3,1:3-12" is the same as:
  
  ```
  device_map = {0: list(range(0, 3), 1: list(range(3, 12)}
  ```
  
  the keys are the gpu ids.
  
  The number of layers is at the moment just the depth of the encoder stack, so 12 for t5-base, 6 for t5-small, etc.
  
  Later we should have a different way as well, where we define the desired balance, rather than the specific layers.
  
  Since each `t5` model has a different number of blocks, the easiest way is to first run without the device map and then check the logger output which will show you which device map it's using. Then I recommend to re-balance it so that gpu0 has less layers than the remaining gpus.
  
  ## Benchmarks
  
  example for 2x 24GB cards
  ```
  export BS=160 MODEL=t5-base; rm -r output_dir; PYTHONPATH=src USE_TF=0 examples/seq2seq/run_seq2seq.py \
  --model_name_or_path $MODEL --output_dir output_dir --adam_eps 1e-06 --do_eval --do_train --evaluation_strategy=steps \
  --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 \
  --max_target_length 128 --val_max_target_length 128  --num_train_epochs 1 --overwrite_output_dir \
  --per_device_eval_batch_size $BS --per_device_train_batch_size $BS  --eval_steps 25000  --sortish_sampler  \
  --warmup_steps 50 \
  --task translation_en_to_ro --dataset_name wmt16 --dataset_config ro-en --source_prefix "translate English to Romanian: " \
  --max_train_samples 10 --max_val_samples 10   \
  --pipeline "chunks=4 device_map=0:0-3,1:3-12" --dataloader_num_workers 4
  ```
  
  Performance-wise:
  - prediction speed is terrible - just as bad as the naive MP we have in t5 and others
  - training/eval w/o prediction is slightly slower 20-40% than the baseline with just one process - this is primarily due to data copying and the current quite inefficient implementation due to the Pipeline api restrictions.
  - the key is to find the value for `chunks` so that there is enough in the pipe so that the gpus don't idle, but not too big as performance goes down. But I wasn't able to overcome 50%/gpu utilization, so it's not much more different from the naive implementation - don't know yet why - probably data copying takes most of the overhead.
  - I think on 4 gpus it'd be good to try an experiment and put the encoder stack on gpu 0+1 and decoder on gpu 2+3, instead of copying data between 4 devices as it's happening now - this will require a more complex device map, that I designed for the Bart MP, which has separate encoder and decoder sub-maps.  But then it'd affect the pipeline as half the gpus will surely idle while encoder is running - so not great either. We will have to experiment with real data once I have access to a rig with 4 gpus and see. That's why I don't think this is urgent to work on. But such change would be easy to do. We will have to do it anyway for other models whose stacks aren't necessarily symmetrical. 
  
  Here are some stats on 2x 24GB Titan RTX:
  
  Baseline: (1gpu)
  
  ```
  export BS=64 MODEL=t5-base; rm -r output_dir;  CUDA_VISIBLE_DEVICES=0  PYTHONPATH=src USE_TF=0  \
  examples/seq2seq/run_seq2seq.py --model_name_or_path $MODEL --output_dir output_dir --adam_eps 1e-06  --do_eval \
  --do_train --evaluation_strategy=steps  --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 \
   --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir \
  --per_device_eval_batch_size  $BS --per_device_train_batch_size $BS --eval_steps 25000  --sortish_sampler   \
  --val_max_target_length 128 --warmup_steps 50 --max_train_samples 1000 --max_val_samples 1000 \
  --task translation_en_to_ro --dataset_name wmt16 --dataset_config ro-en --source_prefix "translate English to Romanian: "
  
    train_runtime              = 6.9149
    eval_loss                 =  3.5492
    eval_runtime              =  3.2802
  
  ```
  XXX: need to re-test with rebased code-base
  
  Now with pipeline:
  
  - can run much higher batch-size
  - note, that I'm using a user-provided device map that has more layers on gpu 1, since gpu 0 needs much more RAM
  
  ```
  # device_map=0:0-3,1:3-12 - so splitting 1:4
  # {0: [0, 1, 2], 1: [3, 4, 5, 6, 7, 8, 9, 10, 11]}
  ```
  
  ```
  export BS=160 MODEL=t5-base;  rm -r output_dir;  CUDA_VISIBLE_DEVICES=0,1  PYTHONPATH=src USE_TF=0  \
  examples/seq2seq/run_seq2seq.py --model_name_or_path $MODEL --output_dir output_dir --adam_eps 1e-06  --do_eval \
  --do_train --evaluation_strategy=steps  --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 \
   --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir \
  --per_device_eval_batch_size  $BS --per_device_train_batch_size $BS --eval_steps 25000  --sortish_sampler   \
  --val_max_target_length 128 --warmup_steps 50 --max_train_samples 1000 --max_val_samples 1000 \
  --task translation_en_to_ro --dataset_name wmt16 --dataset_config ro-en --source_prefix "translate English to Romanian: " \
  --pipeline "chunks=4 device_map=0:0-3,1:3-12" --dataloader_num_workers 4
  ```
  XXX: need to re-test with rebased code-base
  
  
  ## Future
  
  I'm also trying to instrument this feature with reporting that will help users to finetune chunks/device_map
  This is the `model.pipeline_finalize()` call. Things I'm thinking that would be useful:
  * [ ] gpu utilization stats (average/peak) - probably need to fire off a thread that samples pynvml gpu utilization, then calculates average + peak
  * [ ] peak memory usage per device report that I added seems to be too low - I think it has to do with pipeline threads - need to sort it out
  
  Any other ideas/requests/needs?
  
  @PeterAJansen, please let me know if you managed to run this on your 4x gpu setup.
  
  Next, I think I'm going to scratch the current implementation and try a new one afresh.
  
  Also this PR should be good enough to try to figure out how to use with DeepSpeed, once I get access to 4 gpus (need at least 4 gpus to do 2D parallelism).
  
  I did warn you not to look at the code.
  
  I also removed big chunks of MP code for now as it was getting in the way with the noise, will restore it when I sorted this all out.