[trainer] deepspeed integration
Author: stas00
Pull request: yes
Labels: name: DeepSpeed

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 

Comments:
Thanks @stas00 for putting this together! I think there might be a few things we can do on the deepspeed side to smooth a few pain points out. Most of us our out of office until the new year but will definitely be taking a close look at this soon and help where we can.

Thank you, @jeffra! 

New year sounds perfect as a time for you to make suggestions if any, but meanwhile I think it's coming along nicely. 

OK, so this PR also introduces a concept of `self.wrapped_model` so that we have less confusion about which is which in the trainer and user code. 

* `self.model` is always transformers model - this is the internal model
* `self.wrapped_model` is a wrapped model - always the most external model which could be  DDP(Transformers Model), DDP(Deepspeed(Transformers Model)), etc.

It's not documented yet, but @sgugger when you get a chance could you please check that it looks correct what I did here https://github.com/huggingface/transformers/pull/9211/commits/1510444399a4a088e6b14076bd857a47c3f2b1eb

Questions: 

1. is it correct that I set it to `None` if there is no wrapped model?
2. would it be better to call it `model_wrapped` - so the two better align side by side during debug or IDE prediction completion engines?
3. I'm not sure where to document this? And we should probably add a public API accessor?
4. We can now probably remove and refactor the following code, as we now have a simpler way to get the internal model -
https://github.com/huggingface/transformers/blob/cbe63949d76efd153a1f389f38fe9ce1287e06b0/src/transformers/trainer.py#L1635-L1651
 or is it used in some deep code where there is no trainer object? In which case this code won't work as it needs double unwrapping under deepspeed - `model.module.module` (since we have DDP too).
 
 I see it's used only in `floating_point_ops` which does have access to `self` so I'm not sure why it was needed in first place. Also in 2 tests, but that could be moved into the tests if need be.
 
If we want a general unwrap function it needs to do it recursively until there is no more `.module`.

Let me think.
For 1, I think the `wrapped_model` should be the model, in this case, just to avoid the inconvenience of testing if `None`.
For 2, I have no strong opinion, so you can pick the version you prefer.
For 3, none of the attributes of the `Trainer` are properly documented yet. This could be added in the main docstring
For 4, yes, absolutely. This was a quick fix that was merged when I didn't get much time to do a nice solution, I though I had removed all uses of that function.

Could you add the `wrapped_model` or `model_wrapped` in a separate PR? This would be easier to follow and not highjack the discussion on the deepspeed integration. We can rebase this when that PR is merged.

> For 1, I think the `wraped_model` should be the model, in this case, just to avoid the inconvenience of testing if `None`.

Then we somewhat lose information - `None` is telling us that notihng is wrapping the model. But I suppose we could achieve the same by `self.model == self.wrapped_mode` - OK, that works!

Thank you for the rest of the answers, @sgugger. Will integrate those and make a separate PR with wrapped_model.

@stas00 Should it be mentioned in some README/documentation that folks can only use DeepSpeed with the PyTorch `Trainer` and not the TF `Trainer`? There's a hard dependency of using `torch.distributed` with the NCCL backend to use DeepSpeed.

