Export, Quantize, and Serve the LLM

Every downstream path consumes the ExportedProgram, so it comes first. The 2.8 export docs call the API under active development and warn that saved files may not load in newer versions, so record the writing PyTorch version beside a .pt2.

idx = torch.randint(0, 512, (2, 16))          # example (B, T), inside both ranges
ep = torch.export.export(model.eval(), (idx,), dynamic_shapes={
    "idx": {0: Dim("B", min=1, max=64), 1: Dim("T", min=2, max=128)}})

min=2. Tracing specializes sizes 0 and 1, because size 1 broadcasts. A T = 1 example fails with "You marked T as dynamic but your code specialized it to be a constant (1)". Module 28 set the min=2 convention. Here min=1 with a length-16 example also exported, since nothing broadcasts over T.

max=128. wpe has 128 rows, and forward raises ValueError past block_size. Export turns that check into a guard, so asking for max=256 fails with "Suggested fixes: T = Dim('T', max=128)".

Batch. A first probe with a (1, 16) example and only T dynamic rejected batch 2: "Expected input at *args[0].shape[0] to be equal to 1". Unlisted dimensions are static.

One local run (M2 Pro, 4 threads, PyTorch 2.8.0 CPU):

The tied weight is listed twice and stored once. graph_signature.parameters names 53 tensors, including both wte.weight and lm_head.weight, against the model's 52. Placeholder p_wte_weight has no users, and p_lm_head_weight feeds both aten.embedding and the final aten.linear. The tie survives as shared storage (ep.state_dict["wte.weight"] is ep.state_dict["lm_head.weight"]), including after save and load. The .pt2 stores the matrix once, and ep.module() has 52 parameters. Any pass that rewrites parameters by name, a quantizer for example, must handle that alias.

ep.module() returns (logits, None), like GPT.forward. The dropout nodes are there because 2.8 exports a training-level IR; decomposition removes them.

The difference is exactly zero because the program calls the same CPU kernels in the same order as eager. After decomposition or AOTInductor, fusion reorders sums, so check with atol=1e-5 and identical greedy tokens.

The T = 1 row is real 2.8.0 behaviour. The module's input check enforces the upper bound but not the lower one, despite VR[2, 128] (module 28 saw the same lost 0/1 guard). The program was traced assuming T ≥ 2, so ExportedLM reads the range from ep.range_constraints and raises ValueError outside it.

ExportedLM also overrides train(), because ep.module().eval() raises NotImplementedError in 2.8 and module 34's evaluate_nll calls .eval(). It exposes .config so module 35's generate(..., use_cache=False) can crop to block_size. The generation check above ran through that code.

The cached forward from module 35 is the one worth serving, and it does not export, for three reasons:

1. Python state. cache.length is an int that forward reads and writes. Export bakes it in as a constant, so the graph is correct only at the traced position.
2. Data-dependent slicing. cache.k[layer, :, :, :end] has a value-dependent length. With a tensor length it becomes an unbacked symbol needing torch._check hints.
3. Mutation outside the module. KVCache is not an nn.Module, so writes to its tensors are invisible side effects.

There are three ways out:

• (a) Functional cache. Export prefill and decode as two programs. The cache tensors and a position tensor go in, and the updated cache comes out.
• (b) Buffer mutation. Register the cache as buffers and write with index_copy_. Export functionalizes the write and lists it in graph_signature.buffers_to_mutate.
• (c) torch.cond keeps the prefill and decode branches in one program (module 28).

The fragment below is the core of (a). One local run exported it in 0.57 s (177 nodes), fed "Sherlock Holmes" one token at a time, and matched the KVCache path within 2.9e-6.

