LLM 性能优化核心概念与工程实践

二、MHA 工程优化（效果无损）

KV Cache

因为 Decoder only 的特性，每次前向完，把 KV 都保留下来，用于之后计算。

# q、k、v 当前 timestep 的 query, key, value
# K_prev, V_prev 之前所有 timestep 的 key 和 value
for _ in range(time_step):
    ... 
    K = torch.cat([K_prev, k], dim=-2) # [b, h, n, d]
    V = torch.cat([V_prev, v], dim=-2) # [b, h, n, d]
    
    logits = torch.einsum('bhd,bhnd->bhn', q, K)
    weights = torch.softmax(logits/math.sqrt(d), dim=-1)
    outs = torch.einsum('bhn,bhnd->bhd', weights, V)
    ... 
    K_prev, V_prev = K, V

Online Softmax

Safe softmax 和 online softmax：参考 Flash attention && flash decoding。

Flash Attention

背景：一旦模型规模很大长度很长时，QK 根本就存不进缓存。将 QK 两个大的矩阵乘法，拆解为多次运算（平铺、重计算等），放入 SRAM，减少 HBM 访问次数，利用 SRAM 的速度优势，显著提高计算速度。 比如 Llama 7B 模型，hidden size 是 4096，那么每个 timestep 需缓存参数量为 4096232=262144，假设半精度保存就是 512KB，1024 长度那就要 512MB。而现在英伟达最好的卡 H100 的 SRAM 缓存大概是 50MB，而 A100 则是 40MB。而 7B 模型都这样，175B 模型就更不用说了。

Flash Attention 的主要改进点是： 发现 Transformer 的计算瓶颈不在运算能力，而在读写速度上，因此着手降低了对显存数据的访问次数。 传统 attention 流程如下： 从显存中取 QK 计算 -> 将结果 S 写回显存 -> 从显存读 S 计算 softmax -> 将结果 P 写回显存 -> 从显存读取 P 和 V 进行计算 -> 将结果 O 写回显存。

因此想办法进行分块计算，拆到足够小，就能全塞到 L1 缓存上（比如说 A100 的 L1 只有 192KB）进行计算了，不需要将这些参数从显存反复的读入读出，只需要读 L1 缓存，就实现了加速。但是 softmax 是需要知道全局信息的，所以分块计算后，需要一些技巧对结果进行融合。

• FlashAttention 是一种 IO-aware 算法，它通过 tiling 来减少对 HBM 的访存量，从而提高性能。 • FlashAttention 避免了从 HBM 读写一些中间结果，比如 QK 得到的相似度矩阵，以及基于相似度矩阵计算 softmax 得到的概率矩阵。

FlashAttention 对 Transformer 的加速原理简单，但因早期硬件限制未能及时出现，直到 A100 GPU 架构问世。 • 1.大幅度高速 HBM2 显存 • 2.新的异步拷贝指令，可以直接从 HBM 拷贝到 SRAM 大幅度提高的显存、和显存的拷贝效率，使得 FlashAttention 的优势得以大幅发挥。因此 flashAttention 也依赖于 GPU 架构（A100 以上）。

Page Attention

• 每个 block 类比于虚拟内存中的一个 page。每个 block 的大小是固定的，在 vLLM 中默认大小为 16，即可装 16 个 token 的 K/V 值。 • Shared prefix：在某些大模型中，所有请求可能都会共享一个前置信息（比如 system message: '假设你是一个有帮助的 AI 助手…'），这些前置信息没有必要重复存储 KV cache。 • Beam Search、并行采样（Parallel Sampling）中有大量的 KV cache 是重复的。内存使用量降低 55%。 • 对物理块的引用计数进行跟踪，并实现写时复制（Copy-on-Write）机制。

vLLM 主要用于快速 LLM 推理和服务，其核心是 PagedAttention，它将在操作系统的虚拟内存中分页的经典思想引入到 LLM 服务中。在无需任何模型架构修改的情况下，可以做到比 HuggingFace Transformers 提供高达 24 倍的 Throughput。而 PagedAttention 核心则是 attention_ops.single_query_cached_kv_attention。

