已开启

【社区跟随】memory snapshot数据中增加事件时间信息的采集 #5311

xuyun15创建于 19 天前

19 天前创建

https://gitcode.com/Ascend/pytorch/issues/5310 的子任务

NPU Memory Snapshot 时间戳缺失根因分析与补齐方案

1. 概述

本文档针对 ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp 中 memory snapshot 特性的"只有时序、缺少内存事件绝对时间信息"问题，给出根因分析、工作量评估以及补齐方案所需的代码修改片段。

所有路径均使用相对路径（相对仓库根目录），不使用别名与 markdown 链接语法。所有结论均基于代码证据，并提供上下文调用栈。

2. 背景与现象

2.1 现象

NPU 的 memory snapshot 输出 pickle 中，trace 事件（ALLOC / FREE_REQUESTED / SEGMENT_ALLOC 等）只能通过顺序索引推断时序关系，缺少事件发生的绝对时间戳字段 time_us，无法对接 pytorch.org/memory_viz 的时间轴可视化，也无法与 profiler、trace 等其他带时间戳的事件流做时间对齐。

2.2 涉及的接口

Python 用户侧两个接口：

ascend_pytorch/pytorch/torch_npu/npu/memory.py:851 def _record_memory_history(enabled="all", *args, **kwargs)
ascend_pytorch/pytorch/torch_npu/npu/memory.py:994 def _dump_snapshot(filename="dump_snapshot.pickle", augment_with_fx_traces=False)

底层 C++ 入口：

ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:1274 THNPModule_memorySnapshot（CPython 风格，返回 py::dict）
ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:173 _npu_record_memory_history（pybind m.def）
ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp:37 torch_npu::_record_memory_history
ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp:145 _memory_snapshot_pickled（IValue pickle 路径，当前未绑定到 Python，为死代码）

3. 根因分析

3.1 现象确认：NPU TraceEntry 完全没有时间字段

NPU 的 TraceEntry 结构体定义于 ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h:117-159：

struct TraceEntry {
  enum Action {
    ALLOC, FREE_REQUESTED, FREE_COMPLETED,
    SEGMENT_ALLOC, SEGMENT_FREE, SEGMENT_MAP, SEGMENT_UNMAP,
    SNAPSHOT, OOM, WORKSPACE_SNAPSHOT
  };
  TraceEntry(
      Action action, int device, int64_t addr, size_t size,
      aclrtStream stream, MempoolId_t mempool = {0, 0},
      std::shared_ptr<c10::GatheredContext> context = nullptr)
      : action_(action), device_(device), addr_(addr),
        context_(std::move(context)), stream_(stream),
        size_(size), mempool_(std::move(mempool)) {}
  Action action_;
  int device_;
  int64_t addr_;
  std::shared_ptr<c10::GatheredContext> context_;
  aclrtStream stream_;
  int64_t size_;
  MempoolId_t mempool_;
};

成员字段：action_、device_、addr_、context_、stream_、size_、mempool_，没有任何 time_ 字段。

对照 CUDA 基类 pytorch/c10/core/CachingDeviceAllocator.h:123-186：

union trace_time_ {
  time_t t_;            // 已转换的 unix 时间（微秒 since epoch）
  approx_time_t approx_t_;  // 原始 TSC tick
};

struct TraceEntry {
  enum Action { ALLOC, FREE_REQUESTED, ..., SNAPSHOT, OOM, ANNOTATE };
  TraceEntry(
      Action action, c10::DeviceIndex device, size_t addr, size_t size,
      void* stream, MempoolId_t mempool,
      approx_time_t time,                            // <-- 关键
      std::shared_ptr<GatheredContext> context = nullptr,
      std::string compile_context = "",
      std::string user_metadata = "")
      : ... {
    time_.approx_t_ = time;
  }
  ...
  trace_time_ time_{};                                // <-- 关键
  std::string compile_context_;
  std::string user_metadata_;
  std::string internal_metadata_;
};

CUDA 的 TraceEntry 含 trace_time_ time_{} 字段，构造时写入 TSC tick，snapshot 时覆盖为微秒。

3.2 根因：本身缺失该特性（未与 CUDA 原生 snapshot 对齐）

NPU 的 TraceEntry 是自定义版本，并非继承自 c10::CachingDeviceAllocator::TraceEntry。原因是类型适配的"副作用"：

NPU 用 aclrtStream（ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h:142,156）
CUDA 用 void*（pytorch/c10/core/CachingDeviceAllocator.h:152,176）

为适配 aclrtStream 类型，NPU 重写了整个 TraceEntry，但同步丢失了以下 CUDA 基类的能力：

