As model complexity increases exponentially, so too does the need for effective MLOps practices. This post acts as a transparent write-up of all the MLOps frustrations I’ve experienced in the last few days. By sharing my challenges and insights, I hope to contribute to a community that openly discusses and shares solutions for MLOps challenges.

My goal was to improve Inference latency of few of the current state-of-the-art LLMs.

Unfortunately, simply downloading trained model weights & existing code doesn’t solve this problem.

The Promise of Faster Inference

My first target here was Llama 2. I wanted to convert it into ONNX format, which could then be converted to TensorRT, and finally served using Triton Inference Server.

TensorRT optimizes the model network by combining layers and optimizing kernel selection for improved latency, throughput, power efficiency and memory consumption. If the application specifies, it will additionally optimize the network to run in lower precision, further increasing performance and reducing memory requirements.

From online benchmarks [1, 2] it seems possible to achieve a 2~3x boost to latency (by reducing precision without hurting quality much). But the workings for these kind of format conversions feel super flaky, things break too often (without any solution to be found online). Yes, it’s somewhat expected since these models are so new, with different architectures using different (not yet widely-supported) layers and operators.

Model Conversion Errors

Let’s start with Llama 2 7B chat,

1. Firstly I’ve downloaded Llama-2-7B-Chat weights from Meta’s Official repository here after requesting.
2. Convert raw weights to huggingface format using this script by Huggingface. Let’s say we save it under llama-2-7b-chat-hf directory locally.

Now I considered two options for converting Huggingface models to ONNX format:

torch.onnx.export gibberish text​

Let’s write an export_to_onnx function which will load the tokenizer & model, and export it into ONNX format:

We can also check if the exported & original models’ outputs are similar:

Assuming we’ve saved the ONNX model in ./llama-2-7b-onnx/, we can now run inference using onnxruntime:

😖 On my machine, this generates really funky outputs:

… which is mostly due to missing a proper decoding strategy (greedy, beam, etc.) while generating tokens.

optimum-cli gibberish text and tensorrt slowness

To solve the problem above, we can try a different exporter which includes decoding strategies.

Using the Optimum ONNX exporter instead (assuming the original model is in ./llama-2-7b-chat-hf/), we can do:

⌛ This takes a few minutes to generate. If you don’t has a GPU for this conversion, then remove --device cuda from the above command.

The result is:

Now when I try to do inference using optimum.onnxruntime.ORTModelForCausalLM, things work fine (though slowly) using the CPUExecutionProvider:

After waiting a long time, we get a result:

But when switching to the faster CUDAExecutionProvider, I get gibberish text on inference:

Even with different temperature and other parameter values, it always yields unintelligible outputs, as reported in optimum#1248.

🎉 Update: after about a week this issue seemed to magically disappear — possibly due to a new version of llama-2-7b-chat-hf being released.

Using the new model with max_length=128, :

1. Prompt: Why should one run Machine learning model on-premises?
   1. ONNX inference latency: 2.31s
   2. HuggingFace version latency: 3s

🚀 The ONNX model is ~23% faster than the HuggingFace variant!

⚠️ However, while both CPU and CUDA providers work, there now seems to be a bug when trying TensorrtExecutionProvider — reported in optimum#1278.

optimum-cli segfaults

Next let’s try with the Dolly-v2 7B from Databricks. The equivalent optimum-cli command for ONNX conversion would be:

😢 It uses around 17GB of my GPU RAM, seemingly working fine but finally ending with a segmentation fault:

Confusingly, despite this error, all model files seem to be converted and saved to disk. Other people have reported similar segfault issues while exporting (transformers#21360, optimum#798).

Results using the Dolly v2 model:

1. Prompt: Why should one run Machine learning model on-premises?
   1. ONNX inference latency: 8.2s
   2. HuggingFace version latency: 5.2s

😠 The ONNX model is actually ~58% slower than the HuggingFace variant!

To make things faster, we can try to optimize the model:

The different optimization levels are:

1. -O1: basic general optimizations.
2. -O2: basic and extended general optimizations, transformers-specific fusions.
3. -O3: same as O2 with GELU approximation.
4. -O4: same as O3 with mixed precision (fp16, GPU-only).

We still get the same segfault error for all of the levels.

For -O1, the model gets saved but there’s no noticeable performance change. For -O2 it gets killed (even though I have 40GB A100 GPU + 80GB CPU RAM). Meanwhile for -O3 & -O4 it gives seg-fault (above) while only partially saving the model files.

torch.onnx.export gibberish images

Moving on from text-based models, let’s now look at an image generator. We can try to speed up the Stable Diffusion 2.1 model. In an IPython shell:

The latency (as measured by the %time magic) is 3.25 s.

To convert to ONNX format, we can use this script:

The result is:

There’s a separate ONNX model for each Stable Diffusion subcomponent model.

Now to benchmark this similarly we can do the following:

The overall performance results look great, at ~59% faster! We also didn’t see any noticeable quality difference between the models.

1. Prompt: Iron man laughing
   1. ONNX inference latency: 1.34s
   2. HuggingFace version latency: 3.25s

Since we know that the unet model is the bottleneck, taking ~90% of the compute time, we can focus on it for further optimization. We try to serialize the ONNX version of the UNet to a TensorRT engine-compatible format. When building the engine, the builder object selects the most optimized kernels for the chosen platform and configuration. Building the engine from a network definition file can be time-consuming, and should not be repeated each time we need to perform inference unless the model/platform/configuration changes. You can transform the format of the engine after generation and save it to disk for later reuse (known as serializing the engine). Deserializing occurs when you load the engine from disk into memory:

To setup TensorRT properly, follow this support table. It’s a bit painful, and (similar to cuda/cudnn) if you just want a quick solution you can use NVIDIA’s tensorrt:22.12-py3 docker image as a base:

We can then use the following script for serialization:

Now let’s move to deserializing unet.trt for inference. We’ll use the TRTModel class from x-stable-diffusion’s trt_model:

The above script runs, but the generated output looks like this:

Something’s going wrong, and changing to different tensor shapes (defined above) also doesn’t help fix the generation of blank/noisy images.

I don’t know how to make Stable Diffusion 2.1 work with TensorRT, though it’s proved possible for other Stable Diffusion variants in AUTOMATIC1111/stable-diffusion-webui. Others reporting similar issues in stable-diffusion-webui#5503 have suggested:

1. Use more than 16-bits: I did, but it didn’t help.
2. Use xformers: For our model we need pytorch‘s recently added scaled_dot_product_attention operator.

Other Frustrations

Maybe the code above is partially in my control, but there are also other issues that have nothing to do with my code:

1. Licences: Text Generation Inference recently they came up with a new license which is more restrictive for newer versions. I can only use old releases (up to v0.9).
2. Lack of GPU support: GGML doesn’t currently support GPU inference, so I can’t use it if I want very low latency.
3. Quality: I’ve heard from peers that saw a big decrease in output quality vLLM. I’d like to explore this in future.

Conclusion

I’ve listed my recent errors and frustrations. I need more time to dig deeper and solve them, but if you think you can help please do reply in any of the issues linked above! By sharing my experiences and challenges, I hope this can spark lots of discussions and new ideas. Maybe you’ve faced something similar?

While the world likes showcasing the latest advancements and shiny results, it’s important to also acknowledge and address the underlying complexities that come with deploying & maintaining ML models. There’s a scarcity of documentation/resources for these problems in the ML community. As the field continues to rapidly evolve, there is a need for more in-depth discussions and solutions to these technical hurdles.