NLP——LLM内存优化技术总结

本文主要介绍内存优化相关技术

LLM内存优化技术总结

  • LLM（大型语言模型）在内存优化方面采用了多种技术，常见的方法包括

梯度检查点（Gradient Checkpointing）

  • 在前向传播时只保存部分激活值，其余的在反向传播时重新计算，主要用在预训练阶段
  • 特点：显著减少内存占用，但会增加计算量

混合精度训练（Mixed Precision Training）

  • 使用16位浮点数（FP16）代替32位浮点数（FP32）进行计算和存储，通过自动混合精度（AMP）工具实现，关键部分（如梯度更新）仍使用FP32以保证数值稳定性
  • 特点：减少内存使用并提升计算速度，但对显卡有要求，适用于支持FP16的硬件（如NVIDIA Tensor Core GPU）

模型并行（Model Parallelism）

  • 将模型的不同层分配到多个设备上，减少单个设备的内存负担
  • 特点：支持更大模型的训练，但增加了通信开销

数据并行（Data Parallelism）

  • 将数据批次分配到多个设备上，每个设备拥有完整的模型副本
  • 特点：通过增加设备数量来分摊内存压力

梯度累积（Gradient Accumulation）

  • 在多个小批次上累积梯度后再更新模型参数，多次累计梯度后一次更新参数，能够用小内存实现大梯度更新参数
  • 特点：减少单次内存需求，支持更大的批次训练

参数卸载（Parameter Offloading）

  • 将部分模型参数存储在CPU或磁盘上，需要时再加载到GPU，可通过框架（如DeepSpeed的ZeRO-Offload）自动管理参数的加载和卸载
  • 特点：减少GPU内存占用，但可能增加I/O开销

稀疏注意力机制（Sparse Attention）

  • 只计算输入序列中部分位置的注意力权重，通过设计稀疏模式（如局部窗口、随机采样）减少计算量
    • 稀疏注意力机制（Sparse Attention）是一种优化 Transformer 模型中注意力计算的技术，旨在减少计算复杂度和内存占用。它通过限制每个输入位置只与部分其他位置进行注意力计算，而不是与所有位置进行全连接计算，从而实现高效的计算和内存管理
  • 特点：降低内存和计算复杂度、支持更长序列等，对效果是有损的

量化（Quantization）

  • 将模型参数从高精度（如FP32）转换为低精度（如INT8），比如 QLoRA 等可以做到更低的量化，大幅降低模型微调和推理的内存
  • 特点：减少内存占用和计算量，对效果可能有损

知识蒸馏（Knowledge Distillation）

  • 用大模型训练小模型，使其性能接近大模型，小模型模仿大模型的输出
  • 特点：减少内存和计算资源需求，可加速模型的部署和推理

内存高效优化器（Memory-Efficient Optimizers）

  • 使用如Adafactor等优化器，减少存储优化状态的内存
    • AdaFactor是一种优化算法，旨在减少在训练深度学习模型时的内存占用，同时保持或提高模型性能。它是由Google的研究人员提出的一种自适应学习率优化器，其特点是显存成本（Memory Cost）是次线性的（Sublinear），意味着随着参数数量的增长，所需的额外内存不会线性增长
  • 特点：降低训练时的内存使用

分层训练（Layer-wise Training）

  • 逐层训练模型，每次只加载当前层的参数和梯度
  • 特点：减少内存需求，但可能影响模型性能

内存池（Memory Pooling）

  • 预先分配并复用内存块，减少频繁分配和释放的开销
  • 特点：提高内存使用效率

模型剪枝（Model Pruning）

  • 通过删减网络结构，移除不重要的神经元或连接，实现模型压缩
  • 特点：减少模型大小和内存占用，但这种方式不常用

低秩分解（Low-Rank Factorization）

  • 将大矩阵分解为多个小矩阵，常见的方式就是LoRA相关的技术，也可以和量化结合，如QLoRA等
  • 特点：减少内存和计算需求，适用于各种微调场景

vLLM推理框架

  • vLLM 是一个专注于高效推理的框架，要用于推理阶段 ，通过 PagedAttention、连续批处理、量化等技术优化内存和计算效率，显著提升吞吐量和响应速度
  • PagedAttention
    • 这是 vLLM 的核心技术，灵感来源于操作系统的虚拟内存分页机制。它通过分页管理注意力机制中的键值（KV）缓存，显著减少了内存浪费，并支持动态调整缓存大小，从而提高了吞吐量和内存利用率
  • 连续批处理（Continuous Batching）
    • vLLM 支持将多个请求批量处理，通过共享计算资源（如 KV 缓存）来减少重复计算，从而提高吞吐量
  • 量化技术
    • 支持多种量化方法（如 GPTQ、AWQ、FP8 KV Cache 等），通过降低模型参数的精度来减少内存占用和计算开销
  • 张量并行（Tensor Parallelism）
    • 支持将模型分布到多个 GPU 上运行，通过并行计算加速推理过程
  • 推测解码（Speculative Decoding） ，也称为 投机采样
    • 使用较小的模型预测词元，再用大模型验证结果，从而加速文本生成
  • Flash Attention
    • 优化 Transformer 模型的注意力计算，减少计算复杂度和内存占用
  • OpenAI 兼容 API
    • 提供与 OpenAI API 兼容的接口，便于集成到现有应用中
  • 多 LoRA 支持
    • 支持多 LoRA（低秩适应）模型，允许在同一框架下运行多个微调模型

