PyTorch CUDA 内存可视化与快照分析实战:用 memory_viz 定位显存泄漏与 OOM

发布时间:2026/11/8 22:48:29

PyTorch CUDA 内存可视化与快照分析实战:用 memory_viz 定位显存泄漏与 OOM PyTorch CUDA 内存可视化与快照分析实战用 memory_viz 定位显存泄漏与 OOM【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读本文围绕 PyTorch 官方文档 docs/source/torch_cuda_memory.md 展开系统讲解如何使用torch.cuda.memory._record_memory_history()、_snapshot()与_dump_snapshot()生成 CUDA 内存快照并借助交互式可视化器Active Memory Timeline、Allocator State History逐块还原张量的分配与释放过程。读完本文你将掌握一套可落地的显存排查方法论既能定位“内存被谁占用”也能识别非 PyTorch 分配器管理的显存还能通过快照 API 的底层字段理解分配器 Segment / Block 的真实行为。一、理解 PyTorch CUDA 内存分析的基本原理1.1 快照能记录什么PyTorch 提供了一套内存快照机制在任意时刻记录 CUDA 已分配内存的状态并可选择性地记录导致该状态的分配事件历史。生成的快照可以拖拽到交互式查看器pytorch.org/memory_viz中浏览查看器是运行在本地的 JavaScript 应用不会上传任何快照数据适合在含有敏感模型信息的场景下使用。需要特别明确快照的可见范围来自原文档的 note 及源码实现默认情况下内存分析器只能看到经PyTorch 分配器分配和管理的 CUDA 设备内存例如torch.empty(..., devicecuda)或通过torch.cuda.memory.CUDAPluggableAllocator分配的内存CPUPinned 内存主机锁页内存可以通过向_record_memory_history传入record_pinned_host_memoryTrue选择性纳入直接在 C 中调用 CUDA API如cudaMalloc、cuMemCreate分配的内存或通过第三方 Python 绑定如 cuda-python分配的内存不会出现在 PyTorch 内存分析器中NCCL是典型代表——它用于 CUDA 设备上的分布式通信但其 GPU 内存分配绕过了 PyTorch 分配器因此不会被快照统计。1.2 快照字典的底层结构从源码 torch/cuda/memory.py 可以看到_snapshot()返回的字典具有如下 TypedDict 结构理解它是读懂后续可视化视图的前提segmentsSegment列表每个 Segment 是一次cudaMalloc返回的内存所有 Segment 大小之和即为“保留内存reserved memory”总量Segment 会被缓存复用若复用尺寸小于 Segment则会被拆分为多个 Blocktorch.cuda.empty_cache()只会释放完全空闲的 Segmentdevice_tracesTraceEntry列表仅在开启内存历史记录时存在逐条记录分配器动作Segment关键字段address、total_sizecudaMalloc 分配大小、stream、segment_typesmall或large1MB 为 large、segment_pool_id、allocated_size、active_size、blocksBlock关键字段size、requested_size申请大小可能因对齐小于size、address、stateactive_allocated/active_awaiting_free/inactive、frames分配发生处的栈回溯TraceEntry关键字段actionalloc、free_requested、free_completed、segment_alloc、segment_free、oom、snapshot、annotate、addr、frames、size、stream、pool_idOOM 事件还包含device_freeOOM 发生时 CUDA 仍报告的空闲量。这一结构意味着你不仅能看到“谁占用了内存”还能精确区分“保留但未使用”inactive与“真正被张量持有”active_allocated的内存。二、生成一份内存快照三行代码的标准流程原文档给出了生成快照的经典模式——先开启内存历史再运行待观察代码最后把快照 pickle 保存到文件# 开启内存历史会为快照附加 traceback 与事件历史 torch.cuda.memory._record_memory_history() run_your_code() torch.cuda.memory._dump_snapshot(my_snapshot.pickle)保存后打开pytorch.org/memory_viz把 pickle 文件拖拽进查看器即可开始分析。2.1_record_memory_history的参数全景该函数在源码中的完整签名torch/cuda/memory.py远不止enabled一个参数以下是各参数含义与默认值参数可选值 / 默认值作用enabledNone/state/all默认allNone关闭记录state仅保留当前已分配内存的信息all额外保留全部 alloc/free 调用历史contextNone/state/alloc/all默认allNone不记录任何 tracebackstate记录当前已分配内存的 tracebackalloc额外记录 alloc 调用的 tracebackall再额外记录 free 调用的 tracebackstackspython/all默认allpython仅包含 Python、TorchScript 与 inductor 帧all额外包含 C 帧max_entriesint默认sys.maxsize历史中最多保留的 alloc/free 事件数长时运行任务务必设置合理上限每条记录约几 KB过大将导致内存爆炸详见下文环形缓冲clear_historybool默认False开启时是否清空历史skip_actionslist[str]跳过记录的动作类型用于降低内存开销可选alloc、free_requested、free_completed、segment_alloc、segment_free、oom、snapshotrecord_pinned_host_memorybool默认False同时记录 CPU 锁页内存分配器的历史数据位于快照的host_segments与host_traces键record_cudabool默认True是否记录 CUDA 设备分配器置False可与record_pinned_host_memoryTrue组合实现只记录主机锁页内存2.2 环形缓冲与开销生产环境开启的取舍源码注释揭示了历史记录的环形缓冲实现torch/cuda/memory.pyif (record_history) { if (alloc_trace-size() alloc_trace_max_entries_) { alloc_trace-emplace_back(te); } else { (*alloc_trace)[alloc_trace_next] te; if (alloc_trace_next alloc_trace_max_entries_) { alloc_trace_next 0; } } }当达到max_entries上限后新事件会覆盖最旧事件因此长时间运行的任务只保留最近一段历史延迟开销较低Python 回溯收集约 2µs/条C 回溯收集约 50ns/帧多数程序折算约 2µs/条因此官方建议即使生产任务也可开启以备日后排查快照文件体积_dump_snapshot的 docstring 明确指出文件大小随max_entries与每条栈深增长每条记录数 KB长时运行 大max_entries时文件可达 GB 级。2.3 配套 API_snapshot与_dump_snapshot_snapshot(deviceNone, augment_with_fx_tracesFalse)返回内存快照字典。device为None时捕获当前设备augment_with_fx_tracesTrue时会为栈帧附加 FX 调试信息fx_node_op、fx_node_name、fx_original_trace把生成的 FX 代码映射回原始模型源码_dump_snapshot(filenamedump_snapshot.pickle, augment_with_fx_tracesFalse)把_snapshot()结果 pickle 到文件供pytorch.org/memory_viz打开torch/cuda/memory.py。三、包含 Pinned主机内存锁页内存分析默认快照只含 CUDA 设备内存。若要一并捕获 CPU 锁页内存分配例如pin_memoryTrue创建的张量传入record_pinned_host_memoryTruetorch.cuda.memory._record_memory_history(record_pinned_host_memoryTrue) run_your_code() snapshot torch.cuda.memory._snapshot() # 主机分配器数据位于 snapshot[host_segments] 与 snapshot[host_traces] torch.cuda.memory._dump_snapshot(my_snapshot.pickle)若只想记录主机锁页内存、完全跳过 CUDA 设备分配器可将record_pinned_host_memoryTrue与record_cudaFalse组合torch.cuda.memory._record_memory_history(record_pinned_host_memoryTrue, record_cudaFalse)注意当前pytorch.org/memory_viz可视化器尚不支持展示主机内存数据需要以编程方式从快照字典中检查host_segments与host_traces键对应源码 torch/cuda/memory.py 的参数说明。四、Active Memory Timeline查看张量生命周期Active Memory Timeline活跃内存时间线展示快照内某个 GPU 上所有存活张量随时间的变化是定位“某块显存什么时候被谁持有”的第一站可对绘图区域进行 Pan/Zoom聚焦较小的分配块鼠标悬停已分配块可查看该块被分配时的栈回溯以及地址等细节数据量较大时可调整 detail 滑块减少渲染的分配数量以提升性能。下图即该视图的典型形态图片来自 docs/source/_static/img/torch_cuda_memory/active_memory_timeline.png五、Allocator State History还原每个分配事件的分配器状态Allocator State History分配器状态历史在左侧以时间线列出单个分配器事件。选中某个事件后右侧会给出该事件发生时分配器状态的视觉汇总汇总展示每个从cudaMalloc返回的 Segment以及它如何被拆分成“单个分配 Block”或“空闲空间”悬停 Segment / Block 可查看内存分配时的栈回溯悬停事件如张量被释放可查看事件发生时的栈回溯OOM 事件会被标记出来。当一次分配失败但保留内存reserved memory仍然充足时查看 OOM 时刻的内存状态往往能揭示失败根因——例如碎片化导致无法满足大块连续分配。下图即该视图的典型形态图片来自 docs/source/_static/img/torch_cuda_memory/allocator_state_history.png5.1 地址标识字符串的跨视图追踪栈回溯信息还会报告分配发生的地址。形如b7f064c000000_0的字符串含义是block地址为7f064c000000且这是该地址第_0第 0次被分配。这个唯一字符串可以被在 Active Memory Timeline 中检索查看张量存活区间在 Active State History 中搜索检查张量被分配或释放时的内存状态。这构成了“时间线定位分配 → 状态历史还原现场”的完整排查闭环。六、识别非 PyTorch 分配对比 device_memory_used 与快照如果怀疑 CUDA 内存被分配在 PyTorch 分配器之外典型如 NCCL 通信缓冲、cuda-python 绑定分配可以通过pynvml 包收集裸 CUDA 分配信息再与 PyTorch 报告的分配量对比从而确认差异来源。使用device_memory_used获取 PyTorch 之外的原始显存占用import torch device_idx ... print(torch.cuda.device_memory_used(device_idx))从源码实现看torch/cuda/init.py在 CUDA 平台上该函数通过 pynvml 的nvmlDeviceGetMemoryInfo(handle).used返回全局已用显存字节数即整个进程含非 PyTorch 分配的真实占用在 ROCm/HIP 平台torch.version.hip非空上则改用 amdsmi 的amdsmi_get_gpu_vram_usage换算字节数。因此对比公式为device_memory_used()全局真实占用− 快照中 PyTorch 分配的合计 ≈ 非 PyTorch 分配 / 其他进程占用。若差值显著应重点排查 NCCL、自定义 CUDA 扩展与第三方库的直接 CUDA API 调用。七、完整排查工作流从怀疑到定位综合以上能力一套可复用的显存排查流程如下开启历史torch.cuda.memory._record_memory_history()生产长任务建议显式设置max_entries并视需要设置stackspython降低体积运行并转储运行待观察代码段torch.cuda.memory._dump_snapshot(my_snapshot.pickle)时间线初筛在 Active Memory Timeline 中定位“峰值时刻”与异常长寿命分配块状态历史深挖在 Allocator State History 中检索该块的地址标识字符串如b7f064c000000_0查看其分配 / 释放时刻的分配器全局状态若存在 OOM 事件检查该时刻保留内存与空闲块的分布以判断是否碎片化排除外部分配用torch.cuda.device_memory_used()与快照总量对比识别 NCCL 等非 PyTorch 分配必要时补查主机内存以record_pinned_host_memoryTrue重新记录编程检查host_segments/host_traces程序化分析对于超大规模快照可直接在 Python 中操作_snapshot()返回的字典segments / blocks / device_traces自行统计各栈帧累计分配量而不必依赖可视化器。仓库中 torch/cuda/_memory_viz.py 还提供了基于快照生成 flamegraphsegments、memory、compare等函数的命令行辅助工具可用于批量对比前后两次快照的差异适合在 CI 或脚本化巡检中落地。八、相关 API 速查以下 API 均在torch.cuda.memory模块下torch/cuda/memory.py_record_memory_history(enabledall, contextall, stacksall, max_entriessys.maxsize, deviceNone, clear_historyFalse, skip_actionsNone, record_pinned_host_memoryFalse, record_cudaTrue)开启 / 关闭内存历史与栈回溯记录_snapshot(deviceNone, augment_with_fx_tracesFalse)返回当前 CUDA 内存状态字典segments / device_traces以及可选的 host_segments / host_traces_dump_snapshot(filenamedump_snapshot.pickle, augment_with_fx_tracesFalse)将快照 pickle 落盘供pytorch.org/memory_viz加载torch.cuda.device_memory_used(deviceNone)位于 torch/cuda/init.py返回设备全局已用显存NVML 或 AMDSMI用于对比识别非 PyTorch 分配。小结通过快照 可视化器PyTorch 把 CUDA 分配器内部的 Segment / Block 结构与每次 alloc / free 事件完整暴露给开发者Active Memory Timeline 回答“内存何时被占”Allocator State History 回答“分配失败时内存长什么样”device_memory_used对比则划清了 PyTorch 与外部分配的边界。结合本文对底层快照字典结构与各参数开销的说明你可以依据任务规模单次训练 / 长时服务 / 分布式选择合适的记录粒度把显存排查从“盲猜”升级为“有据可查”。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
延伸阅读

