虽然CPU和GPU的算力是主要考量,但可用内存量通常成为处理大型模型时的第一个实际瓶颈。如果处理器无法获取其计算所需的数据和模型参数,它将是无用的。因此,了解内存需求对于训练和部署现代人工智能系统都非常必要。在人工智能硬件背景下,我们主要指两种类型的内存:系统内存 (随机存取内存): 这是服务器的主内存,CPU可以直接访问。它用于存储操作系统、数据预处理脚本以及等待被送入加速器的整个数据集。GPU显存 (视频内存): 这是直接位于GPU卡上的高速内存。它比系统内存明显更快,是所有加速深度学习活动发生的地方。GPU只能处理存在于其自身显存中的数据。对于GPU加速工作负载,显存几乎总是更重要的限制。如果训练任务的组成部分无法放入单个GPU的显存中,或者无法跨多GPU配置的汇总显存,该任务将失败。训练期间显存被什么占用?在神经网络训练期间,多个组成部分必须同时驻留在显存中。总内存占用远大于模型本身。digraph G { rankdir=TB; bgcolor="transparent"; node [shape=box, style="filled,rounded", fontname="sans-serif", margin="0.2,0.1"]; edge [fontname="sans-serif"]; subgraph cluster_vram { label = "训练期间GPU显存分配"; style="filled,rounded"; color="#e9ecef"; fontcolor="#495057"; bgcolor="#f8f9fa"; weights [label="模型参数\n(权重与偏差)", fillcolor="#a5d8ff"]; gradients [label="梯度", fillcolor="#ffc9c9"]; optimizer [label="优化器状态\n(例如Adam的动量、方差)", fillcolor="#d8f5a2"]; activations [label="激活值\n(前向传播缓存)", fillcolor="#ffd8a8"]; data [label="输入数据批次", fillcolor="#eebefa"]; weights -> gradients [style=invis]; gradients -> optimizer [style=invis]; optimizer -> activations [style=invis]; activations -> data [style=invis]; } }模型训练周期中GPU显存的主要占用者。所需的总内存显著超出仅模型参数本身的大小。我们来分解这些组成部分:模型参数: 这些是神经网络的权重和偏差。其大小取决于参数数量和存储它们所用的精度(例如,32位浮点数,即FP32)。梯度: 在反向传播期间,必须为模型中的每个参数计算梯度。因此,梯度所需的内存通常等于模型参数所需的内存。优化器状态: 像Adam这样的现代优化器会维护状态,以调整每个参数的学习率。例如,Adam为每个参数存储两个移动平均值(动量和方差),这有效地使模型参数的内存占用再次翻倍。激活值: 这些是前向传播期间计算的每个层的中间输出。它们需要存储在内存中,以便在随后的反向传播期间用于梯度计算。激活值的大小是动态的,取决于输入数据的批次大小和序列长度。输入数据: 当前正在处理的数据批次也必须加载到显存中。实际内存计算为了使其具体化,我们来估计训练一个使用Adam优化器和标准32位精度(FP32)的70亿参数语言模型所需的显存,其中每个参数需要4字节。模型参数: $$ 70 \text{ 亿参数} \times 4 \text{ 字节/参数} = 28 \text{ GB} $$梯度: 梯度的尺寸与参数相同。 $$ 70 \text{ 亿参数} \times 4 \text{ 字节/参数} = 28 \text{ GB} $$优化器状态 (Adam): Adam为每个参数存储两个状态。 $$ 2 \times 70 \text{ 亿参数} \times 4 \text{ 字节/参数} = 56 \text{ GB} $$仅这三个组成部分所需的总显存为: $$ 28 + 28 + 56 = 112 \text{ GB} $$此计算甚至不包括激活值内存或CUDA内核开销,但它已经超出NVIDIA A100 80GB等高端GPU的容量。这就是为什么你经常听到“内存不足”(OOM)错误的原因。一个OOM错误简单来说意味着你尝试分配给GPU显存的数据量超出了其可用量。这可能是因为模型过大、batch_size过高(这会增加激活内存),或者多种因素组合造成的。推理内存当模型部署用于推理(即进行预测)时,内存需求要低得多。你不再需要存储梯度或优化器状态。主要内存占用者是模型参数本身,加上当前输入的激活值。对于我们的70亿参数模型,推理的基础内存将是模型权重所需的28 GB,再加上单个推理请求的激活值内存。这更易于管理,并解释了为什么可以在一个太小而无法从头训练模型的GPU上运行该模型的推理。了解这种分解是优化的前提。混合精度训练(使用16位浮点数)或选择不同的优化器等技术可以大幅减少这种内存占用,这些我们将在后续章节中讨论。目前,主要收获是:在规划基础设施时,模型大小和训练配置直接决定了你的最小内存需求。