大模型炼丹记：训练微调与推理显存占用深度对比及优化实战

内容摘要：大模型炼丹记：训练微调与推理显存占用深度对比及优化实战,

大模型，尤其是近几年涌现的 Transformer 架构模型，在 NLP、CV 等多个领域展现出强大的能力。然而，动辄数十亿甚至数千亿参数的模型，对硬件资源，尤其是 GPU 显存提出了极高的要求。本文将深入剖析大模型训练、微调和推理阶段的显存占用情况，并结合实际案例，分享一些优化策略。

1. 显存占用：问题场景重现

相信很多开发者都遇到过类似的问题：明明在 CPU 上跑得很流畅的代码，一放到 GPU 上就 OOM (Out of Memory) 了。尤其是在进行大模型训练和微调时，这个问题更加突出。我们先来重现一个常见的 OOM 场景。

假设我们使用 PyTorch 加载一个预训练的 BERT 模型，并尝试在一个较小的 batch size 下进行微调：

import torch
from transformers import BertForSequenceClassification, BertTokenizer

model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name)

model.cuda() # 将模型加载到 GPU

# 模拟数据
batch_size = 8
sequence_length = 128
input_ids = torch.randint(0, tokenizer.vocab_size, (batch_size, sequence_length)).cuda()
attention_mask = torch.ones((batch_size, sequence_length)).cuda()
labels = torch.randint(0, 2, (batch_size,)).cuda()

# 前向传播
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss

# 反向传播
loss.backward()

print("训练完成！")

即使 batch size 设得很小，也可能因为显存不足而导致程序崩溃。 错误信息类似 torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate xxx bytes. GPU ...

2. 显存占用：底层原理深度剖析

要解决 OOM 问题，首先需要理解显存占用主要由哪些部分组成。

• 模型参数 (Model Parameters): 这是显存占用的主要部分，尤其是大模型，参数量巨大。模型参数包括权重 (weights) 和偏置 (biases)，通常以 FP32 (单精度浮点数) 或 FP16 (半精度浮点数) 存储。
• 激活值 (Activations): 在前向传播过程中，每一层的输出（激活值）都需要存储在显存中，以便在反向传播时计算梯度。Transformer 模型的自注意力机制会产生大量的激活值。
• 梯度 (Gradients): 在反向传播过程中，每一层的参数梯度也需要存储在显存中。
• 优化器状态 (Optimizer States): 优化器（如 Adam, SGD）需要维护一些状态信息，例如动量 (momentum) 和方差 (variance)，这些状态信息也会占用显存。 特别是 Adam 优化器，它需要为每个参数保存两个额外的状态变量，因此会占用更多的显存。
• 临时变量 (Temporary Variables): 在计算过程中，会产生一些临时变量，例如中间结果，这些变量也会占用显存。

训练 vs. 推理：显存差异

训练阶段需要存储激活值、梯度和优化器状态，因此显存占用远大于推理阶段。推理阶段只需要加载模型参数，并进行前向传播，不需要进行反向传播，因此不需要存储激活值和梯度。 这也是为什么我们通常可以部署更大参数量的模型用于推理，而训练则需要更强的硬件配置。

微调的影响：

微调本质上也是训练，所以也会遇到与训练相似的显存问题。但是，由于微调通常是在预训练模型的基础上进行，参数更新的幅度相对较小，因此可以采用一些优化策略来减少显存占用。比如：梯度累积（Gradient Accumulation）。

3. 显存优化：具体的代码/配置解决方案

以下是一些常用的显存优化策略：

• 混合精度训练 (Mixed Precision Training): 使用 FP16 (半精度浮点数) 代替 FP32 (单精度浮点数) 来存储模型参数、激活值和梯度。FP16 可以减少显存占用，并提高计算速度。PyTorch 提供了 torch.cuda.amp 模块来实现混合精度训练。

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler() # 创建 GradScaler 对象

