Speeding up T5 inference :rocket:
seq2seq decoding is inherently slow and usingonnxis one obvious  solution to speed it up. Theonnxt5package already provides one way to useonnxfort5.But if we export the complete T5 model toonnx, then we can’t use thepast_key_valuesfor decoding since for the first decoding steppast_key_valueswill beNoneandonnxdoesn’t acceptNoneinput. Withoutpast_key_valuesonnxwon’t give any speed-up over torch for beam search.One other solution is to export theencoderandlm_headtoonnxand keep thedecoderintorch, this way thedecodercan use thepast_key_values.I’ve written a proof-of-concept script which does exactly this and also makes it compatible with thegeneratemethodgist.github.comhttps://gist.github.com/patil-suraj/09244978af5f7598dd30fb9b4f54fe29onnx_t5.pyimport inspect
import logging
import os
from pathlib import Path

import torch
from psutil import cpu_count
from transformers import T5Config, T5ForConditionalGeneration, T5Tokenizer
from transformers.generation_utils import GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, Seq2SeqLMOutputThis file has been truncated.show originalrequirements.txt--find-links https://download.pytorch.org/whl/torch_stable.html
torch==1.6.0+cpu
transformers>=3.1.0
onnxruntime>=1.4.0
onnxruntime-tools>=1.4.2
psutilWith this you canenc = tokenizer("translate English to French: This is cool!", return_tensors="pt")
onnx_model = OnnxT5(model_name_or_path="t5-small", onnx_path="onnx_models")
tokens = onnx_model.generate(**enc, num_beams=2, use_cache=True) # same HF's generate method
tokenizer.batch_decode(tokens)In my experiments this gave ~1.4-1.6x speed-up with beam search.The first time you callOnnxT5it’ll load the model from the hub, export it toonnxas described above and save the exported graphs atonnx_path. So loading will be  slower the first time.Now to gain further speed-up we could distill the model and use less decoder layers.onnx + distillation should give even more speed-up with minimal drop in accuracy.@sshleiferhas just published awesomeseq2seq distillation paperwhich can be used to distillT5model  as well.I’ll be sharing T5 distillation results soon!now this is a very hacky solution, so feel free to suggest feedback and improvements or any other method that can help speed things upcc.@abel,@patrickvonplaten,@sshleifer
2020-11-01T18:06:01Z
Here’s a self-containednotebookand small benchmark for summarization and translation tasks.summ_benchmark862×682 8.7 KBtranslation_benchmark862×682 8.97 KB
2020-11-02T06:50:27Z
A question on twittercan’t we import T5Model from transformer and then may be pass it through a wrapper or decorator for onnx optimization. Or do we really need to change some ops in torch so that they can be onnx compatible?We could export the T5 as it is but then we can’t usepast_key_valuesbecause they will beNonefor first decoding step andonnxdoesn’t acceptNoneinput.  So here we are exporting individual components to be able to usepast_key_valuesHere you could also export the encoder and lm_head directly instead of creating classes for it. I did it this way to make it more readable. And we need to write the extra code inOnnxT5class to make it compatible withgenerate.Feel free to suggest changes if you think it can be improved
2020-11-02T08:04:03Z
Thanks@valhalla. The reason of the doubt was, we have to make changes to some parts of the code to make it work as it is. And this sane issue is there in tensorflow too. Very nice post and super useful. Thanks for sharing:-)
2020-11-02T18:15:11Z
What is causing the Onnx T5 model to be faster than the regular PyTorch T5 model?
2020-11-03T16:56:54Z
Because ofonnxruntime, which does several graph optimizations to speed-up inference. You could find more info in onnx runtime docs
2020-12-05T09:29:30Z
you can quantize the model. As mentioned in the PyTorchdocPyTorch supports INT8 quantization compared to typical FP32 models allowing for a 4x reduction in the model size and a 4x reduction in memory bandwidth requirements. Hardware support for INT8 computations is typically 2 to 4 times faster compared to FP32 compute.thisnotebookshows the benchmarks of using quantized models with onnx.download1306×990 20.1 KB
2020-12-21T06:53:43Z
hey@valhalla, does this onnx runtime solution work with mT5 also?
2020-12-21T07:11:53Z
might not work as it is, but should be easy to adapt. I think replacing T5 related classes with MT5 should make it work, for example,T5ForConditionalGenerationwithMT5ForConditionalGenerationetc
2021-01-16T17:37:36Z
Hey@valhallaHope you’re well. Firstly, I want to thank you for sharing this.I was trying to run thenotebookyou had shared to recreate the results you reported but kept running into an error. Would you be able to help me fix it? The only change to the notebook I had made was installed the sentencepiece library as it was required by T5Tokenizer.Screen Shot 2021-01-17 at 4.36.46 AM2626×818 246 KBAlso, did you get a chance to distill the T5 model and produce some results? Looking forward to thoseThanks!
2021-01-19T06:01:57Z
Hi@rigoh5,what transformer version are you using ? could you try with 3.1.0, I haven’t tested it with the latest version.
2021-03-15T07:44:27Z
hey,@valhallaI’ve createdfastT5library. it exports T5 model to onnx withpast_key_values.  the exported model supportsgenerate()method.for onnx models I was getting the inference speed boost of up to5xforgreedyand3-4xforbeam search. for more details check out thefastT5 repo.currently, transformers library does not support exporting of t5 to onnx withpast_key_values, you can fix this issue by following the guide in thisnotebook. created PR for this supporthere
2021-03-15T08:00:17Z
This is awesome, great job! You should open a new thread to share this so people will know this.I’ll take a look at the PR
2021-03-15T09:29:00Z
thank you!@valhalla. created a new threadhere.
2021-09-25T10:42:24Z
How can we do this with Pegasus.I was trying to modify the script for Pegasus just by loading Pegasus model and i was able to load the encoder and decoder with:encoder = model.model.encoder
decoder = model.model.decoder
lm_head = model.lm_headbut i get get error even when i try to build the original T5. It does generate the t5 onxx model, but i get errors when trying to use it.
2022-02-21T18:15:11Z
Hi@kira,Thanks forfastT5.Do you plan to update it as thelast version(0.1.2, 01/06/2022) isincompatible with transformers>=4.16?Note: for people who want to use the last fastT5 version (0.1.2),install fastT5 from sourceby running in a notebook cell the following code:!git clone https://github.com/Ki6an/fastT5
%cd fastT5Download the filesetup.pyand change the line"transformers>4.6.1"by"transformers==4.15.0"as the transformers versions >= 4.16.0 do not work with the last fastT5 version (0.1.2).Save and upload the filesetup.pybefore to launch the following cell:pip3 install -e .Et voilà.
2022-08-26T07:27:46Z
if you are using PegasusForConditionalGenerationmodel = PegasusForConditionalGeneration.from_pretrained(‘model-path-from-huggingface’)encoder = model.get_encoderdecoder = model.get_decoderlm_head = model.lm_head
2022-08-26T07:30:47Z
if you are using PegasusModel class from transformersmodel  = PegasusModel.from_pretrained(‘model-path-from-huggingface’)encoder = model.encoderdecoder = model.decoderbut you can’t getmodel.lm_headbecause it’s not part ofPegasusModel.