已开启

【社区跟随】_dump_snapshot特性补齐 #5310

xuyun15创建于 21 天前

xuyun15成员

21 天前创建

长期任务

PyTorch Feature 设计与实现分析

模块设计目标与背景

为什么需要这个 feature

_dump_snapshot 是 PyTorch 内存调试体系中的核心功能，用于将当前设备内存分配器的完整状态（包括 segment 布局、block 状态、分配/释放历史 trace）序列化保存到文件。该功能主要服务于：

OOM 调试：当发生 OOM 时，自动触发 dump，保留现场快照供事后分析
内存泄漏排查：开发者可主动调用，获取内存分配的完整调用栈
碎片化分析：通过 segment/block 布局分析内存碎片情况
可视化：生成的 pickle 文件可上传至 pytorch.org/memory_viz 交互式查看

不使用它会怎样

没有 _dump_snapshot，OOM 发生时只能看到简单的错误信息，无法回溯哪些 tensor 占用了内存、哪些调用栈分配了内存，排查效率极低。

在整体架构中的位置

_dump_snapshot 位于 torch.cuda.memory / torch.npu.memory 模块的底层，是内存调试工具链（_record_memory_history → _snapshot → _dump_snapshot → 可视化）的输出端。

长期任务

整体设计架构

核心组件说明

组件	职责	文件路径
`_dump_snapshot` (Python)	入口函数，调用 `_snapshot()` 获取快照并 pickle 写入文件	npu/memory.py / cuda/memory.py
`_snapshot` (Python)	获取内存快照 dict 对象，可选 FX trace 增强	npu/memory.py / cuda/memory.py
`_record_memory_history` (Python)	启用/禁用内存历史记录，控制 trace 采集策略	npu/memory.py / cuda/memory.py
`THNPModule_memorySnapshot` (C++)	Python 绑定，调用 allocator 的 snapshot() 并构建 py::dict 返回	torch_npu/csrc/npu/Module.cpp
`_memory_snapshot_pickled` (C++)	C++ 内部路径，构建 IValue dict 后 pickle 序列化为 string	torch_npu/csrc/npu/memory_snapshot.cpp
`_record_memory_history` (C++)	C++ 实现，配置 traceback 采集器和记录策略	torch_npu/csrc/npu/memory_snapshot.cpp
`NPUCachingAllocator::snapshot()` (C++)	核心数据源，遍历所有 segment/block 生成 SnapshotInfo	torch_npu/csrc/core/npu/NPUCachingAllocator.h
`NPUWorkspaceAllocator::snapshot()` (C++)	workspace allocator 的快照数据	torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h
`OptionsManager::IsOomSnapshotEnable()` (C++)	OOM 时自动触发 dump 的环境变量配置	torch_npu/csrc/core/npu/register/OptionsManager.cpp

整体执行流程

flowchart TD
    A[用户调用 _dump_snapshot] --> B[_snapshot 获取快照]
    B --> C[torch_npu._C._npu_memorySnapshot]
    C --> D[THNPModule_memorySnapshot C++]
    D --> E[NPUCachingAllocator::snapshot]
    D --> F[NPUWorkspaceAllocator::snapshot]
    E --> G[构建 segments + device_traces]
    F --> G
    G --> H[py_symbolize 符号化调用栈]
    H --> I[返回 py::dict]
    I --> J[pickle.dump 写入文件]
    J --> K[_npu_saveDevMemUsageInfo NPU特有]
    K --> L{保存成功?}
    L -->|是| M[完成]
    L -->|否| N[启动 profiler 采集 NPU特有]
    N --> M

入口分析

入口1: `torch.npu.memory._dump_snapshot`

文件路径：torch_npu/npu/memory.py:994

入口函数：_dump_snapshot(filename="dump_snapshot.pickle", augment_with_fx_traces=False)

职责：将内存快照 pickle 序列化写入文件，可选 FX trace 增强，并额外保存设备内存使用信息

调用场景：

开发者主动调用进行内存调试
OOM 发生时由 oom_observer 自动调用

入口2: `torch.cuda.memory._dump_snapshot`（原生实现）

文件路径：torch/cuda/memory.py:1156

入口函数：_dump_snapshot(filename="dump_snapshot.pickle", augment_with_fx_traces=False)

职责：将内存快照 pickle 序列化写入文件，可选增强 FX trace 信息

调用场景：同上

接口使用说明

CUDA 原生用法：

