【深度学习新浪潮】如何估算大模型的训练和推理内存需求？

要估算大模型的训练和推理内存需求，需要考虑模型参数、数据类型、中间激活值、优化器状态等关键因素。下面以开源模型 Llama 2 7B 为例，详细展示计算过程，并提供辅助计算的代码工具。

一、内存需求核心组成部分

1. 训练阶段内存消耗

• 模型参数（Parameters）：模型本身的权重
• 梯度（Gradients）：反向传播时产生的参数梯度
• 优化器状态（Optimizer States）：如Adam优化器需要存储动量和方差
• 中间激活值（Activations）：前向传播中产生的临时计算结果
• 输入/输出数据：批量数据存储

2. 推理阶段内存消耗

• 模型参数：仅需存储权重（无需梯度和优化器状态）
• 中间激活值：推理时的临时计算结果
• 输入/输出数据：批量数据存储