I'm planning a comprehensive guide on quantizing a ViT model with torchao, covering its installation, the core `quantize_` API, and the different config types like weight-only and dynamic activation quantization at int8/int4 precision. I need to address ViT-specific considerations too: the model's linear-heavy structure, how the patch embedding conv layer needs separate handling, using filter_fn to exclude certain modules, hardware/kernel dimension constraints for int4 and fp8, and compute capability requirements. I also want to include practical examples with timm/torchvision/HuggingFace ViT models, benchmarking and accuracy evaluation, model saving/loading, and QAT and calibrated PTQ approaches. I should note the current torchao API uses Config classes like Int4WeightOnlyConfig and Float8DynamicActivationFloat8WeightConfig passed to quantize_, while older versions used function-style calls like int8_weight_only() — worth mentioning both since users may encounter either in the wild, along with the filter_fn signature for quantize_. I'm also recalling the granularity imports from torchao's quantization module, and thinking about the ViT quantization example from the torchao repo that loads a vision transformer model and applies quantization to it. I'm thinking through the ordering nuance between quantize_ and torch.compile — since compile is lazy, calling quantize before or after compile shouldn't matter as long as the model hasn't run yet, though the safer documented pattern is to quantize first then compile. Without compilation, weight-only quantization can actually be slower due to per-op dequant overhead, so I want to stress that compile isn't optional for real speedups. I should also flag some inductor config flags that tune fused int matmul behavior, note that int4 tinygemm requires bf16 weights while int8 works with bf16/fp16, and account for ViT-B/16's specific dimensions (768 embed dim, 3072 MLP, 2304 combined qkv) when reasoning about shape constraints. I'm checking whether the 1000-class classifier output satisfies tinygemm's divisibility rules -- it needs n divisible by 8, which 1000 satisfies, and k divisible by group_size multiples, which 768 and 3072 handle fine at group size 128. I'm also trying to recall the exact shape constraints for int8 dynamic quantization using cuBLAS's int8 matmul, since those kernels have their own alignment requirements on k and n. For unsupported shapes, the practical solution is using filter_fn to skip those layers so torchao gracefully falls back rather than erroring. I should also mention unwrap_tensor_subclass() for compatibility with PyTorch versions before 2.5 during export or compile, and note the separate PT2E quantization flow using XNNPACKQuantizer for executorch/XNNPACK deployment, which follows a prepare/convert pattern distinct from the eager-mode API. For QAT, I'll present the newer unified API in torchao 0.13+, which uses QATConfig wrapping a base config through prepare and convert steps, while mentioning the older Int8DynActInt4WeightQATQuantizer for context. I should also briefly note calibration and static quant tools like the observer utilities and the newer Float8StaticActivationFloat8WeightConfig, pointing to the static_quant tutorial. For accuracy, I want to reference typical degradation patterns on ImageNet: int8 weight-only is nearly lossless, int8 dynamic loses under 0.5%, while int4 group-wise can drop 1-3% on ViT, sometimes requiring HQQ, AWQ, or skipping sensitive layers. For ViT specifically, attention stays in bf16 since SDPA isn't quantized by these APIs, so speedup is capped by the linear layer share of FLOPs (roughly 70-80% in ViT-B at 224 resolution). Batch size matters too: weight-only int8 shines at small batch sizes since it's memory-bound, while dynamic int8/fp8 wins at larger batches where compute dominates. I also want to mention the autoquant API, which benchmarks different quantization strategies automatically when wrapped around a compiled model and run on sample input. For saving and loading quantized models, the recommended pattern involves creating the model on a meta device, applying the same quantization config to set up the tensor subclass structure, then loading the saved state dict with `assign=True` so the quantized weights get properly placed without needing to reconstruct them from scratch. I also want to cover the HuggingFace transformers integration path, where you can pass a `TorchAoConfig` wrapping a torchao quantization config directly into `from_pretrained`, letting HF handle the quantization automatically when loading a model like ViT with a specified dtype and device map. I should mention the CPU-specific quantization option too, since dynamic activation quantization has a distinct config for CPU inference. There's also int4 weight-only quantization with a CPU-specific layout, and an int8-dynamic/intx-weight config aimed at ARM and mobile deployment via KleidiAI. Separately, torchao's float8 module is for training rather than inference — worth distinguishing. On the sparsity side, semi-structured 2:4 sparsity combined with int8 quantization gave a notable speedup in the SAM/ViT benchmarks, which is worth citing. Now I'm thinking through how someone would verify quantization actually applied — printing the model to see the weight tensor type, inspecting a specific layer's weight, and checking memory via CUDA's max allocated stat. I want to structure the full answer around a mental model of what torchao does, setup requirements, a decision table matching recipes to hardware/goals, a minimal working example, and a breakdown of a ViT's architecture showing what gets quantized where. Beyond that I'm planning to cover each recipe in depth (int8 weight-only, int8 dynamic, fp8, int4 variants, sparse+int8), torch.compile/inductor details, benchmarking and accuracy harnesses, autoquant, serialization, framework integration, QAT, and edge deployment, plus a troubleshooting table and expected performance numbers — while keeping the code examples realistic and flagging where APIs might have shifted between versions. I'm working out the input/output shapes for each layer of a torchvision ViT-B/16 to ground the module-filtering example. I catch an important subtlety: torchvision's ViT uses `nn.MultiheadAttention` internally, so the QKV projections are stored as a raw weight parameter rather than an `nn.Linear` submodule, meaning torchao's default `quantize_` filter will silently skip them since it only targets `nn.Linear` instances. Only the `out_proj` (a `NonDynamicallyQuantizableLinear`, itself a `Linear` subclass) would actually get picked up by the standard `_is_linear` check, which is a gotcha worth calling out explicitly. This means timm's ViT, which implements QKV as an explicit `nn.Linear(768, 2304)`, will quantize far more of the attention FLOPs than torchvision's implementation — so for benchmarking purposes I should recommend timm or HF over torchvision. I'm now mapping out the equivalent module structure for HF's `ViTModel` (separate query/key/value Linear layers, output dense, intermediate dense, patch embedding conv, classifier head) so I can compare it against timm's structure next. I'm sketching a `filter_fn` example that selectively skips quantizing certain linear layers — excluding the classification head to preserve precision and skipping small linear layers that wouldn't benefit from quantization. I also need to check torchao's internal helpers like `_is_linear` and `_replace_with_custom_fn_if_matches_filter`, plus verify group_size divisibility for int4 quantization — 768 and 3072 divide evenly by 128, and I'm checking whether the head's 1000 output features satisfy the tinygemm requirement of being divisible by 8. Checking ViT-L and ViT-H dimensions too — all divide cleanly by 128, and the inner_k_tiles=8 constraint requiring k % 128 == 0 also holds. Now I'm mapping out hardware requirements: int8 weight-only quantization works broadly across CUDA and CPU since it mainly benefits memory bandwidth, while int8 dynamic activation quantization needs int8 tensor core support on newer CUDA architectures. I'm also noting that int4 tinygemm needs bf16 with newer CUDA support, fp8 requires specific architectures (L4/L40S/4090/H100), and the newest mxfp8/mxfp4/nvfp4 formats need Blackwell-class GPUs. I'll briefly mention the relevant torchao config classes for these newer formats, then move into writing a thorough, well-organized answer with headers, code blocks, and tables while keeping it as tight as possible given the request for thoroughness. I want to avoid fabricating specific benchmark numbers since I'm not certain of exact figures — I'll give qualitative expectations with ranges, note them as approximate, and recommend the user measure their own results rather than citing precise claims I can't verify. Since batch size varies but sequence length doesn't for ViT, I should note that static shape compilation works well, but variable batch sizes cause recompilation unless the batch dimension is marked dynamic—also worth mentioning the dynamo cache size limit config. For the static quantization example, I'm sketching out the calibration/observer flow using torchao's float8 static activation and weight quantization config with per-tensor granularity. That requires computing a scale via observers, so I'll reference the calibration tutorial at a high level rather than risk inaccurate API details, pointing to the relevant torchao tutorial folder. Now I'm laying out the full answer structure: quickstart, mental model of torchao quantization, install/hardware matrix, what parts of a ViT get quantized, a full end-to-end timm ViT-B/16 example, recipe-by-recipe breakdown, filter_fn control, why torch.compile is essential, benchmarking/accuracy harness, autoquant, save/load, and HF transformers integration. I'm continuing the outline with QAT, combining sparsity, CPU/edge deployment, troubleshooting, and finally a decision-path recommendation before diving into the full write-up. # Quantizing a Vision Transformer with torchao — Complete Guide --- ## 0. TL;DR quickstart ```python import torch, timm from torchao.quantization import quantize_, Int8DynamicActivationInt8WeightConfig model = timm.create_model("vit_base_patch16_224", pretrained=True) model = model.eval().cuda().to(torch.bfloat16) # 1) swap weights for quantized tensor subclasses (in-place, mutates `model`) quantize_(model, Int8DynamicActivationInt8WeightConfig()) # 2) compile — this is where the actual speedup comes from model = torch.compile(model, mode="max-autotune", fullgraph=True) x = torch.randn(32, 3, 224, 224, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): y = model(x) # first call = warmup/autotune ``` Everything else in this document is about *why* those three lines are what they are, how to choose the config, and how to not shoot yourself in the foot. --- ## 1. Mental model: what torchao actually does torchao's inference quantization is **not** a graph rewrite and **not** a custom runtime. It works like this: 1. `quantize_(model, config)` walks the module tree, and for every module matching a filter (default: `isinstance(m, nn.Linear)`), it **replaces `module.weight` (an `nn.Parameter` wrapping a plain `torch.Tensor`) with a tensor subclass** — e.g. `AffineQuantizedTensor` / `Int4Tensor` / `Float8Tensor`. The subclass holds `int_data`, `scale`, `zero_point`, and a layout. 2. The subclass overrides `__torch_function__` / `__torch_dispatch__`, so when `F.linear(x, weight)` is called, it dispatches to a quantized kernel (or to a dequant-then-matmul fallback). 3. `torch.compile` traces through the subclass, sees the dequant/quant math as ordinary ATen ops, and **fuses them** (e.g. fuses dequantization into the matmul epilogue, fuses activation quantization into the preceding op). **Without `torch.compile`, weight-only quantization is usually slower than bf16** because you pay a full materialized dequant on every forward. Consequences worth internalizing: - The model is **still an `nn.Module` with the same structure**. `print(model)` looks almost identical. `model(x)` works. Autograd mostly works (for QAT paths). - `quantize_` is **in-place** and **returns `None`**. `model = quantize_(model, cfg)` is a classic bug that gives you `None`. - Only ops the subclass knows about get accelerated. **Attention (SDPA), LayerNorm, GELU, and the patch-embedding Conv2d stay in bf16 by default.** - Because it's a tensor subclass, saving/loading and `torch.export` need a little care (§10). --- ## 2. Install and hardware matrix ```bash pip install torch --index-url https://download.pytorch.org/whl/cu124 # 2.6+ strongly recommended pip install torchao pip install timm # or transformers, or torchvision ``` Use **PyTorch ≥ 2.5** (2.6/2.7+ preferred). On PyTorch < 2.5 you must call `torchao.utils.unwrap_tensor_subclass(model)` before `torch.compile`/`torch.export`; on 2.5+ that's unnecessary and deprecated. | Recipe | Config class | Min GPU arch | Compute dtype | Primary win | |---|---|---|---|---| | int8 weight-only | `Int8WeightOnlyConfig()` | sm70+ (works on CPU too) | bf16/fp16 | 2× weight memory, latency at **small batch** | | int8 dynamic act + int8 weight | `Int8DynamicActivationInt8WeightConfig()` | sm75/sm80+ (`torch._int_mm`) | bf16/fp16 | throughput at **large batch** | | int4 weight-only (group-wise) | `Int4WeightOnlyConfig(group_size=128)` | sm80+ | **bf16 only** | 4× weight memory, small-batch latency | | fp8 dynamic act + fp8 weight | `Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())` | **sm89+** (L40S/4090/H100) | bf16 | best accuracy/speed tradeoff at large batch | | fp8 weight-only | `Float8WeightOnlyConfig()` | sm89+ | bf16 | memory | | mxfp8 / nvfp4 | `MXFPInferenceConfig()`, `NVFP4InferenceConfig()` | **sm100** (Blackwell) | bf16 | next-gen | | int8 dyn act + 2:4 sparse weight | `Int8DynamicActivationInt8SemiSparseWeightConfig()` | sm80+ | bf16/fp16 | extra ~1.1–1.2× on top of int8 | | intx / int8-dyn-act-int4-weight | `Int8DynamicActivationInt4WeightConfig()`, `Int8DynamicActivationIntxWeightConfig()` | CPU / ARM (KleidiAI) | fp32/bf16 | edge deployment | > **API note on versions.** Older torchao (≤0.7) used *functions*: `int8_weight_only()`, `int4_weight_only()`, `int8_dynamic_activation_int8_weight()`, `float8_dynamic_activation_float8_weight()`. These still work but are deprecated aliases for the `*Config` dataclasses shown above. Newer versions also added a `version=` / `int4_packing_format=` argument to `Int4WeightOnlyConfig`. Always check `torchao.__version__` and `help(Int4WeightOnlyConfig)` if a kwarg is rejected. --- ## 3. Know your ViT: what will and won't get quantized This is the single most ViT-specific thing in this guide, and it's where people silently lose half their speedup. `quantize_`'s default filter matches **only `nn.Linear`**. So the layout of your particular ViT implementation matters enormously. ### timm (`vit_base_patch16_224`) — ✅ ideal ``` patch_embed.proj Conv2d(3, 768, k=16, s=16) ← NOT quantized (conv) blocks.{0..11}.attn.qkv Linear(768 -> 2304) ← quantized blocks.{0..11}.attn.proj Linear(768 -> 768) ← quantized blocks.{0..11}.mlp.fc1 Linear(768 -> 3072) ← quantized blocks.{0..11}.mlp.fc2 Linear(3072 -> 768) ← quantized head Linear(768 -> 1000) ← quantized (consider skipping) ``` All the heavy GEMMs are real `nn.Linear`. ### HuggingFace `transformers` ViT — ✅ good ``` vit.embeddings.patch_embeddings.projection Conv2d ← not quantized vit.encoder.layer.{i}.attention.attention.query Linear(768,768) .key Linear(768,768) .value Linear(768,768) vit.encoder.layer.{i}.attention.output.dense Linear(768,768) vit.encoder.layer.{i}.intermediate.dense Linear(768,3072) vit.encoder.layer.{i}.output.dense Linear(3072,768) classifier Linear(768, num_labels) ``` Note Q/K/V are three separate Linears — slightly less efficient than timm's fused `qkv`, but fully quantizable. ### torchvision `vit_b_16` — ⚠️ **gotcha** torchvision builds attention with `nn.MultiheadAttention`, whose input projection is a **raw `Parameter` (`in_proj_weight`), not an `nn.Linear` module**. `quantize_` cannot see it. You'll quantize `out_proj`, `mlp.0`, `mlp.3` and the head, but **the entire QKV projection (~1/3 of attention GEMM FLOPs) stays bf16**. If you must use torchvision, either (a) accept the loss, (b) monkey-patch attention into explicit Linears, or (c) just use timm/HF. I recommend (c). ### FLOP budget (why ViT quantization pays off at all) For ViT-B/16 @ 224 (197 tokens), the encoder is roughly: - MLP GEMMs: ~50% of FLOPs - QKV + out_proj GEMMs: ~25% - Attention `QK^T` and `AV` (SDPA): ~10% - Norms / GELU / softmax / patch conv: the rest So `nn.Linear` quantization covers **~75% of the math**. With Amdahl's law, an ideal 2× on those gives ~1.6× end-to-end at best — and in practice less, because small batches are bandwidth-bound and quantization overhead is real. **Do not expect 2×.** Realistic: 1.2–1.5× latency at large batch, plus a solid ~2× (int8) or ~4× (int4) reduction in weight memory. Larger models (ViT-L/H/g) and higher resolutions benefit more, because the Linear fraction and the memory pressure both grow. --- ## 4. End-to-end example (timm ViT-B/16, int8 dynamic) ```python import time import torch import torch.nn as nn import timm from torchao.quantization import quantize_, Int8DynamicActivationInt8WeightConfig torch.set_float32_matmul_precision("high") DEV, DTYPE = "cuda", torch.bfloat16 def build(): m = timm.create_model("vit_base_patch16_224", pretrained=True) return m.eval().to(device=DEV, dtype=DTYPE) # ---------- filter: which Linears to touch ---------- def vit_filter(module: nn.Module, fqn: str) -> bool: if not isinstance(module, nn.Linear): return False if fqn in ("head",): # keep classifier in bf16: cheap + accuracy-sensitive return False if min(module.weight.shape) < 256: # tiny layers: overhead > benefit return False return True # ---------- quantize ---------- model = build() quantize_(model, Int8DynamicActivationInt8WeightConfig(), filter_fn=vit_filter) # verify the swap actually happened print(type(model.blocks[0].mlp.fc1.weight)) # -> (or LinearActivationQuantizedTensor) print(type(model.head.weight)) # -> (skipped, as intended) # ---------- compile ---------- model = torch.compile(model, mode="max-autotune", fullgraph=True) # ---------- run ---------- x = torch.randn(64, 3, 224, 224, device=DEV, dtype=DTYPE) with torch.no_grad(): for _ in range(5): model(x) # warmup + autotune torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(50): model(x) torch.cuda.synchronize() print(f"{(time.perf_counter()-t0)/50*1e3:.2f} ms/iter") print(f"peak mem: {torch.cuda.max_memory_allocated()/2**30:.2f} GiB") ``` **Ordering note:** `torch.compile` is lazy, so `quantize_` → `compile` and `compile` → `quantize_` both work *as long as you haven't run the model yet*. Prefer **quantize first, then compile** — it's less surprising and avoids recompiles if you ever run the fp model. --- ## 5. Recipe-by-recipe deep dive ### 5.1 int8 weight-only — the safe default ```python from torchao.quantization import quantize_, Int8WeightOnlyConfig quantize_(model, Int8WeightOnlyConfig(), filter_fn=vit_filter) ``` - Per-output-channel (row-wise) symmetric int8 weights; activations stay bf16. - **Accuracy: essentially free.** For ViT-B/L/H I'd expect ≤0.1% top-1 delta. This is the one recipe you can ship without an eval, though you should still eval. - **Speed:** wins when you're **memory-bandwidth bound**, i.e. batch size 1–8, or when the model is huge (ViT-H/g). At large batch the GEMM becomes compute-bound and there's no win (activations are still bf16, so the matmul is still a bf16 matmul after dequant) — sometimes a small *loss*. - Halves weight memory. Great for fitting a big ViT on a small GPU. ### 5.2 int8 dynamic activation + int8 weight — best int8 throughput ```python from torchao.quantization import quantize_, Int8DynamicActivationInt8WeightConfig quantize_(model, Int8DynamicActivationInt8WeightConfig(), filter_fn=vit_filter) ``` - Activations quantized **per-token** (per-row) on the fly, weights per-channel; the GEMM is a true int8 tensor-core matmul (`torch._int_mm`) with a fused rescale epilogue. - **This is the recipe that helps at large batch / high throughput serving.** - **Accuracy:** typically 0.1–0.6% top-1 drop on ViT. The risk is activation outliers — ViTs are known to have large-magnitude activations in the post-LayerNorm residual stream and in `mlp.fc2` inputs (post-GELU). If you see a big drop, the fix is almost always "skip `fc2`" or "skip the first and last block" (see §6.3). - Requires reasonable shapes for the int8 kernel. ViT-B/L/H hidden dims (768/1024/1280, 3072/4096/5120) are all fine. Odd projection dims can fall back to a slow path. ### 5.3 fp8 dynamic — best overall on Hopper/Ada ```python from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig from torchao.quantization import PerRow, PerTensor # also torchao.quantization.granularity quantize_( model, Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()), filter_fn=vit_filter, ) ``` - `PerRow()` (per-token activation scale, per-channel weight scale) is much more robust than `PerTensor()` and is what you want for ViT. `PerTensor()` is slightly faster but more fragile. - Uses `torch._scaled_mm` → fp8 tensor cores. **Needs sm89+** (L4, L40S, RTX 4090, H100). On sm80 (A100) there are no fp8 tensor cores; use int8 instead. - **Accuracy: usually better than int8 dynamic** (fp8e4m3 has a much wider dynamic range than int8, which is exactly the ViT activation-outlier problem), often within noise of bf16. - Also check `torchao.quantization.Float8WeightOnlyConfig()` for the weight-only fp8 variant. If you're on H100 and want both speed and accuracy with minimal fuss: **this is the recipe.** ### 5.4 int4 weight-only — maximum compression ```python from torchao.quantization import quantize_, Int4WeightOnlyConfig model = model.to(torch.bfloat16) # MANDATORY: int4 tinygemm path is bf16-only cfg = Int4WeightOnlyConfig(group_size=128) # 128 | 64 | 32 ; smaller = better accuracy, more scales quantize_(model, cfg, filter_fn=vit_filter) ``` Constraints to respect: - **Weights must be bfloat16** before quantizing (fp16/fp32 will error or silently take a slow path). - `in_features % group_size == 0`. ViT-B: 768 and 3072 are both divisible by 128. ✅ ViT-L: 1024/4096 ✅. But a custom model with, say, `in_features=1000` will fail — filter it out or use a smaller `group_size`. - `out_features` must be divisible by 8 for the tinygemm packing. - sm80+. Accuracy: **this is where ViTs get unhappy.** ViT-B/16 int4 group-128 commonly loses 1–4% top-1. Mitigations, in order of effort: 1. **Smaller group size**: `group_size=64` or `32`. 2. **HQQ** (Half-Quadratic Quantization) — a better weight rounding, free at quantize time: ```python Int4WeightOnlyConfig(group_size=64, use_hqq=True) ``` 3. **Skip sensitive layers** — keep `patch_embed`, `head`, and blocks 0 and −1 in int8 or bf16 (§6.3). Mixed precision is usually the highest-ROI fix. 4. **AWQ** (activation-aware weight scaling, needs calibration data): ```python from torchao.prototype.awq import AWQConfig # 1) insert observers, 2) run N calibration batches, 3) convert ``` API lives under `torchao.prototype.awq` and has moved around between versions — read the docstrings in your installed version. 5. **QAT** (§12) if you need int4 and can afford a fine-tune. Also available: `UIntXWeightOnlyConfig(dtype=torch.uint3..uint7, group_size=...)` and `FPXWeightOnlyConfig(ebits, mbits)` for sub-4-bit / fp6 experiments — prototype quality. ### 5.5 int8 + 2:4 semi-structured sparsity ```python from torchao.quantization import quantize_, Int8DynamicActivationInt8SemiSparseWeightConfig quantize_(model, Int8DynamicActivationInt8SemiSparseWeightConfig(), filter_fn=vit_filter) ``` This assumes the weights are *already* 2:4 sparse (or you accept the accuracy hit of magnitude-pruning them in one shot). For real use, prune + fine-tune with `torchao.sparsity` first. Good for a further ~10–20% on top of int8 on Ampere+. This is the combination used in the PyTorch "Accelerating Segment Anything" work, which is a ViT-H backbone. ### 5.6 Static activation quantization (calibrated) Dynamic activation quantization costs a small per-op reduction to compute scales. If you want to eliminate that, use static scales derived from calibration: ```python from torchao.quantization import quantize_, Float8StaticActivationFloat8WeightConfig, PerTensor # You must supply the activation scale, obtained by observing calibration data. quantize_(model, Float8StaticActivationFloat8WeightConfig(scale=my_scale, granularity=PerTensor())) ``` The full observer-based calibration flow (`insert_observers_`, `AffineQuantizedObserver`, min/max or moving-average observers) is documented in the torchao repo under `tutorials/calibration_flow/static_quant.py`. For ViT, static activation quant is noticeably riskier than dynamic because of outlier tokens (CLS token in particular behaves very differently from patch tokens). **Use dynamic unless you've measured that the scale-computation overhead matters.** --- ## 6. `filter_fn`: surgical control Signature: `filter_fn(module: nn.Module, fqn: str) -> bool`, where `fqn` is the fully-qualified name (`"blocks.3.mlp.fc1"`). Return `True` to quantize. ### 6.1 Inspect first ```python import torch.nn as nn for fqn, m in model.named_modules(): if isinstance(m, (nn.Linear, nn.Conv2d)): print(f"{fqn:50s} {type(m).__name__:8s} {tuple(m.weight.shape)}") ``` ### 6.2 Reusable helpers ```python import re import torch.nn as nn def only(*patterns): pats = [re.compile(p) for p in patterns] def f(m, fqn): return isinstance(m, nn.Linear) and any(p.search(fqn) for p in pats) return f def exclude(*patterns): pats = [re.compile(p) for p in patterns] def f(m, fqn): return isinstance(m, nn.Linear) and not any(p.search(fqn) for p in pats) return f # e.g. MLPs only (the biggest GEMMs, usually the most robust to quantization) quantize_(model, Int8DynamicActivationInt8WeightConfig(), filter_fn=only(r"\.mlp\.")) ``` ### 6.3 Mixed precision by layer sensitivity (very effective for ViT) `quantize_` can be called repeatedly with different configs and disjoint filters. Order matters only in that a Linear whose weight is already a subclass will be skipped by the default `_is_linear` check. ```python from torchao.quantization import ( quantize_, Int4WeightOnlyConfig, Int8WeightOnlyConfig, ) NUM_BLOCKS = len(model.blocks) SENSITIVE = {0, 1, NUM_BLOCKS - 1} # first + last blocks are usually the most sensitive def is_sensitive(m, fqn): if not isinstance(m, nn.Linear): return False mo = re.match(r"blocks\.(\d+)\.", fqn) if mo and int(mo.group(1)) in SENSITIVE: return True return fqn in ("head",) # and the classifier def is_robust(m, fqn): return isinstance(m, nn.Linear) and not is_sensitive(m, fqn) quantize_(model, Int8WeightOnlyConfig(), filter_fn=is_sensitive) # 8-bit for the fragile ones quantize_(model, Int4WeightOnlyConfig(group_size=64), filter_fn=is_robust) ``` **A practical sensitivity-ranking loop** (worth the 20 minutes): ```python def per_layer_sensitivity(build_fn, config, eval_fn, layer_fqns): baseline = eval_fn(build_fn()) out = {} for fqn in layer_fqns: m = build_fn() quantize_(m, config, filter_fn=lambda mod, name, t=fqn: name == t) out[fqn] = baseline - eval_fn(torch.compile(m)) return dict(sorted(out.items(), key=lambda kv: -kv[1])) ``` Quantize one layer at a time, measure the drop, then keep the worst offenders in higher precision. On ViTs the usual suspects are `patch_embed`/first block, `mlp.fc2` (post-GELU inputs have heavy tails), and the head. ### 6.4 Quantizing the patch-embedding Conv2d `quantize_` targets Linear only. A ViT's `patch_embed.proj` is a `Conv2d(3, D, k=P, s=P)` — a non-overlapping conv, which is **mathematically identical to a Linear on flattened patches**. If you want it quantized, rewrite it: ```python class PatchEmbedLinear(nn.Module): def __init__(self, conv: nn.Conv2d): super().__init__() self.p = conv.kernel_size[0] self.proj = nn.Linear(conv.in_channels * self.p * self.p, conv.out_channels, bias=conv.bias is not None) with torch.no_grad(): self.proj.weight.copy_(conv.weight.reshape(conv.out_channels, -1)) if conv.bias is not None: self.proj.bias.copy_(conv.bias) def forward(self, x): # x: (B, C, H, W) x = torch.nn.functional.unfold(x, self.p, stride=self.p) # (B, C*p*p, L) return self.proj(x.transpose(1, 2)) # (B, L, D) ``` Honestly: **usually not worth it.** For ViT-B/16 the patch embed is ~1% of FLOPs and it's one of the most quantization-sensitive layers. Leave it alone. --- ## 7. `torch.compile`: the non-optional part ```python model = torch.compile(model, mode="max-autotune", fullgraph=True) ``` Useful knobs: ```python import torch torch._inductor.config.coordinate_descent_tuning = True torch._inductor.config.coordinate_descent_check_all_directions = True torch._inductor.config.force_fuse_int_mm_with_mul = True # fuses int8 mm + rescale (older versions) torch._dynamo.config.cache_size_limit = 32 # if you use many batch sizes ``` (Several of these are defaults in recent PyTorch; setting them is harmless.) **Modes** - `mode="default"` — fast to compile, gets most of the fusion benefit. - `mode="max-autotune"` — benchmarks Triton/CUTLASS GEMM configs; can take minutes but gives the best numbers. Use for serving. - `mode="max-autotune-no-cudagraphs"` — if CUDA graphs cause trouble (dynamic control flow, CPU-side syncs). **Shapes.** ViT is beautifully static: fixed 224×224 → fixed 197 tokens. Only the batch dimension varies. - Fixed batch: compile as-is, static shapes, best perf. - Variable batch: either pad to a small set of bucket sizes and accept a handful of recompiles, or mark it dynamic: ```python torch._dynamo.mark_dynamic(x, 0) ``` Dynamic shapes cost a bit of perf and sometimes disable `max-autotune` benefits. Bucketing + padding is usually faster in practice. **`fullgraph=True`** will error on graph breaks instead of silently degrading — use it during development so you find out that, e.g., your custom attention has a `.item()` in it. **Verify the fusion happened.** If quantized is not faster than bf16, dump the generated kernels: ```bash TORCH_LOGS="output_code" python bench.py 2>&1 | tee inductor.log grep -c "_int_mm\|_scaled_mm" inductor.log ``` If you see `to(torch.bfloat16)` on the weight followed by a plain `mm` for every layer, the dequant wasn't fused and you're on the slow fallback. --- ## 8. Benchmark + accuracy harness ### 8.1 Latency / throughput / memory ```python import time, torch def bench(model, batch=64, dtype=torch.bfloat16, iters=50, warmup=10, res=224): x = torch.randn(batch, 3, res, res, device="cuda", dtype=dtype) torch.cuda.reset_peak_memory_stats() with torch.no_grad(): for _ in range(warmup): model(x) torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(iters): model(x) torch.cuda.synchronize() dt = (time.perf_counter() - t0) / iters return { "ms": dt * 1e3, "img/s": batch / dt, "peak_GiB": torch.cuda.max_memory_allocated() / 2**30, } def weight_bytes(model): tot = 0 for p in model.parameters(): # tensor subclasses report the *packed* element size correctly via nbytes in recent torchao; # if not, sum the inner tensors: try: tot += p.nbytes except Exception: tot += p.numel() * p.element_size() return tot / 2**20 # MiB ``` **Always sweep batch size.** The ranking of recipes flips: at batch 1 weight-only int4/int8 wins; at batch 128 fp8/int8-dynamic wins and weight-only can be a regression. ### 8.2 ImageNet top-1 ```python import torch, timm from timm.data import resolve_data_config, create_transform from torchvision.datasets import ImageFolder @torch.no_grad() def eval_top1(model, model_for_cfg, val_dir, batch=64, dtype=torch.bfloat16, limit=None): cfg = resolve_data_config({}, model=model_for_cfg) ds = ImageFolder(val_dir, transform=create_transform(**cfg, is_training=False)) dl = torch.utils.data.DataLoader(ds, batch_size=batch, shuffle=False, num_workers=8, pin_memory=True, drop_last=False) correct = n = 0 for i, (x, y) in enumerate(dl): x = x.to("cuda", dtype=dtype, non_blocking=True) pred = model(x).float().argmax(-1).cpu() correct += (pred == y).sum().item(); n += y.numel() if limit and n >= limit: break return 100.0 * correct / n ``` Pass the *unquantized* model (or just the timm model object) as `model_for_cfg` so `resolve_data_config` can read `default_cfg`. A 5–10k image subset of the val set is plenty for ranking recipes (±0.3% noise); use the full 50k for the final number. ### 8.3 Quick sanity check without a dataset ```python @torch.no_grad() def logit_fidelity(fp_model, q_model, x): a, b = fp_model(x).float(), q_model(x).float() return { "rel_l2": ((a - b).norm() / a.norm()).item(), "top1_agree": (a.argmax(-1) == b.argmax(-1)).float().mean().item(), "kl": torch.nn.functional.kl_div(b.log_softmax(-1), a.softmax(-1), reduction="batchmean").item(), } ``` `rel_l2 < 0.02` and `top1_agree > 0.99` on a few hundred real images ⇒ you almost certainly haven't broken anything. This catches "I quantized the wrong thing" bugs in 10 seconds. --- ## 9. `autoquant`: let torchao pick per layer If you don't want to reason about any of this: ```python import torch, torchao model = build() # bf16, eval, cuda model = torch.compile(model, mode="max-autotune") model = torchao.autoquant(model) # wraps params in AutoQuantizableLinearWeight x = torch.randn(64, 3, 224, 224, device="cuda", dtype=torch.bfloat16) model(x) # <-- benchmarks every candidate kernel for every layer shape; SLOW (minutes) model(x) # now running the chosen quantization ``` Details: - It micro-benchmarks each candidate (bf16, int8-wo, int8-dyn, …) **for the actual shapes it sees**, per layer, and keeps the fastest. This means autoquant is **batch-size-specific** — run it with your production batch size. - Restrict the search space: ```python torchao.autoquant(model, qtensor_class_list=torchao.quantization.DEFAULT_INT4_AUTOQUANT_CLASS_LIST) # also: DEFAULT_AUTOQUANT_CLASS_LIST, OTHER_AUTOQUANT_CLASS_LIST (fp8), ALL_AUTOQUANT_CLASS_LIST ``` - Finalize explicitly with multiple calibration shapes: ```python model = torchao.autoquant(model, manual=True) for bs in (1, 8, 32): model(torch.randn(bs,3,224,224, device="cuda", dtype=torch.bfloat16)) model.finalize_autoquant() ``` - `set_inductor_config=True` (default) mutates global inductor flags. - **Caveat:** autoquant optimizes *latency only*, with no accuracy awareness. Always eval afterwards. For ViT it will happily pick int8-dynamic on `mlp.fc2`, which may be exactly the layer you wanted to protect. --- ## 10. Saving and loading quantized ViTs The state dict contains tensor subclasses, so plain `load_state_dict` into a freshly built fp model won't work (shapes/dtypes differ). The supported pattern: ```python # ---- save ---- quantize_(model, Int8WeightOnlyConfig()) torch.save(model.state_dict(), "vit_b16_int8.pt") # ---- load ---- with torch.device("meta"): model = timm.create_model("vit_base_patch16_224", pretrained=False) model = model.to(torch.bfloat16) sd = torch.load("vit_b16_int8.pt", map_location="cuda", weights_only=False) model.load_state_dict(sd, assign=True) # assign=True is essential model = model.eval() ``` Notes: - **`assign=True`** replaces the parameter objects instead of copying into them — required because the destination is a meta tensor / different dtype. - **`weights_only=False`** is needed because the pickle references torchao classes. To keep `weights_only=True` (safer), register the classes first: ```python import torch.serialization as ts from torchao.dtypes.affine_quantized_tensor import AffineQuantizedTensor from torchao.quantization.quant_primitives import MappingType # ... plus the layout/tensor-impl classes your config uses ts.add_safe_globals([AffineQuantizedTensor, MappingType]) ``` Recent torchao exposes `torchao.utils.get_safe_globals()` or similar helpers; check your version. - The checkpoint is **not portable across torchao versions** for prototype layouts. Pin `torchao` in your serving image, or re-quantize from the fp checkpoint at load time (cheap for a ViT: a few seconds). - Simplest robust alternative: **ship the fp16/bf16 checkpoint and run `quantize_` at startup.** For ViT-scale models the quantization cost is negligible and you sidestep all serialization issues. ### Export / AOTInductor ```python example = (torch.randn(1, 3, 224, 224, device="cuda", dtype=torch.bfloat16),) ep = torch.export.export(model, example) # PyTorch 2.5+: no unwrap needed path = torch._inductor.aoti_compile_and_package(ep) # -> .pt2 for C++/Python-free serving ``` On PyTorch < 2.5, call `torchao.utils.unwrap_tensor_subclass(model)` first. --- ## 11. HuggingFace `transformers` integration If your ViT is a HF checkpoint, you don't even need to call `quantize_` yourself: ```python import torch from transformers import AutoModelForImageClassification, AutoImageProcessor, TorchAoConfig from torchao.quantization import Int8WeightOnlyConfig quant_cfg = TorchAoConfig(quant_type=Int8WeightOnlyConfig()) model = AutoModelForImageClassification.from_pretrained( "google/vit-base-patch16-224", quantization_config=quant_cfg, torch_dtype=torch.bfloat16, device_map="cuda", ) proc = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224") model.forward = torch.compile(model.forward, mode="max-autotune") ``` - Older `transformers` versions take strings: `TorchAoConfig("int8_weight_only")`, `TorchAoConfig("int4_weight_only", group_size=128)`. - `model.push_to_hub(...)` / `save_pretrained(..., safe_serialization=False)` works for torchao-quantized models (safetensors can't represent the subclasses, hence `safe_serialization=False`). - To exclude modules, use `TorchAoConfig(quant_type=..., modules_to_not_convert=["classifier"])` (name may vary by version) or just quantize manually with a `filter_fn`. **timm** has no special integration — and doesn't need any; just `quantize_(timm_model, cfg)`. --- ## 12. Quantization-Aware Training (QAT) for ViT Use QAT when PTQ int4 (or int8-dynamic-act-int4-weight) loses too much accuracy and you can afford a fine-tune. ViTs fine-tune on ImageNet-1k very cheaply (a few epochs at low LR usually recovers most of the gap). Modern unified API: ```python import torch, timm from torchao.quantization import quantize_, Int8DynamicActivationInt4WeightConfig from torchao.quantization.qat import QATConfig model = timm.create_model("vit_base_patch16_224", pretrained=True).cuda().to(torch.bfloat16) base_config = Int8DynamicActivationInt4WeightConfig(group_size=32) # --- step 1: prepare (inserts fake-quant into forward; weights stay high precision) --- quantize_(model, QATConfig(base_config, step="prepare"), filter_fn=vit_filter) # --- step 2: fine-tune normally --- opt = torch.optim.AdamW(model.parameters(), lr=1e-5, weight_decay=0.0) for x, y in train_loader: loss = torch.nn.functional.cross_entropy(model(x.cuda().bfloat16()), y.cuda()) loss.backward(); opt.step(); opt.zero_grad() # --- step 3: convert to real quantized tensors for inference --- quantize_(model, QATConfig(base_config, step="convert"), filter_fn=vit_filter) model = torch.compile(model.eval(), mode="max-autotune") ``` Practical tips: - **Low LR** (1e-5 to 5e-5 for AdamW), short schedule, keep augmentation identical to the original recipe (or milder). You're recovering, not retraining. - Fake-quant makes training ~1.3–2× slower; don't `torch.compile` the prepared model until you've verified it trains (compile + fake-quant is supported but is where bugs live). - Older API, still present: `Int8DynActInt4WeightQATQuantizer(groupsize=32)` with `.prepare(model)` / `.convert(model)`, and the intermediate `IntXQuantizationAwareTrainingConfig` / `FromIntXQuantizationAwareTrainingConfig` pair. Prefer `QATConfig` on new versions. - **The base config used in `prepare` must match the one used in `convert`**, or your fine-tuning was optimizing for a different numeric format than you deploy. ### FP8 *training* is a different thing `torchao.float8` (`convert_to_float8_training(model)`) accelerates *training* by doing the GEMMs in fp8. It's for pretraining/fine-tuning throughput, not inference quantization. Relevant if you're pretraining a large ViT on H100s; orthogonal to everything above. --- ## 13. CPU, ARM, and edge deployment ### x86 CPU ```python from torchao.quantization import quantize_, Int8DynamicActivationInt8WeightConfig, Int4WeightOnlyConfig from torchao.dtypes import Int4CPULayout model = model.eval() # keep float32 or bfloat16 quantize_(model, Int8DynamicActivationInt8WeightConfig()) model = torch.compile(model) # inductor CPU generates the fused int8 kernels ``` For int4 on CPU: `Int4WeightOnlyConfig(group_size=128, layout=Int4CPULayout())`. ### ARM / mobile (int8 activations, low-bit weights, KleidiAI) ```python from torchao.quantization import quantize_, Int8DynamicActivationIntxWeightConfig, PerGroup from torchao.quantization.granularity import PerAxis quantize_(model, Int8DynamicActivationIntxWeightConfig( weight_dtype=torch.int4, granularity=PerGroup(32), )) ``` ### ExecuTorch / XNNPACK — the PT2E flow (different API!) For on-device deployment, torchao also hosts the **PT2 Export quantization** flow, which is graph-based rather than tensor-subclass based: ```python import torch from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import ( XNNPACKQuantizer, get_symmetric_quantization_config, ) example = (torch.randn(1, 3, 224, 224),) ep = torch.export.export_for_training(model.eval(), example) gm = ep.module() quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config(is_qat=False)) gm = prepare_pt2e(gm, quantizer) for x, _ in calib_loader: # 100–500 images is plenty for a ViT gm(x) gm = convert_pt2e(gm) # -> then lower to ExecuTorch / XNNPACK ``` Key difference: PT2E **quantizes the whole graph** including the patch-embedding Conv2d and (depending on the quantizer) matmuls inside attention, using observers + calibration. Use PT2E for ExecuTorch/mobile; use `quantize_` for GPU/server. --- ## 14. Troubleshooting | Symptom | Cause / fix | |---|---| | `model` becomes `None` | `quantize_` returns `None` and mutates in place. Write `quantize_(model, cfg)`, not `model = quantize_(...)`. | | Quantized model is **slower** than bf16 | (a) You didn't `torch.compile`. (b) Batch too large for a weight-only recipe — switch to dynamic/fp8. (c) Dequant didn't fuse — inspect with `TORCH_LOGS=output_code`. (d) You're timing without `torch.cuda.synchronize()`. | | No change in `print(type(layer.weight))` | Your `filter_fn` matched nothing, or the module isn't an `nn.Linear` (torchvision `MultiheadAttention`!). Enumerate modules and check. | | int4 errors about dtype | `Int4WeightOnlyConfig` requires **bfloat16** weights. `model.to(torch.bfloat16)` before `quantize_`. | | int4 errors about shapes | `in_features % group_size != 0` or `out_features % 8 != 0`. Reduce `group_size` or exclude the layer. | | fp8 raises "not supported" / `_scaled_mm` error | GPU is sm < 89. Use int8 on A100/V100. | | Huge accuracy drop (top-1 → near random) | Almost always a *bug*, not quantization: wrong dtype for inputs, model not in `eval()`, or you quantized the head with an incompatible recipe. Run `logit_fidelity` (§8.3) to localize. | | Accuracy drop of 1–5% with int4/int8-dyn | Real quantization error. In order: `use_hqq=True`, smaller group size, exclude `head` + first/last blocks + `mlp.fc2`, move to fp8, then AWQ, then QAT. | | `load_state_dict` size/dtype mismatch | Use the meta-device + `assign=True` pattern (§10). | | `UnpicklingError` / `weights_only` failure on load | `torch.load(..., weights_only=False)` or `add_safe_globals`. | | Constant recompilation | Varying batch size. Bucket + pad, or `torch._dynamo.mark_dynamic(x, 0)`, and raise `torch._dynamo.config.cache_size_limit`. | | Works eager, breaks under `torch.compile` | Try `mode="default"`, then `fullgraph=False`, then `backend="aot_eager"` to isolate; report with the minified repro. | | `unwrap_tensor_subclass` deprecation warning | You're on PyTorch ≥ 2.5 — just delete the call. | --- ## 15. Recommended decision path ``` Deploying on H100 / L40S / RTX 4090 (sm89+)? └─ yes → Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()) Best accuracy/speed. Start here, skip `head`. Deploying on A100 / A10 / T4 (sm75–86)? ├─ throughput-oriented (batch ≥ 16) → Int8DynamicActivationInt8WeightConfig() └─ latency-oriented (batch 1–8) → Int8WeightOnlyConfig() Memory-constrained (must fit ViT-H/g on a small GPU)? └─ Int4WeightOnlyConfig(group_size=64, use_hqq=True) + keep patch_embed / head / block0 / block-last in int8 or bf16 + eval carefully; QAT if the gap is unacceptable CPU server? → Int8DynamicActivationInt8WeightConfig() + torch.compile Mobile/edge? → PT2E + XNNPACKQuantizer → ExecuTorch Don't want to think? → torchao.autoquant(torch.compile(model)) at your production batch size, then eval. ``` **Non-negotiables regardless of recipe:** 1. `model.eval()` and `torch.no_grad()` / `inference_mode()`. 2. `torch.compile` after `quantize_`. 3. Measure latency **at your real batch size** and top-1 **on real data**, against a bf16 baseline compiled the same way (comparing quantized-compiled vs. fp32-eager is how people convince themselves of fake 4× speedups). 4. Use timm or HF ViT, not torchvision's, unless you patch `nn.MultiheadAttention`.

please explain thoroughly how to quantize vit model using to...