import torch
torch.cuda.memory._record_memory_history(enabled="all", context="all", stacks="python")
# ... 执行模型代码 ...
torch.cuda.memory._dump_snapshot("my_snapshot.pickle")
torch.cuda.memory._record_memory_history(enabled=None)  # 关闭记录

NPU 当前用法：

import torch_npu
torch_npu.npu.memory._record_memory_history(enabled="all", context="all", stacks="python")
# ... 执行模型代码 ...
torch_npu.npu.memory._dump_snapshot("my_snapshot.pickle")
torch_npu.npu.memory._record_memory_history(enabled=None)  # 关闭记录

完整调用链分析

阶段1：Python 入口调用

函数：_dump_snapshot(filename="dump_snapshot.pickle", augment_with_fx_traces=False)

文件：torch_npu/npu/memory.py:994

职责：获取快照并写入文件

关键逻辑（torch_npu/npu/memory.py:1006-1022）：

s = _snapshot(augment_with_fx_traces=augment_with_fx_traces)
with open(filename, "wb") as f:
    pickle.dump(s, f)

# NPU 特有：额外保存设备内存使用信息
device = torch_npu.npu.current_device()
save_success = torch_npu._C._npu_saveDevMemUsageInfo(device)
if not save_success:
    # 回退到 profiler 采集
    prof_path = os.path.dirname(os.path.abspath(filename))
    activities = {torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU}
    torch_npu._C._profiler._init_profiler(prof_path, activities)
    prof_config = [prof_path, False, True, False, False, False, torch_npu._C._profiler._ExperimentalConfig()]
    npu_prof_config = torch_npu._C._profiler.NpuProfilerConfig(*tuple(prof_config))
    torch_npu._C._profiler._start_profiler(npu_prof_config, activities)
    torch_npu._C._profiler._stop_profiler()
    torch_npu._C._profiler._finalize_profiler()

对比 CUDA 原生（torch/cuda/memory.py:1172-1175）：

s = _snapshot(augment_with_fx_traces=augment_with_fx_traces)
with open(filename, "wb") as f:
    pickle.dump(s, f)

关键差异：

文件写入方式已对齐：NPU 与 CUDA 均使用 open(filename, "wb")（旧版 NPU 曾用 os.fdopen(os.open(...))，现已统一）
augment_with_fx_traces 参数已对齐：NPU 与 CUDA 均支持该参数
NPU 额外调用 _npu_saveDevMemUsageInfo，CUDA 无此逻辑
NPU 在 saveDevMemUsageInfo 失败时回退到 profiler 采集，CUDA 无此逻辑

↓

阶段2：_snapshot 调用

函数：_snapshot(device=None, augment_with_fx_traces=False)

文件：torch_npu/npu/memory.py:902

职责：调用 C++ 绑定获取完整快照 dict，可选 FX trace 增强

关键逻辑（torch_npu/npu/memory.py:987-991）：

s = torch_npu._C._npu_memorySnapshot(None)
if augment_with_fx_traces:
    from torch._utils import _augment_memory_snapshot_stack_traces
    s = _augment_memory_snapshot_stack_traces(s)
return s

对比 CUDA 原生（torch/cuda/memory.py:1150-1153）：

s = _C._cuda_memorySnapshot(None)
if augment_with_fx_traces:
    s = _augment_memory_snapshot_stack_traces(s)
return s

关键差异：

FX trace 增强已对齐：NPU 与 CUDA 均支持 augment_with_fx_traces，均调用 _augment_memory_snapshot_stack_traces
import 位置差异：NPU 在函数内部 from ... import（延迟导入），CUDA 在文件顶部 import
_npu_memorySnapshot 不接受参数（传 None 占位），CUDA 的 _cuda_memorySnapshot 接受 (mempool_id_first, mempool_id_second, include_traces) 元组
device 参数均未实际使用：两侧声明 device 但均未传给 C++

↓

阶段3：C++ 绑定层 - THNPModule_memorySnapshot

函数：THNPModule_memorySnapshot(PyObject*, PyObject*)

文件：torch_npu/csrc/npu/Module.cpp:1274

职责：调用 allocator snapshot()，构建 Python dict 返回

关键逻辑：

调用 c10_npu::NPUCachingAllocator::snapshot() 获取主分配器快照
调用 c10_npu::NPUWorkspaceAllocator::snapshot() 获取 workspace 分配器快照
合并 workspace segments 和 traces 到主快照
构建 py::dict，包含 segments 和 device_traces
批量符号化 traceback：py_symbolize(to_gather_frames)

