大模型推理服务如何解决显存不足导致的频繁 OOM 问题 - 智猿学院-前后端，数据库，人工智能，云计算等领域前沿技术讲座

大模型推理服务显存优化：从容应对OOM挑战

大家好，今天我们来聊一聊大模型推理服务中一个非常常见，但也相当棘手的问题：显存不足导致的频繁OOM (Out of Memory) 错误。OOM不仅会中断推理任务，还会影响服务的稳定性和用户体验。我们将深入探讨几种有效的显存优化策略，并提供相应的代码示例，帮助大家更好地应对这一挑战。

一、理解OOM的根源：显存需求分析

在深入优化策略之前，我们需要理解OOM的根本原因。大模型推理对显存的需求主要来自以下几个方面：

模型参数： 模型本身的大小，参数越多，占用的显存越大。
激活值： 模型在推理过程中产生的中间结果，例如每一层的输出。激活值的大小与模型结构、输入数据大小和批处理大小密切相关。
其他开销： 包括CUDA上下文、张量缓存、优化器状态等。

了解这些因素有助于我们更有针对性地选择优化方法。我们可以使用PyTorch提供的工具来分析显存使用情况：

import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 创建模型实例
input_size = 1024
hidden_size = 2048
output_size = 10
model = SimpleModel(input_size, hidden_size, output_size).cuda()

# 创建随机输入
batch_size = 32
input_data = torch.randn(batch_size, input_size).cuda()

