Modify BERT/BERT-descendants to be TorchScript-able (not just traceable)Β #5067
Description
π Feature request
Modify BERT models (src/transformers/modeling_bert.py) to conform to TorchScript requirements, so they can be jit.script()
-ed, not just jit.trace()
-ed (as is currently the only supported option)
Note: I have a working version implementing this, which I would like to contribute.
See below.
Motivation
A scriptable model would allow for variable-length input, offering big speedup gains and simplification (no need to create different models for different input lengths).
In addition, it would avoid other potential pitfalls with tracing (e.g., code paths that are input dependent and not covered by the tracing example input).
Related issues:
#2417
#1204
possibly also
#1477
#902
Your contribution
I have a working PR that modifies all the models in src/transformers/modeling_bert.py and makes them TorchScript-able. I have not tested it on other models that use BERT components (e.g., albert), but it should be possible to expand the capability to those, as well.
However, it would require some significant work to make it ready for submission: besides formatting, documentation, testing etc., my current version changes the method signatures, and I would need to avoid that to maintain backward-compatibility.
Before putting in that work, I'd like to make sure that such a PR is something you'd be interested in and would be willing to merge in, assuming it meets the requirements.