更多相关文章

2026/11/8 16:59:49

哈尔滨可信数据交易空间:隐私计算与区块链的创新实践

1. 项目背景与核心价值 哈尔滨可信数据交易空间项目以1.8亿元投资规模引发行业关注,这标志着东北地区首个大型数据要素市场化配置基础设施的落地。作为深耕数据交易领域多年的从业者,我观察到这个项目不同于传统的数据交易平台,其创新性体现在…

2026/11/1 3:25:54

从零散灵感碎片到完整项目的重建方法论

1. 项目概述 作为一名从业多年的技术博主,我经常遇到一个困扰:当灵感突然来临时,却因为各种原因没能及时记录下完整的项目构思。这种情况在创意工作者和技术开发者中尤为常见——我们可能只来得及写下几个关键词或一个模糊的想法,…

2026/11/3 11:33:23

2026年高职物联网专业就业前景与薪资分析

1. 高职物联网专业就业前景概述 2026年高职物联网专业毕业生的薪资水平是当前许多在校生和家长关注的焦点。作为一位在物联网行业摸爬滚打多年的从业者,我亲眼见证了物联网技术从概念到落地的完整发展历程。物联网作为新一代信息技术的重要组成部分,已经…

2026/11/8 22:47:39

郑州电工证考试地址多久拿证别怕挂科,实操考试技巧全在这儿