class DecodeStep(nn.Module):                     # one token in; cache in and out
    def __init__(self, gpt):
        super().__init__()
        self.gpt = gpt

    def forward(self, idx, pos, k_cache, v_cache):  # (B,1), (1,), (L,B,H,S,d_head) x2
        g, B, H, D = self.gpt, idx.size(0), self.gpt.config.n_head, self.gpt.config.d_head
        x = g.wte(idx) + g.wpe(pos)                                   # (B, 1, d_model)
        visible = (torch.arange(k_cache.size(3)) <= pos)[None, :]     # (1, S)
        ks, vs = [], []
        for i, blk in enumerate(g.blocks):
            q, k, v = blk.attn.c_attn(blk.ln_1(x)).split(g.config.d_model, dim=2)
            q, k, v = (t.view(B, 1, H, D).transpose(1, 2) for t in (q, k, v))
            ks.append(k_cache[i].index_copy(2, pos, k))               # functional write
            vs.append(v_cache[i].index_copy(2, pos, v))
            y = F.scaled_dot_product_attention(q, ks[-1], vs[-1], attn_mask=visible)
            x = x + blk.attn.c_proj(y.transpose(1, 2).reshape(B, 1, -1))
            x = x + blk.mlp(blk.ln_2(x))
        return g.lm_head(g.ln_f(x)), torch.stack(ks), torch.stack(vs)

Every step attends over all 128 slots and masks the unwritten ones: the price of static shapes. The payoff on a GPU is one CUDA graph replayed for every step (CUDA graphs). We do not ship it because it duplicates attention outside the frozen llm/model.py.

For each output row \(r\) of \(W \in \mathbb{R}^{d_\text{out} \times d_\text{in}}\):

The error is at most \(s_r/2\), and the symmetric ±127 grid keeps zero exact. Model the error as uniform, with variance \(s_r^2/12\). A Gaussian row of 128–512 weights has absmax ≈ \(3.5\sigma\), so the relative error is about \((3.5\sigma/127)/(\sqrt{12}\sigma) \approx 0.8\%\), matching the measurement below. Scales are per channel because each row is its own dot product, and one tensor-wide scale would let an outlier row set every row's step (fence 3). See module 29, the design space, and Jacob et al..

Int8WeightOnlyLinear holds weight (torch.int8, (d_out, d_in)), scale (fp32, (d_out,)) and bias buffers. Its forward, F.linear(x, weight.float() * scale[:, None], bias), rebuilds the fp32 matrix every call: numerics and storage, not speed. Real kernels dequantize tiles in registers or multiply in int8 (formats and kernels).

The tied head stays fp32. quantize_int8_weight_only deep-copies the model and replaces the 16 block Linears. Replacing lm_head too would give it a private int8 copy: the matrix stored twice, and the lookup reading different numbers from the head.

The Linears shrink 4×, but the fp32 embeddings are now 28 % of the file. In GPT-2 124M, wte alone is 50,257 × 768 = 38.6M parameters. That is why production schemes quantize the embedding too, or use 4 bits for the Linears.

