Extend Transformers Trainer Class to Enable PyTorch Torchscript for Inference

[jianan-gu]

What does this PR do?

This PR intends to extend Transformers Trainer Class with PyTorch Torchscript (torch.jit.trace) to speed up model inference with just-in-time compilation.

Users could simply enable "--jit_mode" Trainer input args to get benifits.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

This may be more optimum's domain than Transformers, what do you think @mfuntowicz @LysandreJik @lewtun ?

[lewtun]

This may be more optimum's domain than Transformers, what do you think @mfuntowicz @LysandreJik @lewtun ?

Thanks for the ping and thank you @jianan-gu for opening the PR!

I agree that inference-based improvements are better suited in optimum, so in this case we'd need to implement a dedicated torchscript subpackage that effectively extends the Trainer in this manner.

Having said that, I wonder if it's overkill to create a new package just to support this feature, which is natively included in PyTorch.

@mfuntowicz @michaelbenayoun do you see any other benefits that could be had by having a dedicated subpackage for torchscript in optimum?

[stas00]

I'd be happy to work on this PR if you give me a green light that you'll accept this feature.

We can mark is as experimental - should you decide that it'd fit better elsewhere down the road. So it would be easier to experiment and try it on.

[sgugger]

Thanks for adding this. We would need some test before we can merge the PR, and the code needs to be adapted to the way the labels are handled in Transformers.

[sgugger]

Thanks for addressing the comments, this will also need a test before we can merge.

[jianan-gu]

Adding UT Tests in test_trainer.py for jit which covers evaluate and predict.
Thanks.

[sgugger]

Thanks for adding the tests! I have a few last small comments.

Do you want to have a second look @stas00 ?

[stas00]

What is the sounds of a tree falling in the forest when there is nobody to hear it?

All these new features need to be documented on the user-side so that they actually get used. So the same as with IPEX let's add user docs and as before ideally a small benchmark to give users an incentive to try the feature.

updade: I added a doc for inference as we haven't created it and then you can push the usage examples there - like we did with IPEX PR. So please fill out the blank and then we are good to go I think.

and added a few nits in the code.

[stas00]

but isn't your PR adding support for inference only?

        jit_mode_eval (`bool`, *optional*, defaults to `False`):
            Whether or not to use PyTorch jit trace for inference.

[jianan-gu]

Yes, it is for inference only.

[stas00]

so why the generic torch_jit_model name? let's name it accordingly? i.e. adding _eval to it

[jianan-gu]

Sure, have made changes accordingly, thanks for your suggestion.

[stas00]

Thank you, @jianan-gu - please ping me when it's ready for a final review.

[jianan-gu]

What is the sounds of a tree falling in the forest when there is nobody to hear it?

All these new features need to be documented on the user-side so that they actually get used. So the same as with IPEX let's add user docs and as before ideally a small benchmark to give users an incentive to try the feature.

updade: I added a doc for inference as we haven't created it and then you can push the usage examples there - like we did with IPEX PR. So please fill out the blank and then we are good to go I think.

and added a few nits in the code.

Though this jit mode works both for CPU and GPU, the current IPEX release covers CPU side optimizations with jit mode for model inference. Therefore we only have added and updated the inference doc for CPU (perf_infer_cpu.mdx), please take a review of the contents @stas00 , thanks. (for the small benchmark, for now, we only have relative performance numbers like shown in #17137, so we will prepare for the numbers with --skip_memory_metrics 0 as a follow-up)

[stas00]

last few small things to adjust and then it should be good.

[stas00]

@sgugger, would you like to have a quick look at the last changes since your review - mainly the newly added doc. Thank you!

otherwise it's good to be merged.

[sgugger]

Some small comments on the new doc, thanks for adding it!