# 4.2 Deep Learning Frameworks and Libraries

_Inference Engineering_ by Philip Kiely. © 2026 Baseten Labs, Inc. All rights reserved. From Chapter 4: Software. [Full book index](https://www.baseten.co/inference-engineering/llms.txt)

Deep learning frameworks and libraries are the bridge between working directly in CUDA and using off-the-shelf inference engines like vLLM. These libraries are used in both training and inference.

Over the past few years, PyTorch has emerged as the clear leader at this level of the stack. There are two other frameworks that, for concision, I will only mention briefly:

- **TensorFlow:** An end-to-end machine learning platform officially supported by Google, TensorFlow was prominent in the 2010s ML era but has fallen out of favor today.
- **JAX:** A research project unofficially associated with Google, JAX presents a simpler interface without as many legacy features and operations. However, as the documentation warns, expect sharp edges.

The remainder of this section focuses on PyTorch and the technologies around and above it in the stack.

## 4.2.1 PyTorch

PyTorch is a Python package for describing tensor operations. Originally created at Meta and now a part of the Linux Foundation, PyTorch is the industry standard technology underlying both training and inference for generative AI models.

Personally, I’ve been a Python programmer for my entire career, and I find writing low-level C++ difficult. With PyTorch, I can write highly performant inference code in Python for both CPUs and GPUs, but I always have the option to dip down into CUDA when necessary by plugging in specific kernels.

PyTorch can train any kind of neural network. The PyTorch documentation shows a basic neural network example:

![Figure 4.3: A basic neural network in PyTorch, adapted from the PyTorch documentation.](https://www.datocms-assets.com/104802/1788123348-inference-engineering-figure-4-3.png)
_Figure 4.3: A basic neural network in PyTorch, adapted from the PyTorch documentation._

While this is a very simple example, you may recognize the linear layers and ReLU activation functions discussed in chapter 2 as key pieces of neural networks.

PyTorch automatically computes gradients for any differentiable function via its `autograd` module. This is what makes PyTorch so powerful for training – you define a computation graph, and you get a gradient to train against.

But PyTorch is special because it isn’t just great for training, it’s also powerful for inference. PyTorch balances built-in functions and automatic performance optimizations with manual control where you need it.

The step that transforms a model from training to inference is compilation. PyTorch compilation (`torch.compile`) targets a specific GPU and performs automatic kernel selection and kernel fusion to ensure optimal performance.

`torch.compile` can’t fuse plugin kernels like DeepGEMM, FlashAttention, or custom kernels. This limits its utility for LLM inference, where most kernels are custom. However, PyTorch compilation is useful for optimizing less common model architectures and compiling long sequences of lightweight kernels. When you are optimizing a model that has a custom or rare architecture, you may have to rewrite functions to be more abstract – especially with respect to Python-specific language features – for compilation to succeed.

PyTorch alone is a powerful and flexible tool for building high-performance inference services. But there is a rich ecosystem on top of PyTorch that makes it faster to implement, compile, and execute optimized code for common model architectures.

## 4.2.2 Model File Formats

The dominant file format for serializing model weights is safetensors, created by Hugging Face.

Safetensors is a replacement for generic formats like bin designed specifically for holding model weights. The safety in safetensors comes from the fact that unlike a general format that can execute arbitrary Python code during deserialization, safetensors only hold tensor data, not executable code.

Generative AI models have hundreds of gigabytes of weights. These weights are split across dozens of safetensors files. The safetensors format uses memory mapping to ensure that the files are loadable without allocating the full memory, which makes loading model weights faster and safer.

Another leading format, ONNX (Open Neural Network Exchange), stores weights along with an execution graph for the model. Where the safetensors format separates the weights from the architecture, ONNX bundles them together.

ONNX files are highly portable. With deep integration into PyTorch and support for multiple hardware options, ONNX is a great alternative to safetensors for when you want to store model graphs, not just weights.

## 4.2.3 ONNX Runtime and TensorRT

ONNX Runtime and TensorRT are high-performance inference runtimes. PyTorch models can be exported to the ONNX format, which ONNX Runtime can execute directly or TensorRT can compile into a highly optimized engine.

| ONNX Runtime                                      | TensorRT                                                |
| :------------------------------------------------ | :------------------------------------------------------ |
| Open source, associated with the Linux Foundation | Mix of proprietary and open components, built by NVIDIA |
| First-class exporter in the PyTorch ecosystem     | Integrated with PyTorch via Torch-TensorRT              |
| Supports many types of GPUs                       | NVIDIA GPUs only                                        |

ONNX Runtime is an open community standard, while TensorRT is specific to NVIDIA GPUs.

The export process looks somewhat like Torch compilation. However, these standards do not support every data structure, type, and operation within PyTorch. The export process can identify these issues in PyTorch code, but this gets tricky with more complex models.

An example is DeepSeek V3, which introduces Multi-Latent Attention (MLA). MLA, as implemented in PyTorch, is difficult to export, but the transformers architecture overall is simple enough that hand-fusing kernels is feasible.

Today, it’s increasingly popular to skip directly from PyTorch to an inference engine like vLLM or TensorRT-LLM for models that these engines support, bypassing the intermediate representation step and exporting weights only as safetensors. ONNX Runtime and TensorRT are still widely used – TensorRT especially for its strong out-of-the-box runtime for image and video models – but the industry is bifurcating between the control of handwritten PyTorch code or the convenience of prebuilt inference engines.

## 4.2.4 Transformers and Diffusers

The `transformers` and `diffusers` libraries by Hugging Face are built on PyTorch but are not designed to run large-scale production inference. Instead, these libraries offer reference implementations of models for inference engineers to learn from and adapt.

While these libraries are toolboxes for tinkering, they do include essential information about models and useful utilities for building model inference servers. The `config.json` file that ships with models implemented for `transformers` and `diffusers` contains essential information, and the libraries’ utilities for Hugging Face operations like downloading model weights are widely used.

You’ll find `transformers` or `diffusers` sample code in the model card for most popular open models on Hugging Face. This sample code is great for understanding the exact input and output spec of a model or for running local inference and notebooks. But for production, you’ll want to either write and compile PyTorch code directly or use a production-ready inference engine.