对比 CUDA 原生 THCPModule_memorySnapshot：

CUDA 接受参数（METH_O），支持 mempool_id 和 include_traces 过滤
CUDA 返回 4 个顶层 key：segments, device_traces, allocator_settings, external_annotations
NPU 返回 2 个顶层 key：segments, device_traces
NPU 缺少 allocator_settings（分配器配置元数据）
NPU 缺少 external_annotations（外部标注数据）
NPU 的 trace entry 缺少 time_us（时间戳）、compile_context、user_metadata、internal_metadata 字段
NPU 有额外的 WORKSPACE_SNAPSHOT trace action，CUDA 无此概念
NPU 有额外的 pool_id 字段（与 CUDA 一致）
CUDA 有 annotate action，NPU 无此概念

↓

阶段4：C++ _record_memory_history

函数：_record_memory_history(enabled, context, stacks, max_entries)

文件：torch_npu/csrc/npu/memory_snapshot.cpp:37

职责：配置内存历史记录策略

关键逻辑：

c10_npu::NPUCachingAllocator::recordHistory(enabled.has_value(), recorder, max_entries, when);
c10_npu::NPUWorkspaceAllocator::recordHistory(enabled.has_value(), recorder, when);

对比 CUDA 原生：

CUDA 支持更多参数：clearHistory, compileContext, globalRecordAnnotations, skip_actions, record_pinned_host_memory, record_cuda
CUDA 有 setRecordFunctionCallbacks 管理 annotation/compile context 回调
NPU 不支持 clearHistory（清除历史记录）
NPU 不支持 compileContext（编译上下文标注）
NPU 不支持 globalRecordAnnotations（全局标注记录）
NPU 不支持 skip_actions（跳过指定 action 类型的记录）
NPU 不支持 record_pinned_host_memory（记录 host pinned 内存）
NPU 不支持 record_cuda（控制是否记录 device 内存）
NPU 额外调用 NPUWorkspaceAllocator::recordHistory，CUDA 无 workspace 概念

↓

阶段5：Python _record_memory_history 入口

函数：_record_memory_history(enabled="all", *args, **kwargs)

文件：torch_npu/npu/memory.py:851 / torch/cuda/memory.py:898

NPU 实现（torch_npu/npu/memory.py:851-888）：

def _record_memory_history(enabled="all", *args, **kwargs):
    return _record_memory_history_impl(enabled, *args, **kwargs)

CUDA 实现（torch/cuda/memory.py:898-996）：

def _record_memory_history(enabled="all", *args, **kwargs):
    if isinstance(enabled, bool):
        return _record_memory_history_legacy(enabled, *args, **kwargs)
    else:
        return _record_memory_history_impl(enabled, *args, **kwargs)

关键差异：

CUDA 有 bool 入参兼容分支：当 enabled 为 bool 时走 _record_memory_history_legacy（torch/cuda/memory.py:867-895），NPU 无此分支
__signature__ 覆盖：CUDA 在 torch/cuda/memory.py:1027 有 _record_memory_history.__signature__ = signature(_record_memory_history_impl)，NPU 无

NPU _record_memory_history_impl（torch_npu/npu/memory.py:891-899）：

def _record_memory_history_impl(
    enabled: Optional[str] = "all",
    context: Optional[str] = "all",
    stacks: str = "all",
    max_entries: int = sys.maxsize,
    device=None,
):
    torch_npu.npu._lazy_init()
    torch_npu._C._npu_record_memory_history(enabled, context, stacks, max_entries)

CUDA _record_memory_history_impl（torch/cuda/memory.py:999-1024）：

def _record_memory_history_impl(
    enabled: str | None = "all",
    context: str | None = "all",
    stacks: str = "all",
    max_entries: int = sys.maxsize,
    device: "Device" = None,
    clear_history: bool = False,
    compile_context: bool = False,
    global_record_annotations: bool = False,
    skip_actions: list[str] | None = None,
    record_pinned_host_memory: bool = False,
    record_cuda: bool = True,
):
    _C._cuda_record_memory_history(
        enabled, context, stacks, max_entries,
        clear_history, compile_context, global_record_annotations,
        skip_actions if skip_actions is not None else [],
        record_pinned_host_memory, record_cuda,
    )

关键差异：