Quality against the fp32 checkpoint on story XI (23,847 targets, stride 128, module 34's evaluate_nll):

Teacher-forced agreement scores the int8 argmax on the fp32 greedy sequence. Free-running comparison is harsher, since one near-tie makes the continuations diverge. Both are perfect here because the layers are well conditioned: the largest row absmax is 1.4–2.2× the median row's. Outlier features near 6.7B parameters forced LLM.int8(), and 4 bits needs error-compensating rounding such as GPTQ. For why weight-only helps bandwidth-bound GPU decoding, see Transformers 14.

torch.ao.quantization. In 2.8.0, calling quantize_dynamic emits DeprecationWarning: torch.ao.quantization is deprecated and will be removed in 2.10, and the message points eager users to torchao's quantize_. The text is DEPRECATION_WARNING in utils.py (line 794), applied by @typing_extensions.deprecated in quantize.py. The package's __init__.py does not warn, so a bare import is silent.

try_ao_dynamic uses it for comparison. It picks an engine from torch.backends.quantized.supported_engines (['qnnpack', 'none'] on this Mac, 'none' active; fbgemm/x86 on Linux), quantizes every Linear except lm_head, measures while the engine is active, and restores it in finally. With no engine it returns None.

Dynamic quantization also rounds activations, hence its roughly 7× larger delta. At these shapes both quantized variants are slower than fp32, because per-call overhead outweighs the saved traffic. A run under background load measured 1.2–2× higher. These timings are a ranking on one machine, and the test asserts none of them.

torchao (v0.13.0, paired with 2.8; not installed, not executed) is the maintained path. quantize_ swaps weights for tensor subclasses (module 20) whose overrides route linear to packed kernels:

from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_
quantize_(model, Int8WeightOnlyConfig())                 # our scheme, real kernels
quantize_(model, Int4WeightOnlyConfig(group_size=32))    # one scale per 32 inputs

PT2E quantizes an exported graph: a quantizer such as XNNPACKQuantizer annotates it, prepare_pt2e inserts observers, calibration runs, and convert_pt2e emits quantize/dequantize pairs for a backend to fuse. The 2.8 copy in torch.ao.quantization.quantize_pt2e is deprecated; torchao v0.13.0 ships torchao.quantization.pt2e.quantize_pt2e.

Leaving Python. These fragments were not executed: Inductor's C++ toolchain failed on the reference Mac (module 23), and ExecuTorch is not installed.

pkg = torch._inductor.aoti_compile_and_package(ep, package_path="gpt_aoti.pt2")  # .so inside
runner = torch._inductor.aoti_load_package(pkg)

from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import to_edge_transform_and_lower                          # v0.7.0
et = to_edge_transform_and_lower(ep, partitioner=[XnnpackPartitioner()]).to_executorch()
open("gpt.pte", "wb").write(et.buffer)

Both start from the 38.2 program (AOTInductor docs, ExecuTorch v0.7.0, whose installer pins torch 2.8.0). ONNX takes the same route via torch.onnx.export(..., dynamo=True), where dynamo defaults to False in 2.8.0 (ONNX docs; not run, onnx is not installed).

A server is a function with a network in front of it, so test the function first:

complete(model, tokenizer, prompt, max_new_tokens=32, temperature=0.0, seed=0) -> dict
# {"text", "n_tokens", "stop_reason": "eot" | "length" | "context", "prompt_tokens", "model_sha256"}

• Validate first. An empty prompt, max_new_tokens <= 0 or non-int, a negative or non-finite temperature, or a prompt at the context limit raises ValueError (HTTP 422 in a server).
• Cap at the context. Asking for 500 tokens after a 5-token prompt returns 123 tokens and "context". The window never slides silently, since sliding recomputes the whole prefix (module 35).
• Own the randomness. Each call seeds its own torch.Generator, so concurrent requests cannot perturb each other. The test checks torch.get_rng_state() is unchanged.
• Stop on EOT and exclude it from the text.
• Report provenance. model_sha256 hashes every state_dict entry, so fp32 and int8 differ.
• Accept any callable (eager, int8, ExportedLM) through module 35's generate(..., use_cache=False).

The HTTP wrapper is a fragment. It was not executed or load-tested:

class CompletionRequest(BaseModel):                    # pydantic
    prompt: str = Field(min_length=1)
    max_new_tokens: int = Field(32, ge=1, le=128)
    temperature: float = Field(0.0, ge=0.0, le=2.0)
    seed: int = 0

app = FastAPI()
MODEL, TOKENIZER, _ = load_pretrained()                # once, at startup
MODEL = quantize_int8_weight_only(MODEL)               # or ExportedLM(load_exported(path), cfg)

@app.post("/v1/complete")
def complete_route(req: CompletionRequest) -> dict:   # sync def runs in FastAPI's threadpool
    with torch.inference_mode():
        return complete(MODEL, TOKENIZER, req.prompt, req.max_new_tokens, req.temperature, req.seed)

Batching, token budgets, streaming and latency SLOs are omitted: see MLOps and serving and the inference course. A real engine replaces the loop inside complete; the contract stays.

Each row is a feature that code in modules 30–38 actually calls:

Read bottom-up, the table is the dependency graph of the course index run in reverse. Tensors (02–06) became token ids, windows and seeded randomness; autograd and nn (07–13) became the model and its loss; training practice (14–19) became a pretraining run that resumes bit for bit; performance and compilers (20–24) and scale and deployment (26–29) became modules 37 and 38. Every capstone module then consumed its predecessor's frozen artifact: corpus and tokenizer hashes (30), token tensors (31), the model interface (32), the committed weights (33), the evaluation functions (34) and the generator (35) that modules 36 and 38 import. Module 01's eager model underlies every row. Three modules have no row of their own: torch.func (09), custom autograd.Function (08) and custom ops/Triton (25). Each would pay off only with a need this small decoder never had.

Next steps follow from the model's limits. Modernize it with RoPE, RMSNorm, SwiGLU and GQA (/courses/transformers/code/modern_decoder.py, Transformers 16); RoPE also removes the 128-position wall. Add data, because one book is about 15 epochs at our budget. Run build-nanogpt on a GPU with bf16, Inductor and the DDP verified in module 37. Each step reuses the contracts drilled here: equal logits, deltas against a frozen reference, hashed artifacts. The PyTorch notes condense the review.

Fence 1: dynamic export round trip. It proves a Dim range survives save/load, in-range outputs equal eager exactly, and the upper bound is enforced. It does not test the lower bound, which 2.8.0 skips (38.3), or any compiled runtime.

import os
import tempfile

import torch
import torch.nn as nn
from torch.export import Dim, export

torch.manual_seed(0)


class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Embedding(32, 16), nn.Linear(16, 32))

    def forward(self, idx):  # idx (B, T) int64 -> logits (B, T, 32)
        return self.net(idx)


