ONNX — Module 18.3

How This Lesson Fits the Module & Volume

TensorRT wants a graph. ONNX (Open Neural Network Exchange) is the industry IR that lets you train in PyTorch and run on TensorRT, ONNX Runtime, CoreML, TFLite converters, or a CPU laptop. Volume 06 dtypes still apply inside the graph; Volume 12 LLM serving often skips ONNX (vLLM / TRT-LLM load native weights)—but embeddings, rerankers, vision towers, and edge models live here.

This lecture is the export contract: opset, dynamic axes, unsupported ops. Quantization and Edge AI consume the file you produce.

Learning Objectives

By the end of this lesson, students should be able to:

Define ONNX as a portable computation-graph IR plus a file format (.onnx).
Export a PyTorch module with torch.onnx.export (opset, names, dynamic axes).
Choose an execution provider: CPU, CUDA, TensorRT, CoreML, DirectML.
Explain why decoder LLMs are harder to export than static encoders.
List common export failures: custom ops, data-dependent control flow, wrong opset.
Know when to stay in native PyTorch/vLLM instead of forcing ONNX.

Definition

ONNX is an open graph IR: nodes are operators (MatMul, Conv, Attention, …) with typed tensors on edges. A .onnx file stores the graph plus initializers (weights). ONNX Runtime (ORT) is a separate engine that executes that graph on execution providers (CPU, CUDA, TensorRT, CoreML, …). ONNX is not a training framework and not NVIDIA-only—that is the point. Opset version is a compatibility contract; bumping it without re-exporting consumers is a silent break.

Where ONNX Sits in the Stack

Stage	Tool	Notes
Train / fine-tune	PyTorch, TF, JAX (via converters)	Keep native until the graph is stable
Export	`torch.onnx.export`, dynamo exporter, tf2onnx	Fix opset + dynamic axes + dummy shapes
Optimize	onnxsim, ORT graph opts, quantization tools	Constant folding; fuse; INT8 (next lecture)
Compile / run	ORT EP, TensorRT parser, CoreML, TFLite convert	SKU-specific after this point
Skip ONNX	vLLM, TensorRT-LLM, llama.cpp GGUF	LLM decode + KV cache often stay native

Execution Providers vs “Just CUDA”

ORT CPU

Any laptop; Tier 1-adjacent local small models
No NVIDIA driver drama
Too slow for large LLM decode

ORT CUDA / TensorRT EP

Same ONNX, NVIDIA kernels / TRT tactics
Must match CUDA + TRT versions
Great for encoders at QPS

CoreML / QNN / DirectML

Edge / Windows / Apple paths
Bridge to Edge AI
Op coverage varies—test before promising

Export a Module (the contract)

Dummy input shapes teach the tracer. Dynamic axes are how batch and sequence survive production. If export fails, it is usually a Python if on tensor values, a custom autograd Function, or an op newer than your opset. For LLMs, export the embedding or reranker first—not generate() with KV.

import torch
import torch.nn as nn

class TinyEncoder(nn.Module):
    def __init__(self, d_in=768, d_out=256):
        super().__init__()
        self.ln = nn.LayerNorm(d_in)
        self.proj = nn.Linear(d_in, d_out)

    def forward(self, hidden):
        # hidden: [batch, seq, d_in]  — static graph, no KV decode loop
        return self.proj(self.ln(hidden))

model = TinyEncoder().eval()
dummy = torch.randn(1, 32, 768)

torch.onnx.export(
    model,
    dummy,
    "encoder.onnx",
    input_names=["hidden"],
    output_names=["proj"],
    dynamic_axes={
        "hidden": {0: "batch", 1: "seq"},
        "proj": {0: "batch", 1: "seq"},
    },
    opset_version=17,
    do_constant_folding=True,
)