郑州电工证考试地址多久拿证别怕挂科,实操考试技巧全在这儿 实操考试心里没底怕挂科?你不是一个人。郑州电工证考试地址虽然已经明确,但很多人还是对考试流程、内容、拿证时间、培训质量充满疑虑。特别是对于刚入行的三门峡项目现场管理员来说,如何选择靠…

2026/11/8 22:47:39

大厂县考电工证千万别踩坑 这些细节能帮你省下培训费

大厂县考电工证千万别踩坑 这些细节能帮你省下培训费 怕考不过白交培训费?大厂县考电工证千万别踩坑,这些坑踩了就亏大了。我干这行十多年,见过太多人为了考电工证被坑得团团转,培训费花了不说,还因为资料不全、流程不对、信息不透明,最后考试挂了只能…

2026/11/8 22:47:39

嘉兴电工证办理机构良心建议:实操考试心里没底怕挂科怎么办?

嘉兴电工证办理机构良心建议:实操考试心里没底怕挂科怎么办? 实操考试心里没底怕挂科?你不是一个人。我见过太多人,拿着报名表,坐在考场门口,心跳加速、手心冒汗。特别是像你这样来自驻马店,在职建筑工人,时间紧、任务重,还要应对一场关乎未来职业发…

2026/11/8 22:47:39

富阳区特种工报考避坑指南:郑州劳务班组怎么报才不被坑

富阳区特种工报考避坑指南:郑州劳务班组怎么报才不被坑 你是不是也担心,找不对地方报名,结果被中介骗钱?别急,富阳区特种工报考这事,咱们得说清楚。尤其是郑州劳务班组的负责人,你们天天跟特种作业打交道,最怕的就是因为证件问题耽误项目进度,甚至被…

2026/11/8 22:42:39

合肥电工作业证怎么考 别踩坑才能顺利上岗

合肥电工作业证怎么考 别踩坑才能顺利上岗 你还在为找不到工作发愁吗?其实,很多电工岗位都明确要求持证上岗, 合肥电工作业证怎么考 成了刚需。但很多人因为不了解流程、被虚假信息误导,最终浪费时间又花钱。 千万别踩坑…

2026/11/8 0:00:27

电工证丢失还能补吗?千万别踩坑,丢了证等于失业

电工证丢失还能补吗?千万别踩坑,丢了证等于失业 你正在找工作,但手里的电工证突然丢了,这下怎么办?别急,先别慌, 电工证丢失还能补吗 ?答案是肯定的,但 千万别踩坑 ,补办流程如果不了解,可能会耽误你上岗、甚至影响到你的收入。…

2026/11/8 0:00:27

电工证丢失还能补吗?千万别踩坑,丢了证等于失业

电工证丢失还能补吗?千万别踩坑,丢了证等于失业 你正在找工作,但手里的电工证突然丢了,这下怎么办?别急,先别慌, 电工证丢失还能补吗 ?答案是肯定的,但 千万别踩坑 ,补办流程如果不了解,可能会耽误你上岗、甚至影响到你的收入。…

2026/11/7 14:54:47

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

在智慧养宠时代,基于视觉的宠物身份识别技术正逐步从实验室走向家庭应用。然而,与配合度高的人类不同,宠物是典型的“非合作生物识别对象”——它们不会主动配合摄像头调整姿态,更无法控制光照环境。实际应用中,我们常…

2026/11/7 14:50:12

SSM电子书店源码全解析:结构、原理、部署与排错实战

/* MD / 富文本中的 .toc(含博客园搬家等嵌套结构);.toc-box 在侧栏,不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p(:empty 仍占 20px),一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

还想了解更多?直接咨询顾问

免费诊断 + 免费方案 + 透明报价。

全国咨询热线400-8866-253
免费获取方案
☎咨询二维码 ☎ ↑