LLM Speed and Memory Optimization

Purpose and Scope

Large language model optimization in Transformers is best approached as a budgeting exercise: decide which tensors must live in GPU memory, which operations create temporary peaks, and which training or inference controls reduce pressure without changing the model objective. The official optimization guide frames the main deployment constraints as parameter scale and long input sequences. The memory anatomy guide then breaks GPU usage into model weights, optimizer states, gradients, activations, temporary tensors, and other overhead such as generation caches. This page connects that conceptual model to the repository documentation surface for optimizers, schedules, callbacks, and architecture-facing model utilities.

For training, the most actionable source-backed APIs on this page are the .optimization module and the PyTorch Trainer callback system. The optimization docs expose Adafactor, learning-rate scheduler factories, SchedulerType, and gradient accumulation support. These are not just convenience utilities: they are where a training run trades optimizer-state memory, convergence stability, and step timing. The callback docs explain how training behavior can be inspected and controlled through TrainerCallback, TrainerState, and TrainerControl, which matters when memory or speed tuning requires stopping early, logging peaks, or coordinating with external experiment systems. Sources: docs/source/en/main_classes/optimizer_schedules.md, docs/source/en/main_classes/callback.md

For inference, the most important decisions are usually outside the optimizer loop: numerical precision, attention implementation, model architecture, context length, batch shape, and cache behavior. The official LLM optimization guide calls out lower precision, Flash Attention, and architectural innovations such as ALiBi, rotary embeddings, Multi-Query Attention, and Grouped-Query Attention. The requested repository paths do not include the generation or quantization implementation pages, so this page treats those techniques as part of the reader workflow and maps only the supplied source-backed APIs that participate in training-time control and architecture selection.

Relevant Source Files

• docs/source/en/main_classes/optimizer_schedules.md — English API documentation entry for the .optimization module, including Adafactor, scheduler factories, SchedulerType, GreedyLR, and gradient accumulation framing.
• docs/source/ja/main_classes/optimizer_schedules.md — Japanese localization of the optimization API surface, preserving the same public concepts for optimizer, schedule, and gradient accumulation documentation.
• docs/source/ko/main_classes/optimizer_schedules.md — Korean localization of the optimization API surface, including localized anchors for optimization, Adafactor, scheduler types, and warmup schedule functions.
• docs/source/zh/main_classes/optimizer_schedules.md — Chinese localization of the optimization API surface, confirming the same reader-facing grouping around weight decay, schedules, and gradient accumulation.
• docs/source/en/main_classes/backbones.md — Documentation for AutoBackbone, BackboneMixin, BackboneConfigMixin, TimmBackbone, and TimmBackboneConfig, useful when optimization involves selecting feature-extraction architectures rather than full task heads.
• docs/source/en/main_classes/callback.md — Documentation for TrainerCallback, TrainerState, TrainerControl, default reporting behavior, and built-in callbacks such as early stopping and logging integrations.

Memory Anatomy: What You Are Optimizing

When a large model runs out of GPU memory, the model parameters are only one part of the problem. During training, mixed precision can keep both a low-precision copy for forward and backward computation and a higher-precision copy for stable updates. Adam-style optimizers add momentum and variance state, gradients occupy their own storage, and activations must be retained for the backward pass. Long sequences and large batches can therefore fail even when the raw checkpoint appears to fit. Temporary tensors from matrix multiplication, softmax, attention, and generation can also create short-lived peaks that matter in practice.

The practical response is to identify which memory category dominates the workload. If optimizer states dominate, choose a lower-memory optimizer or reduce trainable parameters. If activations dominate, reduce sequence length or batch size, use gradient accumulation to preserve effective batch size, or use checkpointing techniques covered elsewhere in the training docs. If temporary attention tensors dominate, attention backends such as Flash Attention can improve memory locality and reduce peak allocation. If inference cache or context length dominates, use shorter prompts, smaller batches, cache-aware generation settings, or architectures designed for long-context autoregressive decoding.

Transformers exposes one source-backed optimizer lever directly in the supplied docs: Adafactor. The official API text describes Adafactor as an adaptive optimizer with sublinear memory cost, and the repository optimization page lists it as the optimizer documented under .optimization. That makes it a natural candidate when Adam optimizer states are a major memory cost. The same page also lists scheduler helpers, because changing optimizer memory without controlling learning-rate behavior can destabilize fine-tuning. In practice, optimizer selection and scheduler selection should be treated as a pair rather than independent toggles. Sources: docs/source/en/main_classes/optimizer_schedules.md

Training-Time Optimization Controls

The .optimization module documentation states that Transformers provides an optimizer with fixed weight decay for fine-tuning, several schedule objects inheriting from _LRSchedule, and a gradient accumulation class for accumulating gradients from multiple batches. Those three pieces correspond to three common training constraints. The optimizer determines how much state is stored per parameter and how updates are computed. The schedule determines how aggressively learning rate changes across warmup, decay, restarts, or plateaus. Gradient accumulation lets a run use smaller per-device batches while still approximating a larger effective batch over multiple forward and backward passes. Sources: docs/source/en/main_classes/optimizer_schedules.md

The scheduler catalog is broad enough to cover both ordinary fine-tuning and specialized large-model runs. The English page documents SchedulerType, get_scheduler, constant schedules, constant-with-warmup schedules, cosine schedules, cosine schedules with hard restarts, cosine schedules with a minimum learning rate, GreedyLR, linear warmup schedules, polynomial decay, inverse square root schedules, reduce-on-plateau schedules, and get_wsd_schedule. The Japanese, Korean, and Chinese pages preserve the same organization around optimizer, schedules, and gradient accumulation, which is useful for contributors checking whether public optimization terminology remains consistent across localized docs. 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

