当工具申请GPU显存超过限额时，如何触发OOM优雅降级？ - 问题详情 - 创脉思

解读

在国内大模型 Agent 落地场景中，GPU 显存是硬瓶颈。训练/推理框架（如 PyTorch、TensorRT、vLLM）默认直接抛 CUDA OOM，导致整条 Agent 链路崩溃，用户请求 5xx，业务方投诉。面试官想知道：

你能否在框架层提前感知显存不足，而不是等 CUDA 驱动报错；
能否把“降”做成配置化策略，而非临时写 if/else；
能否保证降级前后状态一致（不脏写、不丢上下文），且可观测（告警、埋点、复盘）。
因此，答案要覆盖感知→决策→执行→观测四段闭环，并给出可在K8s+GPU-Operator环境落地的代码级细节。

知识点

显存预分配与预留池：CUDA 11 起支持 cudaMemGetInfo + cudaMallocAsync/cudaMemPool_t，可在真正申请前拿到剩余显存；
PyTorch 缓存分配器钩子：注册 torch.cuda.memory_stats 与 torch._C._cuda_setAllocatorSettings，在c10::Allocator层插入OOM 回调；
Agent 工具链分级：把工具标为 P0（必须 GPU）/P1（可 CPU 回退）/P2（可异步延后），降级时按优先级剪枝；
动态批处理与 KV-Cache 压缩：当显存水位 > 85% 时，即时缩小 max_batch_size 或 启用 INT4/INT8 KV-Cache；
Checkpoint & 热卸载：对 7B+ 模型采用 ZeRO-Offload 或 DeepSpeed-Inference 把层卸载到 NVMe，不杀进程；
K8s 层治理：通过 DCGM Exporter 暴露 nvidia_gpu_memory_used_bytes，HPA 配置 Memory-based 扩容阈值，防止单卡打满；
可观测：在Prometheus+Grafana大盘中增加 agent_oom_fallback_total 指标，告警链路飞书/钉钉 5 分钟内触达值班。

答案

显存水位探针
在 Agent 的 ToolExecutor 基类初始化阶段启动守护线程，每 200 ms 调用
free, total = torch.cuda.mem_get_info(device)
计算水位 = (total - free) / total。当水位 > 0.9 时，立即置位 OOM_PREDICT=1 并通知主循环。
工具准入控制器
每个工具注册时带 gpu_mem_estimate 字段（MB）。执行前由 QuotaManager 做预算检查：
if OOM_PREDICT or (allocated + estimate > quota): raise GpuQuotaExceeded
该异常被统一异常中间件捕获，进入降级逻辑，不抛 500。
三级降级策略（可配置）
- Level-1 模型压缩：把 LLM 的 torch_dtype 由 FP16→INT8，batch_size 减半，KV-Cache 压缩比 2:1，显存瞬间下降 35%±；
- Level-2 计算回退：对 P1 工具，把 GPU kernel 映射到 CPU fallback（提前用 TorchScript 导出 CPU 版），延迟增加但可用；
- Level-3 异步延后：对 P2 工具，把任务快照写入 Redis List，返回“任务已排队”，由夜间低峰调度器重试，用户侧无错误。
状态一致性保证
降级触发前，Agent 上下文写入 Checkpoint（只含张量句柄+ID，不占显存），使用 torch.cuda.comm.broadcast 把关键张量迁到备用卡；若整节点 OOM，则通过 NCCL 异步 AllGather 把未完成的张量同步到兄弟 Pod，进程不重启，用户会话不丢失。

代码级最小实现（PyTorch）

class CudaOOMFallback:
    def __init__(self, quota_mb=8192):
        self.quota = quota_mb << 20
        torch._C._cuda_setAllocatorSettings(
            max_split_size_mb=128,
            garbage_collection_threshold=0.6
        )
        torch.cuda.empty_cache()
    def check(self, need_mb):
        need = need_mb << 20
        allocated = torch.cuda.memory_allocated()
        if allocated + need > self.quota:
            torch.cuda.empty_cache()
            if torch.cuda.memory_allocated() + need > self.quota:
                return False
        return True

在工具调用入口 with cuda_oom_fallback(): ... 捕获 torch.cuda.OutOfMemoryError，自动重试 Level-1/2/3。

上线流程
- 灰度：先在 10% Pod 开启降级开关，对比 P99 延迟与 OOM 率；
- 复盘：每周拉取 agent_oom_fallback_total 指标，若降级率 > 5% 则触发预算扩容；
- 安全对齐：降级后模型输出置信度可能下降，通过 Reward Model 做二次打分，低于阈值结果打标“AI 回答可能不完整”，符合国内监管要求。

拓展思考

多租户场景：如果 Agent 平台对外提供 SaaS，需把显存配额下沉到 Namespace 级，用 NVIDIA MIG 把 A100 拆成 7 个 5GB 实例，通过 K8s device-plugin 暴露 nvidia.com/mig-5g 资源，实现硬隔离+超卖；
强化学习 Agent 的自适应降级：把“是否降级”作为环境状态，用 PPO 训练策略网络，奖励函数 = −(用户等待时间 + 显存溢出惩罚)，让 Agent 自己学会在延迟与显存间权衡；
国产卡适配：在华为昇腾 910B 环境，CANN 驱动不提供 cudaMemGetInfo 等价物，需通过 ascend-smi --query-memory 解析文本，封装成与 CUDA 同接口的 Python 扩展，保证降级逻辑零改动迁移；
法律合规：降级后若返回“请求被延后”，需在用户协议中明确“资源紧张条款”，并记录日志 180 天，满足《生成式 AI 管理办法》第十一条对服务可用性说明的要求。