ZeRO显存优化技术

  • ZeRO（Zero Redundancy Optimizer）是一种用于训练阶段的显存优化技术，主要用于训练阶段，通过分片存储、通信优化和混合精度训练等技术减少显存占用，支持更大规模的模型训练
  • ZERO 技术最初是微软在 2020 年的论文 ZeRO: Memory Optimization Towards Training Trillion Parameter Models 中被提出的，详细阐述了 ZERO 的三个阶段（ZERO-1、ZERO-2、ZERO-3）及其内存优化原理
  • ZERO 技术也是 DeepSpeed 框架的核心创新之一（注：DeepSpeed 是微软开发的一个用于大规模深度学习训练的优化库）
  • 分片存储（Sharding）
    • ZeRO 将模型参数、梯度和优化器状态分片存储到多个 GPU 上，从而减少单个 GPU 的内存占用。分为三个阶段：
      • ZeRO Stage 1（ZeRO-1） ：仅分片优化器状态
      • ZeRO Stage 2（ZeRO-2） ：分片优化器状态和梯度
      • ZeRO Stage 3（ZeRO-3） ：分片优化器状态、梯度和模型参数
  • 通信优化
    • ZeRO 通过优化 GPU 间的通信（如 All-Reduce 和 Reduce-Scatter 等GPU通信操作），减少分布式训练中的通信开销
  • 混合精度训练
    • 支持 FP16 和 FP8 等低精度训练，减少显存占用并加速计算
  • 重计算（Gradient Checkpointing）
    • 在前向传播时只保存部分激活值，反向传播时重新计算其余部分，从而减少显存占用
  • 负载均衡
    • 在 MoE（Mixture of Experts）模型中，通过优化路由策略和负载分配，避免专家模型之间的负载不均衡
  • 后来微软 DeepSpeed 团队继续对 ZERO 技术进行扩展，退出了 ZeRO-Offload 和 ZeRO-Infinity 等高级技术：
    • ZeRO-Offload 可以将已划分的优化器状态和梯度卸载到 CPU 内存中
    • ZeRO-Infinity 是 ZeRO-3 的扩展，它可以利用 CPU 和 NVMe 内存来进一步扩展 GPU 的内存，支持训练更大型的模型

附录：大模型推理中的模型量化技术总结

  • 大模型推理中，模型量化旨在减少模型的存储和计算需求，同时尽量保持模型的性能

GPTQ（Gradient-based Post-training Quantization）

  • TLDR：基于梯度的后训练量化方法
  • 基本原理：
    • 在模型训练完成后，对模型权重进行量化
    • 通过优化目标函数来最小化量化误差，利用梯度调整量化时的权重误差，使量化后模型与未量化模型的表现尽可能接近
    • 采用误差反馈机制，将量化误差传播到后续层进行补偿，减少累积误差对模型输出的影响
  • 特点：
    • 适用于 8-bit 或更低的量化需求，尤其对大语言模型量化效果好
    • 不需要额外的训练数据，精度损失相对较小，特别适合复杂模型
    • 针对 GPU 使用进行了优化，在 GPU 推理时性能较好，能将权重动态去量化为 float16，提高性能的同时保持低内存占用

AWQ（Activation-aware Quantization）

  • TLDR：关注激活值的量化方法，量化过程中考虑激活值分布对模型性能的影响
  • 基本原理：
    • 分析激活值的分布特性，对激活值进行适应性处理
    • 采用非均匀量化，针对不同的激活值范围选择不同的量化尺度
  • 特点：
    • 精度较高，通过对激活值分布的考虑，能更好地保留模型的性能
    • 计算复杂度较大，因为需要分析激活值分布并进行非均匀量化操作

GGUF（Generalized Global Uniform Quantization Framework）

  • TLDR：一种通用的全局统一量化框架，用于处理大规模神经网络
  • 基本原理：
    • 通常采用全局统一量化策略，假设模型的所有层或某一类参数具有相似的分布，对整个模型的权重或激活值采用相同的量化参数
    • 采用均匀量化，将所有数值线性地映射到一个均匀的范围，并引入缩放因子，在推理阶段重定标量化后的数值，避免数值溢出或精度过低
  • 特点：
    • 简单高效，适用于资源受限的部署场景，如普通 CPU 环境
    • 兼容 Windows 和 Linux 操作系统
    • 提供从 2-bit 到 8-bit 的多级量化选项
    • 由于采用统一量化策略，可能导致某些模型层的精度损失