Change the HFOnnx pipeline to use Hugging Face Optimum rather than onnxruntime directly

[nickchomey]

The HF documentation says that you can now export seq2seq to ONNX with the OnnxSeq2SeqConfigWithPast class.
https://huggingface.co/docs/transformers/v4.23.1/en/main_classes/onnx#onnx-configurations

This was added with this PR in March huggingface/transformers#14700

Perhaps it is sufficient to be incorporated into txtai now? It would be great to be able to use ONNX versions of the various HF models, for their increased performance.

Additionally, it seems to support ViT models, along with other enhancements that have been made since then. Here's the history for that class https://github.com/huggingface/transformers/commits/main/src/transformers/onnx/config.py

[davidmezzetti]

I'll have ONNX be a focus in the 5.2 release. About to release 5.1.

[nickchomey]

I made this comment in #369, but I'll put it here as well since it seems to be more focused on ONNX improvements:

The hfonnx.call() method uses opset=12 as its default value. But ONNX v1.12 has added opset v17. I wonder whether it might be possible/prudent to add some sort of version check to this method that says:

ONNX_OPSET = {
  "1.12" : 17,
  "1.11" : 16,
  "1.10" : 15,
  "1.9" : 14,
  "1.8" : 13
}

ONNX_VERSION = onnx.version (or whatever is needed to find it...)
OPSET = ONNX_OPSET[ONNX_VERSION]

class HFOnnx(Tensors):
    """
    Exports a Hugging Face Transformer model to ONNX.
    """

    def __call__(self, path, task="default", output=None, quantize=False, opset=OPSET):

etc...

Here's the ONNX versions and their respective opsets https://onnxruntime.ai/docs/reference/compatibility.html

[davidmezzetti]

Sounds good. I've taken a preliminary look and I think the transformers.onnx package can replace some but not all of the code in HFOnnx.

For whatever reason, the default opset in transformers is 11. I wonder if the thinking behind this is picking the lowest opset necessary to maximize compatibility (i.e. support more versions).

[nickchomey]

Probably. But if there's a check for the installed version, then that should solve any compatibility issues - the highest supported version will always be used as the default.

I look forward to whatever you're able to put together to make onnx model usage more accessible! From everything I've read about it, it provides a massive performance improvement.

[nickchomey]

Rather than simply implement ONNX for Seq2Seq models, as discussed in this Slack thread, it would be beneficial and prudent to outsource HFOnnx()'s custom onnxruntime implementation to Hugging Face Optimum. HF Optimum has made it very easy to use the full extent of ONNX's capabilities (including the currently missing model Optimization) with most/all of their model types, including Seq2Seq.

Some relevant links:

• Introductory blog article
• Comprehensive Optimum Docs - goes beyond just ONNX, but that's probably all that needs to be implemented for now (though the pytorch BetterTransformer mechanism looks quite interesting as well)
• Blog article about replacing Transformers with Optimum to train models - this could probably have positive implications for the txtai Trainer pipeline as well
• A 3rd party blog of how to use it

Please don't feel any pressure to implement this immediately on account of me! But I do think it would be quite helpful to txtai users to be able to make full use of ONNX, and would lighten the load on you to monitor and implement changes in onnxruntime.

[nickchomey]

@davidmezzetti

I've made some good progress on building a mechanism to seamlessly auto-convert any HF transformers model to ONNX with HF Optimum, which also allows for the user to configure the optimization level and other important parameters. All combinations of optimization level (0-4) and quantize (True/False) would generate its own model which can be saved to a models directory of choice for quicker re-loading.

I'd like to make it available to txtai via PR such that anyone could start reaping the massive performance and resource (RAM usage) improvements from a fully optimized ONNX model in any of their existing txtai pipelines/workflows/applications, with as little friction as possible.

However, I am not sure what the preferred approach would be.

The easiest option, it seems to me, would be to create .../txtai/pipeline/optimum.py which could be run before any other pipeline, and it would return a (model, tokenizer) tuple, which could then be passed into any other pipeline as its path or model parameter. It would probably require some slight tweaking of txtai/txtai/src/python/txtai/pipeline/hfpipeline.py, but not much else. I'm not sure if this would enable it to be used in a workflow or not...

Perhaps a better approach would be creating .../txtai/models/optimum.py (or just overrwrite the .../txtai/models/onnx.py module). And then modify .../txtai/pipeline/hfpipeline.py and .../txtai/models/models.py to generate/return an Optimum model when requested via the parameters. The biggest problem I foresee with this is that it would probably also require modifying all pipelines such that the required args could be passed through to the Optimum module. (I think a separate case could be made that **kwargs could/should be added to all pipelines and related modules such that things could be extended more easily...)

Any thoughts? I'll probably proceed with the first option today as it should be relatively seamless. But I'm happy to discuss and modify things to accommodate whatever approach you think would be best.

[nickchomey]

