AI Inferencing On a Budget
May 29, 2025
The quintessential "AI-based product" has become a fine-tuned consumption of popular AI APIs, however this is an inopportune path to take for many. They might want more control, in a bid to make the product better suited to their use-case. Or perhaps security is a concern.1
As they near the end of their quest of creating the AI model, the realization dawns that serving the model is another uphill task. Because they need to ingest a lot of real-time data. Last month, I found myself as this hypothetical "they".
Escaping the Jupyter Notebook
We had been working on a model that would eventually turn into a middleware for web applications. This process went through multiple iterations over a year and half involving 3 model architectures changes before we decided to give Transformers a shot. It was the shiny new thing that many thought could solve all the world's problems. However, we were quite pessimistic as our solution demanded a low false-positive rate.
Nevertheless, we stuck to it, and trained, re-trained, fine-tuned something that was deemed good enough. We had a safetensors model exported to our local machine. Now what?
Model on a platter, but which platter
As an aside, let's look at what this safetensors thing is. Semantically, it's a combination of safe and tensors (Thanks Captain Obvious!). Tensors are multi-dimensional arrays, at least that's what university physics had say. And here too, that is what is.
Tensors store weights and some metadata (shapes and data types) about our model. What it does not do is store any structural (think graph) information about our model, and this will become important in a while.
Okay, that was tensors, what makes them safe?
The .safetensor file format is written to be very restrictive of what the file inside can do. One can write malicious code in it but it will remain harmless. As Reddit user AdComfortable1544 puts it2, it's like stabbing someone with a pillow.
A Github discussion has an answer by the maintainer of safetensors which compares this format against .pkl or older .ckpt file formats.3
This file format can be executed using the PyTorch library in Python.4 And thus we come to the second issue. The combination of the previous one, lack of structural information, and the fact that we can use this only in Python means we are stuck again inside Jupyter Notebook (or a Python project). But we intended to escape this!
Serve up on a cross-platform platter
Our quest now leads us to find something that can work across environments and programming languages. Bring in ONNX, an "open format built to represent machine learning models", developed by the PyTorch team while at Facebook. The raison d'être for this format is interoperability and cross-platform support. And with it being open-source, there's community contributions which improve it daily, and makes it more trustworthy than a black-box AI elixir which we dislike.
Exporting the mode from PyTorch to ONNX is fairly simple. You just need to be aware of the general shape (input tensor dimensions etc.) and weights (input types etc.) of your model. And then run this:5
onnx_program = torch.onnx.export(torch_model, example_inputs, dynamo=True)
onnx_program.optimize() # Optimize to remove redundancies in your model
The result is stored in torch.onnx.ONNXProgram as a binary protobuf file. This format will become important later. Now you can save it to your machine with:
onnx_program.save("model.onnx")
You will find many stating this format inferences much slower 6 7 on equivalent systems, and may wonder why we chose this when we're already strapped for compute. While there is no true answer, we had a few reasons to do so.
One that you already know, was our need to serve this cross-platform and ONNX seemed the most flexible for our use-case. The other is more nuanced. We felt that we could optimize our architecture, model or inferencing in order to minimize some of ONNX's (that's a mouthful) shortcoming. There are other, potentially better formats for inferencing on cosumer hardware like GGUF if you have a transformer model, albeit comparitively constrained.
And into Deployment?
Betteridge's law of headlines, or rather section headings suggests the answer is no. And very much so. There was a lot we were doing in our Jupyter Notebook, like pre and post-processing of the data before our model was inferencing. Since we exported only our model, we have lost all those necessary embellishments. But could we simply use the lines of Python we wrote earlier in a different environment with the ONNX model? And the answer to this, fortunately, was yes.
It was possible using NVIDIA's Triton Inference Server, but even that wasn't without friction. Our pre and postprocessor Python snippets used dependencies not supported by Python, for which we had to write a custom backend and Dockerfile.8 And the first run was with:
docker run --rm --gpus all \
-p 8000:8000 -p 8001:8001 -p 8002:8002 \
-v $(pwd)/models:/models \
karideph \
tritonserver --model-repository=/models
What do all these flags mean?
--rm: Automatically remove the container when it exits.--gpus all: Enables all GPUs for Triton (requires NVIDIA Container Toolkit).-p 8000:8000: HTTP endpoint for inference.-p 8001:8001: gRPC endpoint. (discussion for a later post, but needed)-p 8002:8002: Metrics endpoint. (not covered, but needed)-v /path/to/model/repository:/models: Mounts the local model repository to the container.karideph: My custom Docker image. If you're curious why I called it that, it was a planet in the Triton Trade Route (Star Wars nerd much?).tritonserver --model-repository=/models: Starts the Triton server with my model repository.
Now, for a little clarity into why pre and post processing were necessary. And what the broad architecture of the model looked like. Preprocessing was necessary because we had our own custom string tokenizer (Here is a primer on tokenization: Practial ML Gitbook by Uday Bhaskara ) for the kind of inputs we were expecting. Post was necessary to send a callback to our API when inferencing was done, if it got delayed (and Triton expects real-time inferencing).
It was necessary to tie in these 3 steps (preprocessing -> inference -> postprocessing) together using a Triton Ensemble model. You need to specify how data is fed into each of these steps, including shapes (hence why I said knowing the shapes is essential).
The above image gives an overview of the architecture that was used. You can see the ensemble model depicted here. What we haven't discussed yet is the load balancers part. It is still abstracted so let us see what we did and what the way forward could be.
Finding the balance to walk (and run)
We used a simple load balancer which utilizes the various machines with GPUs in a Round-robin manner. You might think that this is inefficient as the different machines would have different processing potentials. And you would be right.
Thus, we distributed the fraction of requests made to each machine based on a dynamically assigned weight. This weight depends on the fail timeout, which is the amount of time for which the machine can inference before failure. Since the causative factor for this metric is processing potential, it was deemed fair for our purpose.
This brings us to the end of our little adventure. We tested this approach on an NVIDIA RTX 4090 and an RTX 2070 Super Mobile, and were able to reliably handle 200 RPS (Requests Per Second).
Contrast this to a naive setup involving similar hardware and a model (less than 500M parameters) the size of GPT-2 medium. If it was fine-tuned over our dataset, and where each inference request averages 2k tokens, it would reliably handle around 2-5 RPS.
What next
Allocating resources based on predicted usage would again drastically reduce both token usage and improve accuracy. Such an approach has been posited by AutoThink.9 This is a software-driven solution which leverages further machine learning approaches and optimizations. To summarize, it predicts which tokens would be important for the final classification and only sends those for inferencing, thereby reducing token usage.
Another approach involves hardware-driven optimizations, llm-d. The architecture of llm-d is beyond the scope of this post, but I may explore it later.
Both of these processes are quite new, as recent as a few days before I started writing this, so the ground is still fertile to deploy a performant model on a budget.
Footnotes
1. OWASP GenAI Security Project. "LLM01:2025 Prompt Injection." Accessed May 29, 2025. https://genai.owasp.org/llmrisk/llm01-prompt-injection/ ↩
2. Reddit user. "What makes .safetensors files safe?" Reddit, r/StableDiffusionInfo. Accessed May 29, 2025. https://www.reddit.com/r/StableDiffusionInfo/comments/14hztyb/comment/jpdodd3/ ↩
3. GitHub Discussion. "Spend a couple of words for uneducated folks?" GitHub, huggingface/safetensors, Discussion #111. Accessed May 29, 2025. https://github.com/huggingface/safetensors/discussions/111 ↩
4. Hugging Face. "Safetensors Documentation." Accessed May 29, 2025. https://huggingface.co/docs/safetensors/en/index ↩
5. PyTorch. "Export a PyTorch model to ONNX." PyTorch Tutorials. Accessed May 29, 2025. https://docs.pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html ↩
6. GitHub Issue. "[Performance] nearest neighbor Resize operator is significantly slower than pytorch for 3D tensors." GitHub, microsoft/onnxruntime, Issue #14596. Accessed May 29, 2025. https://github.com/microsoft/onnxruntime/issues/14596 ↩
7. GitHub Issue. "ONNX Runtime much slower than PyTorch (2-3x slower)." GitHub, microsoft/onnxruntime, Issue #12880. Accessed May 29, 2025. https://github.com/microsoft/onnxruntime/issues/12880 ↩
8. NVIDIA. "Customize Triton Container." NVIDIA Deep Learning Documentation. Accessed May 29, 2025. https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/customization_guide/compose.html ↩
9. Sharma, Asankhaya. "AutoThink: efficient inference for reasoning LLMs." May 13, 2025. https://ssrn.com/abstract=5253327 ↩