Secondly, what is the plan in terms of introducing DeepSpeed in the `transformers` [`setup.py`](https://github.com/huggingface/transformers/blob/master/setup.py) and the PyTorch GPU [`Dockerfile`](https://github.com/huggingface/transformers/blob/master/docker/transformers-pytorch-gpu/Dockerfile)?

While DeepSpeed has a pip installable PyPI package, IIRC it is highly recommended that it be installed from source. Also, in order to use certain features in DeepSpeed such as 1-bit Adam, there are certain special installations to be done that do not come with the PyPI package. Will this PR support every underlying DeepSpeed feature? If not, can the scope of the initial DeepSpeed integration be defined clearly in some README/documentation, while allowing for further iterations in future to enable the utilization of more DeepSpeed features with the `transformers` `Trainer`?

> @stas00 Should it be mentioned in some README/documentation that folks can only use DeepSpeed with the PyTorch Trainer and not the TF Trainer? There's a hard dependency of using torch.distributed with the NCCL backend to use DeepSpeed.

Yes, we should definitely be clear about that. thank you!

At the moment the idea is to put all the ZeRO related docs here: https://github.com/huggingface/transformers/pull/9208 (that PR covers fairscale at the moment)

> Secondly, what is the plan in terms of introducing DeepSpeed in the transformers setup.py 

It'll be up to users to install `deepspeed`, just like it's the case with `fairscale` or any other libraries the `transformers` core doesn't require. Currently if you use `--deepspeed` and you don't have it installed the trainer will assert with a suggestion to install that library.

> and the PyTorch GPU Dockerfile?

I have no idea. I don't see any reason why it can't be included. 

Let's do it in baby steps. First, make the support available, test it out, solve initial issues if any. Then worry about everything else?

> While DeepSpeed has a pip installable PyPI package, IIRC it is highly recommended that it be installed from source. Also, in order to use certain features in DeepSpeed such as 1-bit Adam, there are certain special installations to be done that do not come with the PyPI package. Will this PR support every underlying DeepSpeed feature? If not, can the scope of the initial DeepSpeed integration be defined clearly in some README/documentation, while allowing for further iterations in future to enable the utilization of more DeepSpeed features with the transformers Trainer?

As I have shown in the example of the upcoming fairscale-support doc PR (we are waiting for fairscale to make a new pypi release before we merge it), we will document the same for DeepSpeed and address your questions.  Your comments would be super-helpful for that document, so please save them for when we get to write that document. With your permission I can tag you on that future PR. Thank you.

Wrt to specifics let's see what ends up working out of box and what needs to be polished. I think the main issues will be bugs on the DS side. Otherwise there is a ton of features and I have only been testing a few.

If you feel inspired and are already experienced with DS it'd be awesome if you made a checklist of features and then between you and I, and anybody else who wants to contribute, test those features and check what's supported and report back to DS what is not. Since DS does most of the things on its own, I don't think there will be much to change in `transformers` Trainer once this PR is polished. I can be wrong of course.

**edit**: actually there is no point waiting - I started adding notes into  docs/source/training.rst in this PR. So already added a few of your comments - will need to expand those later.

> how is the optimizer/scheduler creation handled? E.g. how is the fact we don't apply weight decay to some parameters or the proper schedule with the number of training steps handled? I don't see it in the current version of the code since deepspeed is responsible for creating the optimizer and scheduler

It wasn't, but it is now. Please have a look at what I just committed https://github.com/huggingface/transformers/pull/9211/commits/869173f62373a068c2cb932680093e7fd1d4ab78

- schedulers - I was able to remap 2 `constant_with_warmup` + `linear`
- optimizers - I think only Adam is possible at the moment with our cl args - but users can do whatever they want in the config file
- the main logic is that if either `optimizer` or `scheduler` is configured in the config file we keep those and don't override the corresponding parts - if they are missing - we build our own config.
- one thing I could use your help with, @sgugger. To use `linear` I needed `num_training_steps` early in `__init__` before distributed is set up. I tried postponing ds init for the time we have  `num_training_steps` available in `train`, but it was too late - must set up dist early in `__init__` - so I added a hacky `get_num_training_steps` - please have a look if you can think of a better way to handle this chicken-n-egg problem. Thank you!

> Left some comments. Not super happy about the `get_num_training_steps` method, but I don't see a way around it either.

I agree. I tried to move `_init_deepspeed` into a later stage when we already have that figured out cleanly, but there is a chicken and egg problem there.

I have made another commit to only run `get_num_training_steps` if absolutely necessary, and not by default for `--deepspeed`

> Last thing I wonder is if it works okay for checkpointing (saving model/optimizer/scheduler after `save_steps` steps) and then resuming training? I.e., does DeepSpeed gets in the way of saving/reloading a model/optimzier state/scheduler state.

It's on the todo list in the OP, where I'm tracking what needs to be done.  It's just a matter of doing it.


@sgugger, as `_init_deepspeed` was born during this PR's evolution and it continues evolving and growing - are you happy with it in the main file or should it be moved into some other file, say `trainer_integrations.py`?

@sgugger, I'd love some guidance from you to that last point
 
We have `deepspeed.DeepSpeedEngine.save_checkpoint` and `deepspeed.DeepSpeedEngine.load_checkpoint` ([doc](https://deepspeed.readthedocs.io/en/latest/model-checkpointing.html))

So in `_save_checkpoint` I will delegate to `deepspeed.DeepSpeedEngine.save_checkpoint` the 2 parts:
```
        # Save model checkpoint
        # Save optimizer and scheduler
```
and leave the rest untouched, right?

Then I see `_tune_save_checkpoint` - looks very similar to `_save_checkpoint` but it's much simpler - I guess doing the same here?  

`deepspeed.DeepSpeedEngine.save_checkpoint` saves/loads everything in one call - doesn't separate them into model/sched/optim -.

And then on loading - there are quite a few places where this is happening. Any suggestions at how to approach this? Just match any place where there is `torch.load`?

Alternatively, I can just give it a try and then you can comment on what I have missed / did wrong - if that sounds like a better use of your time.

Thank you!

For the checkpointing, `_tune_save_checkpoint` is only there for Ray tune saving of checkpoints, so it's a bit different.

I think it's fine to delegate to DS the saving of model/optimizer/scheduler. For the reloading, I'd go for using `torch.load` when reloading the model at the end of training but use deepspeed for loading optimzier/scheduler (it will also reload the model, but it should match the previous one).

@sgugger, what should we do about tests - I only have a basic test that runs a full train/eval with `finetune_trainer.py` - like we have for `sharded_ddp`.
1. do we want to test all the myriad of different config combinations - probably not since we can't control them anyway and perhaps it's best to do it on demand if and when we find bugs in our integration code.
2. I could move it from `examples` to `tests` but then it'd be much more difficult to do a functional test I think - please correct me if I'm wrong. If there is an easy way then by all means - and `sharded_ddp` tests need to go there too then. we can also deal with this separately after concluding this PR as it's a project of its own.
3. In either case - should we give it a dedicated test file/have integrations dedicated file/leave it for now as is? I think the latter.

The thing is we only pass on the configuration to DS and aren't really doing anything other than really getting out of its way, so there isn't much to test. Perhaps checkpointing should be tested as they will be part of the trainer code, which I think is the last big thing I need to integrate.

Thanks.


>  Could we however move it to `integrations` since the trainer code is already quite long? It only seems to use `self` to access the `args` so we could have its signature be `(args, model)`

It needs `self` for the num of training steps calculations, and I tried to only run it if it's needed (linear scheduler), so I can't pre-calculate and pass it to the `_init_deepspeed` method as it'd be inefficient for code that doesn't need it. I'm waiting for the DS team to follow up on my question about the chicken and the egg - perhaps there will be a way to avoid that hack.

Otherwise I'm in total agreement with you wrt to moving it away and having only the args it needs.

> It needs `self` for the num of training steps calculations, and I tried to only run it if it's needed (linear scheduler), so I can't pre-calculate and pass it to the `_init_deepspeed` method as it'd be inefficient for code that doesn't need it.

This is just a few math operations, I'd prefer we pass it either way to avoid the added complexity in the Trainer for everyone. 

> @sgugger, what should we do about tests

For now, a simple integration test seems enough (I imagine in the multi-GPU tests). It's going to be challenging to properly test all those frameworks that require multi GPU but this is a discussion for another PR. 

> , I'd prefer we pass it either way to avoid the added complexity in the Trainer for everyone.

Sorry I don't know what you mean when you say that. Would you please rephrase that sentence with explicit terms - pass what/where - "it" being what? Thank you.

> Sorry I don't know what you mean when you say that. Would you please rephrase that sentence with explicit terms - pass what/where - "it" being what? Thank you.

I was speaking of the number of training steps, always pass it even if we end up not needing it because the scheduler does without.

> I was speaking of the number of training steps, always pass it even if we end up not needing it because the scheduler does without.

Thank you for clarifying. 

I disagree. I imagine most of the time users will run their own ds_config.json with custom config, which would not use that expensively calculated number - so why introduce a repeated wasted overhead at calculating something that will never be used? (or am I exaggerating and it's not that much of an overhead?)

Why is it so bad if we pass the `trainer` object if it helps to optimize speed? I understand that since we are talking about changing `_init_deepspeed` not to be a trainer method it'd be good to make it so - but there is nothing wrong with it taking a `trainer` object as the first arg, no?

I hope this whole thing might be resolved in a different, much cleaner way if I can only get help on it here https://github.com/microsoft/DeepSpeed/issues/633 - this is a hack we are doing that shouldn't be needed in first place...


Oh, if you prefer sending it `self`, I have no objections. We do this too for some of the other integrations method.

So our training steps calculation hack is no longer needed as @jeffra just happened to add `deepspeed.init_distributed` recently which is what we were missing, and adding which lead to a much cleaner code. yay!

OK, this is pretty much done, just waiting for some feedback from the DeepSpeed team on the best way to handle checkpoints, but the code is in place.

So please kindly send me any final comments/suggestions/requests and let's merge it. Please re-invite old reviewers and add new ones if needed.

Thank you!

deepspeed-0.3.10 has been just released by @jeffra on pypi - I verified that it works - so we are ready to merge this whenever you're happy with it.

it'd be great if you tried running it too, since I think it has been only me running it, so my work is only as good as my environment is and I may not know of other culprits - e.g. I can't test with pytorch < pt-nightly since my card doesn't work with those pytorch versions.

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
```


@sgugger, 

1. While working on docs I discovered that DS does its own gradient clipping (that doc was buried and I didn't see it) so I had to undo the code in the trainer that did that on behalf of DS - just skipping it
2.  I did a major rewrite/expansion of the docs (including the fairscale section) - so please kindly have a look. It's mainly mirroring the config logic in the integration code.
3. In the docs I used consistently Trainer (upcase) to refer to HF trainer. I know you didn't like it when I did that for Issue in a different PR, let me know if you prefer it to be a lowercase trainer.

While this PR is perfectly ready for a final review, I need to wait for https://github.com/microsoft/DeepSpeed/pull/656 to be answered before we can merge this as I'm unsure about their defaults for gradient clipping.

Thank you.


> Went through the documentation and left comments. 

Awesome - thank you - all integrated.

> On the optimizer side, it doesn't seem like DeepSpeed supports AdamW from what you're saying, so we should document the default optimizer is changed at the very beginning of the DeepSpeed section. It does change drastically the value of `weight_decay` to use.

I found a way to use AdamW, thank you for catching that, @sgugger. I documented the nuances.


I think the DeepSpeed team is on vacation, as there is no response since several days.  And since I have no way of talking to anyone there, I have no way of knowing when they will be back. So I will go ahead and merge this so that others can start experimenting and then we can fix whatever needs to be fixed when I get the gradient clipping Issue answered.

Amazing work @stas00 !