# 使用torch.profiler分析显存使用情况
with profile(activities=[
            ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    with record_function("model_inference"):
        output = model(input_data)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
# 可以保存分析结果到文件
# prof.export_chrome_trace("trace.json")

这段代码使用torch.profiler来分析模型的显存使用情况。通过查看key_averages()的输出，可以了解每一层操作的显存占用和耗时。导出的trace.json文件可以使用Chrome DevTools打开，进行更详细的分析。

二、模型优化：瘦身和压缩

模型优化是减少显存占用的一个重要手段。主要有以下几种方法：

模型剪枝 (Pruning): 移除模型中不重要的连接或神经元，降低模型复杂度。

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

# 假设已经训练好的模型
model = SimpleModel(1024, 2048, 10)

# 对第一层全连接层进行剪枝，移除50%的权重
module = model.fc1
prune.random_unstructured(module, name="weight", amount=0.5)
prune.remove(module, "weight") #移除剪枝mask
# 查看剪枝后的模型参数量
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters after pruning: {total_params}")

这段代码使用torch.nn.utils.prune模块对模型的第一层全连接层进行了剪枝，移除了50%的权重。剪枝后，模型的参数量会减少，从而降低显存占用。

模型量化 (Quantization): 将模型的权重和激活值从浮点数转换为整数，降低存储空间和计算复杂度。
```
import torch

# 假设已经训练好的模型
model = SimpleModel(1024, 2048, 10).eval() # 量化前需要设置为eval模式

# 动态量化
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# 打印量化后的模型
print(quantized_model)
```
这段代码使用torch.quantization.quantize_dynamic函数对模型进行了动态量化，将模型的权重和激活值转换为8位整数。量化后，模型的显存占用会显著降低。需要注意的是，量化可能会导致一定的精度损失，需要在精度和性能之间进行权衡。
知识蒸馏 (Knowledge Distillation): 使用一个较小的模型（学生模型）来学习一个较大的模型（教师模型）的知识，从而得到一个更轻量级的模型。

知识蒸馏的代码实现较为复杂，需要定义损失函数来衡量学生模型和教师模型之间的差异。这里不再提供具体的代码示例，但可以参考相关的论文和教程。

三、推理优化：精打细算显存使用

除了模型优化，我们还可以通过一些推理技巧来减少显存占用：

批处理 (Batching): 将多个输入数据合并成一个批次进行推理，可以提高GPU的利用率，减少推理延迟。但是，批处理会增加激活值的显存占用，需要根据实际情况进行调整。

import torch

# 假设已经加载的模型
model = SimpleModel(1024, 2048, 10).cuda()
model.eval()

# 创建多个输入数据
batch_size = 32
num_samples = 128
input_data = torch.randn(num_samples, 1024).cuda()

# 批处理推理
with torch.no_grad():
    for i in range(0, num_samples, batch_size):
        batch = input_data[i:i+batch_size]
        output = model(batch)
        # 处理输出
        print(output.shape)

这段代码将128个输入数据分成若干个大小为32的批次进行推理。批处理可以提高GPU的利用率，但需要注意控制批次大小，避免OOM。

梯度累积 (Gradient Accumulation): 在显存有限的情况下，可以使用梯度累积来模拟更大的批次大小。梯度累积将多个小批次的梯度累加起来，然后一次性更新模型参数。

import torch
import torch.optim as optim

# 假设已经加载的模型和优化器
model = SimpleModel(1024, 2048, 10).cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 梯度累积的步数
accumulation_steps = 4
batch_size = 8
num_samples = 128
input_data = torch.randn(num_samples, 1024).cuda()
target_data = torch.randint(0, 10, (num_samples,)).cuda() #假设是分类任务

# 梯度累积训练
model.train()
for i in range(0, num_samples, batch_size):
    batch = input_data[i:i+batch_size]
    target = target_data[i:i+batch_size]

    output = model(batch)
    loss = torch.nn.functional.cross_entropy(output, target)
    loss = loss / accumulation_steps # 梯度累积需要除以累积步数
    loss.backward()

    if (i + batch_size) % (accumulation_steps * batch_size) == 0:
        optimizer.step()
        optimizer.zero_grad()

这段代码使用梯度累积来模拟更大的批次大小。梯度累积可以将多个小批次的梯度累加起来，然后一次性更新模型参数。这可以在显存有限的情况下，达到与更大批次大小相似的训练效果。

混合精度训练 (Mixed Precision Training): 使用半精度浮点数 (FP16) 来存储模型的权重和激活值，可以减少显存占用，并提高计算速度。

import torch
from torch.cuda.amp import autocast, GradScaler

# 假设已经加载的模型和优化器
model = SimpleModel(1024, 2048, 10).cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scaler = GradScaler() # 用于混合精度训练

batch_size = 32
num_samples = 128
input_data = torch.randn(num_samples, 1024).cuda()
target_data = torch.randint(0, 10, (num_samples,)).cuda() #假设是分类任务

# 混合精度训练
model.train()
for i in range(0, num_samples, batch_size):
    batch = input_data[i:i+batch_size]
    target = target_data[i:i+batch_size]

    optimizer.zero_grad()
    with autocast(): # 开启自动混合精度
        output = model(batch)
        loss = torch.nn.functional.cross_entropy(output, target)

    scaler.scale(loss).backward() # 使用scaler进行梯度缩放
    scaler.step(optimizer)
    scaler.update()

这段代码使用torch.cuda.amp模块进行混合精度训练。 autocast上下文管理器可以自动将计算转换为半精度浮点数，GradScaler用于梯度缩放，避免梯度消失。

激活值检查点 (Activation Checkpointing): 在前向传播过程中，只保存一部分激活值，其他的激活值在反向传播时重新计算。这可以显著减少激活值的显存占用，但会增加计算量。

import torch
from torch.utils.checkpoint import checkpoint

class Block(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(Block, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        return x

class CheckpointModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_blocks):
        super(CheckpointModel, self).__init__()
        self.blocks = nn.ModuleList([Block(input_size if i == 0 else hidden_size, hidden_size) for i in range(num_blocks)])
        self.fc_out = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        for block in self.blocks:
            x = checkpoint(block, x) # 使用checkpoint包裹block
        x = self.fc_out(x)
        return x

# 创建模型实例
model = CheckpointModel(1024, 2048, 10, 4).cuda()

# 创建随机输入
batch_size = 32
input_data = torch.randn(batch_size, 1024).cuda()

# 推理
output = model(input_data)
print(output.shape)

这段代码使用torch.utils.checkpoint.checkpoint函数对模型中的每个Block进行了激活值检查点处理。在前向传播过程中，checkpoint函数只会保存Block的输入，在反向传播时重新计算Block的输出。这可以显著减少激活值的显存占用。

四、服务架构优化：多卡并行和模型切分

当单卡显存无法满足需求时，我们可以考虑使用多卡并行或模型切分：

数据并行 (Data Parallelism): 将输入数据分成多个部分，分配到不同的GPU上进行推理，然后将结果合并。这可以提高推理速度，但需要注意数据同步的开销。

PyTorch提供了torch.nn.DataParallel和torch.nn.DistributedDataParallel来实现数据并行。 DistributedDataParallel更适合多机多卡的场景。
模型并行 (Model Parallelism): 将模型分成多个部分，分配到不同的GPU上进行推理。这可以降低单个GPU的显存占用，但需要注意模型切分的策略和GPU之间的通信开销。

模型并行需要根据模型的结构进行精细的设计，将不同的层分配到不同的GPU上。可以使用torch.distributed包进行GPU之间的通信。
流水线并行 (Pipeline Parallelism): 将模型分成多个阶段，每个阶段分配到不同的GPU上进行推理。一个批次的数据按照流水线的方式依次通过各个阶段。这可以提高GPU的利用率，但需要注意流水线的平衡和GPU之间的通信开销。

流水线并行通常需要结合模型并行和数据并行来实现。可以使用torch.distributed.pipeline包进行实现。
模型卸载 (Offloading): 将部分模型参数或激活值卸载到CPU内存或硬盘上，在需要时再加载到GPU显存中。这可以显著减少GPU显存占用，但会增加推理延迟。

可以使用torch.cpu.amp进行模型卸载。

五、动态调整：监控与自适应

以上策略可以单独使用，也可以组合使用。在实际应用中，我们需要根据模型的特点、硬件资源和性能需求，选择合适的优化策略。为了更好地应对OOM问题，我们需要对推理服务进行监控，并根据实际情况动态调整优化策略。

监控： 监控GPU的显存使用率、CPU利用率、推理延迟等指标。可以使用nvidia-smi命令或PyTorch的torch.cuda.memory_allocated()函数来获取GPU显存使用情况。
自适应： 根据监控数据，动态调整批次大小、量化精度、激活值检查点等参数。例如，当GPU显存使用率超过阈值时，可以减小批次大小或启用激活值检查点。

六、代码示例：动态批处理大小调整

以下是一个简单的代码示例，演示如何根据GPU显存使用情况动态调整批处理大小：

import torch
import torch.nn as nn
import time

# 假设已经加载的模型
model = SimpleModel(1024, 2048, 10).cuda()
model.eval()

# 初始批次大小
batch_size = 32
# 目标显存占用率
target_memory_ratio = 0.8
# 调整步长
batch_size_step = 4

def infer(model, input_data, batch_size):
    with torch.no_grad():
        output = model(input_data[:batch_size])
    return output

def adjust_batch_size(model, input_data, target_memory_ratio, batch_size, batch_size_step):
    """
    动态调整批次大小，使其满足目标显存占用率。
    """
    max_batch_size = batch_size
    min_batch_size = 1

    while True:
        torch.cuda.empty_cache()  # 清空缓存
        torch.cuda.reset_peak_memory_stats()

        try:
            # 尝试推理
            output = infer(model, input_data, batch_size)
            # 获取显存使用率
            memory_usage = torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated()

            if memory_usage > target_memory_ratio:
                max_batch_size = batch_size - batch_size_step
                batch_size = max(min_batch_size, max_batch_size)
                print(f"显存占用率过高，调整批次大小为: {batch_size},memory_usage:{memory_usage}")
            else:
                min_batch_size = batch_size + batch_size_step
                batch_size = min(len(input_data), min_batch_size)
                print(f"显存占用率较低，调整批次大小为: {batch_size},memory_usage:{memory_usage}")
            break

        except RuntimeError as e:
            if "out of memory" in str(e):
                max_batch_size = batch_size - batch_size_step
                batch_size = max(min_batch_size, max_batch_size)
                print(f"OOM error, 调整批次大小为: {batch_size}")
                continue
            else:
                raise e
    return batch_size

# 创建随机输入
num_samples = 128
input_data = torch.randn(num_samples, 1024).cuda()

# 动态调整批次大小
batch_size = adjust_batch_size(model, input_data, target_memory_ratio, batch_size, batch_size_step)
print(f"最终批次大小: {batch_size}")

# 推理
with torch.no_grad():
    for i in range(0, num_samples, batch_size):
        batch = input_data[i:i+batch_size]
        start_time = time.time()
        output = model(batch)
        end_time = time.time()
        print(f"Batch {i//batch_size + 1} inference time: {end_time - start_time:.4f} seconds")

这段代码演示了如何根据GPU显存使用情况动态调整批处理大小。在每次推理之前，都会尝试调整批次大小，使其满足目标显存占用率。如果发生OOM错误，则减小批次大小，直到推理成功。

七、总结：多管齐下，应对OOM

总而言之，解决大模型推理服务中显存不足导致的频繁OOM问题，需要综合考虑模型优化、推理优化和服务架构优化等多个方面。没有一种方法可以解决所有问题，需要根据实际情况选择合适的优化策略，并进行动态调整。通过精打细算地利用显存资源，我们可以构建更稳定、更高效的大模型推理服务。

以下表格总结了我们讨论的各种优化策略：

| 优化策略 | 优点 | 缺点 | 适用场景 !

优化策略	优点	缺点	适用场景
模型剪枝	减少模型参数量，降低显存占用，提高推理速度。

大模型推理服务显存优化：从容应对OOM挑战

发表回复 取消回复

发表回复取消回复