text: 
  
  TITLE
  [trainer] deepspeed integration
  
  COMMENTS
  26
  
  REACTIONS
  +1: 1
  -1: 0
  laugh: 0
  hooray: 2
  heart: 6
  rocket: 0
  eyes: 0
  
  BODY
  This PR adds experimental support for Deepspeed <https://github.com/microsoft/deepspeed>, whose main feature is ZeRO covered by the paper [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, by Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He](https://arxiv.org/abs/1910.02054). 
  
  Recently added support for sharded DDP ([fairscale](https://github.com/facebookresearch/fairscale)) also implements parts of ZeRO. Deepspeed implements all of ZeRO.
  
  I haven't experimented enough yet, but it indeed delivers incredible results. 
  
  For example I can get about a 5-8 times bigger batch onto the same hardware as compared to the same code running w/o deepspeed and the speedup is huge too. In the following example I was able to get 4.5x speedup on training, and ~2x on validation/testing:
  
  ```
  # baseline
  export BS=3; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0  python \
  -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path \
  sshleifer/distill-mbart-en-ro-12-4 --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro \
  --do_eval --do_predict --do_train --evaluation_strategy=steps --fp16 --freeze_embeds --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 \
  --predict_with_generate --eval_steps 25000 --save_steps 25000 --sortish_sampler --src_lang en_XX --task translation \
  --test_max_target_length 128 --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 2000 \
  --n_val 2000 --n_test 2000 
  
  2020-12-18 22:31:40 | INFO | __main__ |   train_runtime = 144.9132
  2020-12-18 22:37:10 | INFO | __main__ |   val_runtime = 329.8146
  2020-12-18 22:42:37 | INFO | __main__ |   test_runtime = 326.6212
  
  # deepspeed
  export BS=20; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0 deepspeed  \
  ./finetune_trainer.py --model_name_or_path sshleifer/distill-mbart-en-ro-12-4 --output_dir output_dir --adam_eps 1e-06 \
  --data_dir wmt_en_ro --do_eval --do_predict --do_train --evaluation_strategy=steps  --freeze_embeds --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 \
  --predict_with_generate --eval_steps 25000 --save_steps 25000 --sortish_sampler --src_lang en_XX --task translation \
  --test_max_target_length 128 --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 2000 --n_val 2000 \
  --n_test 2000 --deepspeed  ds_config.json 
  
  2020-12-18 22:51:46 | INFO | __main__ |   train_runtime = 32.6825
  2020-12-18 22:54:47 | INFO | __main__ |   val_runtime = 180.5917
  2020-12-18 22:57:51 | INFO | __main__ |   test_runtime = 183.7731
  
  ```
  
  The bleu eval scores were slightly better than the baseline (~0.5 point higher), but it's not enough to make any conclusions based on a single run.
  
  The cool thing is that deepspeed does everything by itself, even the `--fp16` handling, so really it was all about getting out of its way, thus the main part of the integration was to disable a lot of things the trainer does when `--deepspeed` is enabled.
  
  Note the different invocation pattern. If normally we run distributed as:
  ```
  python -m torch.distributed.launch --nproc_per_node=2 ./program.py args
  ```
  deepspeed performs its own DDP internally, and requires the program to be started with:
  ```
  deepspeed  ./program.py args
  ```
  
  The only thing I'm not sure about with this PR is that deepspeed enables all of its features via a json config file, so I'm not sure where to stash a sample one. I guess I will just add it to the documentation. Currently I put one under `examples/seq2seq/ds_config.json` since that's where the test that needs it lives. 
  
  But once this is merged all interested parties can start experimenting with various features, and it won't impact `transformers` code. They will just need to tweak `ds_config.json`. And we convert many trainer cl args into DS config on the fly.
  
  There surely will be competition betweeen fairscale and deepspeed integrations. So far from the few experiments I did `deepspeed` allows for a bigger batch size than fairscale.
  
  To install deepspeed you can just do `pip install deepspeed` - I'm not sure if all bug fixes are in the release. We can make a request to release a new version when this is merged.
  
  If the build fails I recommend pre-compiling its CUDA extensions (otherwise they get built at run time via PTX) via master:
  ```
  git clone https://github.com/microsoft/deepspeed
  cd deepspeed
  DS_BUILD_OPS=1 pip install --no-cache -v --disable-pip-version-check -e . 2>&1 | tee build.log
  ```
  If you want a faster build add an env var `TORCH_CUDA_ARCH_LIST` with the cuda compute capabilities you need, e.g. I do:
  ```
  TORCH_CUDA_ARCH_LIST="6.1;8.6" DS_BUILD_OPS=1 pip install --no-clean --no-cache -v --disable-pip-version-check -e . 2>&1 | tee build.log
  ```
  
  It was awesome that @sgugger has just added `fairscale` support, so it was much easier to do the same for deepspeed seeing how `fairscale` was integrated, so I'm appreciating the work you have done, Sylvain.
  
  ----------------------
  
  Do try it so we get better testing!
  
  You will need 2+ gpus to use it
  
  First install it:
  ```
  pip install deepspeed
  ```
  
  At the very least do the test:
  ```
  cd examples/seq2seq
  pytest -sv test_finetune_trainer.py -k deepspeed
  ```
  
  Or if you want to fiddle with the normal run, here is what I have been using.
  ```
  cd examples/seq2seq
  wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz
  tar -xzvf wmt_en_ro.tar.gz
  export BS=20; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0 deepspeed --num_gpus=2 ./finetune_trainer.py --model_name_or_path t5-small --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_eval --do_predict --do_train --evaluation_strategy=steps --freeze_embeds --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 --predict_with_generate --eval_steps 25000  --sortish_sampler --task translation_en_to_ro --test_max_target_length 128 --val_max_target_length 128 --warmup_steps 5 --n_train 100 --n_val 100 --n_test 100 --deepspeed ds_config.json --fp16 --save_steps 1
  ```
  
  
  ----------------------
  
  Questions that need to be addressed so that all Trainer features continue to work under deepspeed.
  
  * [ ] a notebook with benchmarks was requested
  
  Probably at a later time, my uneven gpu-sized setup doesn't lend to impressive benchmarks - may be someone will send me another rtx-3090 card ;)
  
  @sgugger, @LysandreJik, @patrickvonplaten 