model = TinyLM().eval()
example = torch.randint(0, 32, (2, 8))
ep = export(model, (example,), dynamic_shapes={"idx": {1: Dim("T", min=2, max=16)}})
bounds = [str(v) for v in ep.range_constraints.values()]
assert bounds == ["VR[2, 16]"], bounds

with tempfile.TemporaryDirectory() as d:
    path = os.path.join(d, "tiny.pt2")
    torch.export.save(ep, path)
    loaded = torch.export.load(path).module()

with torch.no_grad():
    for T in (3, 16):
        x = torch.randint(0, 32, (2, T))
        torch.testing.assert_close(loaded(x), model(x), rtol=0, atol=0)
try:
    loaded(torch.randint(0, 32, (2, 17)))
    rejected = False
except Exception:
    rejected = True
assert rejected, "T=17 must violate Dim('T', max=16)"
print(f"range {bounds[0]}; T=3 and T=16 match eager exactly; T=17 rejected")

Fence 2: the per-channel bound. It proves \(|W - \hat W| \le s_r/2\) elementwise, one byte per stored weight, and that every row reaches ±127. It says nothing about model quality, which 38.5 measures.

import torch

g = torch.Generator().manual_seed(0)
W = torch.randn(64, 32, generator=g) * 0.05          # (out_features, in_features)

scale = W.abs().amax(dim=1) / 127                     # (64,) one scale per output row
q = torch.round(W / scale[:, None]).clamp(-127, 127).to(torch.int8)
deq = q.float() * scale[:, None]

err = (W - deq).abs()
assert bool((err <= scale[:, None] / 2 + 1e-7).all()), "rounding error must be <= scale/2"
assert q.dtype == torch.int8 and q.untyped_storage().nbytes() == W.numel()
assert int(q.abs().amax(dim=1).min()) == 127          # every row uses the full grid
rel = (err.norm() / W.norm()).item()
assert rel < 0.01
print(f"max err / (scale/2) = {(err / (scale[:, None] / 2)).max():.3f}; "
      f"int8 bytes {q.numel()} vs fp32 {W.numel() * 4}; rel Frobenius err {rel:.4f}")

Fence 3: per-tensor vs per-channel. One row scaled ×100 mimics an outlier channel, so a tensor-wide scale gives the other rows a grid 100× too coarse. Locally 79 % of their weights rounded to zero (relative error 0.707 vs 0.0062 per channel). The fence asserts only the 10× gap. The capstone's layers have no such outlier, so per-tensor was only about 1.7× worse there.

import torch

g = torch.Generator().manual_seed(1)
W = torch.randn(16, 64, generator=g) * 0.02
W[0] *= 100                                           # one outlier row, as in real LLM layers


def roundtrip(w, scale):
    return torch.round(w / scale).clamp(-127, 127) * scale


per_tensor = roundtrip(W, W.abs().max() / 127)
per_channel = roundtrip(W, W.abs().amax(dim=1, keepdim=True) / 127)

