LLM 优化面试笔记：训练与推理核心技术

✍ Gauri Gupta🕐 2026-05-06📦 28.2 KB 🟢 已读 𝕏 文章列表

这是一份针对 AI 实验室面试准备的笔记，涵盖了高效训练和部署大语言模型（LLM）的核心优化策略。内容主要分为三部分：内存优化（如 Flash Attention、MQA/GQA、激活检查点）、计算优化（序列打包、高效变体 Transformer）以及推理优化（KV 缓存、投机解码、量化技术）。文章旨在总结在大规模模型开发中应对算力和内存瓶颈的关键技术。

LLM 优化推理训练 Flash Attention KV Cache 量化面试 Transformer 性能优化

# LLM Optimization Interview Notes: Training and Inference **作者**: Gauri Gupta **日期**: 2026-05-06T04:33:11.000Z **来源**: [https://x.com/gauri__gupta/status/2051882947758993815](https://x.com/gauri__gupta/status/2051882947758993815) --- ![Image](/doc/1674/image/img_068cb700fd0e.jpg) Here is a collection of my personal notes from preparing for interviews at several leading AI labs and revisiting the core ideas behind efficient large-scale model training. Along the way, I compiled these notes, part interview preparation, part personal revision, and thought they might be worth sharing. Training and deploying large language models efficiently has become one of the most critical challenges in modern AI. As models scale into the billions of parameters, traditional approaches quickly start to break down. The techniques outlined here highlight some of the essential optimization strategies that have emerged as industry standards for building and deploying large-scale models. These notes aren’t meant to be exhaustive or perfectly structured, but rather a reflection of the key concepts and techniques that repeatedly came up in technical discussions and that form the backbone of large-scale model development. This is also not a deeply detailed discussion of each topic, but rather a high-level overview of the core approaches and concepts providing a concise summary of the most common practices and ideas used in efficient training and inference. I hope they’re useful to anyone looking to dive deeper into these optimization techniques or preparing for similar technical interviews. 1. Memory Optimization Techniques Memory is the biggest bottleneck in LLM training and inference. As models scale to billions of parameters, traditional memory management approaches become insufficient. The techniques in this section focus on reducing memory footprint while maintaining model quality, enabling the training and deployment of larger models on existing hardware. 1.1 Flash Attention The attention mechanism has quadratic time and memory complexity in sequence length, presenting significant runtime and memory challenges for longer sequences. Flash Attention reduces attention memory complexity through tiling and recomputation techniques. Instead of processing entire attention matrices at once, it processes attention in blocks and stores normalization factors instead of full attention matrices. The tiling technique decomposes inputs based on shared memory size, while recomputation stores softmax normalization factors (linear to sequence length) instead of softmax results (quadratic to sequence length). Tiling Technique: Decomposes inputs based on shared memory size and calculates softmax one tile at a time. Instead of working on entire query, key, value tensors at once, it makes several passes and combines results in subsequent steps. Recomputation Technique: Stores softmax normalization factors (linear to sequence length) instead of softmax results (quadratic to sequence length), using these factors to recompute attention scores. This reduces memory requirements and I/O traffic between global and shared memory. Additional Resources: [1] Matrix multiplication tiling [2] Online softmax and tiling 1.2 Multi-Query and Grouped Query Attention - MQA (Multi-Query Attention): Reduces memory by sharing keys and values across attention heads - GQA (Grouped Query Attention): Balances efficiency and quality by grouping queries 1.3 Activation Checkpointing Input activations easily saturate device memory when training LLMs with large sequence lengths or micro-batch sizes. Checkpointing a few activations and recomputing the rest reduces device memory requirements. 2. Compute Optimization Techniques Maximize GPU utilization and reduce computational overhead through smarter data handling and model architectures. 2.1 Sequence Packing A training technique where multiple training sequences are concatenated into one long sequence. This eliminates padding and allows more tokens to be processed per micro-batch, maximizing both GPU compute and memory utilization. 2.2 Efficient Transformers As sequence lengths and dataset sizes grow, standard transformer architectures become prohibitively expensive due to their quadratic time and memory complexity with respect to sequence length. Efficient transformer techniques have been developed to address these limitations, enabling the processing of much longer sequences and larger datasets by reducing computational and memory requirements. - BigBird: Uses a combination of local, random, and global attention patterns to reduce complexity to O(n). - Longformer: Utilizes sliding window (local) attention combined with global attention for improved efficiency. - Low-Rank Approximations: Projects key and value matrices into lower-dimensional spaces. - LongNet: At lower layers, tokens attend to nearby tokens (small dilation). At higher layers, dilation factor grows, allowing tokens to reach further. Scales linearly with sequence length O(Nd). Additional Resources: [1]Scaling Transformers with LongNet 3. Inference Optimization Techniques Inference is where most production costs occur. These techniques dramatically speed up generation while maintaining quality. 3.1 KV Caching KV caching works by storing the key and value tensors computed for each token as the model generates a sequence. In autoregressive generation, at each new step, the model only needs to compute attention for the newly generated token, using the cached keys and values from previous steps instead of recomputing them for the entire sequence. This drastically reduces both computation and memory usage, as the attention mechanism can reuse the stored key-value pairs for all previous tokens. Advanced KV Cache Optimizations: - Grouped Multi Query Attention: Reduces KV cache memory by grouping multiple queries with same keys and values - Multi-head Latent Attention: Projects K, V, Q into lower-dimensional latent space, computing attention in latent space then projecting back - Cross Layer KV-sharing: Ties KV cache across neighboring attention layers - Interleaving Local and Global Attention: Uses global attention in every 4-6 layers Additional Resources: [1] KV Caching Video [2] FLOPS computation efficiency with KV cache 3.2 Stateful Caching Stateful caching stores conversation history using rolling hashes, allowing reuse of overlapping prefixes. For example, if "Hello, how are you?" is cached, it can be reused when the new prefix is "Hello, how are you doing today?" The cache is organized in a tree structure with LRU eviction to manage memory efficiently. For a new query, compute rolling hashes for all its prefixes and find the longest cached match.Load the KV tensors from cache and continue computation only from the new tokens. KV cache organized in a tree structure with LRU (least recently used) eviction, so you can drop old contexts if memory is full. 3.3 Speculative Decoding Speculative decoding uses a smaller draft model to generate responses, then uses the target model to verify them, achieving 2-3x speedup in inference. The draft model must be fast and well-aligned with the target model for this technique to be effective. 3.4 Quantization Techniques Compressing a model by representing weights/activations with fewer bits instead of standard fp32 (32-bit float). Quantization Types: - Min/Max: Simple but susceptible to outliers - MSE: Minimizes MSE between original and quantized values - Cross-entropy: Preserves order of largest values after quantization for softmax; argmin(softmax(v), softmax(v’)) Post-Training Quantization (PTQ): After fully training a model, its weights are converted to lower precision (e.g., int8 or float16) without further training. PTQ is generally inexpensive to implement compared to retraining, but as model sizes scale to billions of parameters, outlier activations of large magnitude can appear in transformer layers, making naive low-bit quantization less effective. To address this, quantization-aware observers are attached to collect statistics (such as mean and standard deviation) on the input data, which are then used to determine quantization parameters. Mixed-Precision Quantization: Instead of quantizing all weights and activations to the same bit width, mixed-precision quantization assigns different precisions to different parts of the model. For example, sensitive layers or activations may use higher precision (e.g., 8 or 16 bits), while less sensitive parts use lower precision. This approach balances memory savings and model accuracy, and is especially useful for large models where uniform low-bit quantization would degrade performance. Mixed-precision quantization can be applied separately to weights and activations, optimizing for both efficiency and quality. Quantization-Aware Training (QAT): Quantization is applied during pre-training or further fine-tuning. Simulate quantization and de-quantization in the forward pass. This simulates the quantization error and acts as a regularizer to make the model robust to it. Back-propagation: Quantization is not differentiable, so gradients are approximated using the straight-through estimator (STE), which sets the gradient to 1 within the quantization range (alpha, beta) and 0 outside. Additional Resources: [1] Quantization Video [2] Lilian Weng's Inference Optimization 4. Training Optimization Training large models requires sophisticated parallelism strategies. Know the different approaches and their trade-offs. 4.1 Mixed Precision Training Mixed precision training leverages lower-precision number formats, most notably bfloat16 (Brain Floating Point 16), to reduce memory usage and accelerate training. Bfloat16 has the same exponent range as standard 32-bit floats (fp32), but with fewer mantissa bits, allowing it to represent very large and very small numbers, which helps preserve the dynamic range needed for deep learning. However, because bfloat16 and fp16 have reduced precision, gradients can underflow or overflow during back-propagation, leading to instability. To address this, loss scaling is used: the loss is multiplied by a large constant before back-propagation, and gradients are later rescaled down, preventing small gradients from vanishing due to limited precision. Careful management of loss scaling is essential to maintain numerical stability and fully realize the benefits of mixed precision training. 4.2 Data Parallelism DataParallel: Single-process, multi-threaded approach that works when the model fits on a single GPU. Each GPU keeps a copy of the model, processes different micro-batches, then averages gradients across GPUs. The main bottleneck is that it relies on single-process, multi-threaded communication, leading to inefficient inter-GPU communication and potential slowdowns due to CPU overhead. DataParallel runs in a single process with multiple threads, so it suffers from GIL contention. Synchronization Approaches: At the end of each minibatch, workers need to synchronize gradients or weights to avoid staleness. There are two main synchronization approaches: - Bulk Synchronous Parallel (BSP): Workers synchronize at the end of every minibatch. It prevents model weights staleness and provides good learning efficiency, but each machine has to halt and wait for others to send gradients. - Asynchronous Parallel (ASP): Every GPU worker processes the data asynchronously with no waiting or stalling. However, it can easily lead to stale weights being used and thus lower the statistical learning efficiency. Even though it increases the computation time, it may not speed up training time to convergence. Distributed Data Parallel (DDP): Each GPU has its own process and can work on multiple nodes/machines. Uses Ring All-Reduce algorithm to avoid central bottlenecks. DDP has lower communication overhead compared to DataParallel. ZeRO (Zero Redundancy Optimizer): Not just the model parameters and gradients, but the optimizer state (including Adam momentum, variance) also takes a lot of memory. ZeRO-DP has three main optimization stages: - Optimizer State Partitioning: 4x memory reduction, same communication volume as DP. Gradient computation can be done independently for each GPU. When it is being done for parameters that are not on the current GPU, we incur a communication cost but this is the same as gradient averaging in DP. So always do this ZeRO stage-1. - Gradient Partitioning: 8x memory reduction, same communication volume as DP. This is similar to optimizer state partitioning in practice. Optimizer states are calculated per parameter anyway so this doesn't incur any extra cost. - Parameter Partitioning: Memory reduction is linear with DP degree. For example, splitting across 64 GPUs will yield a 64x memory reduction. There is a modest 50% increase in communication volume. This works because at any time doing forward or backward, only a subset of parameters (in a layer) are required for the operation. At best you're gonna need memory equivalent to a layer size. The model parameters can be sliced in any manner (vertically or horizontally). The way it's different from tensor parallelism or pipeline parallelism is that every computation still happens on each GPU using full tensors, just that the parameters are not all on single GPU. Communication Primitives: - All-Reduce: Each process starts with its own data and ends up with the sum (or other reduction) of all data across processes. Commonly used for synchronizing gradients. Can be implemented as a combination of reduce-scatter followed by all-gather. - Ring All-Reduce: Each GPU sends and receives data in a ring topology, minimizing bandwidth bottlenecks. Communication overhead: 2 × (N-1) × X/N bytes, where N is the number of GPUs and X is the data size. - Reduce-Scatter: Each process reduces (e.g., sums) a chunk of data across all processes and keeps only its own chunk. Used as the first step in optimized all-reduce. - All-Gather: Each process gathers data chunks from all other processes so that everyone ends up with the complete data. Used as the second step in optimized all-reduce. - Broadcast: One process sends data to all others (e.g., distributing model weights at initialization). - Reduce: Data from all processes is reduced (e.g., summed) and the result is sent to a single process. - Scatter: One process splits data and sends different chunks to each process. - Gather: Each process sends data to a single process, which collects all the data. These primitives are the building blocks for distributed training and are used to synchronize parameters, gradients, and optimizer states efficiently across multiple GPUs or nodes. All-reduce = reduce-scatter + all-gather. Ring-reduce overhead: 2 × (N-1) × X/N bytes Additional Resources: [1] Scaling ML Models [2] Training Optimization [3] Understanding data parallelism, ZeRO, FSDP [4] Communication overhead slides 4.3 Pipeline Parallelism Naive Model Parallel: Partition the model by layers and put each partition on a separate GPU. The main deficiency and why this one is called "naive" MP, is that all but one GPU is idle at any given moment. GPipe: Pipeline parallelism (PP) combines model parallelism with data parallelism to reduce inefficient time "bubbles". The main idea is to split one mini-batch into multiple micro-batches and enable each stage worker to process one micro-batch simultaneously. Given m evenly split micro-batches and d partitions, the bubble is (d-1)/(m+d-1). Activation Re-computation: Only activations at partition boundaries are saved and communicated between workers. Intermediate activations at intra-partition layers are still needed for computing gradients so they are recomputed during backward passes. With activation re-computation, the memory cost for training M(l) is M(l) = O(l/d) + O(d) = O(√l). PipeDream: It schedules each worker to alternatively process the forward and backward passes. PipeDream does not have an end-of-batch global gradient sync across all the workers. A naive implementation of 1F1B can easily lead to the forward and backward passes of one micro-batch using different versions of model weights, thus lowering the learning efficiency. PipeDream proposed a few designs to tackle this issue: - Weight Stashing: Each worker keeps track of several model versions and makes sure that the same version of weights are used in the forward and backward passes given one data batch. - Vertical Sync: The version of model weights flows between stage workers together with activations and gradients. Then the computation adopts the corresponding stashed version propagated from the previous worker. This process keeps version consistency across workers. - PipeDream-flush: PipeDream-flush adds a globally synchronized pipeline flush periodically, just like GPipe. - PipeDream-2BW: PipeDream-2BW maintains only two versions of model weights, where "2BW" is short for "double-buffered weights". It generates a new model version every k micro-batches and k should be larger than pipeline depth d. A newly updated model version cannot fully replace the old version immediately since some leftover backward passes still depend on the old version. In total only two versions need to be saved so the memory cost is much reduced. Advanced Pipeline Techniques: - Breadth First Pipeline Parallelism: Looped pipeline with the principle of GPipe constitutes breadth first search approach whereas Looped pipeline with principle of 1F1B constitutes a depth first search approach. - Zero Bubble Pipeline Parallelism: Split Backward pass into two: Backward for Input and Backward for weights. Backward for input needs to be done first, backward for weights can be done later.ZB-H1: Bubble reduction is because B is initiated earlier across all workers compared to 1F1B, and the tail-end bubbles are filled by the later-starting W passes. ZB-H2: We introduce more F passes during the warm-up phase to fill the bubble preceding the initial B. We also reorder the W passes at the tail, which changes the layout from trapezoid into a parallelogram, eliminating all the bubbles in the pipeline. - Bypassing optimizer synchronization: Use post-validation strategy to replace optimizer synchronization. - LLaMA-3: Current implementations of pipeline parallelism have batch size constraint, memory imbalance due to embedding layer and warmup micro-batches and computation imbalance due to output & loss calculation making the last stage execution latency bottleneck. They modify the pipeline schedule to run an arbitrary number of micro-batches in each batch. To balance the pipeline, we reduce one Transformer layer each from the first and the last stages, respectively. This means that the first model chunk on the first stage has only the embedding, and the last model chunk on the last stage has only output projection and loss calculation. - DeepSeek-V3: The key idea of DualPipe is to overlap the computation and communication within a pair of individual forward and backward chunks. It employs a bidirectional pipeline scheduling, which feeds micro-batches from both ends of the pipeline simultaneously and a significant portion of communications can be fully overlapped. Additional Resources: [1] GPipe Paper [2] PipeDream Paper [3] Zero Bubble Pipeline 4.4 Tensor Parallelism Tensor Parallelism splits large matrix multiplications across multiple devices, enabling efficient scaling of model size. The two most common approaches are column-wise and row-wise parallelism. (1) Column-wise Parallelism The weight matrix is split along its columns. Each device holds a subset of columns and computes its portion of the output. - If input X and weight A = [A₁, A₂, ..., Aₙ], then Output O = [X @ A₁, X @ A₂, ..., X @ Aₙ] (Each device computes X @ Aᵢ for its assigned columns.) ``` Input X | | ┌─────┬─────┬─────┐ | │ A₁ │ A₂ │ A₃ │ (A split column-wise) | └─────┴─────┴─────┘ | | | | | v v v | Device 1 2 ... n | | | | | [X@A₁] [X@A₂] [X@A₃] |___________|_____|_____| | Concatenate | Output O ``` - After each device computes its partial output (X @ Aᵢ), the results must be gathered and concatenated (usually via an all-gather operation) to form the full output. This step incurs communication overhead proportional to the output size and the number of devices. The protocol is typically an all-gather across devices. (2) Row-wise Parallelism In row-wise (sometimes called output) parallelism, the input X is split column-wise across devices, and the weight matrix A is split row-wise. Each device receives a slice of the input (a subset of columns of X) and a corresponding slice of the weight matrix (a subset of rows of A). Each device computes a partial output of the same shape as the final output, and the final result is obtained by summing (reducing) these partial outputs across all devices. - If input X = [X₁, X₂, ..., Xₙ] (split column-wise) and weight A = [A₁; A₂; ...; Aₙ] (split row-wise), then Each device computes: Oᵢ = Xᵢ @ Aᵢ The final output: O = sum(O₁, O₂, ..., Oₙ) (element-wise addition across devices) ``` Input X split by columns: [X₁ | X₂ | X₃] | | | | | v v v | Device 1 Device 2 Device 3 | | | | | ┌─────┐ ┌─────┐ ┌─────┐ | │ A₁ │ │ A₂ │ │ A₃ │ (A split row-wise) | └─────┘ └─────┘ └─────┘ | | | | | [X₁@A₁] [X₂@A₂] [X₃@A₃] |_____________|______|______| | Reduce (sum) | Output O ``` - After each device computes its partial output, the results must be summed across devices to obtain the final output. This introduces communication overhead proportional to the output size and the number of devices. The protocol is typically an all-reduce across devices. Implementation: Megatron-LM provides open-source tensor parallelism implementation. For distributed attention in transformers, Megatron splits the Q, K, and V linear projections across devices, computes local attention scores and softmax independently, and then uses collective communication (such as all-reduce) to aggregate the partial attention outputs across devices for the final result. Additional Resources: [1] Megatron-LM Paper [2] Megatron-LM GitHub 4.5 Context Parallelism Context Parallelism is about how to parallelize the sequence length into multiple GPUs. During forward propagation, each GPU handles a segment of the sequence, storing only the necessary Key and Value (KV) pairs. In the backward pass, these KV pairs are reassembled across GPUs using advanced communication schemes like all-gather and reduce-scatter transformed into point-to-point communications in a ring topology. Additional Resources: [1] Context Parallelism Paper [2] Sequence Parallelism 4.6 Expert Parallelism (MoE) Instead of every token being processed by the same dense network, we introduce a set of experts (sub-networks, usually feed-forward MLPs inside Transformer blocks). Each token is assigned (via a gating function) to one or a few experts. Experts are sharded across devices (GPUs/TPUs). Each device hosts a subset of experts, and tokens are routed to whichever device hosts their assigned expert. Mixture of Experts (MoE): Instead of every token being processed by the same dense network, a set of experts (sub-networks, typically feed-forward MLPs) is introduced. A router—usually implemented as a softmax gate—assigns each token to one or more experts based on routing strategies: - Top-1 routing (e.g., Switch Transformer): Each token is sent to the single highest-scoring expert. - Top-k routing (e.g., GShard, GLaM): Each token is sent to the top k experts, and their outputs are combined. Load Balancing Challenges: In practice, load balancing can be a significant challenge in expert parallelism. Some experts, and thus some GPUs, may become overloaded with tokens while others remain underutilized. This imbalance can create bottlenecks during training or inference, reducing overall efficiency. - Device Balance Loss: Add a regularization term ("device balance loss") to the training objective that encourages an even distribution of tokens across devices. - Communication Balance Loss: Additional loss term to balance communication patterns. - Auxiliary Free Load Balancing: Instead of adding a penalty in loss, use architectural or algorithmic tricks to achieve balance automatically:Randomized routing with constraints Capacity-based routing (set a hard cap on tokens per expert) Priority dropping (drop excess tokens if an expert is full) Additional Resources: [1] Switch Transformer Paper [2] GLaM Paper [3] GShard Paper [4] Mixture of Experts Survey 5. Other Key Resources - [1] Stanford CS229s - [2] Stanford CS224n - [3] NVIDIA NeMo Framework - [4] Lilian Weng's Inference Optimization - [5] Scaling ML Models - [6] Training Optimization - [7] Efficient Training and Tradeoffs Conclusion Optimizing large language models requires careful consideration across multiple dimensions. The techniques discussed here represent the current state-of-the-art in LLM optimization, from memory-efficient attention mechanisms to advanced parallelism strategies. As models continue to grow, these optimization techniques become increasingly critical for practical deployment. Thanks for dropping by and reading till the end! Hope this was helpful in your journey! ## 相关链接 - [Gauri Gupta](https://x.com/gauri__gupta) - [@gauri__gupta](https://x.com/gauri__gupta) - [24K](https://x.com/gauri__gupta/status/2051882947758993815/analytics) - [Matrix multiplication tiling](https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html) - [Online softmax and tiling](https://www.youtube.com/watch?v=LKwyHWYEIMQ&t=14s) - [Scaling Transformers with LongNet](https://www.youtube.com/watch?v=nC2nU9j9DVQ) - [KV Caching Video](https://www.youtube.com/watch?v=Mn_9W1nCFLo&t=3869s) - [FLOPS computation efficiency with KV cache](https://docs.google.com/presentation/d/14hK7SmkUNfSEIRGyptFD2bGO7K9sJOTnwjAVg3vgg6g/edit?slide=id.g286de50af37_0_933#slide=id.g286de50af37_0_933) - [Quantization Video](https://www.youtube.com/watch?v=0VdNflU08yA) - [Lilian Weng's Inference Optimization](https://lilianweng.github.io/posts/2023-01-10-inference-optimization/) - [Scaling ML Models](https://www.youtube.com/watch?v=hc0u4avAkuM) - [Training Optimization](https://www.youtube.com/watch?v=toUSzwR0EV8) - [Understanding data parallelism, ZeRO, FSDP](https://www.youtube.com/watch?v=UVX7SYGCKkA) - [Communication overhead slides](https://docs.google.com/presentation/d/14SxjHdkvIw80FCAu5c1NGvFKDVF5DgvD2MJ1OwQ-5Gs/edit?slide=id.g24fe79ce068_0_154#slide=id.g24fe79ce068_0_154) - [GPipe Paper](https://arxiv.org/abs/1811.06965) - [PipeDream Paper](https://arxiv.org/abs/1806.03377) - [Zero Bubble Pipeline](https://arxiv.org/abs/2011.06448) - [Megatron-LM Paper](https://arxiv.org/abs/1909.08053) - [Megatron-LM GitHub](https://github.com/NVIDIA/Megatron-LM) - [Context Parallelism Paper](https://arxiv.org/abs/2105.03824) - [Sequence Parallelism](https://arxiv.org/abs/2104.04473) - [Switch Transformer Paper](https://arxiv.org/abs/2101.03961) - [GLaM Paper](https://arxiv.org/abs/2112.06905) - [GShard Paper](https://arxiv.org/abs/2006.16668) - [Mixture of Experts Survey](https://arxiv.org/abs/2202.08906) - [Stanford CS229s](https://cs229s.stanford.edu/fall2023/calendar/) - [Stanford CS224n](https://web.stanford.edu/class/cs224n/) - [NVIDIA NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/24.07/nemotoolkit/index.html) - [Lilian Weng's Inference Optimization](https://lilianweng.github.io/posts/2023-01-10-inference-optimization/) - [Scaling ML Models](https://www.youtube.com/watch?v=hc0u4avAkuM) - [Training Optimization](https://www.youtube.com/watch?v=toUSzwR0EV8) - [Efficient Training and Tradeoffs](https://www.youtube.com/watch?v=UVX7SYGCKkA) - [Upgrade to Premium](https://x.com/i/premium_sign_up) - [12:33 PM · May 6, 2026](https://x.com/gauri__gupta/status/2051882947758993815) - [24.4K Views](https://x.com/gauri__gupta/status/2051882947758993815/analytics) - [View quotes](https://x.com/gauri__gupta/status/2051882947758993815/quotes) --- *导出时间: 2026/5/6 16:20:50* --- ## 中文翻译 # LLM 优化面试笔记：训练与推理 **作者**: Gauri Gupta **日期**: 2026-05-06T04:33:11.000Z **来源**: [https://x.com/gauri__gupta/status/2051882947758993815](https://x.com/gauri__gupta/status/2051882947758993815) --- ![图片](/doc/1674/image/img_068cb700fd0e.jpg) 这是我个人的一些笔记，整理自为几家顶尖 AI 实验室面试做准备的过程，同时也回顾了高效大规模模型训练背后的核心理念。在此过程中，我汇编了这些笔记，既是面试准备的一部分，也是个人的复习，我想它们或许值得分享。高效地训练和部署大语言模型已成为现代 AI 最关键的挑战之一。随着模型扩展到数十亿参数，传统方法很快就开始失效。此处概述的技术突显了一些基本的优化策略，这些策略已作为构建和部署大规模模型的标准在业界普及。这些笔记并非旨在详尽无遗或结构完美，而是反映了在技术讨论中反复出现的关键概念和技术，这些也是大规模模型开发的支柱。这也并非对每个主题的深入探讨，而是对核心方法和概念的高层概述，提供了高效训练和推理中最常见实践和思想的简明总结。我希望它们对那些希望深入研究这些优化技术或正在准备类似技术面试的人有所帮助。 1. 内存优化技术内存是 LLM 训练和推理中最大的瓶颈。随着模型扩展到数十亿参数，传统的内存管理方法变得捉襟见肘。本节的技术侧重于在保持模型质量的同时减少内存占用，从而能够在现有硬件上训练和部署更大的模型。 1.1 Flash Attention 注意力机制的序列时间和空间复杂度呈二次方增长，这对长序列处理带来了显著的运行时和内存挑战。Flash Attention 通过分块和重计算技术降低了注意力的内存复杂度。它不是一次性处理整个注意力矩阵，而是分块处理注意力，并存储归一化因子而不是完整的注意力矩阵。分块技术根据共享内存大小分解输入，而重计算则存储 Softmax 归一化因子（对序列长度线性相关）而不是 Softmax 结果（对序列长度二次方相关）。分块技术：根据共享内存大小分解输入，并一次计算一个分块的 Softmax。不是一次性处理整个查询、键、值张量，而是进行多次传递并在后续步骤中合并结果。重计算技术：存储 Softmax 归一化因子（对序列长度线性相关）而不是 Softmax 结果（对序列长度二次方相关），利用这些因子重计算注意力分数。这减少了内存需求以及全局内存和共享内存之间的 I/O 流量。补充资源：[1] 矩阵乘法分块 [2] 在线 Softmax 和分块 1.2 多查询和分组查询注意力 - MQA (Multi-Query Attention，多查询注意力)：通过在注意力头之间共享键和值来减少内存。 - GQA (Grouped Query Attention，分组查询注意力)：通过分组查询来平衡效率和质量。 1.3 激活检查点在使用大序列长度或小微批次训练 LLM 时，输入激活很容易占满设备内存。检查点少量的激活并重计算其余部分可以减少设备内存需求。 2. 计算优化技术通过更智能的数据处理和模型架构，最大化 GPU 利用率并减少计算开销。 2.1 序列打包一种训练技术，将多个训练序列拼接成一个长序列。这消除了填充，允许每个微批次处理更多的 Token，从而最大化 GPU 计算和内存利用率。 2.2 高效 Transformer 随着序列长度和数据集规模的增加，由于标准 Transformer 架构在序列长度上的时间和内存复杂度呈二次方，其成本变得极高。高效 Transformer 技术应运而生以解决这些限制，通过降低计算和内存需求，使得处理更长的序列和更大的数据集成为可能。 - BigBird：结合使用局部、随机和全局注意力模式，将复杂度降低至 O(n)。 - Longformer：利用滑动窗口（局部）注意力结合全局注意力以提高效率。 - 低秩近似：将键和值矩阵投影到低维空间。 - LongNet：在较低层，Token 关注附近的 Token（小扩张率）。在较高层，扩张因子增大，允许 Token 关注更远的地方。随序列长度线性扩展 O(Nd)。补充资源：[1] 使用 LongNet 扩展 Transformer 3. 推理优化技术推理是产生大部分生产成本的地方。这些技术能在保持质量的同时显著加速生成过程。 3.1 KV 缓存 KV 缓存的工作原理是：在模型生成序列时，存储为每个 Token 计算出的键和值张量。在自回归生成中，在每一步新生成时，模型只需要计算新生成 Token 的注意力，利用之前步骤缓存的键和值，而不是为整个序列重新计算。这极大地减少了计算和内存使用，因为注意力机制可以重用所有先前 Token 存储的键值对。高级 KV 缓存优化： - 分组多查询注意力：通过将使用相同键和值的多个查询分组来减少 KV 缓存内存。 - 多头潜在注意力：将 K、V、Q 投影到低维潜在空间，在潜在空间中计算注意力，然后投影回原空间。 - 跨层 KV 共享：在相邻的注意力层之间绑定 KV 缓存。 - 交错局部和全局注意力：每 4-6 层使用一次全局注意力。补充资源：[1] KV 缓存视频 [2] KV 缓存的 FLOPS 计算效率 3.2 有状态缓存有状态缓存使用滚动哈希来存储对话历史，允许重用重叠的前缀。例如，如果缓存了“Hello, how are you?”，当新前缀是“Hello, how are you doing today?”时就可以重用。缓存以树状结构组织，并使用 LRU 淘汰策略来高效管理内存。对于新查询，计算其所有前缀的滚动哈希并找到最长的缓存匹配。从缓存中加载 KV 张量，并仅从新 Token 开始继续计算。KV 缓存以树状结构组织，并使用 LRU（最近最少使用）淘汰策略，因此如果内存已满，可以丢弃旧上下文。 3.3 投机解码投机解码使用一个较小的草稿模型来生成响应，然后使用目标模型进行验证，从而实现 2-3 倍的推理加速。为了使该技术有效，草稿模型必须速度快且与目标模型高度一致。 3.4 量化技术通过用更少的位表示权重/激活值来压缩模型，而不是标准的 fp32（32 位浮点数）。量化类型： - Min/Max：简单但容易受异常值影响。 - MSE：最小化原始值与量化值之间的均方误差。 - 交叉熵：在量化后保持 Softmax 最大值的顺序；argmin(softmax(v), softmax(v’))。训练后量化 (PTQ)：在完全训练好模型后，将其权重转换为低精度（如 int8 或 float16），无需进一步训练。与重新训练相比，PTQ 实施成本通常较低，但随着模型规模扩展到数十亿参数，Transformer 层中会出现大幅度异常激活值，使得简单的低比特量化效果不佳。为了解决这个问题，会附加量化感知观察器来收集输入数据的统计数据（如均值和标准差），然后用于确定量化参数。混合精度量化：不是将所有权重和激活量化为相同的比特宽度，而是为模型的不同部分分配不同的精度。例如，敏感层或激活可能使用更高的精度（如 8 或 16 位），而不敏感的部分则使用较低的精度。这种方法在内存节省和模型准确性之间取得了平衡，对于统一低比特量化会降低性能的大型模型尤其有用。混合精度量化可以分别应用于权重和激活，同时优化效率和质量。量化感知训练 (QAT)：在预训练或进一步微调期间应用量化。在前向传播中模拟量化和反量化。这模拟了量化误差，并充当正则化器使模型对其具有鲁棒性。反向传播：量化是不可微的，因此使用直通估计器 (STE) 近似梯度，该估计器在量化范围内将梯度设为 1，在范围外设为 0。补充资源：[1] 量化视频 [2] Lilian Weng 的推理优化 4. 训练优化训练大模型需要复杂的并行策略。了解不同的方法及其权衡。 4.1 混合精度训练混合精度训练利用较低精度的数字格式，最著名的是 bfloat16（脑浮点数 16），来减少内存使用并加速训练。Bfloat16 具有与标准 32 位浮点数相同的指数范围，但尾数位更少，允许它表示非常大和非常小的数字，这有助于保持深度学习所需的动态范围。然而，由于 bfloat16 和 fp16 的精度降低，梯度在反向传播期间可能会下溢或上溢，从而导致不稳定。为了解决这个问题，使用了损失缩放：在反向传播之前将损失乘以一个大常数，然后按比例缩小梯度，防止小梯度因精度有限而消失。谨慎管理损失缩放对于保持数值稳定性和充分实现混合精度训练的好处至关重要。 4.2 数据并行 DataParallel：单进程、多线程方法，适用于模型适合单个 GPU 的情况。每个 GPU 保留一份模型副本，处理不同的微批次，然后在 GPU 之间平均梯度。主要的瓶颈是它依赖于单进程、多线程通信，导致 GPU 间通信效率低下，并可能因 CPU 开销而导致速度下降。DataParallel 在单个具有多个线程的进程中运行，因此会遭受 GIL 争用的影响。同步方法：在每个小批次的结束时，工作节点需要同步梯度或权重以避免陈旧。主要有两种同步方法： - 整体同步并行 (BSP)：工作节点在每个小批次的结束时同步。这可以防止模型权重陈旧并提供良好的学习效率，但每台机器都必须停止并等待其他机器发送梯度。 - 异步并行 (ASP)：每个 GPU 工作节点异步处理数据，无需等待或停止。然而，这很容易导致使用陈旧的权重，从而降低统计学习效率。即使它增加了计算时间，但可能不会加速收敛时间。分布式数据并行 (DDP)：每个 GPU 都有自己的进程，并且可以在多个节点/机器上工作。使用环形 All-Reduce 算法来避免中心瓶颈。与 DataParallel 相比，DDP 具有更低的通信开销。 ZeRO (零冗余优化器)：不仅仅是模型参数和梯度，优化器状态（包括 Adam 的动量和方差）也占用大量内存。ZeRO-DP 有三个主要的优化阶段： - 优化器状态分片：内存减少 4 倍，通信量与 DP 相同。梯度计算可以独立于每个 GPU 进行。当针对不在当前 GPU 上的参数进行计算时，会产生通信成本，但这与 DP 中的梯度平均相同。所以始终执行 ZeRO 阶段-1。 - 梯度分片：内存减少 8 倍，通信量与 DP 相同。这在实践中类似于优化器状态分片。优化器状态本来就是按每个参数计算的，所以这不会产生任何额外的成本。 - 参数分片：内存减少与 DP 度数呈线性关系。例如，跨 64 个 GPU 分割将产生 64 倍的内存减少。通信量会有适度的 50% 增加。这是有效的，因为在任何时候进行前向或反向传播，只需要参数的一个子集（在一层中）。最多只需要相当于一层大小的内存。模型参数可以以任何方式（垂直或水平）切片。它与张量并行或流水线并行的不同之处在于，每次计算仍然在每个 GPU 上使用完整的张量进行，只是参数并不都在单个 GPU 上。通信原语： - All-Reduce：每个进程从自己的数据开始，最后得到所有进程中数据的总和（或其他归约结果）。通常用于同步梯度。可以结合 reduce-scatter 和 all-gather 来实现。 - Ring All-Reduce：每个 GPU 在环形拓扑中发送和接收数据，最大限度地减少带宽瓶颈。通信开销：2 × (N-1) × X/N 字节，其中 N 是 GPU 数量，X 是数据大小。 - Reduce-Scatter：每个进程减少（例如求和）所有进程中的数据块，并只保留自己的块。用作优化 all-reduce 的第一步。 - All-Gather：每个进程从所有其他进程收集数据块，以便每个人最终都拥有完整的数据。用作优化 all-reduce 的第二步。 - Broadcast：一个进程向所有其他进程发送数据（例如，在初始化时分发模型权重）。 - Reduce：来自所有进程的数据被归约（例如求和），结果被发送到一个进程。 - Scatter：一个进程分割数据并将不同的块发送给每个进程。 - Gather：每个进程向一个进程发送数据，该进程收集所有数据。这些原语是分布式训练的构建块，用于在多个 GPU 或节点之间高效同步参数、梯度和优化器状态。All-reduce = reduce-scatter + all-gather。环形 reduce 开销：2 × (N-1) × X/N 字节补充资源：[1] 扩展 ML 模型 [2] 训练优化 [3] 理解数据并行、ZeRO、FSDP [4] 通信开销幻灯片 4.3 流水线并行朴素模型并行：按层对模型进行分区，并将每个分区放在单独的 GPU 上。主要的缺陷以及为什么它被称为“朴素”MP，是因为在任何给定时刻，除一个 GPU 外的所有 GPU 都处于空闲状态。 GPipe：流水线并行 (PP) 结合了模型并行和数据并行，以减少低效的“气泡”时间。主要思想是将一个小批次分成多个微批次，并使每个阶段工作节点能够同时处理一个微批次。给定 m 个均匀分割的微批次和 d 个分区，气泡比例为。激活重计算：仅保存分区边界处的激活并在工作节点之间通信。分区内部层的中间激活仍然需要计算梯度，因此它们在反向传播期间被重新计算。通过激活重计算，训练 M(l) 的内存成本为 M(l) = O(l/d) + O(d) = O(√l)。 PipeDream：它调度每个工作节点交替地进行 pr

🔗 相关文档