Inference Optimization

Purpose and Scope

Inference optimization in Transformers means reducing latency, memory use, or serving cost while preserving the model behavior required by an application. The most visible optimizations happen during generation and serving: assisted decoding can reduce the number of expensive forward passes, optimized attention backends can improve kernel efficiency, and multi-device placement can make large checkpoints usable. This page explains those inference-facing ideas while also mapping them to the documentation surface present in this repository shard, especially the optimization, backbone, and callback reference pages that define adjacent public APIs and clarify where training-time optimization ends and runtime inference optimization begins.

Sources: docs/source/en/main_classes/optimizer_schedules.md, docs/source/en/main_classes/backbones.md, docs/source/en/main_classes/callback.md

The important distinction is that not every module named “optimization” is an inference optimizer. The documented .optimization module provides fine-tuning tools: a weight-decay-aware optimizer, learning-rate schedules inheriting from _LRSchedule, and a gradient accumulation class for multiple batches. Those APIs matter when preparing a model that will later be served, but they are not a substitute for generation strategies, attention kernels, batching, or serving architecture. Treat training optimization as a way to produce a better checkpoint, and inference optimization as the set of choices made when loading, decoding, batching, and executing that checkpoint.

Sources: docs/source/en/main_classes/optimizer_schedules.md, docs/source/ja/main_classes/optimizer_schedules.md, docs/source/ko/main_classes/optimizer_schedules.md, docs/source/zh/main_classes/optimizer_schedules.md

Relevant Source Files

• docs/source/en/main_classes/optimizer_schedules.md — English reference entry for .optimization, Adafactor, SchedulerType, and scheduler factory functions such as get_scheduler, get_linear_schedule_with_warmup, and get_wsd_schedule.
• docs/source/ja/main_classes/optimizer_schedules.md — Japanese localization of the optimization reference, preserving the same high-level contract for optimizers, schedules, and gradient accumulation.
• docs/source/ko/main_classes/optimizer_schedules.md — Korean localization of the optimization reference, including localized anchors for optimization, Adafactor, and schedules.
• docs/source/zh/main_classes/optimizer_schedules.md — Chinese localization of the optimization reference, documenting the same .optimization module responsibilities.
• docs/source/en/main_classes/backbones.md — Reference page for AutoBackbone, BackboneMixin, BackboneConfigMixin, TimmBackbone, and TimmBackboneConfig, which are relevant when inference workloads reuse feature extractors for vision tasks.
• docs/source/en/main_classes/callback.md — Reference page for TrainerCallback, TrainerState, and TrainerControl; it clarifies that callbacks customize the PyTorch Trainer loop and are read-only except through returned control objects.

System-to-Code Mapping

The repository documentation separates training-loop optimization from inference architecture. In the optimization reference, Adafactor is documented alongside SchedulerType and scheduler factories including constant, warmup, cosine, hard-restart cosine, minimum-learning-rate cosine, greedy, linear, polynomial decay, inverse-square-root, reduce-on-plateau, and WSD schedules. These names are part of the public training API. They help when a checkpoint is being fine-tuned for later deployment, but they do not directly determine how generate() schedules tokens or how a server batches requests.

Sources: docs/source/en/main_classes/optimizer_schedules.md

The backbone reference gives a second inference-relevant axis: model structure. A backbone is defined as a model used for feature extraction for higher-level computer-vision tasks such as object detection and image classification. AutoBackbone initializes a Transformers backbone from pretrained weights, while BackboneMixin and BackboneConfigMixin manage output features and feature indices. For inference, this matters because many vision systems are built as staged pipelines: an image processor prepares inputs, a backbone extracts reusable feature maps, and a task head or downstream model consumes those features.

Sources: docs/source/en/main_classes/backbones.md

Callbacks form a boundary rather than an inference mechanism. The callback reference says callbacks customize the PyTorch Trainer training loop, inspect state for progress reporting or logging, and may take decisions such as early stopping. It also states they are read-only apart from the TrainerControl object they return. That means callbacks are appropriate for monitoring or controlling training and evaluation jobs, but production inference should usually be optimized through model loading options, decoding arguments, attention implementations, batching, device placement, or a serving runtime rather than TrainerCallback hooks.

Sources: docs/source/en/main_classes/callback.md

Generation-Time Strategies

