我的 LLM 优化面试笔记
以下是我为几家领先 AI 实验室面试做准备时整理的个人笔记,同时也重新复习了高效大规模模型训练的核心思想。在这个过程中,我把这些笔记汇总起来——既是面试准备,也是个人复习——觉得值得分享。
高效训练和部署大型语言模型已成为现代 AI 中最具挑战性的问题之一。随着模型参数规模达到数十亿,传统方法很快就会失效。这里介绍的技术突出了构建和部署大规模模型的一些关键优化策略,这些策略已成为行业标准。
这些笔记并非详尽无遗或结构完美,而是对技术讨论中反复出现的关键概念和技术的总结,也是大规模模型开发的基础。这不是对每个主题的深入讨论,而是一个高层概述,提供高效训练和推理中最常用实践和思想的简洁总结。希望对想要深入了解这些优化技术或准备类似技术面试的人有所帮助。
1. 内存优化技术
内存是 LLM 训练和推理的最大瓶颈。随着模型扩展到数十亿参数,传统内存管理方法变得不足。本节技术重点在于减少内存占用,同时保持模型质量,从而能在现有硬件上训练和部署更大模型。
1.1 Flash Attention
注意力机制在序列长度上具有二次方的计算和内存复杂度,对长序列带来了显著的运行时和内存挑战。Flash Attention 通过分块(tiling)和重计算技术降低了注意力的内存复杂度。它不再一次性处理整个注意力矩阵,而是分块处理,并存储归一化因子而不是完整的注意力矩阵。
- 分块技术:根据共享内存大小分解输入,每次计算一个 tile 的 softmax,而不是一次性处理整个 Query、Key、Value 张量,通过多次传递并合并结果。
- 重计算技术:存储 softmax 归一化因子(与序列长度线性相关),而不是 softmax 结果(二次方相关),利用这些因子重新计算注意力分数,从而减少内存需求和全局内存与共享内存之间的 I/O 流量。
1.2 多查询与分组查询注意力
- MQA (Multi-Query Attention):通过在注意力头之间共享 Key 和 Value 来减少内存。
- GQA (Grouped Query Attention):通过对 Query 进行分组,在效率和质量之间取得平衡。
1.3 激活检查点(Activation Checkpointing)
在训练具有大序列长度或微批次大小的 LLM 时,输入激活很容易耗尽设备内存。仅对少量激活做检查点,剩余部分在需要时重新计算,从而显著降低设备内存需求。
2. 计算优化技术
通过更智能的数据处理和模型架构,最大化 GPU 利用率并减少计算开销。
2.1 序列打包(Sequence Packing)
将多个训练序列拼接成一个长序列的训练技术。它消除了填充(padding),允许每个微批次处理更多 token,从而最大化 GPU 计算和内存利用率。
2.2 高效 Transformer
随着序列长度和数据集规模增长,标准 Transformer 因序列长度的二次方复杂度而变得极其昂贵。高效 Transformer 技术旨在解决这些限制。
- BigBird:结合局部、随机和全局注意力模式,将复杂度降至 O(n)。
- Longformer:使用滑动窗口(局部)注意力结合全局注意力。
- Low-Rank Approximations:将 Key 和 Value 矩阵投影到低维空间。
- LongNet:低层使用小扩张率(dilation)关注附近 token,高层扩张率增大,实现线性 O(Nd) 复杂度。
3. 推理优化技术
推理是生产环境中成本最高的部分。这些技术能在保持质量的同时显著加速生成。
3.1 KV Caching
在自回归生成过程中,缓存每个 token 计算得到的 Key 和 Value 张量。新生成 token 时,只需计算新 token 的注意力,并复用之前缓存的 KV,从而大幅减少计算和内存开销。
高级 KV Cache 优化:
- Grouped Multi-Query Attention
- Multi-head Latent Attention
- Cross Layer KV-sharing
- 交替使用局部和全局注意力(每 4-6 层使用全局注意力)
3.2 有状态缓存(Stateful Caching)
使用滚动哈希存储对话历史,支持重叠前缀的复用。缓存以树形结构组织,并采用 LRU 淘汰策略管理内存。
3.3 推测解码(Speculative Decoding)
使用较小的 draft 模型生成候选序列,再由目标模型验证,可实现 2-3 倍推理加速。
3.4 量化技术
将权重/激活用更少的比特表示,从而压缩模型。
- 量化类型:Min/Max、MSE、Cross-entropy 等。
- PTQ(训练后量化):训练完成后直接转为低精度。
- 混合精度量化:模型不同部分使用不同比特宽度。
- QAT(量化感知训练):在训练过程中模拟量化,使用 Straight-Through Estimator (STE) 处理梯度。
4. 训练优化
4.1 混合精度训练
使用 bfloat16 等低精度格式减少内存并加速训练。通过 Loss Scaling 解决梯度下溢/上溢问题。
4.2 数据并行(Data Parallelism)
- DataParallel vs Distributed Data Parallel (DDP)
- ZeRO(Zero Redundancy Optimizer):
- Stage 1:Optimizer State Partitioning(4x 内存节省)
- Stage 2:Gradient Partitioning(8x 内存节省)
- Stage 3:Parameter Partitioning(线性内存节省)
4.3 流水线并行(Pipeline Parallelism)
- GPipe、PipeDream、Zero Bubble Pipeline Parallelism 等
- LLaMA-3 和 DeepSeek-V3 的最新优化实践
4.4 张量并行(Tensor Parallelism)
- 列并行(Column-wise)
- 行并行(Row-wise)
Megatron-LM 是典型实现。
4.5 上下文并行(Context Parallelism)
将序列长度维度并行化到多个 GPU。
4.6 专家并行(MoE - Mixture of Experts)
通过路由机制将 token 分配给不同专家子网络,实现稀疏激活。需解决负载均衡问题。
5. 其他关键资源
- Stanford CS229s / CS224n
- NVIDIA NeMo Framework
- Lilian Weng 的推理优化笔记等
6、结束语
优化大型语言模型需要在多个维度进行权衡。这里讨论的技术代表了当前 LLM 优化的前沿,从内存高效的注意力机制到先进的并行策略。随着模型持续变大,这些优化技术对于实际部署将越来越关键。
原文链接:LLM Optimization Interview Notes: Training and Inference
汇智网翻译整理,转载请标明出处