AI & ML / Inference & serving / 08_model_export_and_runtimes.md

Model export and inference runtimes

Updated 5 interview angles 5 min read source
On this page6
  1. Why not just ship the checkpoint
  2. The formats
  3. Runtimes and containers
  4. The integration contract
  5. Related
  6. Interview angle

Model export and inference runtimes

Training gives you a checkpoint. Production needs an artefact that runs without your training code, on hardware you did not train on, at a latency someone promised. The gap between those two is export, and it is where most “the model works on my machine” incidents live.

Why not just ship the checkpoint

A PyTorch .pt file is pickled Python. Loading it needs the model class, the same library versions, and the same custom layers importable on the path. That couples your serving container to your training repository, and it means a refactor of the training code can break production without touching the model.

Export breaks that coupling by writing the computation graph to a portable format:

python 
torch.onnx.export(
    model, sample_input, "model.onnx",
    input_names=["input"], output_names=["logits"],
    dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
    opset_version=17,
)

dynamic_axes is the line people forget. Without it the batch dimension is frozen at whatever sample_input had, and the served model accepts batches of exactly that size — which works in testing at batch 1 and fails the moment you batch requests.

The formats

Format What it buys
ONNX Portability across runtimes and hardware
TensorRT Maximum NVIDIA GPU throughput
TorchScript Staying in the PyTorch runtime, no Python

ONNX is the interchange format: one graph, many runtimes (ONNX Runtime, TensorRT, OpenVINO, CoreML). Export once, deploy to CPU, GPU or edge. The cost is operator coverage — a model using an exotic op or Python control flow may not export cleanly, and the error arrives at export time as an unsupported operator rather than as a wrong answer.

TensorRT compiles a graph for a specific GPU architecture, fusing layers and selecting kernels. It is the fastest option on NVIDIA and the least portable: an engine built for an A100 does not run on an L4. Build the engine in the deployment image, on the target device class, not on a developer laptop.

TorchScript keeps you in PyTorch but removes the Python dependency. torch.jit.trace records one execution path — so any if depending on input values is silently baked to the branch your sample took. torch.jit.script compiles the control flow properly and is the safer default when the model branches.

Always verify numerically

Export is a translation, and translations drift. Assert it before you ship:

python 
torch_out = model(sample_input).detach().numpy()
sess = onnxruntime.InferenceSession("model.onnx")
onnx_out = sess.run(None, {"input": sample_input.numpy()})[0]

np.testing.assert_allclose(torch_out, onnx_out, rtol=1e-3, atol=1e-5)

That test belongs in CI, not in a notebook. A silently wrong export produces plausible predictions that are subtly worse, which no error rate alarm catches until someone audits the model weeks later.

Runtimes and containers

The serving container should hold the artefact, the runtime, and nothing else — no training dependencies, no CUDA toolkit if the runtime ships its own.

Triton Inference Server is the usual answer when you serve several models: it handles dynamic batching, multiple frameworks in one process, model versioning and GPU sharing. ONNX Runtime or TorchServe behind FastAPI is simpler and right for one model.

The decision that actually sets your latency is dynamic batching: the server waits a few milliseconds to group concurrent requests into one GPU call. It trades a little p50 for a large throughput gain, and the wait window is the tuning knob. See Batching and serving engines.

The integration contract

When another team consumes the model, the handoff is a contract, and writing it down is the part that prevents a month of confusion.

It has to state, precisely:

  • Input shape, dtype and layout — float32[N,3,224,224], NCHW, and whether values are 0-255 or already normalised.
  • The preprocessing that is not in the graph. Resize interpolation, mean and standard deviation, tokenizer version. This is the most common source of train/serve skew: the backend resizes with a different filter than training did, accuracy drops a few points, and nothing errors.
  • Output semantics — logits or probabilities, class index order, and the threshold if there is one. “Index 0 is negative” belongs in the contract, not in someone’s memory.
  • Latency and batch expectations, so the caller knows whether to batch.
  • Version, pinned. The model artefact and the contract version together.

The strongest move is to push preprocessing into the exported graph where you can, so there is less contract to get wrong. What cannot go in the graph goes in a shared library the backend imports, not in prose the backend reimplements.

Interview angle 5

  • “Why export a model instead of loading the checkpoint?” - a checkpoint is pickled Python and needs the training code, its classes and its library versions importable at serving time. Export writes the computation graph to a portable artefact, so the serving container stops depending on the training repository.
  • “ONNX or TensorRT?” - ONNX for portability across runtimes and hardware; TensorRT when you need maximum throughput on NVIDIA and can accept that the compiled engine is tied to a GPU architecture. Build TensorRT engines on the target device class, not on a laptop.
  • “What breaks most often in export?” - a frozen batch dimension from omitting dynamic_axes, unsupported operators, and torch.jit.trace baking in a branch because it only records the path the sample took. Use script rather than trace when the model has data-dependent control flow.
  • “How do you know the exported model is correct?” - run both and assert numerical closeness with a tolerance, in CI. A bad export produces plausible but subtly worse predictions, which no error-rate alarm catches.
  • “What goes in the contract with the backend team?” - input shape, dtype and layout; the preprocessing that is not in the graph, including interpolation and normalisation constants; output semantics and class order; latency and batching expectations; and a pinned version. Preprocessing mismatch is the classic train/serve skew, and it fails silently.
Previous
Azure and Vertex