字段 / 能力	CUDA 基类位置	NPU 状态
`trace_time_ time_`	`pytorch/c10/core/CachingDeviceAllocator.h:179`	缺失
`compile_context_`	`pytorch/c10/core/CachingDeviceAllocator.h:180`	缺失
`user_metadata_`	`pytorch/c10/core/CachingDeviceAllocator.h:181`	缺失
`internal_metadata_`	`pytorch/c10/core/CachingDeviceAllocator.h:185`	缺失
`ANNOTATE` action	`pytorch/c10/core/CachingDeviceAllocator.h:144`	缺失（被 `WORKSPACE_SNAPSHOT` 替代）

3.3 证据链：四段链路全部缺失

链路 1：record_trace 不取时间

ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:3502-3537：

void record_trace(
    TraceEntry::Action action, int64_t addr, size_t size,
    aclrtStream stream, int device,
    MempoolId_t mempool_id = {0, 0},
    std::shared_ptr<c10::GatheredContext> context = nullptr) {
  if (!record_history && trace_trackers_.empty()) {
    return;
  }
  auto te = TraceEntry(
      action, device, addr, size, stream, mempool_id,
      record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);
  // ↑ 只传 7 个参数，未调用 getApproximateTime()
  ...
}

对照 CUDA pytorch/c10/cuda/CUDACachingAllocator.cpp:4626-4637：

TraceEntry te(
    action, device, addr, size,
    reinterpret_cast<void*>(stream), mempool_id,
    getApproximateTime(),   // <-- 关键差异
    record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr,
    compile_string,
    metadata_override ? std::move(*metadata_override) : user_metadata);

链路 2：trace() 不做 TSC→μs 转换

ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:2167-2175：

std::vector<TraceEntry> trace() {
  std::lock_guard<std::recursive_mutex> lock(mutex);
  std::vector<TraceEntry> result;
  result.reserve(alloc_trace->size());
  result.insert(result.end(), alloc_trace->begin() + alloc_trace_next, alloc_trace->end());
  result.insert(result.end(), alloc_trace->begin(), alloc_trace->begin() + alloc_trace_next);
  return result;
}

对照 CUDA pytorch/c10/cuda/CUDACachingAllocator.cpp:3175-3186：

std::vector<TraceEntry> trace(
    const std::function<time_t(approx_time_t)>& tsc_to_us) const {
  std::lock_guard<std::recursive_mutex> lock(mutex);
  std::vector<TraceEntry> result;
  alloc_buffer.getEntries(result);
  for (auto& te : result) {
    te.time_.t_ = tsc_to_us(te.time_.approx_t_);   // <-- 关键：就地转换
  }
  return result;
}

链路 3：snapshot wrapper 缺少时钟转换器

ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:3926-3935：

SnapshotInfo snapshot(MempoolId_t mempool_id) override {
  SnapshotInfo result;
  int count = static_cast<int>(device_allocator.size());
  for (int i = 0; i < count; i++) {
    result.device_traces.emplace_back(device_allocator[i]->trace());
    auto snap = device_allocator[i]->snapshot(mempool_id);
    result.segments.insert(result.segments.end(), snap.begin(), snap.end());
  }
  return result;
}

对照 CUDA pytorch/c10/cuda/CUDACachingAllocator.cpp:5049-5077：

SnapshotInfo snapshot(MempoolId_t mempool_id, bool include_traces) override {
  SnapshotInfo result;
  if (include_traces) {
    auto tsc_to_ns = clock_converter.makeConverter();        // <-- 关键
    auto tsc_to_us = [=](approx_time_t t_approx) {
      return tsc_to_ns(t_approx) / 1000;
    };
    for (auto& da : device_allocator) {
      result.device_traces.emplace_back(da->trace(tsc_to_us));
      auto snap = da->snapshot(mempool_id);
      result.segments.insert(result.segments.end(), snap.begin(), snap.end());
    }
  } else { ... }
  return result;
}

CUDA 的 clock_converter 是 NativeCachingAllocator 的成员（pytorch/c10/cuda/CUDACachingAllocator.cpp:4736）：

c10::ApproximateClockToUnixTimeConverter clock_converter;

NPU 侧完全没有 ApproximateClock 相关代码，文件头 include 列表（ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:1-38）也没有 <c10/util/ApproximateClock.h>。

链路 4：序列化层缺 time_us 键

Python dict 路径（实际生效路径）ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:1428-1446：

for (const auto& traceInfo : snapshot.device_traces) {
  py::list trace;
  for (const auto& te : traceInfo) {
    py::dict trace_entry;
    ...
    trace_entry[action_s] = action_to_str(te.action_);
    trace_entry[te.action_ == TraceEntry::OOM ? device_free_s : addr_s] = te.addr_;
    trace_entry[size_s] = te.size_;
    trace_entry[stream_s] = int64_t(te.stream_);
    trace_entry[pool_id_s] = te.mempool_;
    // 缺少 time_us 键
    trace.append(trace_entry);
  }
  traces.append(trace);
}

