LLM训练与推理的内存需求计算
Github: https://github.com/manuelescobar-dev/LLM-System-Requirements 网站: https://llm-system-requirements.streamlit.app
大型语言模型（LLM）通常需要耗费巨大的算力资源，往往包含数十亿个参数，并且需要在数TB的数据上进行训练。近十年来，得益于计算能力的显著提升，以及新型优化技术和架构的出现，开发如此庞大的模型才成为可能。尽管在这些方面已取得了重大进展，但由于算力资源的限制和模型的专有性，LLM的广泛应用仍面临诸多挑战。举例来说，据估计，仅训练GPT-4一个模型的成本就高达约一亿美元。
不过，随着Llama 3等开源模型的出现，各类企业和个人现在可以使用并定制这些模型。这些模型有不同规模的版本，甚至最小型的变体也能适用于移动应用。虽然微调一个70亿参数的模型仍需要相当可观的算力资源，但成本已大幅降低，使得像我这样的技术爱好者也能在本地运行其中的一些模型。
推理：参数数量 * 精度（通常为2或4字节） 训练：推理资源的 4-6 倍
免责声明：本文仅供估算内存需求参考，实际所需内存可能因具体配置和所用框架而异。因此，准确评估内存需求的唯一方法是进行实际测试。
推理（Inference）
在进行推理时，需要加载模型权重，并为KV缓存和激活值分配内存空间。由于推理阶段不涉及模型训练，因此不需要存储梯度和优化器状态。
公式：
推理总内存 = 模型大小 + KV 缓存 + 激活
模型权重（Model Weights）
内存需求中最主要且最重要的部分是加载模型所需的内存。这主要取决于模型的参数数量和选择的精度。一种常用的优化技术称为“量化”（Quantization）。量化是指以较低精度加载模型权重的方法。虽然这可能会略微影响模型性能，但通常影响不大，而且相比选择参数更少的高精度模型，使用量化后的大模型往往能获得更好的效果。
公式：
模型大小 = 参数数量 * 精度
精度：
4 字节: FP32 / 全精度 / float32 / 32 位
2 字节: FP16 / float16 / bfloat16 / 16 位
1 字节: int8 / 8 位
0.5 字节: int4 / 4 位
进一步内存优化：双重量化
KV缓存（KV Cache）
在Transformer模型中，解码过程是逐步进行的。每个时间步骤会生成一个新的token，这个过程依赖于之前生成的所有token信息。为了提高效率，避免重复计算，这些先前计算得到的token相关张量会被缓存在GPU内存中。
公式：
KV 缓存 = 2 * 批大小 * 序列长度 * 层数 * 隐藏层大小 * 精度
进一步内存优化：分页注意力
激活（Activations）
在模型的前向传播过程中，需要存储各层的中间激活值。这些激活值代表了神经网络每一层在前向传播中产生的输出数据。为了防止数值溢出并确保模型的稳定收敛，这些激活值通常需要以32位浮点数（FP32）格式保存。
公式：
激活内存 = 批大小 * 序列长度 * 隐藏层大小 *（34 +（5 * 序列长度 * 注意力头数）/（隐藏层大小））
分页注意力 序列并行化 激活重计算
训练（Training）
相比推理过程，模型训练需要更多的计算资源，这主要是因为训练过程中需要额外存储优化器状态和梯度信息。这些额外的数据在模型训练时是必不可少的，但同时也显著增加了内存占用。
公式：
总内存 = 模型大小 + KV 缓存 + 激活 +（优化器状态 + 梯度) * 可训练参数数量）
模型参数（权重）：4字节 Adam优化器（2个状态）：8字节 梯度：4字节 激活和临时内存：8字节（高端估计，大小可变） 总计：4 + 20字节/参数
微调（Fine-Tuning）
由于需要计算优化器状态和梯度，训练时需要额外的内存资源。通常采用参数高效微调（PEFT）技术，例如低秩适配（LoRA）和量化低秩适配（QLoRA），以减少可训练参数的数量。
优化器状态（Optimizer States）
优化算法需要资源来存储参数和辅助变量。这些变量包括优化算法如 Adam（2 个状态）或 SGD（1 个状态）所用的动量和方差。具体取决于优化状态的数量及其精度。
AdamW（2 个状态）：每个参数8字节 AdamW（bitsandbytes 量化）：每个参数2字节 SGD（1 个状态）：每个参数4字节
替代优化器（如 NVIDIA/apex, Adafactor, 量化Adam等） 分页优化器
梯度（Gradients）
在模型的反向传播过程中计算梯度值。它们表示损失函数相对于每个模型参数的变化率，对于在优化过程中更新参数至关重要。为了保持数值稳定性，梯度必须存储在FP32。
公式：每个参数4字节
梯度累积 梯度检查点
参考资料
[1] https://huggingface.co/docs/transformers/model_memory_anatomy
[2] https://huggingface.co/docs/transformers/perf_train_gpu_one
[3] https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/
[4] https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html
[5] https://www.oreilly.com/library/view/generative-ai-on/9781098159214/ch04.html
--【本文完】---
近期受欢迎的文章：
更多交流，可添加本人微信
（请附姓名/单位/关注领域）