NPU 传 4 个参数给 C++，CUDA 传 10 个
NPU 缺 clear_history、compile_context、global_record_annotations、skip_actions、record_pinned_host_memory、record_cuda
NPU 多了 torch_npu.npu._lazy_init() 调用，CUDA 无

↓

阶段6：OOM 自动 dump

函数：oom_observer(device, allocated, device_total, device_free)

文件：torch_npu/csrc/core/npu/register/OptionsManager.cpp

职责：OOM 发生时自动调用 Python _dump_snapshot

关键逻辑：

auto module = THPObjectPtr(PyImport_ImportModule("torch_npu.npu.memory"));
PyObject* p_func = PyObject_GetAttrString(module, "_dump_snapshot");
PyObject_CallObject(p_func, p_args);

配置入口：环境变量 OOM_SNAPSHOT_ENABLE

0：不启用
2：启用 state 级别记录
其他：启用 all 级别记录

完整调用链总结

sequenceDiagram
    participant User as 用户代码
    participant PyMem as npu/memory.py
    participant CBind as Module.cpp (C++绑定)
    participant Alloc as NPUCachingAllocator
    participant WS as NPUWorkspaceAllocator
    participant Sym as py_symbolize

    User->>PyMem: _dump_snapshot(filename)
    PyMem->>PyMem: _snapshot(augment_with_fx_traces)
    PyMem->>CBind: _npu_memorySnapshot()
    CBind->>Alloc: snapshot()
    Alloc-->>CBind: SnapshotInfo (segments, traces)
    CBind->>WS: snapshot()
    WS-->>CBind: SnapshotInfo (workspace segments, traces)
    CBind->>CBind: 合并 workspace 数据
    CBind->>Sym: py_symbolize(tracebacks)
    Sym-->>CBind: 符号化后的 frames
    CBind-->>PyMem: py::dict {segments, device_traces}
    PyMem->>PyMem: pickle.dump 写入文件
    PyMem->>Alloc: _npu_saveDevMemUsageInfo(device)
    alt 保存失败
        PyMem->>PyMem: 启动 profiler 采集
    end

扩展点分析

NPU 与 CUDA 功能差异全景

注：本表反映当前最新代码状态（NPU torch_npu/npu/memory.py / CUDA torch/cuda/memory.py）。

维度	CUDA 原生	NPU 当前	差异影响
`_dump_snapshot` 参数	`filename`, `augment_with_fx_traces`	`filename`, `augment_with_fx_traces`	已对齐
`_dump_snapshot` 文件写入	`open(filename, "wb")`	`open(filename, "wb")`	已对齐（旧版 NPU 曾用 `os.fdopen(os.open(...))`）
`_dump_snapshot` 兜底逻辑	无	`_npu_saveDevMemUsageInfo` + profiler 回退	NPU 额外采集设备内存信息
`_snapshot` 参数	`device`, `augment_with_fx_traces`	`device`, `augment_with_fx_traces`	已对齐（`device` 两侧均未实际使用）
`_snapshot` FX 增强	调用 `_augment_memory_snapshot_stack_traces`	调用 `_augment_memory_snapshot_stack_traces`	已对齐（NPU 延迟 import，CUDA 顶部 import）
`memory_snapshot` 参数	`mempool_id`, `include_traces`	`mempool_id`	NPU 缺少 `include_traces`（无法轻量快照）
`_npu_memorySnapshot` C++ 参数	`mempool_id`, `include_traces`	无参数（传 `None` 占位）	无法按 pool 过滤、无法轻量快照
返回值顶层 key	`segments`, `device_traces`, `allocator_settings`, `external_annotations`	`segments`, `device_traces`	缺少配置元数据和标注
trace entry 字段	`action`, `addr`/`device_free`, `size`, `stream`, `time_us`, `pool_id`, `frames`, `compile_context`, `user_metadata`, `internal_metadata`	`action`, `addr`/`device_free`, `size`, `stream`, `pool_id`, `frames`	缺少时间戳、编译上下文、用户元数据、内部元数据
trace action 类型	`alloc`, `free_requested`, `free_completed`, `segment_alloc`, `segment_free`, `oom`, `snapshot`, `annotate`	`alloc`, `free_requested`, `free_completed`, `segment_alloc`, `segment_free`, `oom`, `snapshot`, `WORKSPACE_SNAPSHOT`	NPU 缺 `annotate`，多 `WORKSPACE_SNAPSHOT`
`_record_memory_history` 参数	`enabled`, `context`, `stacks`, `max_entries`, `clear_history`, `compile_context`, `global_record_annotations`, `skip_actions`, `record_pinned_host_memory`, `record_cuda`	`enabled`, `context`, `stacks`, `max_entries`	缺少 6 个控制参数
`_record_memory_history` bool 兼容	有（`_record_memory_history_legacy`）	无	缺少向后兼容
`_record_memory_history.__signature__`	有	无	IDE 显示签名不完整
`_allocation_traceback`	支持（`torch/cuda/memory.py:1030`）	不支持	缺少按 data_ptr 查询分配栈
`_set_memory_metadata` / `_get_memory_metadata`	支持（`torch/cuda/memory.py:1193,1231`）	不支持	缺少自定义元数据标注
`_memory_metadata_supported`	支持（`torch/cuda/memory.py:1178`）	不支持	缺少元数据能力探测
`_record_memory_history_legacy`	支持（`torch/cuda/memory.py:867`）	不支持	缺少 bool 入参兼容路径
`caching_allocator_enable` / `caching_allocator_disabled`	支持（`torch/cuda/memory.py:157,162`）	不支持	缺少分配器启停控制
`list_gpu_processes`	支持（`torch/cuda/memory.py:782`）	不支持	缺少进程级内存查看
`mem_get_info`	支持（`torch/cuda/memory.py:848`）	不支持	缺少全局空闲/总量查询
NPU 特有逻辑	无	`_npu_saveDevMemUsageInfo` + profiler 回退	NPU 额外采集设备内存信息
NPU 特有 action	无	`WORKSPACE_SNAPSHOT`	workspace 分配器特有

可扩展点

扩展目标	修改位置	推荐方式
增加 `memory_snapshot` 的 `include_traces` 参数	`torch_npu/npu/memory.py::memory_snapshot`	传递 `include_traces` 给 C++，需 C++ 绑定支持
增加 `mempool_id` / `include_traces` 参数（C++ 绑定）	`torch_npu/csrc/npu/Module.cpp::THNPModule_memorySnapshot`	修改 C++ 绑定签名从 METH_NOARGS 到 METH_O，解析参数传给 allocator
增加 `allocator_settings` 返回字段	`torch_npu/csrc/npu/Module.cpp::THNPModule_memorySnapshot`	参照 CUDA 实现构建配置元数据 dict
增加 `external_annotations` 返回字段	`torch_npu/csrc/npu/Module.cpp::THNPModule_memorySnapshot`	需先在 allocator 侧实现 annotation buffer
增加 trace entry 的 `time_us` 字段	`torch_npu/csrc/npu/Module.cpp::THNPModule_memorySnapshot`, `NPUCachingAllocator.h/cpp`	需 allocator 侧在 TraceEntry 中记录时间戳（详见 `npu_memory_snapshot_timestamp_gap_analysis.md`）
增加 trace entry 的 `compile_context` 字段	`Module.cpp`, allocator 侧	需 TraceEntry 加字段 + record_trace 采集
增加 trace entry 的 `user_metadata` / `internal_metadata` 字段	`Module.cpp`, allocator 侧	需 TraceEntry 加字段 + 元数据回写机制
增加 `annotate` action	`NPUCachingAllocator.h` TraceEntry::Action 枚举	参照 CUDA 实现 annotation 回调
增加 `_record_memory_history` 的 `clear_history` 参数	`memory_snapshot.cpp::_record_memory_history`, Python 层	传递给 `NPUCachingAllocator::recordHistory`
增加 `_record_memory_history` 的 `skip_actions` 参数	`memory_snapshot.cpp::_record_memory_history`, Python 层	传递给 `NPUCachingAllocator::recordHistory`
增加 `_record_memory_history` 的 `compile_context` / `global_record_annotations`	`memory_snapshot.cpp`, 需实现 `setRecordFunctionCallbacks`	参照 CUDA 实现 annotation/compile context 回调机制
增加 `_record_memory_history` 的 `record_pinned_host_memory` / `record_cuda`	`memory_snapshot.cpp`, allocator 侧	参照 CUDA 实现 host 内存记录
增加 `_record_memory_history_legacy` 兼容路径	`torch_npu/npu/memory.py`	参照 CUDA 实现 bool 入参分支
增加 `_record_memory_history.__signature__` 覆盖	`torch_npu/npu/memory.py`	`from inspect import signature` 后覆盖
增加 `_allocation_traceback` 函数	`torch_npu/npu/memory.py`, C++ 绑定	绑定 `NPUCachingAllocator::allocationTraceback`
增加 `_set_memory_metadata` / `_get_memory_metadata` / `_memory_metadata_supported`	`torch_npu/npu/memory.py`, C++ 绑定	绑定 allocator 的元数据接口
增加 `caching_allocator_enable` / `caching_allocator_disabled`	`torch_npu/npu/memory.py`, C++ 绑定	绑定 allocator 启停接口
增加 `list_gpu_processes`	`torch_npu/npu/memory.py`, C++ 绑定	绑定进程级内存查看接口
增加 `mem_get_info`	`torch_npu/npu/memory.py`, C++ 绑定	绑定全局空闲/总量查询接口