IValue pickle 路径（死代码）ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp:268-285：

for (const auto& traceInfo : snapshot.device_traces) {
  auto trace = new_list();
  for (const auto& te : traceInfo) {
    auto trace_entry = new_dict();
    trace_entry.insert(action_s, action_to_str(te.action_));
    trace_entry.insert(te.action_ == TraceEntry::OOM ? device_free_s : addr_s, te.addr_);
    trace_entry.insert(size_s, (int64_t)te.size_);
    trace_entry.insert(stream_s, int64_t(te.stream_));
    trace_entry.insert(pool_id_s, std::tuple<int64_t, int64_t>(te.mempool_));
    // 缺少 time_us 键
    ...
  }
}

对照 CUDA pytorch/torch/csrc/cuda/Module.cpp:900-920：

auto traceEntryToDict = [&](const TraceEntry& te) {
  py::dict trace_entry;
  ...
  trace_entry[action_s] = action_to_str(te.action_);
  trace_entry[TraceEntry::OOM == te.action_ ? device_free_s : addr_s] = te.addr_;
  trace_entry[size_s] = te.size_;
  trace_entry[stream_s] = int64_t(te.stream_);
  trace_entry[time_us_s] = te.time_.t_;           // <-- 关键
  trace_entry[compile_context_s] = te.compile_context_;
  trace_entry[user_metadata_s] = te.user_metadata_;
  ...
  trace_entry[pool_id_s] = te.mempool_;
  return trace_entry;
};

3.4 结论

根因为本身缺失该特性——NPU fork 时为适配 aclrtStream 自定义了 TraceEntry，但未同步移植 CUDA 原生的时间戳机制（trace_time_ union + getApproximateTime() + ApproximateClockToUnixTimeConverter 三件套），导致 snapshot 只有事件序列、没有事件绝对时间。

3.5 完整调用栈

NPU snapshot 调用链：

[Python 用户层]
torch.npu.memory._dump_snapshot(filename)                          # ascend_pytorch/pytorch/torch_npu/npu/memory.py:994
  └─ torch.npu.memory._snapshot()                                  # ascend_pytorch/pytorch/torch_npu/npu/memory.py:902
       └─ torch_npu._C._npu_memorySnapshot(None)                   # ascend_pytorch/pytorch/torch_npu/npu/memory.py:987
[C++ pybind 入口]
       └─ THNPModule_memorySnapshot                                # ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:1274
            └─ c10_npu::NPUCachingAllocator::snapshot()            # 多设备 wrapper
                 ├─ DeviceAlloc::trace()                           # ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:2167
                 └─ DeviceAlloc::snapshot()                        # ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:2096
                      └─ DeviceAlloc::record_trace(SNAPSHOT, ...)  # ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:3502
[Python 落盘]
pickle.dump(s, f)                                                  # ascend_pytorch/pytorch/torch_npu/npu/memory.py:1008

对照 CUDA 完整链路：

torch.cuda.memory._dump_snapshot(filename)                         # pytorch/torch/cuda/memory.py:1156
  └─ torch.cuda.memory._snapshot()                                 # pytorch/torch/cuda/memory.py:1100
       └─ torch._C._cuda_memorySnapshot(None)
            └─ THCPModule_memorySnapshot                           # pytorch/torch/csrc/cuda/Module.cpp:732
                 └─ c10::cuda::CUDACachingAllocator::snapshot(mempool, include_traces=true)
                      ├─ clock_converter.makeConverter()           # pytorch/c10/cuda/CUDACachingAllocator.cpp:5054
                      ├─ annotation_buffer 转换 time_
                      └─ DeviceAlloc::trace(tsc_to_us)             # pytorch/c10/cuda/CUDACachingAllocator.cpp:3175

NPU 链路里没有 clock_converter 这一步，这是缺失的核心环节。

4. 工作量评估

4.1 评估表（单位：人天）

略

4.2 关键风险与缓冲

aarch64 TSC 路径：若 C10_ARMTSC 未启用，record_trace 会回退到 clock_gettime syscall，性能下降明显。需在第 1 阶段尽早验证，若需补 TSC 读取汇编，工作量 +1.5 人天。
pickle schema 破坏性变更：若下游 memory_viz 工具对 NPU snapshot 有硬编码字段集，加 time_us 可能导致解析异常，需联调验证，最坏 +1 人天。
_dump_snapshot 兜底逻辑：NPU 独有的 _npu_saveDevMemUsageInfo + profiler 回退路径与 time_us 补齐无直接耦合，但需回归测试保护，已计入回归用例。

4.3 总体结论

略

5. 代码修改方案

补齐需打通"取 TSC → 存 TSC → snapshot 时转换为 μs → 序列化 time_us"四段链路。修改涉及 4 层（C++ 结构体 / C++ 算法 / C++ 序列化 / Python 接口）。本节仅展示修改片段，不修改文件。

