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

发布时间:2026/9/10 21:59:31

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/9/10 21:59:31

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

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

2026/9/10 21:59:31

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

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

2026/9/10 21:59:31

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

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

2026/9/10 22:54:37

基于 GB/T 3836 标准的油气站防爆对讲机选型与通信安全方案

文章定位:技术方案类,面向油气站安全管理人员、通信运维人员,梳理油气站防爆通信的标准要求、设备选型要点、组网方案、施工验收与运维保障。 核心内容:防爆标准要求 选型技术参数 组网方案 施工验收 运维培训 应用案例 4 个…

2026/9/10 22:54:37

CANN/ge图引擎节点构建API

Build 【免费下载链接】ge GE(Graph Engine)是面向昇腾的图编译器和执行器,提供了计算图优化、多流并行、内存复用和模型下沉等技术手段,加速模型执行效率,减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友…

2026/9/10 22:54:37

如何10分钟跑通CVAT标注工具:从部署到标注的完整指南

如何10分钟跑通CVAT标注工具:从部署到标注的完整指南 【免费下载链接】cvat Computer Vision Annotation Tool (CVAT) is a leading platform for building high-quality visual datasets for vision AI. It offers open-source, cloud, and enterprise products, a…

2026/9/10 16:39:38

超人会飞不算本事:系统稳定依赖清晰规则与边界设计

开头先不绕弯子。“#斯坦李吐槽dc 所以超人是无缘无故会飞的嘛哈哈哈哈哈哈哈锤哥真是技术人才啊!#雷神 #复联”这类调侃式短标题,第一波冲击力在于它把两个宇宙的角色塞进同一个吐槽箱里,但细想一下就能发现,它真正碰到的根本不是…

2026/9/10 11:16:38

超人VS蜘蛛侠:拆解超级IP的影响力与传播方法论

把“蜘蛛侠 vs 超人”放在 CSDN 上聊,可能很多人第一反应是走错片场了。但如果把这两个角色看成“两个持续运营了 80 多年的文化产品”,你会发现,这场比较本质上是两个不同 IP 策略的长期结果对比:超人赢在定义了整个超级英雄题材…

2026/9/9 16:31:09

基于CNN的调制信号识别:MATLAB实现时频图分类实战

简介:本资源是一套面向通信工程与信号处理方向学习者、研究者的深度学习实践方案,聚焦调制信号自动检测与识别这一典型无线通信任务,解决传统方法依赖人工特征、低信噪比下性能下降等痛点。压缩包共12个文件(10.73MB)&…

2026/9/10 0:00:55

目录对比去重实战:用哈希算法精准清理重复文件

我电脑里现在还有一块换了三次机的“数据墓地”硬盘,里面存着2016年以前所有旧笔记本的完整备份。平时不觉得有什么,直到前阵子想把它整理归档,发现同一个安装包、同一批照片、同一份论文草稿,在几个不同的备份目录里反复出现。更…

2026/9/10 0:00:55

Leaflet离线地图完整Demo合集:内网部署与坐标纠偏实战

简介:这是一份面向Web GIS开发者的LeafLet离线地图示例合集,帮助开发者快速掌握离线地图从搭建到交互的完整流程。压缩包共723个文件,大小14.06MB,以319个js脚本、175个html页面和29个css样式文件为主体,配合png/svg图…

2026/9/10 0:00:55

MATLAB读取Rinex 3.02观测文件:多系统GNSS数据解析实战

简介:基于MATLAB开发的Rinex3.02版观测文件(o文件)读取代码包,面向卫星定位导航方向的学习者与研究人员,用于解决新版观测文件的数据解析、历元提取与时间转换问题。压缩包共4个文件,包含两个m脚本、一个19…

2026/9/10 12:32:02

USB Type-C PCB布局分区设计:电源、高速信号与PD协议全攻略

做硬件这行,Type-C接口算是典型的“看着简单,做起来全坑”的东西。光引脚就24个,高低速信号、电源、控制线全部塞在一个小小的连接器里,如果PCB布局不做规划,打样回来基本就是“插上没反应”、“高速掉线”、“静电一打…

2026/9/10 15:19:50

系统编程学习原型如何补齐稳定性边界

系统编程学习原型如何补齐稳定性边界预算有限时&#xff0c;我先优化明显多余的复制&#xff0c;而不是猜测性地换容器。用借用传递只读数据通常就能减少分配&#xff1a; fn parse(line: &str) -> Result<Item, Error> { /* ... */ }用基准确认热点确实在分配&am…

2026/9/10 15:49:53

雨花区哪家财务公司代理记账比较好?

在雨花区&#xff0c;企业处理财税事务常常面临诸多挑战&#xff0c;选择一家靠谱的财务公司至关重要。湖南巨勤财务管理咨询有限公司就是本地正规实体财税服务机构&#xff0c;深耕本地工商财税行业多年&#xff0c;熟悉当地工商局、税务局最新政策与申报流程。主营公司注册、…

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

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

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