A memory-conscious fine-tuning loop typically starts with the smallest per-device batch that keeps the accelerator busy, then uses gradient accumulation to reach the target effective batch size. Next, choose an optimizer whose state cost fits the model and hardware, then choose a warmup and decay schedule that matches the expected number of optimizer steps after accumulation. Finally, monitor the run: if loss is unstable, the memory savings may have made the update rule too aggressive; if throughput is low, accumulation may be hiding underutilization. These tradeoffs are workload-specific, so the library exposes multiple schedules rather than a single preferred policy.

from transformers import Adafactor, get_scheduler
 
optimizer = Adafactor(
    model.parameters(),
    scale_parameter=False,
    relative_step=False,
    warmup_init=False,
    lr=1e-4,
）
 
lr_scheduler = get_scheduler(
    name="linear",
    optimizer=optimizer,
    num_warmup_steps=500,
    num_training_steps=10_000,
）

The example shows the shape of the decision rather than a universal recipe. Adafactor is the source-backed optimizer name exposed by the optimization page, and get_scheduler is the generic entry point for selecting one of the documented learning-rate policies. The important optimization habit is to count actual optimizer steps after gradient accumulation and distributed training, then pass that count to the scheduler. A schedule configured for raw dataloader batches can decay too quickly or warm up too slowly, which looks like an optimization bug even though the memory settings are technically valid.

Runtime Monitoring with Trainer Callbacks

Optimization changes should be observable. The callback documentation defines callbacks as objects that customize the PyTorch Trainer training loop by inspecting state for progress reporting, logging, or decisions such as early stopping. It also states an important boundary: callbacks are read-only except for the TrainerControl object they return. If the optimization requires changing the training loop itself, the docs direct users to subclass Trainer and override the needed methods. That distinction prevents callback-based memory monitoring from turning into hidden training-loop mutation. Sources: docs/source/en/main_classes/callback.md

For LLM speed and memory work, callbacks are useful for guardrails. EarlyStoppingCallback can stop unproductive runs before they consume more GPU time, while logging callbacks for TensorBoard, Weights & Biases, MLflow, CodeCarbon, and other integrations help correlate training loss, evaluation metrics, and runtime cost. TrainerState gives a callback access to the loop state, and TrainerControl is the documented channel for taking actions. The docs also note that TrainingArguments.report_to defaults to "none", so reporting integrations must be intentionally enabled rather than assumed. Sources: docs/source/en/main_classes/callback.md

A minimal memory-debugging callback can log at phase boundaries without modifying the optimizer or dataloader. Use this pattern when comparing batch size, accumulation steps, scheduler choice, precision, or model architecture. If a callback only prints or reports state, it stays within the documented callback contract. If it must alter forward behavior, replace modules, or change batch construction, move that logic into a custom Trainer, data collator, model wrapper, or lower-level training loop instead of overloading callback hooks.

from transformers import TrainerCallback
 
class MemoryTraceCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, **kwargs):
        print("Starting memory-sensitive training run")
 
trainer.add_callback(MemoryTraceCallback())

Architecture and Inference Considerations

The supplied backbone documentation is not an LLM-specific optimization page, but it demonstrates a recurring Transformers design pattern: separate a reusable feature-extraction body from task-specific heads, expose an auto class for loading from pretrained weights, and make output feature selection configurable. The page defines a backbone as a model used for feature extraction in higher-level computer vision tasks and documents AutoBackbone, BackboneMixin, BackboneConfigMixin, TimmBackbone, and TimmBackboneConfig. That is relevant when speed or memory optimization starts with choosing a smaller or more appropriate architecture instead of only tuning runtime flags. Sources: docs/source/en/main_classes/backbones.md

For LLM inference, the analogous architectural questions are whether the checkpoint uses positional encodings and attention variants that are efficient for the intended context length and decoding pattern. The official LLM optimization guide highlights ALiBi, rotary embeddings, Multi-Query Attention, and Grouped-Query Attention because autoregressive inference repeatedly extends a sequence and benefits from architectures designed for that use. Transformers users usually encounter these choices through model selection and configuration rather than by rewriting attention modules manually. Select the checkpoint architecture first, then tune precision, batching, attention backend, and generation settings around it.

Lower precision and Flash Attention are also inference-first levers. Lower precision reduces memory per parameter and can increase throughput on hardware with optimized low-precision kernels. Flash Attention changes how attention computation uses GPU memory, reducing materialization of large intermediate tensors and improving memory locality. These techniques interact with model family, hardware, and backend support, so validate with representative prompt lengths and batch sizes. A configuration that is fastest for short prompts may not be best for long-context generation, and a setting that fits one GPU may regress on another.

Compact Reference

Practical Workflow and Next Steps

Start optimization by reproducing the workload with realistic sequence lengths, batch sizes, and generation settings. Record whether the failure is an out-of-memory error, slow tokens per second, unstable loss, or underutilized hardware. For training, reduce per-device batch size, add gradient accumulation, choose an optimizer such as Adafactor when optimizer states dominate, and configure a scheduler using the actual optimizer-step count. Add callbacks for logging and early stopping before running long experiments. For inference, evaluate lower precision, attention backend choices, and model architectures designed for long-context autoregressive generation.

Next, read the quantization, inference optimization, generation API, Trainer, and distributed training pages for the lower-level mechanisms that are outside the supplied source set for this page. Use this page as the bridge between memory anatomy and the concrete Transformers controls that are documented here: optimizer state, schedules, gradient accumulation, callback observability, and architecture selection. The fastest successful configuration is usually the one that removes the dominant bottleneck first, then validates quality and throughput under the same conditions that production or fine-tuning will actually use.