LLM Training Optimization Techniques
-
by David Spuler, Ph.D.
LLM Training Optimization
LLM training optimization is the set of techniques to speed up the training phases. Training has a forward pass and a backward pass, each with lots of matrix multiplications, and both need to be optimized. Efficiency of training is one of the largest bottlenecks for training larger frontier models.
This is a companion article to the list of 500+ LLM Inference Optimization Techniques. There is some overlap with the inference optimization techniques, but there's also a whole swathe of issues that don't arise in inference, such as an entire "backward pass" of parameter updates. The network is important, too, but for different reasons to inference, because there's no KV caching in training, but there is outward transmission of training data and the inter-GPU transmission of the gradients and parameter updates.
Training versus Inference Optimization: Some of the key training optimization concepts:
- LLM training runs compute through all the Transformer's layers and sub-components (and then back again!)
- No decoding phase in training. No tokens are "output" during training.
- No KV caches are needed in training.
- Training's forward pass is prefill-like (token-wise parallel), but doesn't store any KV cache, and does many more things than prefill (to prepare for the backprop).
- Training's backward pass updating parameters via "gradients" has no equivalent in inference.
- Backward pass is more expensive than the forward pass in training.
- GEMM is the main training compute cost, mostly in the backward propagation phase.
- The gradient optimizer (AdamW) is memory-bound.
- GPU sharding is key for distributed multi-GPU cluster training. GPUs each train different parts of the model.
- Accuracy is often more important than speed in LLM training, which is why many lossy optimizations are shunned.
- Many important LLM Inference Optimization methods have limited applicability to training (e.g., speculative decoding, KV cache optimizations like Paged Attention or Radix Attention).
- Flash Attention is still a good one for attention in training! (just don't store the KV cache).
- A large model stays distributed across GPUs for the most part during training, only centralized into the full model at particular "consolidated checkpoints" (and at the very end).
Free AI C++ books: for more about LLM optimization, read books online or download a PDF:
- Generative AI in C++, David Spuler, March 2024, full text online, free PDF, bonus materials, source code
- CUDA C++ Optimization, David Spuler, June 2024, full text online, bonus materials, free PDF
- C++ Ultra Low Latency, David Spuler, July 2025, full text online, free PDF
Popular inference optimization articles: additional research articles on faster LLM inference:
- 500+ LLM Inference Optimization Methods
- Promising LLM Inference Optimization Techniques (Sep 2025)
- Hot LLM inference optimization research (August 2024)
- LLM reasoning model efficiency
More lists: lots of general efficiency optimization information:
- 100 CUDA Optimization Techniques
- List of 100+ AI Smartness Techniques
- 600+ low latency C++ techniques
LLM Training Optimizations List
Here's the training list!
-
Major LLM Intelligence Improvement Methods:
- Training
- Fine-Tuning
- RAG & Prompt Augmentation methods
- Tool usage (function calling)
- Plugins
- Knowledge distillation
- Harness and ecosystem
- Next-gen models
- AGI
Prompt Augmentation Methods: - RAG architectures
- RALM
- TALM
- Small reasoning model augmentation
- Reasoning scaffold augmentation
- Context engineering
- Plugins
- Internet search plugin
- Knowledge graph
- Taxonomy/ontology graph
Prompt Preprocessing Methods: - Prompt preprocessing ("hooks")
- Prompt shielding
- Refusal modules
- Prompt compression
- Context compression
- Automatic prompt optimization
- Heuristic prompt preprocessing
Next-Generation Models: - State Space Models (SSMs)
- Mamba
- Hyena
- Hybrid Transformer-SSM Architectures
- World models
- Embodied AI
- Symbolic reasoning
- Large Concept Models
Previous-Generation Classic ML Models: still going strong! - CNNs
- Hybrid CNN-Transformers
- RNNs
- Hybrid RNN-Transformers
- Diffusion models
- Hybrid Diffusion-Transformers
- GAN
Harness Improvement Methods: - Harness engineering
- Prompting techniques
- System prompt
Overall Training Optimization Categories: - Training data loading and preprocessing
- Loss calculations
- Back propagation
- Gradient optimizers
- Attention module optimizations
- Long context training optimizations
- FFN optimizations
- Matrix multiplication optimizations (GEMM/MatMul)
- Network transmission optimizations (both directions)
- Resilience and failure mitigation
High-Level Training Methods: - Pre-Training
- Supervised
- Unsupervised
- RLHF
- Reinforcement Learning
- Reward modeling
- Post-Training Optimizations
Training Data: - Labeled training data
- Unlabeled training data
- Synthetic data
- Open-source training data
- Commercially-licensed training data
- Proprietary training data
- Overrepresentation of common domains (needs downsampling)
- Underrepresentation of common domains (needs upsampling)
Model Media Types: - Text models
- Code generation models
- Multimodal models
- Image models
- OCR
- Audio
- Music
- Speech
- Voice
- Video models
- Computer Vision models
Model Overarching Goals: - Reasoning vs non-reasoning
- Agentic vs non-acting
- Chatbot vs batched models
Model Platforms: - Data Center LLMs
- — GPU
- — TPU
- — CPU
- — NPU
- Edge AI
- — Hybrid cloud-on-device inference
- On-device
- — AI Phones
- — AI PCs (desktops/laptops)
- — IoT devices
- New AI form factors
- — AI gadgets
- — AI glasses
- — AI pendants
- — AI rings
Model Architecture Decisions: - Single Model vs Ensemble (Multi-Model)
- Decoder-only vs encoder-decoder vs encoder-only
- Model dimensions
- Total parameter count
- Layers
- Small models vs large models
- Dense model vs MoE (sparse experts)
- Data size (e.g., FP32 vs INT4)
- Tokenizer
- Vocabulary (size and composition)
- Pre-Norm vs Post-Norm
- Embeddings method
- FFN vs GLU
- NAS analysis
- Autoregressive Decoding vs Parallel decoding
Model Architecture Major Components: - Mixture-of-Experts (MoE)
- Attention algorithm
- FFN
- GLU
- Positional Encoding algorithm
- Activation function
- Normalization algorithm
- Decoding algorithms
- Tokenization
- Embedding matrix
- Unembedding
- Detokenization
Other Model Architectures: - Slimmable
- Bulging (attention or FFN)
- Per-Layer Embeddings (PLE)
- Mixture-of-Attention (MoA)
Reasoning Models: - One-step reasoning model ("long answers")
- Multi-step reasoning model
- Reasoning Harness
- Reasoning System Prompts
Training Compute Optimizations: - Mixed-Precision Training
- — FP16
- — BF16
- Batching of training data (for efficiency & gradient stability)
- Per-batch gradient updates (not per-query or per token)
Open-Source LLM Training Engines: - PyTorch
- TensorFlow
- JAX/Flax (using XLA)
- DeepSpeed (Microsoft)
- Megatron‑LM (NVIDIA)
- Colossal‑AI
- Hugging Face Accelerate
- Determined AI
- Kubeflow
Open-Source LLM Inference Engines: - llama.cpp (known for CPU and Apple platform inference)
- vLLM (mainstream usage data center inference; known for Paged Attention)
- SGLang (mainstream; known for Radix Attention)
- LMCache (known for KV caching optimizations)
Gradient Optimizers: - SGD (Stochastic Gradient Descent): original optimizer:
- — Vanilla SGD (old-style, too unstable)
- — QSGD (Quantized SGD)
- — Batch SGD
- — Mini‑batch SGD
- — SGD with momentum
- — Nesterov Accelerated Gradient (NAG)
- — DeepSpeed QSGD (1‑bit Adam, 1‑bit LAMB)
- Adam (adaptive SGD with momentum; name is wordplay based on "Adaptive Moment Estimation"; legacy production optimizer superceded by AdamW)
- AdamW (Adam + Weight decay; widely used for frontier training; also for fine-tuning; SGD with momentum and modifications; Adam + decoupled weight decay; warmup cosine decay)
- AdamW variants (mostly experimental, not yet widely used instead of AdamW)
- — RAdamW (Rectified AdamW)
- — AdaBeliefW
- — AdamP/AdamPW
- — AdaFactorW
- — LionW
- — K‑FAC‑AdamW hybrids
- Other production usage optimizers:
- — Adafactor (from Google, used for TPU, memory-efficiency)
- — Lion (from Google; looks only at the sign bit of the gradient; simpler than AdamW, used for fine-tuning, esp. LoRA/QLoRA fine-tuning)
- — Sophia/Sophia-G (second-order optimizer; experimental/emerging into production for fine-tuning)
- — Muon (Momentum Orthogonalized by Polar decomposition) (emerging 2026 optimizer)
- Other legacy optimizers: formerly used in production:
- — AdaGrad (legacy optimizers; still somewhat used)
- — RMSProp (legacy optimizer, superceded by AdamW)
- — AdaDelta (modification of AdaGrad; legacy production usage)
- Other experimental gradient optimizers: mostly research or experimental usage:
- — Shampoo / Block‑Shampoo
- — LARS (Layerwise Adaptive Rate Scaling)
- — LAMB (Layerwise Adaptive Moments)
- — Newton’s Method (Hessian)
- — AdaNorm/AdaNormW
- — L‑BFGS (quasi‑Newton)
- — K‑FAC (Kronecker‑factored curvature)
- — AdaHessian
- — RAdam (Rectified Adam)
- — Lookahead
- — SAM (Sharpness‑Aware Minimization)
Gradient Optimizers for RLHF: - PPO (Proximal Policy Optimization)
- TRPO (Trust Region Policy Optimization)
- A2C
- A3C
SOTA Training Setup: - AdamW gradient optimizer
- Low base learning rate
- Momentum: β₁ ≈ 0.9, β₂ between 0.95–0.999
- Epsilon: ε ≈ 1e‑8,
- Weight decay ≈ 0.01 (or 0.02-0.05)
- Warmup phase: 1–3% of total steps
- Decay algorithm: cosine or linear
Random Noise Injection: to reduce overfitting and increase generalization; aimed at accuracy, not for efficiency improvement. - Dropout (randomly zero a percentage of activations)
- Structured dropout (apply Dropout to particular structures)
- Attention Dropout (applies before Softmax)
- DropHead (attention head dropout)
- LayerDrop (randomly skip entire layers)
- Stochastic Depth (token-wise variant of LayerDrop)
- DropPath (in Vision Transformers)
- Token Dropout (randomly mask or prune input tokens in training)
- Randomized Positional Dropout (RoPE dropout, ALiBi dropout)
- Drop Embeddings
- Many more Dropout variants: see this 2022 paper: https://arxiv.org/pdf/2204.02027
Network Optimizations for Training: - Training data outward transmission
- Gradients and updates (returned)
- Gradient compression methods
- Burst management
- Network bandwidth management
LLM Training Major Settings: - Learning rate settings
- Learning Rate warmup (LR warmup)
LLM Training Problems: - Model training fails overall
- Model evaluation failure (bad model trained)
- — Misalignment
- — Emergent misalignment
- — Catastrophic forgetting
- — Overfitting
- — Generalization failures
- — Misgeneralization
- — Degenerate solution learning
- — Reasoning failures
- — Brittle reasoning (narrow learning)
- — Reward hacking (learned shortcuts)
- — Spurious correlations learned
- — Arithmetic failures
- Data problems
- — Data quality
- — Noisy data
- — Skewed data sets
- — Mixed data sets (e.g. text vs code)
- Optimizer instability
- — Vanishing gradient (loss plateaus)
- — Exploding gradients (loss spikes)
- — Floating-point overflow (NaN/Inf)
- — Floating-point underflow (zero rounding)
- — Divergence
- — Loss spikes
- — Loss oscillations
- — Learning rate incorrect
- — LayerNorm value spikes
- — Batch size too small
- — Optimizer state failures
- — Optimizer state drift
- — Distributed gradient corruption
- — Non-deterministic training
- — Long-context gradient problems
- Warmup failures
- — LR warmup issues
- — Checkpoint restore errors
- Hardware faults
- — Transient GPU faults
- — GPU burnouts
- Network faults
- — Network communication failure
- — Gradient communication faults
- — Gradient synchronization failures
- — Network latency spikes
- — Network bandwidth issues
Numerical stability mechanisms: - LR warmup (the learning rate is low to start and increased gradually)
- LR warmdown or "decay phase" (learning rate reduced near the end of training)
- AdamW bias correction
- RMS-scaled gradients (avoid floating-point underflow)
- Gradient clipping (stopping gradients above a threshold; global norm clipping of loss spikes to avoid exploding gradients)
- Norm tracking
- Cosine decay
- Activation scaling (uP, DeepNorm, RMSNorm)
Distributed gradient corruption: - all‑reduce desynchronization
- Silent bit flips (SDC, transient errors)
- Optimizer state inconsistent (across workers)
- AdamW moment vectors inconsistent
- Gradient compression errors
- Over-quantization in gradient compression
- Over-sparsification in gradient compression
- Mixed-precision desynchronization (desync)
- Checkpoint corruption
Resilience Optimizations for Training: avoiding training failures is a big part of efficiency, because restarts are costly. - Checkpointing optimizations
- In-memory checkpointing
- Asynchronous checkpointing
- Silent Data Corruption (SDC) mitigation
- GPU failure mitigation
- GPU burnout mitigation
- Stragglers (one slow GPU or slow network connection forces everyone else to wait)
- Hangs (if one GPU fails)
- NaN/Inf detection ("
std::isnan" CPU/GPU; use "x!=x" trick; use CUDA__float_as_uint intrinsic)
Pruning during Training: can be used for pre-training; rarely for fine-tuning. - Unstructured pruning (overview)
- — Magnitude pruning
- — First-order pruning
- — Movement pruning
- — Second-order pruning
- — LTH (Lottery Ticket Hypothesis) (sparse unstructured pruning)
- — SNIP (Single‑Shot Network Pruning) (unstructured sparsity at initialization)
- — GraSP (Gradient Signal Preservation) (unstructured, initialization-based)
- Dynamic Sparse Training (DST)
- — Unstructured DST
- — RigL (Rigged Lottery) (fixed sparsity budget)
- — SET (Sparse Evolutionary Training) (mostly research usage)
- — NOTE: Unstructured pruning does not reduce FLOPs (only structured pruning does)
- Structured pruning (overview)
- Structured DST
- — 2:4 sparsity (prunes 50%)
- — 4:8 sparsity (prunes 50%)
- — Block sparsity (block-sparse DST)
- — Block-sparse RigL
- — Expert-level DST for MoE training (used by frontier MoE model training)
- Pruning approaches:
- — Phased sparsity increases pruning during warmup
- — Sparsity from initialization (during whole warmup)
Quantization during Training: - Quantization-Aware Training (QAT)
- Post-Training Quantization (PTQ)
Attention optimization subtypes: - Attention optimizations (overview)
- — Multi-Head Attention (MHA)
- — Group Query Attention (GQA)
- — Multi-Query Attention (MQA)
- — Sparse attention
- — Local attention
- — Memory-efficient attention algorithms
- — Flash Attention
- — Paged Attention
- — Linear attention
- — Cross attention
- — Tree attention
- — Sliding window attention
- — Approximate attention heads
- — Attention alternatives/replacements
- — Fused MHA
- — Low-rank matrix attention
- — Medusa attention
- — Block attention
- — Cross attention
- — Fused head attention
- — Hybrid local-global attention
- — FFT attention
- — Additive attention
- — Multiplicative attention
- — Graph attention
- — Attention sink
- — Attention steering
- — Bilinear attention
- — Attention-free methods
- — Star attention
- — Ring attention
- — Flex attention
- — Razor attention
- — Contiguous QKV tensor
- — Relative Attention Bias (RAB)
- — Lightning attention
- — Multihead Latent Attention (MLA (DeepSeek)
- — FFT attention
- — Round attention
- Delta attention
- Gated attention
- KIVI attention
- K=V (KV compute sharing)
- Bulging attention (per-layer attention module size increases)
Attention compute optimizations: - — Chunked attention
- — QKV computation optimizations
- — Mixture-of-Heads (MOH) Attention (MoE+MHA)
- Mixture-of-Attention (MoA) (MoE attention)
Long context optimizations (attention): - — Long context models
- — Length generalization
- — Quadratic attention complexity
- — Long RAG
RAG Architecture Optimizations: - RAG architectures (overview)
- — RAG cache
- — RAG optimizations
- — RAG retriever datastore indexing
- — Advanced RAG
- — Speculative RAG
- — Reranker in RAG
- — Chunk-specific global KV caching
- — Chunk-specific prefix KV caching
- — RAG Knowledge Graph
- — RAG Ontologies/Taxonomies
- — RAG fusion
- — Mini-RAG (single-document RAG)
Non-Multiplication AI Models: - Zero-Multiplication Models (overview)
- — Binary quantization
- — Ternary quantization
- — 2-bit quantization (INT2)
- — Adder networks
- — Bitshift-add networks
- — Bitshift power-of-2 quantization (logarithmic quantization)
- — Double bitshift quantization
- — Add-as-integer networks
- — Logarithmic Models
- — Bitwise neural networks
- — Diff-squared networks
- — Log-sum-exp (LSE) networks
- — Max-Plus networks
- — Min-Max-Plus networks
- — Morphological networks
- — Trigonometric approximate inference
- — Weightless Neural Networks (WNNs)
- — XNOR networks
- — Hadamard elementwise matrix multiplication models
- — Other addition-related zero-multiplication networks
- — Table lookups replace multiplication
- — Other multiplication-free neural networks
Advanced Number System optimizations: - Advanced Number Systems (overview)
- — Posit number system (PNS)
- — Residue number system (RNS)
- — Dyadic numbers
- — Double-base number system (DBNS)
- — Dynamic number systems
- — Hybrid number systems
- — Tropical algebra (max-plus)
- — MiniMax algebra
- — Multi-dimensional logarithmic number system (MDLNS)
- — Multiple-Base Number System (MBNS)
- — Semi-Logarithmic Number System (SLNS)
- — Lattice algebra
Logarithmic Number System optimizations: - Logarithmic number system (LNS) (overview)
- — End-to-end LNS logarithmic model
- — LNS addition and subtraction
- — LNS in AI models
- — LNS Hardware Acceleration
- — LNS mathematical and algorithmic theory
- — LNS algebra
- — LNS extensions
Parameter Efficient Fine-Tuning (PEFT) subtypes: - PEFT (overview)
- — LoRA
- — Multi-LoRA inference
- — QLoRa (Quantized Low-Rank Adapters)
- — LoRA inference optimizations (load/unload)
- — Prompt Tuning (Extended Vocabulary PEFT)
- — Prefix Tuning
Ensemble multi-LLM subtypes: - Ensemble inference (overview of multi-model AI engines)
- — Model selection algorithms
- — Big-little architectures
- — Cascades
- — Collaborative inference
- — Consensus decoding
- — Swarm ensemble architectures
- — Committee ensemble architectures
- — Ensemble averaging
- — Easy-hard queries
- — Submodels (Many-Models-in-One)
- — Distributed Inference
Tool Integration Optimizations: LLMs using tools has gone mainstream, and there is also newer research on speeding it up: - Tool optimizations
- — Tool execution pipelining (overlap with prefill or decode)
- — Speculative tool execution
- — Tool token reduction
- — Concise tool output
- — Disaggregated tool execution
- — Multi-tool parallel execution
Knowledge distillation subtypes: - Knowledge Distillation (overview)
- — Ensemble Distillation
- — Unnatural instructions (data sets)
- — Dataset Distillation
- — Black Box Distillation
- — White Box Distillation
Overall summaries of AI optimizations: - — Deslugging AI engines
- — Accuracy-degrading optimizations
- — Accuracy-retaining optimizations
- — Uncommon inference optimizations
Not Enough?
More inference optimization resources:
- What's Hot in Inference Optimization?
- Inference optimization research overview
- Research Blog
- Patents in inference optimization
Free AI and C++ Books
Generative AI programming books:
- The Sweetest Lesson: Your Brain Versus AI, November 2025: full text online, free PDF available
- RAG Optimization: Accurate and Efficient LLM Applications, June 2025: full text online, free PDF available
- Generative AI Applications: Planning, Design and Implementation, November 2024: full text online, free PDF available
- Generative AI in C++ (Spuler, March 2024): full text online, free PDF available, table of contents, bonus materials, reference lists, source code
CUDA C++ GPU Programming Books:
- CUDA C++ Optimization: Coding Faster GPU Kernels, July 2024: full text online, bonus materials, free PDF available
- CUDA C++ Debugging: Safer GPU Kernel Programming, July 2024: full text online, free PDF available
Modern C++ Programming Books
- C++ AVX Optimization: CPU SIMD Vectorization, 2025: full text online, free PDF available
- C++ Ultra-Low Latency: Multithreading and Low-Level Optimizations, 2025: full text online, free PDF available
- Advanced C++ Memory Techniques: Efficiency and Safety, 2025: full text online, free PDF available
- Efficient C++ Multithreading: Modern Concurrency Optimization, 2025: free PDF available
- Efficient Modern C++ Data Structures: Container and Algorithm Optimizations, 2025: free PDF available
- C++ Low Latency: Multithreading and Hotpath Optimizations, 2025: free PDF available
- Safe C++: Fixing Memory Safety Issues, Oct 2024: full text online, free PDF available
More AI Research Topics
Read more about: