资讯动态

GPU推理显存溢出?用信号量控制并发,告别CUDA OOM

发布时间:2026/10/27 15:43:45 来源:尧图企业网站定制
1. 显存溢出的直接原因模型的显存账本远比你想的复杂做模型部署的人大概都遇到过这种情况本地单测好好的模型加载完nvidia-smi一看显存占用几个G心里还挺踏实。结果接口一上线压测工具并发打个二三十个请求服务直接报CUDA out of memory进程挂掉容器直接被 OOM Killer 干掉。更头疼的是重启之后单独请求又正常了看起来就像“玄学”。先把这个账算清楚。很多人以为模型显存占用量等于模型权重大小这是一个特别危险的误解。实际上一次 GPU 推理的显存开销由三部分组成模型权重与优化器/缓存结构、KV cache如果做生成式推理、前向传播过程的中间激活值activation。举个实际例子你用 PyTorch 加载一个 7B 参数的模型FP16 权重大约占 14GB看起来一张 24GB 的卡完全够用。但单条 prompt 推理时中间激活值可能额外吃掉 2–4GB如果输出序列很长KV cache 又是随 token 数线性增长的。你以为的“24GB 够用了”其实安全余量没剩下多少。而一旦并发请求同时进来问题就从“线性”变成“乘法”。每个并发请求都会在显存里创建独立的中间激活张量夸张点说10 个并发请求可能瞬间多出 20–30GB 的临时显存需求。PyTorch 的 CUDA caching allocator 会提前向驱动申请大块显存并缓存释放块当多个请求交替分配和释放不同形状的张量时还会加剧显存碎片化明明看着总量够却分配不出一个连续大块。这就是为什么单测永远没问题、一上并发就爆炸的根本原因——你事先没有给“并发×峰值显存”的乘积留预算。再补一个底层知识点CUDA context上下文是进程级的第一次调用 CUDA API 时就会初始化每个进程大概固定占用几百 MB 到 1GB 左右的显存取决于驱动和 CUDA 版本这部分是没法省的。所以如果你用多进程方式部署每个 worker 都会有一份独立的 CUDA context 叠加在显存上。这点在后面的 worker 数量设计上很关键先记住。另外显存还有一种“假性占用”。模型跑完之后nvidia-smi里看到的显存占用往往不会立刻降下来因为 PyTorch 的缓存分配器会把释放的显存留在自己的缓存池里方便后续复用。这不代表你的服务还在泄漏只代表可用显存余量没有你以为的那么多。后面我会提到用torch.cuda.memory_stats()看真实活跃分配而不是只看nvidia-smi。2. 控制并发的三种常用方案与最终选型既然根因是“并发请求带来的显存需求叠加”最直接的解法就是限制同一时刻进入 GPU 推理的请求数量。实现思路基本有三条路我挨个说下实际用下来的感受。2.1 不加限制依赖单卡显存硬扛很多框架自带的批处理如 vLLM 的 continuous batching会在推理引擎内部做动态批处理你只管往里塞请求引擎自己调度。但如果你只是用原生 PyTorch 或 HuggingFace Transformers 包一层 FastAPI模型本身不具备排队和动态合批能力请求并发进来就是各算各的显存开销直接线性叠加。这一条只适合显存极富余、单次推理显存占用很小的场景比如一个小分类模型跑在 A100 80G 上那确实不用怎么操心。但对于大模型推理这不是方案是赌运气。2.2 用队列Queue把请求串行化思路是把所有推理请求丢进一个队列后端单线程 worker 依次消费天然实现“同一时间只有一个请求在跑”。好处是显存占用可控到最低坏处是吞吐上限被死死压在单请求延迟上。比如一次生成式推理耗时 2 秒那系统吞吐最高就是 0.5 QPS再多的并发也只能排队。对离线任务这没问题但如果你做的是在线服务用户拖个滑块用一次模型就要等前面几十人排完体验基本是灾难。队列串行的本质是拿吞吐换显存安全适合离线批量任务不适合在线交互式服务。2.3 用信号量Semaphore限制并发数最终我采用的是信号量方案。思路是在 FastAPI 的推理接口入口设置一个全局并发闸门最多允许 N 个请求同时进入 GPU 推理阶段多余的请求在await处挂起等待。这和队列不一样的地方在于它不是把请求串成一个队而是允许少量请求并行比如 4 个让它们共享 GPU 的计算能力。如果单 batch 的 batch size 是 1那 4 个并发其实就是 4 个独立 batch 一直在 GPU 上跑吞吐接近串行的 4 倍同时显存预算可控。这里有个非常关键的工程细节FastAPI 的异步事件循环里不能直接跑同步的 GPU 推理代码否则一个推理请求的 CPU 部分会把整个事件循环卡住其他所有请求全部阻塞。正确做法是把推理函数放进线程池执行同时信号量要选对类型如果推理函数是同步的、跑在线程池里用threading.Semaphore线程级信号量因为多个线程同时竞争这个闸门。如果用asyncio.to_thread或run_in_executor包裹推理还是threading.Semaphore。千万不能用asyncio.Semaphore直接包同步推理代码它只会限制异步任务的调度阻止不了线程池里多个线程真的抢到 GPU 执行权。我见过有人把asyncio.Semaphore写在同步推理函数外面结果该爆显存还是爆。原因就是同步代码跑在线程池中asyncio 的信号量管不到线程池里的并发。这个坑不踩一次真的容易想当然。考虑到我自己的生产环境里还有可能把推理丢给外部推理引擎比如 FastAPI 作为网关去请求 Ollama 或 LocalAI 推理引擎这种场景下并发控制其实是由引擎自己负责的FastAPI 侧只需要做连接池和超时控制不需要动信号量。但如果推理引擎是外部进程你同样会遇到上游并发打满、显存被撑爆的问题这时候要么信任引擎内部的调度器要么在你的网关上限制并发放量逻辑和进程内是一样的。最终我的选择是普通模块里定义一个全局的BoundedSemaphore推理接口用asyncio.to_thread包装同步推理函数信号量在函数内部获取。下面一节具体说实现。3. 信号量控制的推理接口完整实现与并发上限计算3.1 并发上限怎么算先算预算再写代码很多人上来就拍脑袋定并发数等于 4 或者 8结果要么浪费显存要么跑两天就 OOM。正确做法是先算一笔账。假设你的 GPU 总显存是M_totalPyTorch 加载模型并做完预热后的稳态显存占用包括权重、CUDA context、缓存池是M_base单次请求推理时显存峰值是M_peak这个值要通过实际测试拿到测的方法是并发 1 个请求时盯着torch.cuda.max_memory_allocated()看。那么理论上最大并发数N_max int((M_total - M_base - M_reserve) / M_peak)其中M_reserve是你手动留出的安全余量至少预留总显存的 10%–15%。为什么一定要留因为生成式推理的M_peak和输入输出长度强相关你测试时用的是 128 token 的输出线上用户可能一次生成 1024 token激活值和 KV cache 直接翻好几倍。用“最大可能输出长度”的峰值来算并发数不要用平均值。举个例子一块 24GB 的 4090加载 7B 模型 FP16 后M_base约 15GB数值因实现而异预留 3GB剩下 6GB。如果单请求峰值约 2GB那N_max就是 3。听起来有点少但这意味着任意时刻最多有 3 个请求在 GPU 上跑每个请求的峰值即使涨到 6GB累计 18GB加上基础占用 15GB 正好 33GB超过 24GB 还是会炸。所以更稳妥的做法是取N_max 2甚至动态降级。算完之后你才会理解同时并发数往往比你想的小得多这不是耻辱这是硬件边界。3.2 代码骨架BoundedSemaphore asyncio.to_thread核心代码逻辑如下注意信号量定义的位置和获取方式import asyncio import threading import torch from fastapi import FastAPI, HTTPException from pydantic import BaseModel app FastAPI() # 全局信号量限制同一时刻进入GPU推理的线程数 GPU_SEMAPHORE threading.BoundedSemaphore(2) # 根据上面的预算计算得出 model None tokenizer None def load_model(): global model, tokenizer # 读入模型并切换到 eval 模式 model ... # 你的模型加载逻辑 model.eval() model.to(cuda) # 预热跑一次极小输入触发CUDA初始化并建立缓存池 with torch.no_grad(): dummy_input ... model(dummy_input) torch.cuda.synchronize() print(model loaded and warmed up) app.on_event(startup) async def startup_event(): # 注意模型加载放线程池避免阻塞事件循环 await asyncio.to_thread(load_model) class PredictRequest(BaseModel): prompt: str max_tokens: int 256 def _infer(prompt: str, max_tokens: int) - str: 同步推理函数由线程池执行入口用信号量控制 with GPU_SEMAPHORE: # 这里做真正的推理 input_ids tokenizer(prompt, return_tensorspt).to(cuda) with torch.no_grad(): output model.generate(**input_ids, max_new_tokensmax_tokens) return tokenizer.decode(output[0], skip_special_tokensTrue) app.post(/predict) async def predict(req: PredictRequest): try: # to_thread 把同步推理放到线程池线程池内由 BoundedSemaphore 限流 result await asyncio.to_thread(_infer, req.prompt, req.max_tokens) return {result: result} except RuntimeError as e: if out of memory in str(e): # 触发OOM时清理缓存并返回错误码 torch.cuda.empty_cache() raise HTTPException(status_code503, detailGPU busy, please retry) raise几个细节我强调下信号量用BoundedSemaphore而不是Semaphore。BoundedSemaphore在释放次数超过获取次数时会抛异常能帮你及时发现代码里的acquire/release配对错误别问我是怎么知道的。模型加载放到asyncio.to_thread里。否则启动时模型加载那几十秒会直接卡住健康检查K8s 探针一失败就疯狂重启。OOM 异常捕获不能只放在接口层。如果你的推理函数内部循环里也分配中间张量建议在引擎层包一层装饰器统一处理 OOM这样不仅限于当前请求也能记录当时的显存状态方便排查。3.3 流式推理怎么和信号量共存很多生成式服务需要流式输出客户端要拿到 token 一个一个蹦出来这跟“信号量统一进口”天然有冲突——因为在一个流式响应过程中GPU 推理任务始终没结束信号量必须一直被占着。实现上我建议这样拆_infer_stream仍然是一个同步生成器函数在with GPU_SEMAPHORE:内部一个 token 一个 token 地yield结果。FastAPI 端用StreamingResponse包住它但注意要把生成器放到asyncio.to_thread执行然后通过asyncio.Queue把线程池里的数据传到StreamingResponse的异步生成器里。否则生成器在事件循环里跑同步操作又会把事件循环堵死。简化一下流式接口的伪代码长这样app.post(/stream) async def stream_predict(req: PredictRequest): queue: asyncio.Queue asyncio.Queue() async def producer(): # 在线程池中执行同步生成器 def sync_gen(): return _infer_stream(req.prompt, req.max_tokens) gen await asyncio.to_thread(sync_gen) for chunk in gen: await queue.put(chunk) await queue.put(None) task asyncio.create_task(producer()) async def consumer(): while True: chunk await queue.get() if chunk is None: break yield chunk return StreamingResponse(consumer(), media_typetext/plain)这段代码的关键点是_infer_stream内部带着with GPU_SEMAPHORE:并且在整个生成器结束之后才释放信号量这样才是“一个流式请求等于一个并发名额”的语义不会出现一个流式请求还没结束下一个请求又挤进来一起抢显存。4. 显存之外的救火措施半精度、动态批处理与失败重试信号量只解决“同一时刻最多多少个请求进 GPU”的问题但它不负责“单次请求最高峰不要太高”。所以除了闸门还要把单请求的显存峰值压下去以及尽量摊薄固定开销这样才能真正提高少量并发下的吞吐。这一节讲几个我实测有效的组合拳。4.1 半精度是底线不是可选项能开 FP16/BF16 就开这能直接把权重和大部分激活值的显存砍半。注意两点一是model.half()之后一定要重新做预热因为 FP16 下的激活值分布和 FP32 不一样缓存分配器行为也会变化预热后会重新建立适合半精度推理的缓存池二是用 BF16 的时候要确认你的卡支持Ampere 架构及以上基本都支持否则会回退到 FP32 导致速度反降。如果你用的是量化版本——比如 GPTQ、AWQ、bitsandbytes 的 int4——显存占用就更友好但要注意量化模型的首次推理有反量化开销并发时会明显拉高 CPU 负载。CPU 负载一高线程池排队疯狂延迟就会变差。所以量化部署要同时评估 CPU 核数和线程池大小别只盯着显存。4.2 动态批处理是吞吐跃升的关键严格说这不是标准信号量方案的替代品而是补强。既然我们已经允许最多 N 个请求并发那能不能把这 N 个请求合并成一个 batch 喂给模型利用 GPU 的并行计算能力答案是能但这需要改造推理函数。实现上我会在服务里维护一个“攒批队列”请求进来先不立即推理而是放进一个带超时窗口的缓冲池比如等 50ms在这 50ms 内到达的请求合并成一个 batch 送入模型50ms 一到即使只有 1 个请求也直接推理。合并后单次推理的算力利用率更高每个请求拿到的显存增量也更小——因为合并 batch 时权重还是同一份只是激活值按 batch 维度增加。代价是延迟多等一个窗口的排队时间。50ms 对大多数在线服务是可以接受的但对单请求延迟敏感的接口比如交互式一键出图就不合适。我自己的经验是生成式模型一次输出几十上百个 token合批收益很大判别式模型单次推理只有几十毫秒合批反而因等待窗口拖慢整体延迟不如直接让它们并发跑。4.3 KV cache 释放策略和显存碎片Transformers 库做generate时会在编码器和解码器内部维护 KV cache不同长度的输出会导致 cache 形状差异很大大量交替分配释放后显存碎片化很严重。我的建议是每个请求结束后不定时执行torch.cuda.empty_cache()而是只在显存占用超过安全水位时触发。空缓存会把所有缓存块释放回 CUDA代价是下一次请求分配时重新向驱动申请会有一次额外延迟。更细的做法是通过torch.cuda.memory_stats()监控active_bytes真实活跃分配的显存一旦活跃显存超过阈值触发一次empty_cache()。另外生成式推理用torch.inference_mode()代替torch.no_grad()前者会跳过一部分历史梯度和张量记录的额外开销显存占用和速度都更好。4.4 失败重试要带指数退避OOM 捕获到之后直接返回 503 是一种选择但用户体验很差。更稳的做法是客户端拿到 503 后做指数退避重试第一次 200ms第二次 400ms第三次 800ms最多 5 次。这个策略可以在服务端 FastAPI 接口里做也可以让上游网关做。我不建议在服务端内部同步做重试因为如果在请求高峰期触发重试请求量不降反增信号量闸门外堆积的任务更多反而加剧延迟。重试应该放在被动的客户端侧或网关侧别放在服务内部。5. 压测结果、uvicorn 多 worker 配置与日志排查5.1 压测要测什么并发数、延迟、吞吐、显存四维一起看写完接口后我一般用locust或wrk做压测但不是只盯着“平均响应时间”。真正要记录的是四个指标在一张表里并发请求数平均延迟(ms)吞吐(QPS)峰值显存占用(GB)失败率13203.118.20%23355.920.10%438010.223.60%85109.8超过24OOM15%这是我模拟出来的一组比较典型的数据实际会根据模型和卡型变动。注意几个规律并发从 1 加到 2吞吐几乎翻倍延迟只涨一点点。并发 4 时吞吐达到峰值延迟开始明显上升。并发 8 时显存顶到硬件极限吞吐不升反降因为部分请求 OOM 后疯狂重试资源被浪费在失败和重试上。压测结束后用nvidia-smi --query-gpumemory.used --formatcsv -l 1蹲几分钟看看显存在请求高峰过去后会不会回落。如果它一直不回落PyTorch 缓存池特性别慌如果它持续上升并最终触发 OOM那才是真正的泄漏。区分二者最可靠的方法是看torch.cuda.memory_stats()里active_bytes是否随着请求结束而降回基线而不是看nvidia-smi的 raw 值。5.2 uvicorn worker 数别乱开每个 worker 都有一份模型这是多进程部署时最坑的一点。很多人习惯性给 uvicorn 开 4 个 worker然后发现显存直接翻了 4 倍——因为每个 worker 进程独立加载了一份 CUDA context 和一份模型权重。你要么开单 worker 线程池要么开多 worker 但每个 worker 的并发数按“总显存 / worker 数”重新均分。以 24GB 卡、15GB 模型基础占用为例开 1 个 worker信号量并发数可以到 3开 2 个 worker每份模型基础占用变成 15GB两张共享一块卡时总基础占用 30GB已经超了根本跑不动开 2 个 worker只能用两张卡分别部署一个进程一张卡每张卡上信号量并发数 3总并发 6。所以我的结论是在单机单卡场景优先单 worker 多线程并发 信号量多卡场景优先多 worker 且每个 worker 绑定一张卡CUDA_VISIBLE_DEVICES0和1再在每个 worker 内做信号量限制。千万不要多 worker 共享一块卡那是自找 OOM。5.3 uvicorn 日志丢失问题与推理统计热词里提到的“uvicorn/fastapi 日志丢失问题”我在多 worker 部署时确实遇到过。原因是默认日志配置里多个 worker 进程同时往同一个文件句柄写日志时会出现被截断或被互相覆盖的情况而且异步日志队列在进程退出时丢失缓冲直接表现就是“有一批请求的记录消失了”。我的解决方案是给 uvicorn 传--log-config指定一个dictConfig核心是把disable_existing_loggers设为False并给不同 worker 加不同的日志文件通过环境变量注入 worker id。额外用 Sentry 或 Prometheus 记录请求数和推理耗时不要把业务统计完全押在日志文件上。你后面排功耗、做容量规划时可靠的指标比日志文本有用得多。uvicorn main:app --host 0.0.0.0 --port 8000 --workers 1 --log-config log_config.yamlworker 数设置为 1 时业务日志最省心如果实在要多 worker就靠上面的方法分流日志文件。6. 长期运行时的显存监控与异常自动恢复部署完光测一轮压测是远远不够的线上跑几天后显存、缓存池、碎片率都会发生漂移。这一节讲我在长期运维时盯的几个指标和处理手段。6.1 用 pynvml 监控真实显存水位手动跑nvidia-smi只在排查时用生产环境我会写一个小后台任务用pynvml库每秒采样一次显存使用量和利用率推到 Prometheus。监控指标除了总量更推荐加一个torch.cuda.memory_reserved()和torch.cuda.memory_allocated()的自定义 metric前者看 PyTorch 从驱动手里占了多少后者看真实活跃张量占了多少。两者差值越大说明缓存池越大碎片风险越高。当memory_reserved()长时间超过你设定的安全阈值比如可用显存的 85%时触发一次torch.cuda.empty_cache()。如果memory_allocated()也超阈值说明真实活跃显存异常需要检查是否有请求长时间占着信号量没释放、或者生成任务真的把输出长度拉到了上限。6.2 CUDA 上下文异常与自动恢复长时间运行后偶尔会遇到CUDA error: an illegal memory access was encountered或CUDA context lost。这类错误比 OOM 严重因为整个进程的 CUDA context 已经坏了empty_cache()救不回来必须重启进程。我的做法是在 FastAPI 里挂一个全局异常处理器单独捕获这类 CUDA 错误记录现场后调用os._exit(1)让进程退出依赖容器的自动重启策略拉起新进程。别纠结怎么恢复现场CUDA context 损坏后没有任何安全恢复路径尽早自杀重启才是对用户最负责的做法。唯一要做的是给重启策略加一点退避避免启动即崩溃时无限循环重启。6.3 健康检查接口建议做成“真实推理检查”服务存活探针如果只返回 200那只能证明进程还活着证明不了模型可服务。我在/health接口里除了检查进程状态还会执行一次极小输入的推理比如只生成 1 个 token并检测是否有 OOM 异常标记。这样做的代价是每次健康检查都消耗一次 GPU 计算但对在线服务来说宁可少接一点流量也不要带着一个已经接近 OOM 的服务继续收请求——否则一次请求高峰就是一次雪崩。这个检查的频率我设置为每 30 秒一次由负载均衡器触发不影响正常推理。实际线上跑下来救回过好几次因为显存碎片问题导致“能加载但推理极慢”的僵尸状态。关于显存碎片的真正解法我再提一句如果业务形态确实是大量短请求、输出 token 长度波动很大长期跑下来碎片化几乎是必然的。除了定期empty_cache()也可以考虑每运行 24 小时主动重启一次 worker虽然粗暴但能把很多隐性问题一起带走。这个频率要根据你的请求量来调请求量大的服务可以缩短到 12 小时请求量小的服务一周一次就够了。毕竟稳定运行比“零重启记录”这种虚荣指标重要得多。我自己在实际运维这套体系时最深的一个体会是GPU 推理服务和普通 Web 服务最大的区别在于普通 Web 服务加机器就能加并发GPU 服务却会被单卡显存这个物理上限钳死。信号量、动态批处理、半精度这些说到底都是为了“在显存这个硬约束下尽量提高有效并发”方向对了剩下的就是不断根据线上数据调整那几个参数。希望这篇实战记录能让你少走几步弯路。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