# 前向传播
with autocast(): # 使用 autocast 上下文管理器
    outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
    loss = outputs.loss

# 反向传播
scaler.scale(loss).backward() # 使用 scaler.scale 缩放 loss
scaler.step(optimizer) # 使用 scaler.step 更新参数
scaler.update() # 更新 scaler

• 梯度累积 (Gradient Accumulation): 将多个 mini-batch 的梯度累积起来，然后再进行一次参数更新。这可以模拟更大的 batch size，而无需增加显存占用。

accumulation_steps = 4 # 梯度累积步数

# 前向传播
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
loss = loss / accumulation_steps # 归一化 loss

# 反向传播
loss.backward()

# 每 accumulation_steps 步更新一次参数
if (batch_idx + 1) % accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()

• 梯度检查点 (Gradient Checkpointing): 在前向传播过程中，只保存部分层的激活值，在反向传播时，重新计算其他层的激活值。这可以减少显存占用，但会增加计算时间。 PyTorch 提供了 torch.utils.checkpoint 模块来实现梯度检查点。

from torch.utils.checkpoint import checkpoint

def my_model(input_ids, attention_mask):
    x = layer1(input_ids)
    x = checkpoint(layer2, x)
    x = layer3(x)
    return x

• 模型并行 (Model Parallelism): 将模型拆分到多个 GPU 上进行训练。这可以解决单卡显存不足的问题。常用的模型并行策略包括数据并行 (Data Parallelism) 和张量并行 (Tensor Parallelism)。 数据并行相对简单，可以使用 torch.nn.DataParallel 或 torch.nn.DistributedDataParallel 实现。 张量并行则需要更复杂的代码实现，通常需要修改模型的结构。

• 减少 Batch Size: 这是最简单直接的方法。减小 batch size 可以减少激活值的数量，从而减少显存占用。如果以上方法都无法解决 OOM 问题，可以尝试进一步减小 batch size。

• 使用更小的模型: 如果任务允许，可以考虑使用参数量更小的模型。例如，可以使用 DistilBERT 代替 BERT，或者使用 MobileBERT 代替 BERT。 这些模型在保持一定性能的同时，可以显著减少显存占用。

  

4. 实战避坑经验总结

• 监控显存占用: 在训练和推理过程中，需要时刻监控显存占用情况。可以使用 torch.cuda.memory_allocated() 和 torch.cuda.memory_cached() 函数来获取显存占用信息。
• 避免显存碎片: 频繁地分配和释放显存会导致显存碎片，从而降低显存利用率。可以使用 torch.cuda.empty_cache() 函数来释放未使用的显存。
• 选择合适的优化器: 不同的优化器对显存的占用不同。Adam 优化器会占用较多的显存，可以尝试使用其他的优化器，例如 SGD。
• 代码审查: 仔细检查代码，避免不必要的显存占用。例如，避免在循环中创建大量的临时变量。
• 利用NVIDIA提供的工具: 使用诸如 NVIDIA Nsight Systems 和 NVIDIA Nsight Compute 等工具，可以深入分析 GPU 的性能瓶颈，找出显存占用过高的原因。

Nginx 反向代理与大模型推理：

在实际部署大模型推理服务时，通常会使用 Nginx 作为反向代理服务器，用于负载均衡和流量转发。Nginx 可以将用户的请求分发到多个 GPU 服务器上，从而提高推理服务的吞吐量。 需要注意的是，Nginx 的并发连接数需要根据 GPU 服务器的性能进行调整，以避免服务器过载。同时，也要注意 Nginx 的缓存配置，合理利用缓存可以减少 GPU 服务器的压力。 甚至可以考虑使用宝塔面板来简化 Nginx 的配置和管理。

总之，解决大模型训练和推理阶段的显存问题需要综合考虑多个方面。 通过选择合适的优化策略，并结合实际情况进行调整，我们可以有效地减少显存占用，并提高模型的训练和推理效率。