Model serving

Problem

Training ends with a state_dict and the code that built the model. A serving system has different goals from a training loop:

• Latency. A user waits for each answer. What matters is not the mean but the tail: the 90th or 99th percentile request.
• Throughput and cost. Hardware is billed by the hour, so the number of requests answered per second per device decides what serving costs.
• Portability. The serving process may not be the training environment: a different machine, a C++ binary, a phone. The model has to survive leaving its Python source behind.
• Safety and stability. Loading a file must not run arbitrary code, and the same input must give the same output today and next month.

The tension is between the first two. Sixteen requests arrive over 22 ms at a model that takes 5 ms of fixed overhead plus 1 ms per item. Run them one at a time and the queue grows: the mean request waits 40.25 ms and the last finishes at 96 ms. Group them into batches and the same hardware averages 15.31 ms. Batching is the main lever in serving, and choosing how to batch is most of the design.

Intuition

There are two things to ship: the weights (tensors) and the program (what to compute with them). PyTorch offers three ways to carry the program:

• Eager. Ship the Python class, load the weights into it, call it. Every operation is dispatched one at a time from Python. Simple, flexible, and the slowest per call.
• Compiled. Still Python, but torch.compile captures the model into graphs at run time and generates fused kernels for them (see GPU execution). Faster, but the first calls pay for compilation and new input shapes can recompile.
• Exported. torch.export traces the model ahead of time into a single graph of ATen operators with no Python control flow left in it. The result is a file that can be loaded without the model's source, and AOTInductor can compile it into a shared library that runs without Python at all.

The second idea is batching. A GPU computing one row costs nearly the same as computing a dozen, because most of the time goes to fixed overhead: launching kernels and reading the weights from memory. So the server holds requests briefly and runs them together. A dynamic batcher collects requests until the batch is full or the oldest has waited long enough. For text generation, where a request is many decode steps long and requests finish at different times, a continuous batcher lets a new request take a finished one's place at the next step instead of waiting for the whole batch to finish.

Mechanics

Weights: saving and loading safely

torch.save(obj, f) writes a zip archive: tensor storages as raw byte files, and a pickle that describes how to rebuild obj around them. The usual thing to save is model.state_dict(), a dict from parameter names to tensors, rather than the model object, so loading doesn't depend on the class being importable under the same name.

Pickle is a small program: unpickling can call any importable function, so loading an untrusted pickle can run arbitrary code. torch.load therefore defaults to weights_only=True (when no custom pickle_module is passed), which uses a restricted unpickler that only rebuilds tensors, primitive types and an allowlist of safe classes. Two other arguments matter for serving: map_location moves storages to a different device as they load (a GPU checkpoint onto a CPU server), and mmap=True maps the file into memory instead of reading it, so a large checkpoint can be loaded without holding two copies.

Eager, compiled, exported

In eager mode, Python runs forward line by line and each operator goes through the dispatcher separately. For a large model on a big batch, the kernels dominate and the overhead is noise; for a small model or batch size 1, the Python and launch overhead can be most of the time.

