引言
大语言模型(LLM)如GPT-3、LLaMA等参数量达到数十亿甚至数千亿,训练和推理都面临巨大挑战。本文系统介绍LLM的优化技术,帮助您高效地训练和部署大模型。
一、训练优化
1.1 分布式训练策略
数据并行(Data Parallelism):
- 每个GPU保存完整模型副本
- 不同GPU处理不同数据批次
- 梯度同步后更新模型
- 适用于模型较小的场景
模型并行(Model Parallelism):
- 将模型参数分布到多个GPU
- 适用于单GPU装不下模型的情况
- 类型:
- 张量并行(Tensor Parallelism): 单层内切分
- 流水线并行(Pipeline Parallelism): 按层切分
ZeRO优化器:
- ZeRO-1: 切分优化器状态
- ZeRO-2: 切分优化器状态 + 梯度
- ZeRO-3: 切分优化器状态 + 梯度 + 模型参数
- 显著降低显存占用
1.2 混合精度训练
FP16训练:
- 使用16位浮点数代替32位
- 速度提升2-3倍,显存节省一半
- 挑战: 数值下溢,梯度消失
解决方案:
- Loss Scaling: 放大损失避免下溢
- FP32主权重: 保持FP32版本用于更新
- 动态Loss Scaling: 自动调整缩放因子
BF16(Brain Float 16):
- 与FP32相同的指数位,更大的数值范围
- 不需要Loss Scaling
- 需要硬件支持(如A100、H100)
1.3 梯度累积和检查点
梯度累积(Gradient Accumulation):
- 多个小batch累积梯度后再更新
- 模拟大batch训练
- 解决显存限制
梯度检查点(Gradient Checkpointing):
- 只保存部分激活值
- 需要时重新计算
- 用时间换空间
1.4 高效注意力
Flash Attention:
- 优化attention的显存访问模式
- 速度提升2-4倍
- 显存节省5-20倍
Multi-Query Attention(MQA):
- 多个Query共享单个Key和Value
- 减少KV cache大小
- 提升推理速度
Grouped-Query Attention(GQA):
- MQA和传统Attention的折中
- 多个Query共享一组KV
- 平衡性能和效率
二、推理优化
2.1 模型量化
量化类型:
- INT8量化: 8位整数,通用方案
- INT4量化: 4位整数,极致压缩
- 混合精度量化: 敏感层保持FP16
量化方法:
- Post-Training Quantization(PTQ): 训练后量化,简单快速
- Quantization-Aware Training(QAT): 训练时考虑量化,精度更高
热门工具:
- GPTQ: 高质量PTQ方案
- AWQ: 激活感知量化
- SmoothQuant: 平滑激活分布
2.2 KV Cache优化
KV Cache机制:
- 缓存历史Token的Key和Value
- 避免重复计算
- 显存占用随序列长度线性增长
优化策略:
- PagedAttention: 分页管理KV Cache
- Multi-Query Attention: 减少KV维度
- 量化KV Cache: INT8存储KV
2.3 批处理优化
Dynamic Batching:
- 动态组合多个请求
- 提高GPU利用率
- 需要处理变长序列
Continuous Batching:
- 完成的序列立即替换为新请求
- 不等整个batch完成
- 进一步提升吞吐量
2.4 推理框架
vLLM:
- PagedAttention内存管理
- Continuous Batching
- 高吞吐量,低延迟
TensorRT-LLM:
- NVIDIA官方推理引擎
- 深度优化,性能顶级
- 支持多种量化
Text Generation Inference(TGI):
- Hugging Face推出
- 易用性好,集成度高
- 支持流式输出
三、显存优化技巧
3.1 计算显存需求
模型参数:
- FP32: 参数量 × 4 bytes
- FP16: 参数量 × 2 bytes
- INT8: 参数量 × 1 byte
训练额外开销:
- 梯度: 与参数相同大小
- 优化器状态(Adam): 参数 × 8 bytes
- 激活值: 取决于batch size和序列长度
3.2 节省显存策略
- 使用混合精度训练
- 开启梯度检查点
- 使用梯度累积减小batch size
- 使用ZeRO优化器
- 使用Flash Attention
- 减少序列长度或模型深度
四、性能监控与调优
4.1 关键指标
训练阶段:
- Throughput: 每秒处理的样本数或Token数
- GPU利用率: 计算资源使用情况
- 显存使用率: 内存占用情况
- 通信开销: 多GPU同步时间
推理阶段:
- Latency: 单次推理延迟
- Throughput: QPS(每秒查询数)
- Time to First Token(TTFT): 首Token延迟
- Tokens per Second: 生成速度
4.2 性能分析工具
- nvidia-smi: 基础GPU监控
- PyTorch Profiler: 详细性能分析
- DeepSpeed Flops Profiler: 计算效率分析
- Weights & Biases: 训练监控和可视化
4.3 调优流程
- 1. 建立基准,记录初始性能
- 2. 识别瓶颈(计算、显存、通信、I/O)
- 3. 应用针对性优化
- 4. 测量效果,对比改进
- 5. 迭代优化,找到最佳配置
五、实践建议
5.1 训练配置建议
- 7B模型: 8× A100(80GB)+ DeepSpeed ZeRO-2
- 13B模型: 16× A100(80GB)+ DeepSpeed ZeRO-2
- 70B模型: 64× A100(80GB)+ DeepSpeed ZeRO-3
- 开启Flash Attention和混合精度
- 使用梯度检查点节省显存
5.2 推理配置建议
- 7B模型: 单张A100可部署,考虑INT8量化
- 13B模型: 单张A100(80GB)或双卡INT8
- 70B模型: 4-8张A100或使用INT4量化
- 使用vLLM或TensorRT-LLM框架
- 开启KV Cache和Continuous Batching
5.3 成本优化
- 训练: 使用Spot实例,设置检查点恢复
- 推理: 根据负载动态伸缩
- 考虑模型蒸馏,使用小模型
- 混合部署: 简单任务用小模型,复杂任务用大模型
总结
LLM的优化是系统工程,需要在精度、速度、成本间权衡。理解底层原理,合理选择工具和策略,才能高效地训练和部署大语言模型。