Fused optimizer and gradient clipper using apex
Author: slayton58
Pull request: yes
Labels: name: wontfix

Significant (40ms / iter for XLNet squad finetuning) performance increase.

Also adds fused grad clipping, gives further ~30ms / iter saving in the same XLNet squad case. 

Redefines the `AdamW` implementation such that the existing code will be used if apex's multi_tensor_apply code isn't available, and should drop-in speedup all existing scripts using `AdamW`. Also abstracts the gradient clipping (in order to keep run scripts concise and move apex-specific logic into `optimizations.py`)

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=h1) Report
> Merging [#1516](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/1d4d07025635c998acf8c7abab426b013e87206c?src=pr&el=desc) will **decrease** coverage by `0.18%`.
> The diff coverage is `25%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/1516/graphs/tree.svg?width=650&token=9qOlN6Hb1c&height=150&src=pr)](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1516      +/-   ##
==========================================
- Coverage   85.17%   84.98%   -0.19%     
==========================================
  Files          94       94              
  Lines       13920    13953      +33     
==========================================
+ Hits        11856    11858       +2     
- Misses       2064     2095      +31
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [transformers/tests/optimization\_test.py](https://codecov.io/gh/huggingface/transformers/pull/1516/diff?src=pr&el=tree#diff-dHJhbnNmb3JtZXJzL3Rlc3RzL29wdGltaXphdGlvbl90ZXN0LnB5) | `99.02% <100%> (ø)` | :arrow_up: |
| [transformers/optimization.py](https://codecov.io/gh/huggingface/transformers/pull/1516/diff?src=pr&el=tree#diff-dHJhbnNmb3JtZXJzL29wdGltaXphdGlvbi5weQ==) | `75.2% <18.18%> (-21.43%)` | :arrow_down: |
| [transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/1516/diff?src=pr&el=tree#diff-dHJhbnNmb3JtZXJzL2ZpbGVfdXRpbHMucHk=) | `74.17% <0%> (-2.2%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=footer). Last update [1d4d070...a359214](https://codecov.io/gh/huggingface/transformers/pull/1516?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Does this apply that apex's FusedAdam is a fused implementation of AdamW rather than Adam (vanilla)?

It might be nice to first try to import [torch's native AdamW](https://pytorch.org/docs/stable/_modules/torch/optim/adamw.html#AdamW) (from 1.2), and if not available fallback to the transformers implementation. Cf. https://github.com/huggingface/transformers/pull/1593

https://github.com/huggingface/transformers/pull/1516/files#diff-59de7b854fbd60c6ba87f68027a2db36R208 enables AdamW support in the `FusedAdam` optimizer. 

I see you've done the work to check for native PyT implementation & defining if necessary, it should be an easy rebase & resolve if/when these individual PRs get merged.

Oh, my bad. I hadn't noticed this `adam_w_mode=True` in Apex's Adam before. Good to know!

Updated to clip gradient outside of gradient accumulation inner loop as we are trying to do now.

We should update the other training scripts as well (`run_glue` for instance). I think it may be soon time to refactor and gather the common portions of the examples (which are numerous) so we spend less time synchronizing them.

What do you think @LysandreJik?

I agree that some scripts definitely need some refactoring and that having shared pieces of code like that gradient clipping seems like the way to go.

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.