Separate implementation for Torch-Scriptable BERT model
Author: sbrody18
Pull request: yes
Labels: name: wontfix

<!-- This line specifies which issue to close after the pull request is merged. -->
Fixes #5067 

Separate (re-)implementation of BERT such that it can be used with TorchScript's script() method (not just trace). This allows for better/more reliable handling of internal model logic and removes the requirement for fixed input size, resulting in large speedups when average input size is << than max input size.

Tests duplicate all the ones on the original BERT model, with the only fundamental difference being the test_torchscript* tests, which now use script() rather than trace().




Comments:
Thanks a lot for doing this! It's a great way for us to realize the needed changes to get fully torch-scriptable models. Aside from the tests (we can help fix them once the design is approved), I'd love to see what parts we can reuse from bert (with potential non-harmful modifications) and what parts need to be rewritten because they're not compatible with the rest of our API.

For instance, the change in the embeddings layer is just a type annotation which we could do in bert (it would be a nice addition) and then import that layer. On the other hand, the whole parts with `return_dict` are probably fully incompatible with scripting.

I guess in an ideal world, we would reuse the same internal layers from bert and only change the full models if that is possible.

As you can see in a [previous comment on the thread](https://github.com/huggingface/transformers/issues/5067#issuecomment-662586999) my initial implementation tried to go the minimal-duplication route. I modified the original models to be scriptable, and then had a thin wrapper around them to transform the output into dictionary form.
So basically, you had BertScriptableModel returning a tuple of fixed size, and BertModel who's forward just ran BertScriptableModel and put the output in a dictionary, to keep the interface.
The main issue with that was that the code kept changing. Other than that, it should be doable.


90% of the changes were type annotations, and assertions about Nullity (which would improve the code quality regardless).
The added bonus of the minimal duplication route is that it makes it easier to convert other models that use BERT components, e.g., Albert.

# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6846?src=pr&el=h1) Report
> Merging [#6846](https://codecov.io/gh/huggingface/transformers/pull/6846?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/2de7ee0385bee4134ca894a208fa3a2aaf7d5371?el=desc) will **decrease** coverage by `0.85%`.
> The diff coverage is `18.92%`.

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

```diff
@@            Coverage Diff             @@
##           master    #6846      +/-   ##
==========================================
- Coverage   80.20%   79.35%   -0.86%     
==========================================
  Files         157      158       +1     
  Lines       28734    29257     +523     
==========================================
+ Hits        23047    23216     +169     
- Misses       5687     6041     +354     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6846?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_scriptable\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19zY3JpcHRhYmxlX2JlcnQucHk=) | `18.92% <18.92%> (ø)` | |
| [src/transformers/tokenization\_t5.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdDUucHk=) | `59.43% <0.00%> (-35.85%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.66% <0.00%> (+0.25%)` | :arrow_up: |
| [src/transformers/tokenization\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHMucHk=) | `89.84% <0.00%> (+0.39%)` | :arrow_up: |
| [src/transformers/data/data\_collator.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL2RhdGFfY29sbGF0b3IucHk=) | `92.30% <0.00%> (+0.40%)` | :arrow_up: |
| [src/transformers/tokenization\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYmVydC5weQ==) | `92.41% <0.00%> (+0.44%)` | :arrow_up: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `88.05% <0.00%> (+0.55%)` | :arrow_up: |
| [src/transformers/configuration\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3V0aWxzLnB5) | `96.66% <0.00%> (+0.66%)` | :arrow_up: |
| [src/transformers/tokenization\_utils\_fast.py](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfZmFzdC5weQ==) | `94.28% <0.00%> (+1.42%)` | :arrow_up: |
| ... and [5 more](https://codecov.io/gh/huggingface/transformers/pull/6846/diff?src=pr&el=tree-more) | |

------

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


Yes, I can see that clearly now. Sorry for going back and forth with you on this. We definitely want the type annotations in the main bert file, and I think the first implementation is better on that regard. It just misses the `return_dict` argument, which is easy to add with the way you designed things (happy to do it myself if you give me access to a branch).

The previous implementation is at https://github.com/sbrody18/transformers/tree/scripting
As mentioned, it is a behind head, and still needs some work.
I sent an invite for access to my repo. Let me know if there's a better way to share the branch.

Yes, saw the invite and accepted it. I have some stuff to finish but it's on my TODO and I hope to be able to add the missing stuff before the end of the week. Do you prefer a PR or can I directly commit on this branch?

No rush on my side.
A PR might be better, to make it easier to comment, but direct is fine if that's too much trouble.

Super cool PR!
I can tweak our benchmarking tools a bit to get some numbers on speed improvements using your scriptable Bert model tomorrow

@patrickvonplaten That would be great!
The major improvement is expected when running a large set of inputs with varying lengths, individually or in small batches (that's where not having to pad to max_length would come into play)

> @patrickvonplaten That would be great!
> The major improvement is expected when running a large set of inputs with varying lengths, individually or in small batches (that's where not having to pad to max_length would come into play)

Got it!

This is great, looking forward to this PR!

Okey I did some benchmarking, which can be seen here: https://github.com/huggingface/transformers/pull/6907. 

@sbrody18 - it would be awesome if you could take a look if I am using the function correctly.

Ok, after reviewing this PR and the other design in [this diff](https://github.com/huggingface/transformers/compare/clean_scripting?expand=1), along with @patrickvonplaten benchmark results in #6907 we've come to the conclusion that adding scriptable layers is a bit too much for almost no gain, since `script` and `trace` now have the same speed in PyTorch.

All type annotations and asserts are welcome additions on the other hand, if you want to suggest a PR with just those changes.

Sure. Makes sense. I'll see if I can put one together, but other things might take priority.
Thanks for all the work you've put in to look into this.

@sbrody18 - Thanks a lot for making us aware of this issue! I learned a lot about the differences between `torch.jit.trace` and `torch.jit.script` thanks to you!

Yes thanks a lot for all your work on this, I learned a lot on scriptable pytorch modules thanks to the PR!

I just wanted to point out that, IIUC, a big benefit of making everything scriptable is free reuse from languages other than Python (for example, from the C++ frontend). I know that the prescribed setup is to train in python, trace, then deploy at runtime with a traced TorchScript, but the freedom to train from C++, or even the JVM with a few extra bindings, is a pretty big win. 

This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.