生产环境部署BERT推理服务时，有没有遇到过明明单请求显存占用看起来还好，突然就爆出“CUDA out of memory”的错误？而且这个错误还很随机，有时候测着没事，上线就崩，大概率是动态图的内存碎片在搞鬼。接下来就一步步说怎么从内存碎片化排查根因，再给长期解决的办法。
一、搞懂：BERT推理爆显存的“奇怪”表象
1.1 为什么动态图容易攒出内存碎片
可以把GPU显存比作小区的快递驿站，动态图就像每次寄快递都临时裁纸做盒子，不同大小的快递对应不同尺寸的纸盒，处理完快递后，纸盒不会完整退回驿站，而是拆成各种小纸片散在角落。当后面来一个大包裹需要整张大纸时，就会发现“总纸张够，但没整块的大纸”，这就是内存碎片。BERT推理用的动态图（比如PyTorch默认模式），每次处理不同长度的文本，都会动态创建对应大小的“张量盒子”，处理完后不会立刻释放整块显存，攒多了就没大空间了。
1.2 怎么初步判断是碎片导致的爆显存
简单测个小实验：在测试环境，单独跑一个1000词的长文本请求，显存只用到3G；再跑10个10词的短请求，总显存才2.5G，第8个短请求就爆显存——这就是典型的碎片问题，不是总显存不够，是没有连续的大显存块可用。
二、从动态图内存碎片化排查根因
2.1 用PyTorch内置工具找碎片的痕迹
这里用实际代码模拟生产环境的混合长度请求，定位碎片来源，技术栈统一为PyTorch 1.13 + Transformers 4.28：
import torch
from transformers import BertTokenizer, BertForSequenceClassification
# 技术栈说明：PyTorch 1.13框架 + Hugging Face Transformers库，推理环境为NVIDIA GPU
# 加载预训练BERT情感分类模型，放到GPU上
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
model = BertForSequenceClassification.from_pretrained("bert-base-chinese")
model.to("cuda")
def simulate_inference(text_batch):
# 模拟生产环境批量处理不同长度的请求
for idx, text in enumerate(text_batch):
# 分词并转成GPU张量
inputs = tokenizer(text, return_tensors="pt").to("cuda")
# 模型推理
outputs = model(**inputs)
# 打印当前显存状态：已分配（实际用的）、预留（显卡预留给程序的）
allocated = torch.cuda.memory_allocated() / (1024**2) # 转成MB
reserved = torch.cuda.memory_reserved() / (1024**2)
print(f"请求{idx+1}：已分配显存={allocated:.2f}MB，预留显存={reserved:.2f}MB")
# 这里故意不手动释放（生产环境如果忘释放就会攒碎片）
del inputs, outputs
# 模拟生产环境的混合长度请求：3个短请求+2个长请求
mixed_texts = [
"今天天气好", "我想吃米饭", "人工智能很厉害",
"人工智能是未来的技术，会改变生活，带来便利，还能帮人解决很多复杂问题",
"我昨天去了公园，看到粉色的荷花、黄色的向日葵，还喂了鸽子，心情特别轻松"
]
simulate_inference(mixed_texts)
跑这段代码时会发现：每处理一个请求，预留显存（reserved）会持续上涨，但已分配显存（allocated）不会降回初始值，说明中间攒了很多小碎片，等长请求来时，整块显存不够用就报错了。
2.2 从请求链路找碎片来源
生产环境里，除了动态图本身的特性，还有两个常见场景会加重碎片：一是请求队列里短请求和长请求交替进来，比如先处理100个短请求，再处理一个长请求，短请求占的小碎片会把大空间隔开；二是并发线程太多，每个线程申请自己的小显存，最后散成几百个小碎片，没大空间。
三、长期有效治理策略
3.1 把动态图转成静态图，从根源减少碎片化
动态图是临时裁纸，静态图就是提前剪好固定大小的纸板，拼的时候不会浪费边角料。PyTorch的jit工具可以把BERT转成静态图，固定计算图的内存结构，避免动态申请：
import torch
from transformers import BertTokenizer, BertForSequenceClassification
# 技术栈：PyTorch 1.13，静态图推理优化
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
model = BertForSequenceClassification.from_pretrained("bert-base-chinese")
model.to("cuda")
# 用示例输入追踪并转换为静态图，固定输入长度范围（比如最大128词）
example_input = {
"input_ids": torch.randint(0, 1000, (1, 128), dtype=torch.long).to("cuda"),
"attention_mask": torch.ones((1,128), dtype=torch.long).to("cuda")
}
static_model = torch.jit.trace(model, example_input)
# 保存静态模型，线上直接加载推理
torch.jit.save(static_model, "bert_static_128.pt")
# 线上推理时的显存表现：内存结构固定，碎片大幅减少
def static_infer(text):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to("cuda")
with torch.no_grad():
outputs = static_model(**inputs)
allocated = torch.cuda.memory_allocated() / (1024**2)
print(f"静态图推理显存：{allocated:.2f}MB")
静态图的优点是推理速度快30%左右，碎片几乎消失，缺点是灵活性稍差，输入长度不能超设定的最大值（比如这里最多128词）。
3.2 请求队列的流量调度，避免碎片积累
并发太高、请求长度混排会加重碎片，所以要做流量限流和批量处理同长度请求，比如用FastAPI的限流控制并发：
# 技术栈：FastAPI 0.100，生产API服务，先安装依赖
pip install fastapi uvicorn slowapi
from fastapi import FastAPI
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
import torch
from transformers import BertTokenizer
app = FastAPI()
# 限流配置：每秒最多处理5个请求，避免并发太高
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
# 加载预训练的静态BERT模型
static_model = torch.jit.load("bert_static_128.pt", map_location="cuda")
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
# 推理接口
@app.post("/bert/predict")
@limiter.limit("5/second")
async def predict(text: str):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to("cuda")
with torch.no_grad():
outputs = static_model(**inputs)
return {"sentiment": outputs.logits.argmax(dim=1).item()}
这个方案的优点是容易实现，不需要修改模型；缺点是限流阈值要根据GPU显存调整，避免影响请求延迟。
3.3 显存预分配与碎片整理的配合
在服务启动时预分配一部分显存池，给推理用，避免动态申请小内存，配合按需释放减少碎片：
import torch
from transformers import BertTokenizer, BertForSequenceClassification
# 技术栈：PyTorch 1.13，显存预分配优化
# 启动时设置：用80%的GPU显存，预留20%给系统用（避免其他程序崩溃）
torch.cuda.set_per_process_memory_fraction(0.8, 0)
# 预分配1GB的显存池，用于推理的固定内存块
pre_pool = torch.empty((1024*1024*1024//4,), dtype=torch.float32, device="cuda") # 1GB显存
# 加载模型
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
model = BertForSequenceClassification.from_pretrained("bert-base-chinese").to("cuda")
# 推理时优先用预分配的内存，推理后手动整理碎片
def pool_infer(text):
inputs = tokenizer(text, return_tensors="pt").to("cuda")
with torch.no_grad():
outputs = model(**inputs)
# 推理结束后，释放没用的临时内存，保留预分配池
torch.cuda.empty_cache()
allocated = torch.cuda.memory_allocated() / (1024**2)
print(f"显存池推理显存：{allocated:.2f}MB")
注意：empty_cache不要每个请求都调，会增加延迟，最好在请求批量处理完后调。
四、方案的优缺点与生产注意事项
4.1 各方案的优缺点
静态图方案：优点是碎片少、速度快，适合长期稳定的服务；缺点是固定输入长度，不支持超长文本推理，需要提前设定最大长度。 流量调度方案：优点是实现简单，对现有代码改动小；缺点是可能延迟部分请求，限流阈值要根据GPU性能调整。 显存预分配方案：优点是内存稳定，基本不会出现碎片；缺点是会浪费部分预分配的显存（如果用不完），需要根据服务并发量调整预分配大小。
4.2 生产环境的注意事项
转静态图时，最大输入长度要比业务里的最长文本多留10%的余量，避免边界情况报错；限流阈值要做压测，比如16G的GPU每秒最多处理8-10个BERT请求，根据显存占用调整；预分配显存比例不要超过80%，留20%给系统和其他进程，防止崩溃。
五、总结
生产环境BERT推理的显存溢出，很多时候不是真的显存不够，是动态图的内存碎片化导致的。排查时先通过显存分配数据判断是不是块不够，再用静态图替换动态图做根源优化，配合流量调度和显存预分配，就能长期稳定运行，不会再随机爆显存。
评论
围绕“生产环境部署BERT推理服务时频繁遭遇显存溢出，如何从Transformer动态图内存碎片化角度排查根因并制定长期有效治理策略”参与讨论