修改已有行为

改 _dump_snapshot 使其与 CUDA 一致：

Python 层 (torch_npu/npu/memory.py)：
- _dump_snapshot：augment_with_fx_traces 参数与 FX 增强已对齐，无需改动
- _snapshot：FX 增强已对齐，无需改动
- _npu_saveDevMemUsageInfo 和 profiler 回退逻辑：NPU 独有扩展，建议保留（或作为可选扩展）
- 文件写入方式：已对齐为 open(filename, "wb")
C++ 绑定层 (torch_npu/csrc/npu/Module.cpp)：
- THNPModule_memorySnapshot 从 METH_NOARGS 改为 METH_O
- 解析 (mempool_id_0, mempool_id_1, include_traces) 参数元组
- 传递给 NPUCachingAllocator::snapshot(mempool_id, include_traces)
- 增加 allocator_settings 和 external_annotations 构建逻辑
C++ 实现层 (torch_npu/csrc/npu/memory_snapshot.cpp)：
- _record_memory_history 增加 clearHistory, compileContext, globalRecordAnnotations, skip_actions, record_pinned_host_memory, record_cuda 参数
- 实现 setRecordFunctionCallbacks 机制

哪些逻辑不能乱改：

NPUCachingAllocator::snapshot() 的数据结构 SnapshotInfo 是核心数据源，修改需谨慎
py_symbolize 的符号化逻辑与 traceback 采集紧密耦合，不可随意修改
OOM observer 的 GIL 获取逻辑 (pybind11::gil_scoped_acquire) 必须保留
_npu_saveDevMemUsageInfo + profiler 回退是 NPU 独有扩展，与 time_us 补齐无耦合，建议保留

常见扩展模式

新增 trace action 类型：在 TraceEntry::Action 枚举中增加，然后在 action_to_str lambda 中增加映射
新增 trace entry 字段：在 THNPModule_memorySnapshot 构建 trace_entry dict 时增加 key-value，同时需在 TraceEntry 结构体中加字段
新增顶层返回字段：在 result dict 中增加 key，如 allocator_settings
新增 Python API：在 memory.py 中定义函数，在 __all__ 中导出，必要时增加 C++ 绑定

容易踩坑的地方

C++ 绑定签名不一致：NPU 的 _npu_memorySnapshot 是 METH_NOARGS，CUDA 是 METH_O。修改时需同步更新 PyMethodDef 表和函数签名，否则会 segfault
workspace allocator 数据合并顺序：workspace traces 被插入到主 traces 的开头（begin()），修改合并逻辑需注意顺序对可视化工具的影响
GIL 与 OOM observer：OOM observer 在 C++ 侧获取 GIL 后调用 Python _dump_snapshot，如果 _dump_snapshot 内部逻辑变更导致需要额外 GIL 操作，可能死锁
_npu_saveDevMemUsageInfo 的 profiler 回退：当前回退逻辑硬编码了 profiler 配置参数（torch_npu/npu/memory.py:1014-1022），修改需确保与 profiler 模块兼容
pickle 兼容性：_dump_snapshot 生成的 pickle 文件需与 pytorch.org/memory_viz 可视化工具兼容，新增字段不应破坏已有字段的格式
_memory_snapshot_pickled 与 THNPModule_memorySnapshot 双实现：两者逻辑需保持同步，修改一处需同步修改另一处（注意前者当前是死代码，未绑定到 Python）
aarch64 vs x86_64 分支：gather(), gather_with_cpp(), py_symbolize 在不同平台调用不同实现，修改需覆盖两个平台
time_us 补齐的 aarch64 TSC 风险：若 C10_ARMTSC 未启用，record_trace 会回退到 clock_gettime syscall，热路径性能下降（详见 npu_memory_snapshot_timestamp_gap_analysis.md）