For text generation workloads, the fastest useful optimization is often to change decoding work rather than model weights. The official assisted decoding guide describes speculative decoding: a smaller assistant model drafts candidate tokens, and the main model verifies those candidates in one forward pass. This can replace many expensive forward passes by the main model when the assistant is much cheaper and shares the tokenizer. In application terms, speculative decoding is best suited to latency-sensitive generation where a large model is the quality authority and a smaller model is accurate enough to propose likely continuations.

A minimal assisted-decoding call uses the same tokenizer for the main and assistant checkpoints and passes the assistant to generate(). The same concept is available through the Pipeline API with an assistant_model argument. The practical constraints are important: the assistant should be significantly smaller, tokenizer compatibility matters, and the official guide notes support for greedy search and sampling while excluding batched inputs for that method. If a service relies heavily on batching for throughput, compare speculative decoding against continuous batching or a specialized inference server before standardizing on it.

from transformers import AutoModelForCausalLM, AutoTokenizer
 
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-1.7B")
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-1.7B", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")
 
outputs = model.generate(**inputs, assistant_model=assistant_model)

Attention Backends and Model Compatibility

Optimized attention is a model-implementation concern as well as a runtime concern. The official backend-compatibility guidance describes AttentionInterface, ALL_ATTENTION_FUNCTIONS, and the _supports_attention_backend flag on PreTrainedModel as the route for custom and optimized attention functions. In practice, this means a model implementation should route attention through the shared interface and propagate **kwargs from the model forward path into the attention layers. That pattern lets inference engines use their performance features without reimplementing each model family from scratch.

This design also explains why model authors should think about inference while adding architectures. If a model hard-codes attention behavior and drops backend-specific keyword arguments, downstream engines cannot easily switch to optimized attention implementations. If it advertises support through the common interface, the same model checkpoint can be loaded in Transformers and reused by engines such as vLLM or SGLang when those engines support the needed backend contract. The repository shard here does not include the implementation files for AttentionInterface, so use the backend guide together with the model implementation and generation API references when changing code.

Multi-GPU and Vision Feature Extraction

Multi-GPU inference is usually a placement and execution strategy, not a scheduler setting. Large decoder models may need tensor parallelism, pipeline parallelism, device maps, or an external inference engine; vision pipelines may instead benefit from separating feature extraction and task-specific heads. The backbone reference is relevant to the latter pattern because it documents the supported backbone abstraction and its configuration mixins. When a computer-vision application repeatedly needs intermediate features, using a supported backbone contract makes it easier to standardize which feature maps and indices are returned.

Sources: docs/source/en/main_classes/backbones.md

The supported backbone catalog in the reference includes BEiT, BiT, ConvNext, ConvNextV2, DiNAT, DINOV2, FocalNet, MaskFormer, NAT, ResNet, Swin Transformer, Swin Transformer v2, and ViTDet, plus timm models through TimmBackbone and TimmBackboneConfig. For inference design, that list is a signal that the library treats feature extraction as a first-class reusable interface across multiple model families. If you are optimizing a vision service, start by identifying whether the bottleneck is preprocessing, backbone execution, task-head execution, or postprocessing before changing the model family.

Sources: docs/source/en/main_classes/backbones.md

Compact Reference

Practical Workflow

Start by naming the workload: single-request low-latency generation, high-throughput batched serving, vision feature extraction, or fine-tuned model preparation. For generation latency, evaluate decoding choices first: greedy or sampling settings, cache behavior, and assisted decoding with a compatible smaller assistant. For backend performance, verify that the target model supports the shared attention backend interface. For large checkpoints, choose a placement strategy or inference engine before tuning minor Python-level settings. For vision workloads, decide whether AutoBackbone and the backbone feature-index contract can isolate expensive feature extraction from downstream task heads.

Finally, keep training controls and inference controls separate in your mental model. Adafactor, scheduler factories, gradient accumulation, and TrainerCallback are important when producing and validating a checkpoint. Inference optimization begins after that checkpoint is loaded: it is about decoding work, attention execution, batching, device placement, and reusable feature extraction. Next, read the Text Generation, Caching and KV Cache, Continuous Batching, Serve CLI, Quantization Overview, and Distributed and Parallel Training pages for the runtime-specific knobs that complement the APIs summarized here.