Ring Attention

旨在解决处理长序列时面临内存限制问题。我们只需要把 seq_len 分为卡数那么多份 (n = num_gpu)，每张卡计算一个 block，只存储一份 Qi,Ki,Vi，通过跨卡的 p2p 通信互相传递 K，V，来实现迭代计算，就可以实现多卡的超长 context length。

Striped Attention

Striped Attention 是 Ring Attention 的一个简单扩展，它通过改变设备间分配工作的方式来解决 ring attention 的工作负载不平衡的问题。

三、FFN 部分的优化

MoE（Mixture of Experts）

参数量方面：近 2/3 的参数集中在 FFN 结构中。计算量方面：如果不是超长序列，也是 FFN 结构占大头，序列越短，FFN 计算量的占比越大。通常认为 FFN 中的 MLP 压缩了大量的知识，有一些观点将这个 MLP 看成存储了大量具体知识的 Key-Value 存储器，那么也有利用让模型学习到在不同的 context 中访问不同的知识。MLP 相对于 transformer 中的其他结构来讲，也更容易做稀疏化。因此有充分的动机对 FFN 中的 MLP 进行稀疏化。

四、微调

有多种微调方式。Freeze-tuning，Adapter Tuning，Prefix-Tuning，P-Tuning，LoRA 等。lora 用的比较多。比如 72B 微调，可以选择量化 4bit、lora_dim = 64，具有较高的性价比。

五、训练相关

混合精度

直接使用 float16 的问题：

• 精度溢出：gradient×lr 超出 float16 的精度，为 0。 • 舍入误差：权重和梯度差异大，相加的时候被舍弃。 • 原因是：由于浮点数的特性，FP16 在两个相邻的，能够被 FP16 表达的数值之间存在一定的间隔，当计算数值存在于间隔之中时，运算将会出现舍入误差。 • 具体例子：在 FP16 中 1.0 与 1.0000001 完全一样，就是因为 1.0 的最小间隔为 0.0000001，因此 1.0000001 将在这次相加中丢失（未被丢失）。

混合精度训练：

在传导过程中使用 FP16（一份权重一份梯度，即 2 份 FP16），然后使用 FP32 接受更新的梯度以及保存模型（即优化器参数为 FP32。对于 adam 来说，保存 1 份权重 +2 份辅助变量，即 3 份 FP32）。混合精度训练能够极大的提高模型训练速度，同时保留几乎 99% 的训练精度。

具体过程：

使用 float16 权重进行前向传播、并反向传播得到 float16 的梯度； 通过优化器计算出 float32 精度的权重更新量并更新 float32 权重； 将 float32 权重转换为 float16；

细节说明：

前向传播时，数据精度是 fp16。但根据 Hugging Face 源码、LLaMA 官方实现，在自注意力层有一个细节：算 softmax 之前，需要把数据精度转换成 fp32；softmax 算完后再转换回 fp16。

为什么保存两份权重反而显存占用降低？ 训练的时候，前向 + 反向所占用的显存减半了，只是权重更新的时候使用了 FP32，因此，总体上显存占用会显著减小。看下面这张图可以比较清晰，为什么保存了多份权重，训练时候显存占用反而降低。

并行、调度、训练框架

数据并行、模型并行、流水线并行、张量并行。

3D 并行：3D 并行实际上是三种常用并行训练技术的组合，即数据并行、流水线并行和张量并行。 相关的框架：Huggingface Transformer, deepspeed, megatron。

Megatron LM

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 使用的是模型并行（把一个层中间切开）。 与 gpipeline（按层切开）方式是正交、互补的。 优点：在原始代码的基础上修改简单，不需要编译器，只需要改 pytorch 代码。 通用性：只适用于 transformers。 把 MLP 和 attention 拆开。 MLP 输入：b*l

ZeRO

