PyTorch, Model Formats, ONNX, TensorRT, Transformers

Writing CUDA kernels by hand gives full control and takes enormous effort. Using a ready-made inference engine such as vLLM takes little effort and gives less control. Between the two sits a middle layer: deep learning frameworks and libraries. A framework lets you describe a neural network in a high-level language, run it on a CPU or a GPU, and still drop down to custom kernels when needed. Frameworks are used both for training (adjusting a model's weights from data) and for inference (using the trained weights to produce outputs).

This lesson covers PyTorch, the two other frameworks worth knowing by name, the file formats that store model weights, the runtimes that execute exported models, and the Hugging Face libraries that serve as reference implementations. Afterwards you can explain what torch.compile does and does not do, choose between safetensors and ONNX, compare ONNX Runtime and TensorRT, and say where transformers and diffusers fit in production. Terms such as kernel and fusion come from CUDA, kernels, and kernel fusion.

Over the past few years PyTorch has become the clear leader at this level. Two other frameworks deserve a short mention.

The rest of the lesson concentrates on PyTorch and the technologies around and above it.

Example 1: PyTorch and a small network

PyTorch is a Python package for describing tensor operations, where a tensor is a block of numbers with any number of dimensions. It was created at Meta and is now part of the Linux Foundation. It is the industry standard under both training and inference for generative AI models. A Python developer can write fast inference code for CPUs and GPUs in PyTorch and still drop down into CUDA by plugging in specific kernels when needed.

PyTorch can train any kind of neural network. The basic example in its documentation is a network that reads a 28 by 28 image and produces 10 scores. It is built from three linear layers separated by two ReLU activations. A linear layer multiplies its input by a matrix of weights and adds a bias. ReLU replaces every negative number with 0 (max(0, x)). The structure in diagram form:

flowchart LR
  I["Image 28 x 28"] --> F["Flatten to 784 numbers"]
  F --> L1["Linear 784 to 512"] --> R1["ReLU"]
  R1 --> L2["Linear 512 to 512"] --> R2["ReLU"]
  R2 --> L3["Linear 512 to 10"] --> O["10 scores (logits)"]

The size of this model follows from the layer shapes. A linear layer from a inputs to b outputs has a x b weights plus b biases.

Linear 784 -> 512: 784 x 512 + 512 = 401,408 + 512 = 401,920 parameters
Linear 512 -> 512: 512 x 512 + 512 = 262,144 + 512 = 262,656 parameters
Linear 512 -> 10 : 512 x 10  + 10  =   5,120 +  10 =   5,130 parameters
Total: 401,920 + 262,656 + 5,130 = 669,706 parameters

At 4 bytes each (FP32): 669,706 x 4 = 2,678,824 bytes, about 2.7 MB
At 2 bytes each (FP16): 669,706 x 2 = 1,339,412 bytes, about 1.3 MB

The ReLU layers have no parameters. This tiny network has the same building blocks as a large language model, which differs mostly in size and in the attention layers covered in Attention and Mixture of Experts.

Example 2: autograd and the step from training to inference

PyTorch computes gradients automatically for any differentiable function, through its autograd module. A gradient says how much the output changes when a weight changes a little, and training uses it to nudge each weight in the direction that reduces the error. You define a computation graph, and PyTorch hands you the gradient to train against. A hand-worked case of one neuron shows the idea:

model: y = w x + b, with x = 2, w = 3, b = 1, target t = 10, loss = (y - t)^2
forward: y = 3 x 2 + 1 = 7, loss = (7 - 10)^2 = 9
gradients (chain rule): dloss/dy = 2 (y - t) = -6
  dloss/dw = dloss/dy x x = -6 x 2 = -12
  dloss/db = dloss/dy x 1 = -6
update with step size 0.1: w = 3 - 0.1 x (-12) = 4.2,  b = 1 - 0.1 x (-6) = 1.6
forward again: y = 4.2 x 2 + 1.6 = 10.0, loss = 0

Autograd does the gradient arithmetic for graphs with millions of operations. This is what makes PyTorch strong for training. It is also strong for inference, because it combines built-in functions and automatic performance optimizations with manual control where you want it.

Training turns into inference through compilation. PyTorch compilation, torch.compile, targets a specific GPU and performs automatic kernel selection and kernel fusion for good performance. Fusion removes memory round trips between operations, as the previous lesson showed. There are limits:

• torch.compile cannot fuse plugin kernels such as DeepGEMM, FlashAttention, or custom kernels. That limits its use for LLM inference, where most kernels are custom.
• It is useful for less common model architectures and for compiling long sequences of lightweight kernels.
• For a custom or rare architecture, functions may have to be rewritten in a more abstract form, especially to avoid Python-specific language features, before compilation succeeds.

PyTorch alone is a flexible tool for building high-performance inference services. On top of it sits an ecosystem that speeds up implementing, compiling, and running optimized code for common architectures.

Example 3: model file formats, safetensors

Trained weights must be stored in a file, a step called serialization. Loading them back is deserialization. The dominant format for weights is safetensors, created by Hugging Face. It replaces generic formats such as bin with a format designed only for model weights.

The "safe" in the name matters. A general format can run arbitrary Python code while it is loaded, so a downloaded model file could execute a hidden program on your machine. Safetensors holds only tensor data, never executable code. This is like the difference between reading a JSON file and deserializing arbitrary objects with a mechanism that can construct anything.

The layout is simple. The file starts with 8 bytes that hold the length of the header as a little-endian integer. Next comes the header, a JSON text that names each tensor with its type, its shape, and where its bytes start and end. The raw tensor bytes follow. A reader can parse the small header first and then locate any tensor without reading the rest.

[ 8 bytes: header size as a little-endian integer ][ JSON header: tensor names, types, shapes, byte ranges ][ raw tensor bytes ]

Large models are big: generative AI models have hundreds of gigabytes of weights, split across dozens of safetensors files. Safetensors uses memory mapping, which asks the operating system to expose the file as if it were in memory and to load pieces from disk only when they are touched. Loading does not need to allocate memory for the whole file first, which makes loading faster and safer. A worked case:

weights: 140 GB split into 30 files, about 4.7 GB per file (140 / 30)
copy-everything load: needs 140 GB of free host memory before the first byte is used
memory-mapped load: files are mapped without copying, and each tensor is read as it is moved to the GPU

Example 4: ONNX, weights together with the graph

Another leading format is ONNX (Open Neural Network Exchange). It stores the weights together with an execution graph of the model, the list of operations and how data flows between them. Safetensors separates weights from the architecture, and ONNX bundles them. ONNX files are highly portable. With deep integration into PyTorch and support for many hardware options, ONNX is a good alternative to safetensors when you want to store model graphs and not only weights.

Example 5: ONNX Runtime and TensorRT

ONNX Runtime and TensorRT are high-performance inference runtimes. A PyTorch model is exported to ONNX. ONNX Runtime can execute the result directly, and TensorRT can compile it into a highly optimized engine.

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

flowchart LR
  P["PyTorch model"] -- "export" --> O["ONNX file (weights + graph)"]
  O --> OR["ONNX Runtime executes it"]
  O --> TR["TensorRT compiles it into an optimized engine"]
  P -- "weights only, as safetensors" --> E["Inference engine: vLLM or TensorRT-LLM"]

The export step looks somewhat like Torch compilation. Neither standard supports every data structure, type, and operation in PyTorch. The export process can detect these problems in PyTorch code, but it gets difficult with complex models. DeepSeek V3 is an example. It introduces Multi-Latent Attention (MLA), and MLA as implemented in PyTorch is hard to export. The transformer architecture as a whole is simple enough that fusing kernels by hand is feasible.

It is now increasingly common to skip from PyTorch straight to an inference engine such as vLLM or TensorRT-LLM for models that those engines support. That bypasses the intermediate representation and exports weights only, as safetensors. ONNX Runtime and TensorRT remain widely used, TensorRT especially for its strong out-of-the-box runtime for image and video models. The industry is splitting between the control of handwritten PyTorch code and the convenience of prebuilt inference engines (Inference engines).

Example 6: transformers and diffusers

Hugging Face publishes two libraries built on PyTorch, transformers (for language and other models) and diffusers (for image and video generation models). They are not built to run large-scale production inference. They provide reference implementations: readable versions of models that inference engineers study and adapt.

They are still toolboxes with valuable pieces for building model servers:

• The config.json file that ships with models implemented for these libraries holds key facts about the model, such as its architecture settings. A worked use of such a file appears below.
• The libraries' helpers for Hugging Face operations, such as downloading model weights, are widely used.
• The model card (the documentation page) of most popular open models on Hugging Face has transformers or diffusers sample code. The sample shows the exact input and output specification of a model and is good for local inference and notebooks.

For production, either write and compile PyTorch code directly or use a production-ready inference engine.

A config file gives numbers that feed the memory arithmetic of earlier lessons. For instance, a configuration listing 32 layers, 8 key-value heads, and a head size of 128 for a model in FP16 gives the size of the KV cache per token (the saved attention data per token, see KV caching). The numbers here are hypothetical:

KV cache per token = 2 (keys and values) x layers x KV heads x head size x bytes
                   = 2 x 32 x 8 x 128 x 2 = 131,072 bytes = 128 KB
for 8,000 tokens: 8,000 x 128 KB = 1,024,000 KB, about 1 GB

Example 7: choosing a layer

Interview points

• Frameworks sit between CUDA and inference engines. PyTorch leads. TensorFlow (Google, out of favor) and JAX (research, sharp edges) are the others.
• PyTorch was created at Meta and is hosted by the Linux Foundation. Autograd computes gradients, and torch.compile performs kernel selection and fusion for a specific GPU.
• torch.compile cannot fuse plugin kernels (DeepGEMM, FlashAttention, custom kernels), so it helps LLM inference less. It helps rare architectures and long sequences of lightweight kernels, and it may force rewrites that avoid Python-specific features.
• Safetensors stores only tensor data and cannot execute code. Weights are split across dozens of files, and memory mapping makes loading faster and safer.
• ONNX stores weights together with the execution graph and is portable. Safetensors separates weights from architecture.
• ONNX Runtime is open source, a first-class PyTorch exporter, and supports many GPUs. TensorRT is NVIDIA-only, integrated through Torch-TensorRT, and compiles an optimized engine.
• Export can fail on unsupported operations. MLA in DeepSeek V3 is an example.
• Many teams skip ONNX and go from PyTorch to vLLM or TensorRT-LLM with safetensors. TensorRT stays popular for image and video models.
• transformers and diffusers are reference implementations with config.json files and download helpers. They are not for large-scale production.

Summary

• Frameworks bridge raw CUDA and prebuilt engines, and serve training and inference.
• PyTorch: tensors, autograd, torch.compile with selection and fusion, limited for custom kernels.
• Formats: safetensors (weights only, safe, memory-mapped) and ONNX (weights plus graph, portable).
• Runtimes: ONNX Runtime (open, many GPUs) and TensorRT (NVIDIA, optimized engines).
• The industry is divided between handwritten PyTorch and prebuilt inference engines.
• transformers and diffusers are for learning, adapting, and sample code, with config.json and weight download utilities.