Force einsum to run in fp16
Author: slayton58
Pull request: yes

As noted in the comments, this will force `torch.einsum` to run in fp16 for the squad finetuning task (it should be valid for other tasks, but I haven't verified that) when run with `--fp16_opt_level="O1"` which is the default.

Otherwise, `torch.einsum` is treated as a "promote" operation by `apex.amp`, and if any argument is fp32, all arguments will be cast to fp32, and the answer will return in fp32. This will happen at any point when a parameter is used (XLNet in particular suffers here). Given all uses I've seen for einsum are to express gemm, batched-gemm and transpose, operations we'd normally consider to be safe in fp16, this should be a safe change. From a performance standpoint it allows TensorCore usage which can significantly boost achieved performance.

This change doesn't affect accuracy in my testing, and gives ~20-25% higher throughput on XLNet-based finetuning.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/1513?src=pr&el=h1) Report
> Merging [#1513](https://codecov.io/gh/huggingface/transformers/pull/1513?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/f62f992cf7aa7f1e4eb0d1ef912bd06d26c4dd8c?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `n/a`.

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

```diff
@@           Coverage Diff           @@
##           master    #1513   +/-   ##
=======================================
  Coverage   85.98%   85.98%           
=======================================
  Files          91       91           
  Lines       13579    13579           
=======================================
  Hits        11676    11676           
  Misses       1903     1903
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/1513?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/1513?src=pr&el=footer). Last update [f62f992...4e6a557](https://codecov.io/gh/huggingface/transformers/pull/1513?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Great, thanks a lot for your work on that @slayton58!

Why not adopt this change to other finetuning tasks?  Currently, I only see the code snippet in the squad task.

Einsum is tricky because it can express both tasks that are very likely to be good in fp16 (gemm, batch-gemm) and some that are not (large summations).

It could be adopted for other tasks but it needs to be done task-by-task (with testing) to ensure that no problems are caused. 

> On Jul 12, 2021, at 2:01 AM, Gordon Lee ***@***.***> wrote:
> 
> ﻿
> Why not adopt this change to other finetuning tasks? Currently, I only see the code snippet in the squad task.
> 
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub, or unsubscribe.