超线性的加速，100B 模型。 背景：模型大了，一个卡放不下。模型并行（这里特指张量并行）：垂直切开，每个层以内都要做通讯，通讯量太大，单机内多卡还行，多机器就不好（大概做到 8 卡，多了的话计算通讯比差）。内存花在什么地方：1. 参数的值 梯度 优化器的状态（冲量 variance 等）。2. 中间值 临时的 buffer。

• Optimizer->ZeRO1 • 将 optimizer state 分成若干份，每块 GPU 上各自维护一份 • 每块 GPU 上存一份完整的参数 W，做完一轮 forward 和 backward 后，各得一份梯度，对梯度做一次 AllReduce（reduce-scatter + all-gather） ， 得到完整的梯度 G，由于每块 GPU 上只保管部分 optimizer states，因此只能将相应的 W 进行更新，对 W 做一次 All-Gather

• Gradient+Optimzer->ZeRO2 • 每个 GPU 维护一块梯度 • 每块 GPU 上存一份完整的参数 W，做完一轮 forward 和 backward 后， 算得一份完整的梯度，对梯度做一次 Reduce-Scatter，保证每个 GPU 上所维持的那块梯度是聚合梯度，每块 GPU 用自己对应的 O 和 G 去更新相应的 W。更新完毕后，每块 GPU 维持了一块更新完毕的 W。同理，对 W 做一次 All-Gather，将别的 GPU 算好的 W 同步到自己这来

• Parameter+Gradient+Optimizer->ZeRO3 • 每个 GPU 维护一块模型状态 • 每块 GPU 上只保存部分参数 W，做 forward 时，对 W 做一次 All-Gather，取回分布在别的 GPU 上的 W，得到一份完整的 W， forward 做完，立刻把不是自己维护的 W 抛弃，做 backward 时，对 W 做一次 All-Gather，取回完整的 W，backward 做完，立刻把不是自己维护的 W 抛弃。做完 backward，算得一份完整的梯度 G，对 G 做一次 Reduce-Scatter，从别的 GPU 上聚合自己维护的那部分梯度，聚合操作结束后，立刻把不是自己维护的 G 抛弃。用自己维护的 O 和 G，更新 W。由于只维护部分 W，因此无需再对 W 做任何 AllReduce 操作

• ZeRO-Offload • forward 和 backward 计算量高，因此和它们相关的部分，例如参数 W（fp16），activation，就全放入 GPU • update 的部分计算量低，因此和它相关的部分，全部放入 CPU 中。例如 W(fp32)，optimizer states（fp32）和 gradients(fp16) 等 • ZeRO-Offload 分为 Offload Strategy 和 Offload Schedule 两部分，前者解决如何在 GPU 和 CPU 间划分模型的问题，后者解决如何调度计算和通信的问题

offload

• ZeRO-Offload • forward 和 backward 计算量高，因此和它们相关的部分，例如参数 W（fp16），activation，就全放入 GPU • update 的部分计算量低，因此和它相关的部分，全部放入 CPU 中。例如 W(fp32)，optimizer states（fp32）和 gradients(fp16) 等 • ZeRO-Offload 分为 Offload Strategy 和 Offload Schedule 两部分，前者解决如何在 GPU 和 CPU 间划分模型的问题，后者解决如何调度计算和通信的问题

• ZeRO-Infinity • 一是将 offload 和 ZeRO 的结合从 ZeRO-2 延伸到了 ZeRO-3，解决了模型参数受限于单张 GPU 内存的问题 • 二是解决了 ZeRO-Offload 在训练 batch size 较小的时候效率较低的问题 • 三是除 CPU 内存外，进一步尝试利用 NVMe 的空间

六、总结

LLM 的性能优化是一个系统工程，涉及模型架构设计、显存管理、计算加速以及分布式训练策略等多个层面。从 MHA 结构的 MQA/GQA/MLA 改进，到工程层面的 Flash Attention 和 Page Attention，再到训练阶段的混合精度与 ZeRO 并行，每一步都在权衡效果与效率。开发者在实际应用中，应根据具体的硬件资源和业务需求，选择合适的优化组合，以实现最佳的推理速度和训练成本。