On second thought, I'm not sure that the first pipeline approach would be all that useful - you'd have to run a separate one for each txtai pipeline that you want to use, and probably also have to know which task (e.g. zero-shot-classification is being requested.

The only reasonable way to do it would be to allow for setting arguments in any pipeline (e.g. labels = Labels(model_id, onnx=True, optimization_level=3, quantize = true, architecture="avx512_vnni") and it will handle all the model selection, conversion, optimization and quantization behind the scenes... But that would require modifying all existing pipelines to have the required parameters. I don't mind doing it if that's the case, but perhaps, as mentioned, adding **kwargs to each pipeline would allow for more flexibility going forward?

Any thoughts on any of this? I'm happy to do all of the work if you can point me in the right direction, such that all you'd need to do is review the code and modify things to accommodate your style, and be more robust, efficient etc...

[davidmezzetti]

For this change to be most effective, the best path is to replace this class - https://github.com/neuml/txtai/blob/master/src/python/txtai/models/onnx.py with an optimum version. And this method https://github.com/neuml/txtai/blob/master/src/python/txtai/models/models.py#L118 to detect loading an ONNX model.

Last I checked with Optimum it doesn't support loading streaming models. It expects everything to work with files. That will cause issues with some existing functionality.

The other piece would be changing HFOnnx to use Optimum to convert models. From there you shouldn't need to add these extra arguments as that would be done at model conversion time.

[nickchomey]

Thanks! I'll explore those files and submit something for your review when ready.

Yeah, Optimum works off of files. The workflow I've set up is:

1. download the transformers model
2. convert to ONNX and save file
3. optionally optimize the ONNX model and save
4. optionally quantize the raw ONNX model or the optimized ONNX model.

I've built it such that it checks for existing files as early as possible to minimize processing. But can also store many versions depending on what combination of optimization and quantization is desired. Once the initial download is processing is done, then it should all load quickly from the disk.

I don't know anything about streaming - could you please point me to the relevant txtai code where streaming is used, so that I can dig in and see what it does and what options might exist for this implementation?

At the very least, it seems like this could be implemented as a sort of optional alternative. It seems far easier and more future-maintainable to offload all the work to Hugging Face than to build custom mechanisms directly on top of onnxruntime to do optimization and quantization on all the different transformer model types (Seq2Seq, FeatureExtraction etc...).

[davidmezzetti]

Did you try using txtai pipelines with ORTModelForXYZ.from_pretrained?

from optimum.onnxruntime import ORTModelForSequenceClassification, ORTModelForSeq2Seq
from transformers import AutoTokenizer

from txtai.pipeline import Labels

model = ORTModelForSequenceClassification.from_pretrained("path", from_transformers=True)
tokenizer = AutoTokenizer.from_pretrained("path")

labels = Labels((model, tokenizer), gpu=False)
labels("Text to label")

And for Seq2Seq

from optimum.onnxruntime import ORTModelForSeq2SeqLM

from txtai.pipeline import Sequences

model = ORTModelForSeq2SeqLM.from_pretrained("t5-small", from_transformers=True)
tokenizer = AutoTokenizer.from_pretrained("t5-small")

sequences = Sequences((model, tokenizer), gpu=False)
sequences("translate English to French: Hello")

This should all work with quantized models as well.

from optimum.onnxruntime.configuration import AutoQuantizationConfig

quantizer = ORTQuantizer.from_pretrained(model)
qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)

quantizer.quantize(save_dir="quant", quantization_config=qconfig)

So while I like the idea of making this more integrated, it seems like it's pretty straightforward to do run Optimum models with txtai as it stands right now.

[nickchomey]

Yes, I believe I mentioned somewhere above (or in Slack) that it works quite well already by passing a (model, tokenizer) tuple in - precisely how you did.

But, it resulted in a lot of redundant code for me when I started testing it with different pipelines. So, I started refactoring and eventually arrived at something generalized and flexible - along with something that makes use of saving the various converted/optimized/quantized models. So, I figured that it would be worth going the extra mile to make something nice and integrated - both for my sake as well as for everyone else who don't have time/interest to figure all of this out.

If accompanied by a good colab example, I think it would be well received by txtai users - it should make it simple for anyone to experiment with any of this by adding/tweaking some arguments in their existing code. It will also make it easy for people to make use of any hardware-specific acceleration that they have available to them (via the ONNX Execution Providers, such as the OpenVINO EP - which seems to be highly compatible with most systems and even more powerful than raw onnxruntime)

I'm going to build it anyway, so would you be willing to review a PR for this? To start, I'll leave onnx.py alone and create a new module and slightly modify the models.py logic to handle it appropriately - that way nothing at all will be affected. But it would be easy enough to replace onnx.py later if desired.

[davidmezzetti]

If it can be done with changing HFOnnx, models/onnx.py and models/models.py, then yes. I don't want to go down the path of adding a bunch of optimum options to the pipelines. That should be a one time conversion via a process similar to HFOnnx.

Sounds like you're going to work on a lot of this regardless for your own work. But in terms of a PR, that is the type of code above that would make sense for core txtai.