small = slice(1, None)                                # the 15 ordinary rows
rel_t = ((W[small] - per_tensor[small]).norm() / W[small].norm()).item()
rel_c = ((W[small] - per_channel[small]).norm() / W[small].norm()).item()
zeroed = (per_tensor[small] == 0).float().mean().item()
assert rel_c * 10 <= rel_t, (rel_c, rel_t)
print(f"ordinary rows: per-tensor rel err {rel_t:.3f} ({zeroed:.0%} of weights rounded to 0), "
      f"per-channel {rel_c:.4f}")

Fence 4: a quantized Linear forward. This inlines the llm/export.py idea for a c_attn-shaped layer. It proves the output stays within 1 % relative error of fp32 and that the state_dict holds torch.int8. Locally, tensor bytes fell 3.79× rather than 4×, because the fp32 bias and scales remain. It does not measure speed.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


class Int8WeightOnlyLinear(nn.Module):
    def __init__(self, linear: nn.Linear):
        super().__init__()
        w = linear.weight.detach()
        scale = w.abs().amax(dim=1).clamp_min(1e-12) / 127
        self.register_buffer("weight", torch.round(w / scale[:, None]).clamp(-127, 127).to(torch.int8))
        self.register_buffer("scale", scale)
        self.register_buffer("bias", linear.bias.detach().clone())

    def forward(self, x):  # dequantize, then an ordinary fp32 matmul
        return F.linear(x, self.weight.float() * self.scale[:, None], self.bias)


fp32 = nn.Linear(128, 384)
int8 = Int8WeightOnlyLinear(fp32)
x = torch.randn(4, 17, 128)
with torch.no_grad():
    ref, out = fp32(x), int8(x)
rel = ((out - ref).norm() / ref.norm()).item()
sd = int8.state_dict()
assert rel < 1e-2, rel
assert sd["weight"].dtype == torch.int8 and sd["scale"].dtype == torch.float32
fp32_bytes = sum(t.numel() * t.element_size() for t in fp32.state_dict().values())
int8_bytes = sum(t.numel() * t.element_size() for t in sd.values())
print(f"output rel err {rel:.4f}; tensor bytes {fp32_bytes} -> {int8_bytes} "
      f"({fp32_bytes / int8_bytes:.2f}x smaller)")

Tested with Python 3.11 and PyTorch 2.8.0 on CPU.

code/38_export_quantize.py runs export, quantize, serve-check or all (default) on load_pretrained()'s model, writing gpt.pt2 and gpt_int8.pt to --output-dir (default: a deleted temporary directory). --try-ao-dynamic adds the deprecated comparison and --bench the 38.6 timings. The JSON has export, quantize, ao_dynamic (or null), serve and upstream, which records whether module 34's evaluate_nll and module 35's generate were found. If they are missing, llm/export.py uses clearly marked local fallbacks, which the test checks against upstream. Scope: Python dequantization shows size and numerics, not speed; no compiled runtime runs; the HTTP fragment is not load-tested.

python content/courses/pytorch/code/38_export_quantize.py all --output-dir /tmp/gpt-export
python site/test_pytorch_38_export.py

One local run (M2 Pro, 4 threads, with both upstream modules present) produced the numbers in 38.2–38.6. Its greedy continuation of "Sherlock Holmes" was .\n\n“It is a few more thank to the matter. You have the more time of.

The test asserts contracts, not those numbers:

• Export: the .pt2 reloads with ranges VR[1, 64] and VR[2, 128]; logits are within 1e-5 of eager at T ∈ {2, 17, 128}; T = 129 and T = 1 are rejected; 32 greedy tokens equal eager.
• Quantization: int8 bytes < 0.5 × fp32; NLL delta in [−0.02, 0.15); 64-token agreement ≥ 0.75; 16 int8 weights, with lm_head the only nn.Linear, still tied to an unchanged fp32 wte.
• Serving: complete() is deterministic at temperature 0 and under a fixed seed, leaves the global RNG alone, rejects bad arguments, stops on EOT, and reports an allowed stop_reason.
• ao: try_ao_dynamic returns a finite NLL or None and always restores the engine.

The test took 4.5 s wall locally, about 0.7 s of which was import torch.