# Load with ONNX Runtime (CPU example — swap providers=["CUDAExecutionProvider"])
# import onnxruntime as ort
# sess = ort.InferenceSession("encoder.onnx", providers=["CPUExecutionProvider"])
# out = sess.run(["proj"], {"hidden": dummy.numpy()})[0]
# TensorRT lecture: this file is what OnnxParser consumes.

When ONNX Is the Wrong Hammer

Export ONNX when

You need one artifact across GPU, CPU, and edge
The graph is mostly static (encoder, CNN, ASR encoder)
TensorRT / CoreML / ORT is the serve target
Compliance wants a frozen IR, not a Python pickle

Stay native when

Autoregressive LLM + paged KV (vLLM / TRT-LLM / GGUF)
Daily LoRA swaps and Python hooks
Custom ops with no ONNX mapping
Tier 1 API-only: you never load weights locally

Related Lectures

Lecture	Why it sits beside ONNX
TensorRT	Primary NVIDIA consumer of ONNX graphs
Quantization	ORT / ONNX quant tools + bitsandbytes contrast
Edge AI	CoreML / TFLite / ORT mobile EPs
Vol. 06 FP16	Dtype still lives on ONNX tensors
Vol. 12 KV Cache	Why full LLM `generate()` export is painful
18.2 Docker	Pin opset + ORT + CUDA EP together

Common Misconception

“ONNX is a faster PyTorch.” ONNX is an IR. Speed comes from the execution provider (CUDA EP, TensorRT, CoreML), not from the file extension. Second: “If it exports, it matches eager numerics.” Op implementations differ; diff outputs on a holdout. Third: “Export the whole chat LLM to ONNX and you get FlashAttention + paged KV for free.” Those Vol. 12 systems are usually native engines; ONNX is the wrong default for decoder serving.

Knowledge Check

Short Answer: What is ONNX, in one sentence? Answer: A portable computation-graph IR and file format for trained models.
True/False: ONNX Runtime is the same project as the ONNX IR spec. Answer: False—ORT is one executor; the IR is the interchange.
Multiple Choice: dynamic_axes exist so: (a) CSS reflows, (b) batch/seq can vary at run time, (c) gradients flow. Answer: (b).
Short Answer: Name two execution providers. Answer: CPU, CUDA, TensorRT, CoreML, DirectML, QNN (any two).
True/False: A .onnx file is NVIDIA-only. Answer: False—that is the point of the IR.
Multiple Choice: Hardest to export cleanly: (a) Linear+LayerNorm encoder, (b) full autoregressive generate() with KV, (c) a frozen CNN. Answer: (b).
Short Answer: Why pin opset_version? Answer: It is the operator compatibility contract between exporter and runtime/compiler.
True/False: TensorRT’s OnnxParser typically consumes the file torch.onnx.export writes. Answer: True.
Multiple Choice: If a custom autograd Function has no ONNX symbolic: (a) export succeeds magically, (b) export fails or emits unsupported nodes, (c) VRAM doubles. Answer: (b).
Short Answer: When should you skip ONNX for an LLM? Answer: When serving with vLLM / TRT-LLM / llama.cpp that load native or GGUF weights and manage KV themselves.

Key Takeaways

ONNX is the portable graph IR; ORT/TensorRT/CoreML are backends.
Export with named I/O, opset, and dynamic axes; test numerics.
Encoders and edge models shine; decoder LLMs usually stay native (Vol. 12).
ONNX is not automatically faster—the EP is.
Continue with Quantization—shrink dtypes on ONNX or in PyTorch/bitsandbytes.

Trainer’s Guide

Lab: Run the TinyEncoder export. Inspect with Netron (or onnx.helper). Run ORT CPU inference and diff vs PyTorch. Optionally break export by adding an if x.sum() > 0 in forward and read the error.

Discussion: For a RAG stack, which pieces (embedder, reranker, generator) should be ONNX vs vLLM, and how does that map to Module 18.4 hardware tiers?

Recap: ONNX is the portable IR between PyTorch and runtimes like TensorRT and ORT. Continue with Quantization.