5.1 修改点 1：C++ 结构体层

文件：ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h

1.1 文件头部 include 区（约第 1-10 行附近，需新增 include）

修改前：无 ApproximateClock 头文件

修改后（新增一行）：

#include <c10/util/ApproximateClock.h>

1.2 TraceEntry 结构体（第 116-158 行）

修改前（ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h:116-158）：

struct TraceEntry {
  enum Action {
    ALLOC, // API made to the caching allocator for new memory
    FREE_REQUESTED, // API call made to the caching allocator to free memory
    FREE_COMPLETED, // The allocator might have to delay a free because
                    // it is still in use on another stream via
                    // record_stream This event is generated when a free
                    // actually completes.
    SEGMENT_ALLOC, // a call to AclrtMalloc to get more memory from the OS
    SEGMENT_FREE, // a call to aclrtFree to return memory to the OS (e.g. to
                  // defragment or empty_caches)
    SEGMENT_MAP, // a call to AclrtMapMem (used with expandable_segments)
    SEGMENT_UNMAP, // unmap part of a segment (used with expandable
                   // segments)
    SNAPSHOT, // a call to snapshot, used to correlate memory snapshots to
              // trace events
    OOM, // the allocator threw an OutOfMemoryError (addr_ is the amount of
         // free bytes reported by cuda)
    WORKSPACE_SNAPSHOT
  };
  TraceEntry(
      Action action,
      int device,
      int64_t addr,
      size_t size,
      aclrtStream stream,
      MempoolId_t mempool = {0, 0},
      std::shared_ptr<c10::GatheredContext> context = nullptr)
      : action_(action),
        device_(device),
        addr_(addr),
        context_(std::move(context)),
        stream_(stream),
        size_(size),
        mempool_(std::move(mempool)) {}
  Action action_;
  int device_;
  int64_t addr_; // for OOM, this is the amount of free bytes reported by cuda
  std::shared_ptr<c10::GatheredContext> context_;
  aclrtStream stream_;
  int64_t size_;
  MempoolId_t mempool_;
};

修改后（对照 pytorch/c10/core/CachingDeviceAllocator.h:123-186，新增 trace_time_ union + time_ 字段 + 构造参数 + ANNOTATE action）：

// 时间戳存储 union，参考 pytorch/c10/core/CachingDeviceAllocator.h:123-126
// 记录时存 TSC tick (approx_t_)，snapshot 时覆盖为微秒 since epoch (t_)
union trace_time_ {
  c10::time_t t_;
  c10::approx_time_t approx_t_;
};

struct TraceEntry {
  enum Action {
    ALLOC, // API made to the caching allocator for new memory
    FREE_REQUESTED, // API call made to the caching allocator to free memory
    FREE_COMPLETED, // The allocator might have to delay a free because
                    // it is still in use on another stream via
                    // record_stream This event is generated when a free
                    // actually completes.
    SEGMENT_ALLOC, // a call to AclrtMalloc to get more memory from the OS
    SEGMENT_FREE, // a call to aclrtFree to return memory to the OS (e.g. to
                  // defragment or empty_caches)
    SEGMENT_MAP, // a call to AclrtMapMem (used with expandable_segments)
    SEGMENT_UNMAP, // unmap part of a segment (used with expandable
                   // segments)
    SNAPSHOT, // a call to snapshot, used to correlate memory snapshots to
              // trace events
    OOM, // the allocator threw an OutOfMemoryError (addr_ is the amount of
         // free bytes reported by cuda)
    WORKSPACE_SNAPSHOT,
    ANNOTATE // metadata attached post facto to a live allocation (可选，对齐 CUDA)
  };
  TraceEntry(
      Action action,
      int device,
      int64_t addr,
      size_t size,
      aclrtStream stream,
      MempoolId_t mempool = {0, 0},
      c10::approx_time_t time = 0,
      std::shared_ptr<c10::GatheredContext> context = nullptr)
      : action_(action),
        device_(device),
        addr_(addr),
        context_(std::move(context)),
        stream_(stream),
        size_(size),
        mempool_(std::move(mempool)) {
    time_.approx_t_ = time;
  }
  Action action_;
  int device_;
  int64_t addr_; // for OOM, this is the amount of free bytes reported by cuda
  std::shared_ptr<c10::GatheredContext> context_;
  aclrtStream stream_;
  int64_t size_;
  MempoolId_t mempool_;
  trace_time_ time_{};  // 新增：时间戳字段
};

注意：构造参数 time 放在 mempool 之后、context 之前，与 CUDA 基类顺序一致（pytorch/c10/core/CachingDeviceAllocator.h:147-168），避免破坏现有调用方对 context 参数的位置依赖。但 NPU 现有调用方都按位置传 mempool 后再传 context（见 NPUCachingAllocator.cpp:3513-3520），加 time 参数后需同步更新所有调用点。