torch.compile(model) returns a wrapper. On the first call TorchDynamo reads the Python bytecode, captures the tensor operations into a graph (breaking it where it meets something it can't trace), records guards (the input shapes, dtypes and Python values the graph assumed), and hands the graph to a backend, by default Inductor, which generates fused kernels. Later calls check the guards and reuse the compiled code; a failed guard triggers a recompile, and after torch._dynamo.config.recompile_limit recompiles (8 by default) that frame falls back to eager. mode="reduce-overhead" adds CUDA graphs, which replay a recorded sequence of kernel launches with one call, aimed squarely at small-batch serving. dynamic=True asks for shape-generic code up front instead of specializing and recompiling.

torch.export.export(mod, args) instead traces once, ahead of time, and requires the whole model to become one graph: no graph breaks. Data-dependent Python control flow is not allowed (use torch.cond for branches that must stay in the graph). By default (strict=False) it traces by running the Python code with fake tensors; strict=True traces through TorchDynamo instead. The result, an ExportedProgram, holds the graph in the functional ATen operator set, the parameters and buffers, a signature saying which graph inputs are parameters and which are user inputs, and range constraints on symbolic dimensions.

Dynamic shapes in export

By default export specializes on the example input's shapes: export with a batch of 4 and the program accepts only batches of 4. A server needs variable batch size and often variable sequence length, declared with dynamic_shapes: a Dim("batch", min=1, max=64) names a symbolic size with a range, and the hints Dim.AUTO, Dim.DYNAMIC and Dim.STATIC let export infer or fix a dimension. Export then checks the model's code against the declared ranges, and the loaded module checks inputs against them at call time (ExportedProgram.module(check_guards=True)). torch.export.save writes the program to a .pt2 archive; its docstring warns that the format is under active development and saved files may not load in newer PyTorch releases, and torch.export.load uses pickle, so the same rule applies: never load one from an untrusted source.

Inference mode and eval

Two switches belong around every serving call. model.eval() (which is model.train(False)) changes what layers compute: dropout becomes the identity and batch norm uses its running statistics instead of the batch's. torch.inference_mode() changes what autograd records: like no_grad it builds no graph, and it also skips view tracking and version-counter bumps, so tensors cost a little less. Tensors created inside can't be used in autograd later. Forgetting eval() is a correctness bug (outputs change with the batch); forgetting inference mode only wastes memory and time.

Dynamic batching

Model one batch's cost as a fixed part plus a per-item part:

Throughput rises toward 1/c_item as B grows, and the larger c_fixed is compared with c_item, the bigger the batch needed to get there. A dynamic batcher has two knobs: the largest batch it will form, and the longest it will make the oldest queued request wait. When the device frees up it starts a batch as soon as either the queue holds a full batch or the oldest request's deadline passes, whichever comes first, and takes everything that has arrived by then up to the maximum. Clipper (Crankshaw et al., 2017) built a serving system around this idea, tuning the maximum batch size adaptively per model to meet a latency target.

Waiting is a bet. At light load, a short wait collects requests that would otherwise run in near-empty batches. At heavy load the queue fills batches on its own, and any deliberate wait only adds latency. The worked example shows the second case.

Continuous batching for generation

An autoregressive model produces one token per request per forward pass, and requests need different numbers of tokens. With static batching, the batch runs until its longest request finishes; slots whose requests finished early sit idle, and new requests wait for the whole batch. Orca (Yu et al., OSDI 2022) introduced iteration-level scheduling: the scheduler revisits the batch after every decode step, removes finished requests and admits waiting ones into the free slots. This is now called continuous batching. It needs a KV cache that can grow and shrink per request; vLLM's PagedAttention (Kwon et al., SOSP 2023) stores the cache in fixed-size blocks for exactly this reason, so a finished request's memory goes straight back to the pool.

Requests are numbered r0 to r7; each row lists the requests running at that step, packed into the first slots for readability. Real schedulers keep a request in the slot it was admitted to.

CPU versus GPU

The same cost model explains when a CPU is the right server. A CPU has small fixed overhead per call and a large per-item cost; a GPU has a larger fixed cost (kernel launches, a host-to-device copy) and a tiny per-item cost. At batch size 1 the CPU can win outright; the GPU's advantage appears only once batches are large enough to amortize its fixed cost, which in turn needs enough traffic to fill those batches. Small models (linear models, gradient-boosted trees, small MLPs) at modest traffic often serve best on CPU, where instances are cheap and plentiful. Large networks, and anything generating tokens, belong on accelerators, where the batching machinery above is what makes them affordable. Quantization shrinks the per-item cost on both.

A serving checklist

• Artifact. Save a state_dict or an exported program, record the PyTorch version and the preprocessing that produced the inputs, and load with weights_only=True.
• Mode. model.eval() once at load; torch.inference_mode() around every call.
• Shapes. Decide which dimensions vary and declare them (Dim in export, dynamic=True or bucketed sizes for compile). Warm up every shape bucket before taking traffic so no user pays for compilation.
• Parity. Compare served outputs with the eager model on a fixed set of inputs, within a tolerance that allows for kernel and precision differences.
• Batching. Choose a maximum batch from the latency budget and the cost curve; add a wait only if traffic is light enough that batches would otherwise be small.
• Measure the tail. Track p50, p90 and p99 latency, queue length, batch-size distribution and device utilization. Latency and throughput covers the measurement.
• Watch the inputs. Log input statistics and prediction distributions so drift shows up (Monitoring and drift).
• Roll out safely. Version the model, shadow or canary a new version against the old one, and keep the old one loadable for rollback.

Worked example

Implementation

Export a model once with a dynamic batch dimension, save it, and serve it from a process that never imports the model's class. The batcher is a minimal version of the policy above: block for the first request, then keep collecting until the batch is full or the deadline passes.

import queue, threading, time
import torch
from torch.export import Dim

# offline, in the training environment
model.eval()
example = (torch.randn(4, 128),)
ep = torch.export.export(model, example,
                         dynamic_shapes={"x": {0: Dim("batch", min=1, max=64)}})
torch.export.save(ep, "model.pt2")
# optional: a package that runs without Python
# torch._inductor.aoti_compile_and_package(ep, package_path="model_aoti.pt2")

# serving process
served = torch.export.load("model.pt2").module()   # trusted file only: load uses pickle

requests = queue.Queue()                           # items: (input tensor [128], reply queue)

def batcher(max_batch=32, max_wait_s=0.003):
    while True:
        batch = [requests.get()]
        deadline = time.monotonic() + max_wait_s
        while len(batch) < max_batch:
            left = deadline - time.monotonic()
            if left <= 0:
                break
            try:
                batch.append(requests.get(timeout=left))
            except queue.Empty:
                break
        x = torch.stack([item for item, _ in batch])
        with torch.inference_mode():
            y = served(x)
        for (_, reply), out in zip(batch, y):
            reply.put(out)

threading.Thread(target=batcher, daemon=True).start()

def predict(features):
    reply = queue.Queue(maxsize=1)
    requests.put((features, reply))
    return reply.get()

For plain weights the safe load is model.load_state_dict(torch.load("weights.pt", map_location="cpu", weights_only=True)). Production servers (TorchServe, Triton, vLLM, and others) add what this sketch leaves out: per-request timeouts, back-pressure when the queue is too long, multiple model replicas, metrics, and for generation, continuous batching over a paged KV cache.

Tradeoffs

• Eager vs compiled vs exported. Eager is the easiest to debug and change. Compile is a one-line speedup that stays in Python but adds warm-up time and recompiles. Export gives a portable, Python-free artifact but demands a single graph with no data-dependent Python control flow, which can mean rewriting parts of the model.
• Batch size. Larger batches raise throughput and lower cost per request, but every request in a batch waits for the whole batch, so latency rises with it.
• Batching wait. A wait fills batches at light load and only adds latency at heavy load. Many systems keep it at a few milliseconds or tune it adaptively.
• Static vs dynamic shapes. Specialized shapes give the fastest kernels and CUDA graphs but need a graph per shape; dynamic shapes give one graph that is somewhat slower. Padding inputs to a few bucket sizes is a common middle ground.
• CPU vs GPU. CPUs win on small models, low traffic and strict per-request latency; GPUs win when the model is large and traffic fills batches.
• Portability. The exported .pt2 format is still evolving and may not load in a newer PyTorch, so keep the training code and weights to re-export from.

Connections

References

• Crankshaw, Wang, Zhou, Franklin, Gonzalez and Stoica. Clipper: A Low-Latency Online Prediction Serving System. NSDI 2017. usenix.org (Adaptive dynamic batching under a latency target.)
• Yu, Jeong, Kim, Kim and Chun. Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI 2022. usenix.org (Iteration-level scheduling, now called continuous batching.)
• Kwon, Li, Zhuang, Sheng, Zheng, Yu, Gonzalez, Zhang and Stoica. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. arXiv:2309.06180
• Ansel et al. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. ASPLOS 2024. doi:10.1145/3620665.3640366 (TorchDynamo and TorchInductor.)
• PyTorch v2.15.0: torch.export.export, Dim, ExportedProgram.module, torch.compile, aoti_compile_and_package, torch.load, inference_mode.