总结

feature 核心工作机制

_dump_snapshot 通过调用 C++ 绑定 _npu_memorySnapshot 获取内存分配器的完整状态快照（segments 布局 + 分配/释放 trace 历史），将快照 dict pickle 序列化写入文件。NPU 侧额外调用 _npu_saveDevMemUsageInfo 保存设备级内存使用信息，失败时回退到 npu profiler 采集。

核心调用路径

_dump_snapshot → _snapshot(augment_with_fx_traces) → _npu_memorySnapshot (C++) → NPUCachingAllocator::snapshot() + NPUWorkspaceAllocator::snapshot() → py_symbolize → pickle.dump → _npu_saveDevMemUsageInfo (NPU特有)

核心数据结构

SnapshotInfo：包含 segments (vector) 和 device_traces (vector<vector>)
SegmentInfo：segment 地址、大小、pool id、expandable 标记、block 列表
BlockInfo：block 地址、大小、requested_size、状态（active_allocated / active_pending_free / inactive）
TraceEntry：action 类型、地址、大小、stream、调用栈上下文、pool_id（缺 time_us / compile_context / user_metadata / internal_metadata）

当前对齐状态（已对齐项）

_dump_snapshot 参数（filename, augment_with_fx_traces）
_dump_snapshot 文件写入方式（open(filename, "wb")）
_snapshot 参数与 FX trace 增强（_augment_memory_snapshot_stack_traces）
memory_snapshot 的 mempool_id 参数（部分对齐，NPU 缺 include_traces）
trace entry 的 pool_id 字段

当前主要差异（未对齐项）

trace entry 缺 time_us：NPU 的 TraceEntry 无时间字段，无法做时间轴可视化（详见 npu_memory_snapshot_timestamp_gap_analysis.md）
trace entry 缺 compile_context / user_metadata / internal_metadata：NPU 无元数据标注能力
trace action 缺 annotate：NPU 无 annotation 事件类型
返回值缺 allocator_settings / external_annotations：NPU 缺配置元数据和标注
_record_memory_history 缺 6 个参数：clear_history / compile_context / global_record_annotations / skip_actions / record_pinned_host_memory / record_cuda
缺 _record_memory_history_legacy 与 __signature__：无 bool 入参兼容，IDE 签名不完整
缺 _allocation_traceback：无法按 data_ptr 查询分配栈
缺 _set_memory_metadata / _get_memory_metadata / _memory_metadata_supported：无自定义元数据标注
缺 caching_allocator_enable / caching_allocator_disabled：无分配器启停控制
缺 list_gpu_processes / mem_get_info：无进程级内存查看、无全局空闲/总量查询
memory_snapshot 缺 include_traces：无法轻量快照
C++ 绑定缺 mempool_id / include_traces 参数：无法按 pool 过滤

如果要扩展，先看哪些文件

增加参数/字段：先看 torch/cuda/memory.py 对应实现，再改 torch_npu/npu/memory.py 和 torch_npu/csrc/npu/Module.cpp
增加 trace 能力：先看 torch/csrc/cuda/memory_snapshot.cpp，再改 torch_npu/csrc/npu/memory_snapshot.cpp
修改 OOM 行为：看 torch_npu/csrc/core/npu/register/OptionsManager.cpp
补齐 time_us：看 npu_memory_snapshot_timestamp_gap_analysis.md，涉及 4 层联动修改

建议下一步阅读路径

先通读 CUDA 原生 torch/cuda/memory.py 的 _dump_snapshot / _snapshot / _record_memory_history / _record_memory_history_legacy / _allocation_traceback / _set_memory_metadata 完整实现
对比阅读 NPU 侧 torch_npu/npu/memory.py 同名函数，逐项标记缺失功能
深入 C++ 层，对比 THCPModule_memorySnapshot vs THNPModule_memorySnapshot 的返回值结构差异
理解 _record_memory_history 的参数差异，评估哪些参数对 NPU 场景有实际价值
阅读 npu_memory_snapshot_timestamp_gap_analysis.md，理解 time_us 缺失的根因与补齐方案
制定分阶段对齐计划：优先对齐影响可视化兼容性的字段（time_us, allocator_settings），其次对齐参数接口