5.2 修改点 2：C++ 算法层

文件：ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp

2.1 record_trace 函数（第 3502-3537 行）

修改前（ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:3513-3520）：

    auto te = TraceEntry(
        action,
        device,
        addr,
        size,
        stream,
        mempool_id,
        record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);

修改后（新增 c10::getApproximateTime() 参数，对照 pytorch/c10/cuda/CUDACachingAllocator.cpp:4626-4637）：

    auto te = TraceEntry(
        action,
        device,
        addr,
        size,
        stream,
        mempool_id,
        c10::getApproximateTime(),  // 新增：记录 TSC tick
        record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);

2.2 trace() 函数（第 2167-2175 行）

修改前（ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:2167-2175）：

  std::vector<TraceEntry> trace() {
    std::lock_guard<std::recursive_mutex> lock(mutex);
    std::vector<TraceEntry> result;
    result.reserve(alloc_trace->size());
    result.insert(result.end(), alloc_trace->begin() + alloc_trace_next, alloc_trace->end());
    result.insert(result.end(), alloc_trace->begin(), alloc_trace->begin() + alloc_trace_next);

    return result;
  }

修改后（对照 pytorch/c10/cuda/CUDACachingAllocator.cpp:3175-3186，新增 tsc_to_us 转换）：

  std::vector<TraceEntry> trace(
      const std::function<c10::time_t(c10::approx_time_t)>& tsc_to_us) {
    std::lock_guard<std::recursive_mutex> lock(mutex);
    std::vector<TraceEntry> result;
    result.reserve(alloc_trace->size());
    result.insert(result.end(), alloc_trace->begin() + alloc_trace_next, alloc_trace->end());
    result.insert(result.end(), alloc_trace->begin(), alloc_trace->begin() + alloc_trace_next);

    // Convert all the timestamps from tsc to epoch time in microseconds.
    for (auto& te : result) {
      te.time_.t_ = tsc_to_us(te.time_.approx_t_);
    }
    return result;
  }

2.3 多设备 snapshot wrapper（第 3926-3935 行）

修改前（ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp:3926-3935）：

  SnapshotInfo snapshot(MempoolId_t mempool_id) override {
    SnapshotInfo result;
    int count = static_cast<int>(device_allocator.size());
    for (int i = 0; i < count; i++) {
      result.device_traces.emplace_back(device_allocator[i]->trace());
      auto snap = device_allocator[i]->snapshot(mempool_id);
      result.segments.insert(result.segments.end(), snap.begin(), snap.end());
    }
    return result;
  }

修改后（对照 pytorch/c10/cuda/CUDACachingAllocator.cpp:5049-5077，新增 clock_converter + tsc_to_us 转换）：

  // 新增成员：时钟转换器（放在该类的 private 成员区，对照 pytorch/c10/cuda/CUDACachingAllocator.cpp:4736）
  c10::ApproximateClockToUnixTimeConverter clock_converter;

  SnapshotInfo snapshot(MempoolId_t mempool_id) override {
    SnapshotInfo result;
    int count = static_cast<int>(device_allocator.size());

    // Set-up converter to convert timestamps from tsc to microseconds.
    auto tsc_to_ns = clock_converter.makeConverter();
    auto tsc_to_us = [=](c10::approx_time_t t_approx) {
      return tsc_to_ns(t_approx) / 1000;
    };

    for (int i = 0; i < count; i++) {
      result.device_traces.emplace_back(device_allocator[i]->trace(tsc_to_us));
      auto snap = device_allocator[i]->snapshot(mempool_id);
      result.segments.insert(result.segments.end(), snap.begin(), snap.end());
    }
    return result;
  }

注意：clock_converter 需声明为该 wrapper 类的成员变量（非局部变量），因为它需要在构造时测量初始 1001 对时间采样（见 pytorch/c10/util/ApproximateClock.cpp:12-24）。如果作为局部变量每次 snapshot 都重新构造，会增加开销但功能正确；作为成员则只构造一次。建议作为成员，对照 pytorch/c10/cuda/CUDACachingAllocator.cpp:4736。

5.3 修改点 3：C++ 序列化层（IValue pickle 路径，死代码但建议一并修改）

文件：ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp

3.1 trace 序列化循环（第 226-285 行附近）

修改前（ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp:237-238 + 268-285）：

  c10::IValue action_s = "action";
  // ...（其他 action 字符串）
  c10::IValue pool_id_s = "pool_id";

  // ...
  for (const auto& traceInfo : snapshot.device_traces) {
    auto trace = new_list();
    for (const auto& te : traceInfo) {
      auto trace_entry = new_dict();
      trace_entry.insert(action_s, action_to_str(te.action_));
      trace_entry.insert(te.action_ == TraceEntry::OOM ? device_free_s : addr_s, te.addr_);
      trace_entry.insert(size_s, (int64_t)te.size_);
      trace_entry.insert(stream_s, int64_t(te.stream_));
      trace_entry.insert(pool_id_s, std::tuple<int64_t, int64_t>(te.mempool_));
      if (te.context_) {
        auto sc = getFromContext(te.context_);
        frame_tracebacks.push_back(sc);
        frame_dict.push_back(trace_entry);
      }
      trace.push_back(trace_entry);
    }
    traces.push_back(trace);
  }

修改后（新增 time_us_s 键，对照 pytorch/torch/csrc/cuda/memory_snapshot.cpp:335 + 467）：

  c10::IValue action_s = "action";
  // ...（其他 action 字符串）
  c10::IValue pool_id_s = "pool_id";
  c10::IValue time_us_s = "time_us";  // 新增

  // ...
  for (const auto& traceInfo : snapshot.device_traces) {
    auto trace = new_list();
    for (const auto& te : traceInfo) {
      auto trace_entry = new_dict();
      trace_entry.insert(action_s, action_to_str(te.action_));
      trace_entry.insert(te.action_ == TraceEntry::OOM ? device_free_s : addr_s, te.addr_);
      trace_entry.insert(size_s, (int64_t)te.size_);
      trace_entry.insert(stream_s, int64_t(te.stream_));
      trace_entry.insert(time_us_s, te.time_.t_);  // 新增：绝对时间（微秒 since epoch）
      trace_entry.insert(pool_id_s, std::tuple<int64_t, int64_t>(te.mempool_));
      if (te.context_) {
        auto sc = getFromContext(te.context_);
        frame_tracebacks.push_back(sc);
        frame_dict.push_back(trace_entry);
      }
      trace.push_back(trace_entry);
    }
    traces.push_back(trace);
  }

5.4 修改点 4：C++ 序列化层（Python dict 路径，实际生效路径）

文件：ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp

4.1 THNPModule_memorySnapshot 的 trace 序列化（第 1396-1446 行）

修改前（ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:1396-1446）：

  py::str device_free_s = "device_free";
  py::str pool_id_s = "pool_id";

  using namespace c10_npu::NPUCachingAllocator;

  auto action_to_str = [&](TraceEntry::Action action) {
    // ...（switch case）
  };

  for (const auto& traceInfo : snapshot.device_traces) {
    py::list trace;
    for (const auto& te : traceInfo) {
      py::dict trace_entry;
      if (te.context_) {
        auto sc = getFromContext(te.context_);
        to_gather_frames.emplace_back(sc);
        to_gather_dest.emplace_back(trace_entry);
      }
      trace_entry[action_s] = action_to_str(te.action_);
      trace_entry[te.action_ == TraceEntry::OOM ? device_free_s : addr_s] = te.addr_;
      trace_entry[size_s] = te.size_;
      trace_entry[stream_s] = int64_t(te.stream_);
      trace_entry[pool_id_s] = te.mempool_;
      trace.append(trace_entry);
    }
    traces.append(trace);
  }

修改后（新增 time_us 键，对照 pytorch/torch/csrc/cuda/Module.cpp:912）：

  py::str device_free_s = "device_free";
  py::str pool_id_s = "pool_id";
  py::str time_us_s = "time_us";  // 新增

  using namespace c10_npu::NPUCachingAllocator;

  auto action_to_str = [&](TraceEntry::Action action) {
    // ...（switch case，保持不变）
  };

  for (const auto& traceInfo : snapshot.device_traces) {
    py::list trace;
    for (const auto& te : traceInfo) {
      py::dict trace_entry;
      if (te.context_) {
        auto sc = getFromContext(te.context_);
        to_gather_frames.emplace_back(sc);
        to_gather_dest.emplace_back(trace_entry);
      }
      trace_entry[action_s] = action_to_str(te.action_);
      trace_entry[te.action_ == TraceEntry::OOM ? device_free_s : addr_s] = te.addr_;
      trace_entry[size_s] = te.size_;
      trace_entry[stream_s] = int64_t(te.stream_);
      trace_entry[time_us_s] = te.time_.t_;  // 新增：绝对时间（微秒 since epoch）
      trace_entry[pool_id_s] = te.mempool_;
      trace.append(trace_entry);
    }
    traces.append(trace);
  }

4.2 _npu_record_memory_history 绑定（第 173-176 行）—— 可选扩展

修改前（ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:173-176）：

  m.def(
      "_npu_record_memory_history",
      static_cast<void (*)(c10::optional<std::string>, c10::optional<std::string>, std::string, size_t)>(
          torch_npu::_record_memory_history));

修改后（扩展参数到 10 个，对照 pytorch/torch/csrc/cuda/Module.cpp:1239-1251）—— 此为可选扩展，若仅补齐 time_us 不需要改这里：

  m.def(
      "_npu_record_memory_history",
      static_cast<void (*)(c10::optional<std::string>, c10::optional<std::string>, std::string, size_t,
                            bool, bool, bool, std::vector<std::string>, bool, bool)>(
          torch_npu::_record_memory_history));

同时需修改 memory_snapshot.cpp:37-73 的 torch_npu::_record_memory_history 签名与实现，并在 NPUCachingAllocator::recordHistory 中支持 skip_actions / clearHistory 等。此部分为对齐 CUDA 的增强项，非补齐 time_us 的必要项。

5.5 修改点 5：Python 接口层

文件：ascend_pytorch/pytorch/torch_npu/npu/memory.py

5.1 TraceEntry TypedDict（第 947-973 行）

修改前（ascend_pytorch/pytorch/torch_npu/npu/memory.py:947-973，缺 time_us 字段）：

        class TraceEntry(TypedDict):
            # When `torch.npu.memory._record_memory_history()` is enabled,
            # the snapshot will contain TraceEntry objects that record each
            # action the allocator took.
            action: Literal[
            'alloc'  # memory allocated
            'free_requested', # the allocated received a call to free memory
            'free_completed', # the memory that was requested to be freed is now
                            # able to be used in future allocation calls
            'segment_alloc', # the caching allocator ask aclrtMalloc for more memory
                            # and added it as a segment in its cache
            'segment_free',  # the caching allocator called aclrtFree to return memory
                            # to npu possibly trying free up memory to
                            # allocate more segments or because empty_caches was called
            'oom',          # the allocator threw an OOM exception. 'size' is
                            # the requested number of bytes that did not succeed
            'snapshot'      # the allocator generated a memory snapshot
                            # useful to coorelate a previously taken
                            # snapshot with this trace
            ]
            addr: int # not present for OOM
            frames: List[Frame]
            size: int
            stream: int
            pool_id: Tuple[int, int] # id of the memory pool for this entry
            device_free: int # only present for OOM, the amount of
                            # memory npu still reports to be free

修改后（新增 time_us 字段，对照 pytorch/torch/cuda/memory.py 的 TraceEntry TypedDict）：

        class TraceEntry(TypedDict):
            # When `torch.npu.memory._record_memory_history()` is enabled,
            # the snapshot will contain TraceEntry objects that record each
            # action the allocator took.
            action: Literal[
            'alloc'  # memory allocated
            'free_requested', # the allocated received a call to free memory
            'free_completed', # the memory that was requested to be freed is now
                            # able to be used in future allocation calls
            'segment_alloc', # the caching allocator ask aclrtMalloc for more memory
                            # and added it as a segment in its cache
            'segment_free',  # the caching allocator called aclrtFree to return memory
                            # to npu possibly trying free up memory to
                            # allocate more segments or because empty_caches was called
            'oom',          # the allocator threw an OOM exception. 'size' is
                            # the requested number of bytes that did not succeed
            'snapshot'      # the allocator generated a memory snapshot
                            # useful to coorelate a previously taken
                            # snapshot with this trace
            ]
            addr: int # not present for OOM
            frames: List[Frame]
            size: int
            stream: int
            time_us: int  # 新增：事件发生的绝对时间（微秒 since unix epoch）
            pool_id: Tuple[int, int] # id of the memory pool for this entry
            device_free: int # only present for OOM, the amount of
                            # memory npu still reports to be free

5.2 _record_memory_history_impl（第 891-899 行）—— 可选扩展

修改前（ascend_pytorch/pytorch/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)

修改后（对照 pytorch/torch/cuda/memory.py:999-1024，补齐参数——此为对齐 CUDA 的增强项，若仅补齐 time_us 不需要改这里）：

def _record_memory_history_impl(
    enabled: Optional[str] = "all",
    context: Optional[str] = "all",
    stacks: str = "all",
    max_entries: int = sys.maxsize,
    device=None,
    clear_history: bool = False,
    compile_context: bool = False,
    global_record_annotations: bool = False,
    skip_actions: Optional[List[str]] = None,
    record_pinned_host_memory: bool = False,
    record_cuda: bool = True,
):
    torch_npu.npu._lazy_init()
    torch_npu._C._npu_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,
    )


# 新增签名覆盖，对照 pytorch/torch/cuda/memory.py:1027
_record_memory_history.__signature__ = signature(_record_memory_history_impl)

注意：此修改依赖修改点 4.2 的 C++ 绑定扩展，需配套提交。若仅补齐 time_us，可跳过本项。

5.3 _dump_snapshot（第 994-1022 行）—— 无需修改

_dump_snapshot 调用 _snapshot() → torch_npu._C._npu_memorySnapshot(None)，返回的 dict 经修改点 4.1 后已含 time_us 键，pickle 落盘自动包含。NPU 独有的 _npu_saveDevMemUsageInfo + profiler 兜底逻辑与 time_us 无耦合，建议保留不动。

6. 修改点依赖关系总结

6.1 最小化补齐 time_us 的必要修改（7 处）

修改点	文件	行号	必要性
1.1	`ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h`	头部	必要：include ApproximateClock.h
1.2	`ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h`	116-158	必要：TraceEntry 加 `time_` 字段 + 构造参数
2.1	`ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp`	3513-3520	必要：record_trace 调 `getApproximateTime()`
2.2	`ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp`	2167-2175	必要：trace() 加 tsc_to_us 转换
2.3	`ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp`	3926-3935	必要：snapshot wrapper 加 clock_converter
4.1	`ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp`	1396-1446	必要：Python dict 路径加 `time_us` 键（实际生效路径）
5.1	`ascend_pytorch/pytorch/torch_npu/npu/memory.py`	947-973	必要：TypedDict 补 `time_us` 字段（文档/类型对齐）

6.2 可选增强项（对齐 CUDA 完整能力，非 time_us 必需）

修改点	文件	增强内容
3.1	`ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp`	IValue pickle 路径加 time_us（当前是死代码，未绑定 Python）
4.2	`ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp`	`_npu_record_memory_history` 绑定扩展参数
5.2	`ascend_pytorch/pytorch/torch_npu/npu/memory.py`	`_record_memory_history_impl` 补齐 10 参数 + `__signature__`

注：修改点 3.1 的 _memory_snapshot_pickled 函数在 Module.cpp 中未绑定到 Python（死代码），Python 侧实际走修改点 4.1 的 THNPModule_memorySnapshot 返回 dict 路径。若仅补齐 time_us，修改点 3.1 可不改，但建议一并修改以保持两条序列化路径一致。

7. 风险与注意事项

7.1 二进制兼容

TraceEntry 加字段后，pickle 输出 schema 改变，旧版 memory_viz 工具需确认兼容（CUDA 的 time_us 已是标准字段，pytorch/torch/cuda/memory.py 的 TraceEntry TypedDict 已含 time_us，但 NPU 的 memory.py:947-973 TraceEntry TypedDict 缺 time_us，需补上，即修改点 5.1）。

7.2 性能

getApproximateTime() 是 TSC 读取（x86 __rdtsc、aarch64 cntvct_el0），单次 ns 级，对 alloc/free 热路径影响小。makeConverter() 在 snapshot 时做 1001 次采样再算中位数，但只在 dump 时调用，可接受。

7.3 平台

昇腾 NPU 多为 aarch64，需确认 C10_ARMTSC 宏在编译时启用（参考 pytorch/c10/util/ApproximateClock.h:77-83 的 arm64 实现 mrs %0, cntvct_el0），否则会回退到 getTime()（每次 clock_gettime syscall，热路径性能下降）。

7.4 NPU 兜底逻辑

_dump_snapshot 里的 _npu_saveDevMemUsageInfo + npu profiler 兜底是 NPU 独有扩展，与 CUDA 不一致，补齐 time_us 时建议保留（避免破坏现有 OOM 调试链路）。

8. 参考文件清单

NPU 侧（需修改）：

ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.h
ascend_pytorch/pytorch/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.cpp
ascend_pytorch/pytorch/torch_npu/csrc/npu/memory_snapshot.h
ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp
ascend_pytorch/pytorch/torch_npu/npu/memory.py

CUDA 参考实现：

pytorch/c10/core/CachingDeviceAllocator.h
pytorch/c10/cuda/CUDACachingAllocator.h
pytorch/c10/cuda/CUDACachingAllocator.cpp
pytorch/c10/util/ApproximateClock.h
pytorch/c10/util/ApproximateClock.cpp
pytorch/torch/csrc/cuda/memory_snapshot.cpp
pytorch/torch/csrc/cuda/Module.cpp
pytorch/torch/cuda/memory.py

X xuyun15成员

19 天前关联了里程碑：v26.3.0

X xuyun15成员

19 天前添加了label：feature

TorchNPU-Bot成员

19 天前添加了label：triage-reviewmodule: core

TorchNPU-Bot成员

19 天前评论：

无报错，属功能建议/需求（为 memory snapshot 增加时间戳字段），主题为内存快照/核心运行时特性增强
添加分类标签：module: core。
当前分类由社区bot提供，如分类有误请在评论中回复。

C chenqiming10成员

13 天前关联了pull request：Add time_us timestamp collection to memory snapshot trace events

TorchNPU-Bot成员

13 天前添加了label：bot-triaged；删除了label：triage-review

TorchNPU-Bot成员

13 天前评论：

检测到当前 issue 已关联 PR，自动添加标签：bot-triaged

C chenqiming10成员

12 天前关联了pull request：feat: Add time_us timestamp collection to memory snapshot trace events