From b81866c7088facdb70d39d7dd7106fbc0182e56e Mon Sep 17 00:00:00 2001 From: weich97 Date: Sat, 10 Oct 2026 22:57:35 -0400 Subject: [PATCH] feat(agent): memory-analysis skill with attribution read from real snapshots Add `.agent/skills/memory-analysis`: record NPU allocation history, dump snapshots at chosen moments, and break device memory down into what holds it and which line allocated it. The field names, frame shapes and identity relations below were read off snapshots taken on torch_npu 2.9.0.post2 / CANN 9.0.0 / Ascend 910, and the category rules were checked against a run whose composition was known in advance. `breakdown.py` reads a `_snapshot()` dump: * Totals come from the segment fields (`total_size` / `allocated_size` / `active_size`) rather than re-summed blocks; the re-summed values serve the consistency checks instead. * Attribution scans the recorded stack rather than reading `frames[0]`. The stack is innermost-first and its innermost entries are the snapshot machinery itself (`torch_npu::unwind::unwind`, `CapturedTraceback::gather`, `??`), so `frames[0]` is the same for every block and says nothing about the allocation. The call site is the innermost Python frame outside `site-packages`; the first Python frame sat at index 26 of 62 on the collected snapshots, so the position cannot be assumed. * Gradients carry no Python frame at all, because the autograd engine allocates them from C++. They are recognised by `torch::autograd::AccumulateGrad` among the native frames, which is the only available signal. The native fallback also skips `at::` and `c10::` allocator plumbing, or it would report `at::_ops::empty_strided::redispatch` for every such block. * Categories: optimizer state from a `torch/optim` frame (AdamW allocates through `adam.py`, not `adamw.py`), parameters from the module `.to()` conversion, gradients as above, activations when the innermost own-code frame is a `forward`. On a run of three bf16 `Linear(2048, 2048)` layers with AdamW, the split came out as parameters 24.0 MiB, gradients 24.0 MiB, optimizer state 48.0 MiB, input 2.0 MiB and activations 12.0 MiB, summing to the reported `allocated_size` and matching the sizes those tensors must have. A snapshot is one instant and is not the peak, so the report states the snapshot total beside `max_memory_allocated` and flags the gap: on that run a post-step snapshot held 100.01 MiB with no activations at all, a mid-backward snapshot held 112.01 MiB, and the peak was 130.01 MiB. Five invariants are emitted as consistency checks with exit code 2 on failure - per-segment block tiling, the global active total, the two allocator counters, and whether the snapshot is within 2% of the peak. `--compare` diffs two snapshots by category, which is what separates a leak from a one-off allocation; a single snapshot cannot. Across two steady steps every category came out at +0.00 MiB, and against the mid-backward snapshot the growth was attributed to activations at +10.00 MiB. `collect_snapshot.py` exported `MEM_SNAPSHOT_STEPS` and `MEM_SNAPSHOT_PATH` that nothing read, and judged collection by the run's exit code. It now verifies that each dump arrived with its counter sidecar, and the new `mem_hook` is what reads those variables, so the launch side and the training script share one window. `mem_hook.dump(tag)` can be called at any moment, including from an autograd hook while the graph is alive, which is the only way to see activations and the only way to inspect a run that is going badly without waiting for it to end. Each dump also writes the allocator counters, which live only in the process, so the offline report can cross-check its totals and name the peak. `_dump_snapshot` creates the file mode 0200, write-only, so not even its owner can read it back; `mem_hook` restores owner read/write and the parser explains the fix when it meets a dump taken another way. Unit tests cover the attribution order, the native-plumbing skip, each category rule, the segment identity, the peak check in all three states, the comparison in both directions, the unreadable-dump message and the exit codes, against a fixture carrying the real frame shapes (CPU only). --- .agent/skills/memory-analysis/SKILL.md | 102 ++++ .../references/breakdown-dimensions.md | 85 ++++ .../references/snapshot-format.md | 124 +++++ .../memory-analysis/scripts/breakdown.py | 447 ++++++++++++++++++ .../scripts/collect_snapshot.py | 113 +++++ .../memory-analysis/scripts/mem_hook.py | 152 ++++++ .../templates/memory-breakdown.html | 72 +++ AGENTS.md | 1 + tests/ut/tools/test_memory_breakdown.py | 339 +++++++++++++ 9 files changed, 1435 insertions(+) create mode 100644 .agent/skills/memory-analysis/SKILL.md create mode 100644 .agent/skills/memory-analysis/references/breakdown-dimensions.md create mode 100644 .agent/skills/memory-analysis/references/snapshot-format.md create mode 100644 .agent/skills/memory-analysis/scripts/breakdown.py create mode 100644 .agent/skills/memory-analysis/scripts/collect_snapshot.py create mode 100644 .agent/skills/memory-analysis/scripts/mem_hook.py create mode 100644 .agent/skills/memory-analysis/templates/memory-breakdown.html create mode 100644 tests/ut/tools/test_memory_breakdown.py diff --git a/.agent/skills/memory-analysis/SKILL.md b/.agent/skills/memory-analysis/SKILL.md new file mode 100644 index 000000000..2aafa77bd --- /dev/null +++ b/.agent/skills/memory-analysis/SKILL.md @@ -0,0 +1,102 @@ +--- +name: memory-analysis +description: > + NPU memory analysis for training: record allocation history, dump snapshots + at chosen moments (including mid-backward), and break device memory down + into parameters / gradients / optimizer state / activations with each large + block attributed to the call site that allocated it, plus fragmentation, + cross-step growth for leak triage, and a check on whether the snapshot is + the peak at all. Use for 显存分析 / 显存拆解 / 显存峰值 / 显存碎片 / + 显存泄露 / OOM 定位 / memory breakdown. Rule-based, cross-checked against + the allocator's own invariants, and reproducible; emits a visual HTML + report. +--- + +# Memory Analysis + +Record allocation history, dump a snapshot at the moment that matters, and +break device memory down into what is holding it and which line allocated it. +**Rule-based and reproducible** — the same snapshot gives the same breakdown. + +**This file is the index.** The snapshot guide and breakdown dimensions load +on demand. + +> 对齐状态：字段、帧形态与恒等式均在 **torch_npu 2.9.0.post2 + CANN 9.0.0 / +> Ascend 910** 上实采核对，分类规则对着一次已知构成的运行逐项对账（见 +> `references/breakdown-dimensions.md` § 实测对账）。规则版本 v2。 + +## Subcommands + +```bash +python .agent/skills/memory-analysis/scripts/.py +``` + +| Intent | Tool | Purpose | +|---|---|---| +| 采集 | `collect_snapshot.py --run-dir --launch "" [--steps 1-3]` | 驱动被测命令并校验每份 dump 都带计数器旁车 | +| 训练脚本内挂钩 | `mem_hook.start()` / `step()` / `dump(tag)` | 进程内记录与落盘，**任意时刻可 dump** | +| 拆解 + HTML | `breakdown.py --snapshot [--compare ] [--json r.json]` | 解析快照、分类、归属、校验（退出码 0=全通过 / 2=有不通过 / 1=文件缺失）| + +训练脚本接入： + +```python +import sys +sys.path.insert(0, ".agent/skills/memory-analysis/scripts") +import mem_hook + +mem_hook.start() # 开始记录；读 collect_snapshot 导出的 MEM_* 变量 +for step, batch in enumerate(loader): + out = model(batch) + out.register_hook(lambda g: mem_hook.dump("midbwd") and None) # 想看激活就在此处 + loss(out).backward(); opt.step(); opt.zero_grad() + mem_hook.step() # 落在配置窗口内的 step 自动 dump +mem_hook.stop() +``` + +## References + +| When | Read | +|---|---| +| 快照字段、帧形态、落盘坑、实测恒等式 | [references/snapshot-format.md](references/snapshot-format.md) | +| 总量三层、类别判定规则与实测对账、归属规则、校验项 | [references/breakdown-dimensions.md](references/breakdown-dimensions.md) | + +## Report + +`templates/memory-breakdown.html` — 六节：总量与碎片（含与峰值的差距）、 +按类别、最大活跃块及其调用点、跨快照增长、一致性校验、结论。 + +## What it covers + +- **总量与碎片**：Reserved / Allocated / Active 三层与碎片量。 +- **按类别**：参数 / 梯度 / 优化器状态 / 激活 / 其他用户张量 / unclassified， + 判据见 references，已对着已知构成的运行逐项对账。 +- **块级归属**：最大活跃块各自归到**分配它的那行自有代码**（梯度无 Python 帧， + 归到 `AccumulateGrad`）。 +- **跨 step 增长**：`--compare` 比对两份快照，区分泄露与一次性大分配。 +- **是否峰值**：把快照活跃量与 `max_memory_allocated` 对比并报差距—— + 步后快照里激活为 0，实测比峰值低 30 MiB，不标出来就会被误读。 + +### 尚未实现（不要当成已有能力） + +- **Module / Layer 粒度汇总**：当前归属到调用点（文件:行），未按 module 名聚合。 +- **重计算收益估算**：需要 module 粒度激活量，依赖上一项。 +- **峰值时刻构成**：当前只标注快照与峰值的差距；回放 `device_traces` 还原峰值 + 构成的口径 `[待对齐]`。 + +## Hard rules + +- **快照不是峰值。** 报告必须同时给出 `max_memory_allocated`；差距超 2% 即校验 + 不通过。要看激活必须在 backward 期间 dump。 +- **`frames[0]` 不是归属。** 记录栈最内层是快照机制自身，取第 0 帧会让每个块都 + 归到同一处；按 references 的规则扫帧。 +- **梯度按 native 帧的 `AccumulateGrad` 判定**——它没有任何 Python 帧，按 Python + 帧分类会整类漏掉。 +- **Default steps 1-3** 且至少两份以便比对；单份快照分不清泄露与一次性峰值。 +- **校验不通过的数字不采信。** 四条恒等式（逐段与全局、与进程计数器）是解析 + 正确的正面证据。 +- **Reproducible evidence.** 拆解带被采集运行的命令、commit 与环境 + （与 `autoresearch` / `profiling` / `perf-playbook` 共享该纪律）。 +- **换省显存手段必须同时量化性能。** 本 skill 只定位占用来源；手段选择是 + `perf-playbook` 的事，且显存降了但性能掉了不算收益。 +- A classification needing an internal 口径 stays `[待对齐]`; an unparsed field + stays `[待解析]`. Never fill a number that was not computed. diff --git a/.agent/skills/memory-analysis/references/breakdown-dimensions.md b/.agent/skills/memory-analysis/references/breakdown-dimensions.md new file mode 100644 index 000000000..69dbbfd79 --- /dev/null +++ b/.agent/skills/memory-analysis/references/breakdown-dimensions.md @@ -0,0 +1,85 @@ +# 显存拆解维度与分类规则 + +规则版本 **v2**（按 torch_npu 2.9.0.post2 实采字段与帧形态重写）。字段出处与 +实测证据见 [snapshot-format.md](snapshot-format.md)。 + +## 总量三层 + +| 层 | 来源 | 含义 | +|---|---|---| +| Reserved | `Σ segment.total_size` | 分配器向设备保留的总量 | +| Allocated | `Σ segment.allocated_size` | 取样瞬间在用 | +| 碎片 | `Reserved − Allocated` | 保留但未在用；占 Reserved 超 10% 先查分配器配置 | + +段总量**直接取段字段**，不重新累加 block——累加值另算，专用于校验（见下）。 + +峰值 `max_memory_allocated` 不在快照里，由 `mem_hook` 写进 +`.stats.json` 旁车。**快照 ≠ 峰值**，差距作为校验项输出。 + +## 类别判定规则（按记录栈，顺序判定） + +| 顺序 | 类别 | 判据 | 实测验证 | +|---|---|---|---| +| 1 | 优化器状态 | Python 帧路径含 `/optim/` | AdamW 走 `torch/optim/adam.py:177/181 _init_group`（AdamW 复用 adam.py，不要按 `adamw` 匹配） | +| 2 | 参数 | Python 帧含 `_module.py` 且函数名为 `convert`/`to` | `torch_npu/utils/_module.py:76 convert`，由 `module.to(device)` 触发 | +| 3 | 梯度 | **native** 帧含 `AccumulateGrad` | 梯度无任何 Python 帧，这是唯一信号 | +| 4 | 激活 | 最内层自有代码帧的函数名含 `forward` | 仅在 backward 期间的快照里出现 | +| 5 | 其他用户张量 | 有自有代码帧但不匹配上述 | 如输入 batch | +| 6 | unclassified | 无可用帧 | 如实标注，不猜 | + +### 实测对账（一次已知构成的运行） + +模型：3 × `Linear(2048, 2048, bias=False)`，bf16；AdamW；输入 512×2048 bf16。 + +| 类别 | 工具给出 | 地面真值 | 核对 | +|---|---|---|---| +| 参数 | 24.0 MiB | 3 × 2048² × 2B = 24.0 MiB | ✅ | +| 梯度 | 24.0 MiB | 同参数（bf16）= 24.0 MiB | ✅ | +| 优化器状态 | 48.0 MiB | 2 状态 × 3 × 8 MiB = 48.0 MiB | ✅ | +| 激活（backward 中途） | 12.0 MiB | — | 步后快照降至 2.0 MiB | +| 其他用户张量 | 2.0 MiB | 512 × 2048 × 2B = 2.0 MiB（输入） | ✅ | +| 合计 | 112.0 MiB | = `allocated_size` 112.01 MiB | ✅ | + +> 注意 AdamW 的 `exp_avg`/`exp_avg_sq` 走 `zeros_like(param)`，**参数是 bf16 时 +> 状态也是 bf16**（8 MiB/个），不是 fp32 的 16 MiB/个。按 fp32 估算会得出错误的 +> 预期值。 + +## 归属（调用点）规则 + +见 [snapshot-format.md](snapshot-format.md) § frames 的真实形态。要点： + +1. 记录栈最内层优先，最内几帧是快照机制自身（`unwind` / `CapturedTraceback` / + `??`）——**取 `frames[0]` 会让所有块归到同一处**。 +2. 取最内层「不在 site-packages 内」的 `.py` 帧 → 调用方自己的代码位置。 +3. 退到最内层框架 `.py` 帧。 +4. 再退到 native 符号，但**先跳过 `at::` / `c10::` 分配器管道**，否则拿到 + `at::_ops::empty_strided::redispatch` 这类无用帧。 + +## 一致性校验（正面生效标记） + +| 校验 | 依据 | 离线可查 | +|---|---|---| +| 逐段 Σ块 = `total_size`，Σ活跃块 = `allocated_size` | 分配器自洽 | 是 | +| Σ活跃块 = Σ段 `allocated_size` | 同上，全局 | 是 | +| Σ段 `total_size` = `memory_reserved()` | 旁车计数器 | 需旁车 | +| Σ段 `allocated_size` = `memory_allocated()` | 旁车计数器 | 需旁车 | +| 快照是否为峰值（容差 2%） | 旁车 `max_memory_allocated` | 需旁车 | + +前四条是解析正确的正面证据；第五条防止把取样瞬间当峰值读。任一不通过 +→ 退出码 2，报告里标红。 + +## 采集纪律 + +- **默认采 step 1–3**（`--steps`）；跑长不代表构成会变，但峰值可能在后面出现。 +- **要看激活就必须在 backward 期间取样**：`mem_hook.dump(tag)` 可从 autograd + hook 内调用。步后快照里激活为 0。 +- **长跑要能中途取样**：性能/显存异常时当场 dump 当场定位，不要等跑完 + （见仓库纪律：性能是贯穿整个实验的判据）。 +- 省显存的手段（重计算、offload、分配器配置）见 `perf-playbook`；本 skill 只 + 定位占用来源，**换手段必须同时量化性能影响**，不能只报显存降了。 + +## 待对齐 + +- `device_traces` 回放峰值构成的口径 `[待对齐]`（当前只用快照瞬间 + 旁车峰值 + 标注差距）。 +- Module/Layer 粒度归并：需要按模块名正则把调用点聚合，规则随实践补充。 diff --git a/.agent/skills/memory-analysis/references/snapshot-format.md b/.agent/skills/memory-analysis/references/snapshot-format.md new file mode 100644 index 000000000..bcf02d590 --- /dev/null +++ b/.agent/skills/memory-analysis/references/snapshot-format.md @@ -0,0 +1,124 @@ +# 显存快照格式与实测要点 + +基于 `torch_npu.npu.memory`，在 **torch_npu 2.9.0.post2 + CANN 9.0.0 / +Ascend 910** 上实采核对。字段名与行为均为实测，不是推测。 + +## API + +| 调用 | 作用 | +|---|---| +| `_record_memory_history(enabled="all", context="all", stacks="all", max_entries=N)` | 开始记录分配历史。`stacks="all"` 才有 Python 帧，归属全靠它 | +| `_snapshot(device=None)` | 返回快照 dict（进程内） | +| `_dump_snapshot(filename=...)` | 落盘为 pickle | +| `_record_memory_history(enabled=None)` | 停止记录 | + +实测签名：`_record_memory_history_impl(enabled='all', context='all', +stacks='all', max_entries=sys.maxsize, device=None)`。 + +### 两个落盘坑 + +1. **`_dump_snapshot` 写出的文件权限是 `--w-------`（0200，只写）**，连属主自己都 + 读不回来。`mem_hook.dump()` 落盘后补 `chmod 0600`；手工 dump 的要自己 + `chmod 600 `，否则解析会报 PermissionError。 +2. `_dump_snapshot` 还会在当前工作目录生成 `PROF_00000N_*` 与 `FRAMEWORK/` + 目录，属正常副产物。 +3. 记录期间 torch_npu 会打 `AclrtGetMemUsageInfo failed, ret:207000` 的 + **ERROR 日志**，但快照正常产出，属已知噪音，不要据此判定采集失败。 + +## 快照结构 + +```python +{"segments": [...], "device_traces": [...]} +``` + +### segment 字段（实测） + +`active_size` · `address` · `allocated_size` · `blocks` · `device` · `frames` +· `is_expandable` · `requested_size` · `segment_pool_id` · `segment_type` +· `stream` · `total_size` + +| 字段 | 含义 | +|---|---| +| `total_size` | 该段向设备保留的字节数 | +| `allocated_size` | 段内在用字节数 | +| `active_size` | 含 pending-free 的在用字节数；与 `allocated_size` 相等表示无待释放块 | +| `segment_type` | `large` / `small` | +| `frames` | 实测为空 `[]`，归属只能取 block 的 | + +### block 字段（实测） + +`address` · `frames` · `requested_size` · `size` · `state` + +`state` 实测取值：`active_allocated` · `inactive`（`active_pending_free` 亦属 +活跃，按 `active` 前缀判定）。 + +### 实测恒等式 + +**离线可查（只用快照本身）：** + +``` +逐段：Σ(所有 block.size) == segment.total_size +逐段：Σ(活跃 block.size) == segment.allocated_size +``` + +**需要进程内计数器（`mem_hook` 写进 `.stats.json` 旁车文件）：** + +``` +Σ(segment.total_size) == memory_reserved() +Σ(segment.allocated_size) == memory_allocated() +``` + +实测四条全部精确成立（零偏差）。`breakdown.py` 把它们作为校验项输出，不通过 +则退出码 2。 + +## frames 的真实形态（归属的关键） + +记录栈是**最内层优先**，且最内几帧是**快照机制自身**： + +``` +[0] torch_npu::unwind::unwind() ← 取 frames[0] 等于每块都归到这里 +[1] torch_npu::CapturedTraceback::gather(...) +[2..] ?? ← 无符号 +... +[26] 第一个 .py 帧（共 62 帧时） ← 位置不固定，必须扫不能索引 +... +[61] python3 / __libc_start_main +``` + +- **第一个 `.py` 帧的下标不固定**（实测 62 帧里在第 26 位），必须遍历查找。 +- Python 帧内部同样最内层优先：`torch_npu/utils/_module.py:76 convert` → + `torch/nn/modules/module.py:957 _apply` → `train.py:31 main`。 +- **归属取最内层「不在 site-packages/dist-packages 内」的 `.py` 帧**，得到调用方 + 自己的代码位置；没有则退到最内层框架 `.py` 帧；再没有则退到 native 符号。 +- native 退路要先跳过**分配器管道**（`at::` / `c10::` 前缀），否则拿到的是 + `at::_ops::empty_strided::redispatch` 这类对归属无用的帧。跳过后得到 + `torch::autograd::AccumulateGrad::apply`，才是有意义的。 + +### 梯度没有 Python 帧 + +实测：梯度块的 `frames` 里**一个 `.py` 帧都没有**——autograd 引擎在 C++ 侧分配。 +唯一可用信号是 native 帧里的 `torch::autograd::AccumulateGrad::apply`。 +按 Python 帧做分类的实现会把梯度整类漏掉。 + +## device_traces + +每设备一条列表，元素形如 +`{"action": "segment_alloc"|"alloc"|..., "addr":, "size":, "stream":, "frames": [...]}`， +`frames` 同样带 Python 帧。可用于回放峰值时刻的构成（当前 `breakdown.py` 未用， +标 `[待对齐]`：回放口径）。 + +## 快照不等于峰值（最重要的一条） + +实测同一次运行： + +| 取样时刻 | 活跃量 | +|---|---| +| `optimizer.step()` 之后 | 100.01 MiB（**激活为 0**） | +| backward 中途（autograd hook 内） | 112.01 MiB（激活在场） | +| `max_memory_allocated()` | **130.01 MiB** | + +- 步后快照里**没有任何激活**，拿它当峰值构成会严重低估。 +- 即便 backward 中途取样，仍低于真实峰值 18 MiB。 +- 所以：`mem_hook.dump(tag)` 支持**任意时刻**调用（含从 autograd hook 内）， + 并把 `max_memory_allocated` 写入旁车；`breakdown.py` 把「快照是否为峰值」 + 作为一项校验，差距超 2% 即报不通过并在报告里写明。 diff --git a/.agent/skills/memory-analysis/scripts/breakdown.py b/.agent/skills/memory-analysis/scripts/breakdown.py new file mode 100644 index 000000000..3839cb37f --- /dev/null +++ b/.agent/skills/memory-analysis/scripts/breakdown.py @@ -0,0 +1,447 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Parse an NPU memory snapshot and break device memory down into an HTML report. + +Written against snapshots taken with ``torch_npu.npu.memory._snapshot()`` on +torch_npu 2.9.0.post2 / Ascend 910. The snapshot is ``{"segments": [...], +"device_traces": [...]}``; each segment carries ``total_size`` / +``allocated_size`` / ``active_size`` and a ``blocks`` list whose entries have +``size`` / ``requested_size`` / ``state`` / ``frames``. + +Three properties of the real data drive this parser: + +* **A snapshot is one instant, not the peak.** A snapshot taken after + ``optimizer.step()`` holds parameters, gradients and optimizer state but no + activations at all; one taken mid-backward holds activations too and was + still below ``max_memory_allocated``. The report therefore states the + snapshot instant's total next to the recorded peak and flags the gap instead + of letting a snapshot be read as the peak. +* **``frames[0]`` is useless for attribution.** The recorded stack is + innermost-first and its innermost entries are the snapshot machinery itself + (``torch_npu::unwind::unwind``, ``CapturedTraceback::gather``, ``??``), so + every block would be attributed to the same frame. Attribution scans for the + innermost Python frame outside ``site-packages`` for the call site, and uses + the framework frames for the category. +* **Gradients carry no Python frame at all** — the autograd engine allocates + them from C++ — so they are recognised by + ``torch::autograd::AccumulateGrad`` among the native frames. + +Pure stdlib. The snapshot is a file the recording already produced. +""" +from __future__ import annotations + +import argparse +import json +import pickle # nosec B403 - trusted local snapshot produced by our own run +import sys +from pathlib import Path + +MIB = 1024 * 1024 + +# Innermost frames belonging to the capture machinery, not the allocation site. +FRAME_NOISE = ("unwind", "CapturedTraceback", "gather", "??") +# Allocator plumbing that every device allocation goes through; the frame that +# says something about the caller sits just outside it (for a gradient, that is +# `torch::autograd::AccumulateGrad::apply`). +NATIVE_PLUMBING = ("at::", "c10::") +# Installed-library path markers; a frame outside these is the user's own code. +SITE_MARKERS = ("site-packages", "dist-packages") +# Native symbol that marks a gradient buffer allocated by the autograd engine. +GRAD_SYMBOL = "AccumulateGrad" + +ACTIVE_PREFIX = "active" +CATEGORY_OTHER = "其他用户张量" +CATEGORY_UNKNOWN = "unclassified" + + +def load_snapshot(path: Path) -> dict: + """Load a snapshot from pickle or JSON. + + Args: + path: Snapshot file written by the recording hook. + + Returns: + The snapshot dict. + """ + if path.suffix == ".json": + return json.loads(path.read_text(encoding="utf-8")) + try: + with path.open("rb") as handle: + return pickle.load(handle) # nosec B301 - local, self-produced file + except PermissionError as exc: + # `_dump_snapshot` creates the file write-only (mode 0200); a dump taken + # without `mem_hook`, which restores owner read, lands unreadable. + raise PermissionError( + f"{path} is not readable (torch_npu writes snapshots mode 0200); " + f"run `chmod 600 {path}`, or dump through mem_hook which does it" + ) from exc + + +def load_sidecar(path: Path) -> dict: + """Read the allocator counters recorded beside the snapshot, if present. + + ``memory_reserved`` / ``memory_allocated`` / ``max_memory_allocated`` are + live-process values that the snapshot itself does not carry; the recording + hook writes them to ``.stats.json`` so the offline report can + cross-check totals and name the peak. + """ + sidecar = path.with_suffix(path.suffix + ".stats.json") + if not sidecar.is_file(): + return {} + try: + return json.loads(sidecar.read_text(encoding="utf-8")) + except (ValueError, OSError): + return {} + + +def _python_frames(frames: list | None) -> list[dict]: + """Python frames only, in the recorded innermost-first order.""" + return [frame for frame in (frames or []) + if isinstance(frame, dict) and str(frame.get("filename", "")).endswith(".py")] + + +def _named_native(frames: list | None) -> list[str]: + """Native symbol names with the capture machinery filtered out.""" + out = [] + for frame in frames or []: + name = frame.get("name") if isinstance(frame, dict) else None + if name and not any(noise in name for noise in FRAME_NOISE): + out.append(name) + return out + + +def _format_frame(frame: dict) -> str: + """``file.py:line func`` for one recorded frame.""" + filename = str(frame.get("filename", "?")).replace("\\", "/").rsplit("/", 1)[-1] + return f"{filename}:{frame.get('line', '?')} {frame.get('name', '?')}" + + +def call_site(frames: list | None) -> str: + """Attribute a block to the innermost frame that is the caller's own code. + + Falls back to the innermost framework Python frame, then to the innermost + meaningful native symbol, then to ``unclassified``. + """ + pythons = _python_frames(frames) + for frame in pythons: + if not any(marker in str(frame.get("filename", "")) for marker in SITE_MARKERS): + return _format_frame(frame) + if pythons: + return _format_frame(pythons[0]) + native = _named_native(frames) + meaningful = [name for name in native + if not any(name.startswith(prefix) for prefix in NATIVE_PLUMBING)] + chosen = (meaningful or native or [CATEGORY_UNKNOWN])[0] + return chosen.partition("(")[0][:80] + + +def categorize(frames: list | None) -> str: + """Classify a block from its recorded frames. + + Rules, in order, each validated against a run whose composition was known: + optimizer state (a ``torch/optim`` frame), parameter (the module ``.to()`` + conversion), gradient (``AccumulateGrad`` among the native frames, which is + the only signal because gradients carry no Python frame), activation (the + innermost own-code frame is inside a ``forward``). + """ + pythons = _python_frames(frames) + paths = " ".join(str(frame.get("filename", "")).replace("\\", "/") for frame in pythons) + names = " ".join(str(frame.get("name", "")) for frame in pythons) + if "/optim/" in paths: + return "优化器状态" + if "_module.py" in paths and ("convert" in names or " to" in f" {names}"): + return "参数" + if any(GRAD_SYMBOL in symbol for symbol in _named_native(frames)): + return "梯度" + own = [frame for frame in pythons + if not any(marker in str(frame.get("filename", "")) for marker in SITE_MARKERS)] + if own and "forward" in str(own[0].get("name", "")): + return "激活" + if own: + return CATEGORY_OTHER + return CATEGORY_UNKNOWN + + +def _active_blocks(snapshot: dict) -> list[tuple[int, dict]]: + """Every active block as ``(size, block)``.""" + out = [] + for segment in snapshot.get("segments") or []: + for block in segment.get("blocks") or []: + if str(block.get("state", "")).startswith(ACTIVE_PREFIX): + out.append((int(block.get("size", 0) or 0), block)) + return out + + +def summarize(snapshot: dict) -> dict: + """Total reserved / allocated / fragmentation plus the per-category split. + + Segment totals are taken from the segment fields rather than re-summed, and + the re-summed values are kept for the consistency checks. + """ + segments = snapshot.get("segments") or [] + reserved = sum(int(seg.get("total_size", 0) or 0) for seg in segments) + allocated = sum(int(seg.get("allocated_size", 0) or 0) for seg in segments) + active = sum(int(seg.get("active_size", 0) or 0) for seg in segments) + blocks = _active_blocks(snapshot) + categories: dict[str, float] = {} + rows = [] + for size, block in blocks: + category = categorize(block.get("frames")) + categories[category] = categories.get(category, 0.0) + size / MIB + rows.append((size / MIB, category, call_site(block.get("frames")))) + rows.sort(reverse=True) + return { + "segments": len(segments), + "reserved_mib": reserved / MIB, + "allocated_mib": allocated / MIB, + "active_mib": active / MIB, + "fragmentation_mib": max(0.0, (reserved - allocated) / MIB), + "block_sum_mib": sum(size for size, _ in blocks) / MIB, + "categories": categories, + "top_blocks": rows[:20], + "block_count": len(blocks), + } + + +def run_checks(snapshot: dict, stats: dict, sidecar: dict) -> list[tuple[str, bool, str]]: + """Validate the parse against the allocator's own invariants. + + The first two hold inside the snapshot and are always checked; the rest need + the counters the recording hook wrote beside it. + """ + checks: list[tuple[str, bool, str]] = [] + checks.append(_segment_identity(snapshot)) + checks.append(("Σ活跃块 = Σ段 allocated_size", + abs(stats["block_sum_mib"] - stats["allocated_mib"]) < 1e-6, + f"{stats['block_sum_mib']:.2f} vs {stats['allocated_mib']:.2f} MiB")) + for key, label, value in ( + ("memory_reserved", "Σ段 total_size = memory_reserved()", stats["reserved_mib"]), + ("memory_allocated", "Σ段 allocated_size = memory_allocated()", + stats["allocated_mib"])): + if key in sidecar: + recorded = sidecar[key] / MIB + checks.append((label, abs(recorded - value) < 1e-6, + f"{value:.2f} vs {recorded:.2f} MiB")) + checks.append(_peak_check(stats, sidecar)) + return [check for check in checks if check is not None] + + +def _segment_identity(snapshot: dict) -> tuple[str, bool, str]: + """Each segment's blocks must tile it exactly, and match its allocated_size.""" + bad = [] + for index, segment in enumerate(snapshot.get("segments") or []): + blocks = segment.get("blocks") or [] + total = sum(int(block.get("size", 0) or 0) for block in blocks) + live = sum(int(block.get("size", 0) or 0) for block in blocks + if str(block.get("state", "")).startswith(ACTIVE_PREFIX)) + if total != int(segment.get("total_size", 0) or 0): + bad.append(f"seg{index} 块总和≠total_size") + if live != int(segment.get("allocated_size", 0) or 0): + bad.append(f"seg{index} 活跃块≠allocated_size") + return ("逐段：Σ块 = total_size，Σ活跃块 = allocated_size", not bad, + "全部符合" if not bad else "; ".join(bad[:4])) + + +def _peak_check(stats: dict, sidecar: dict) -> tuple[str, bool, str] | None: + """State whether this snapshot is actually the peak, never assume it is.""" + if "max_memory_allocated" not in sidecar: + return ("快照是否为峰值", False, + "无 max_memory_allocated 记录，无法判断；快照仅代表取样瞬间") + peak = sidecar["max_memory_allocated"] / MIB + gap = peak - stats["allocated_mib"] + near = abs(gap) <= peak * 0.02 + detail = (f"快照 {stats['allocated_mib']:.2f} MiB vs 峰值 {peak:.2f} MiB" + f"，差 {gap:.2f} MiB") + return ("快照是否为峰值（容差 2%）", near, + detail + ("" if near else "——本报告的构成是取样瞬间的，不是峰值构成")) + + +def analyse(path: Path) -> dict: + """Parse one snapshot into the full record.""" + snapshot = load_snapshot(path) + sidecar = load_sidecar(path) + stats = summarize(snapshot) + return {"path": path, "stats": stats, "sidecar": sidecar, + "checks": run_checks(snapshot, stats, sidecar)} + + +def compare(baseline: dict, later: dict) -> dict: + """Diff two snapshots' category totals to expose growth across steps. + + A leak shows up as a category that keeps growing between steps, which a + single snapshot cannot distinguish from a large one-off allocation. + + Args: + baseline: ``summarize`` output for the earlier snapshot. + later: ``summarize`` output for the later one. + + Returns: + ``{"allocated_delta_mib": float, "reserved_delta_mib": float, + "categories": {name: delta_mib}}``. + """ + names = set(baseline["categories"]) | set(later["categories"]) + return { + "allocated_delta_mib": later["allocated_mib"] - baseline["allocated_mib"], + "reserved_delta_mib": later["reserved_mib"] - baseline["reserved_mib"], + "categories": {name: later["categories"].get(name, 0.0) + - baseline["categories"].get(name, 0.0) + for name in sorted(names)}, + } + + +def _rows(cells: list[tuple]) -> str: + """Render table rows from a list of cell tuples.""" + return "".join("" + "".join(f"{c}" for c in row) + "" for row in cells) + + +def _total_rows(report: dict) -> str: + """Build the totals table.""" + stats = report["stats"] + sidecar = report["sidecar"] + peak = (f"{sidecar['max_memory_allocated'] / MIB:.1f}" + if "max_memory_allocated" in sidecar else "[待解析]") + return _rows([ + ("Reserved（Σ段 total_size）", f"{stats['reserved_mib']:.1f}", "分配器向设备保留"), + ("Allocated（Σ段 allocated_size）", f"{stats['allocated_mib']:.1f}", "取样瞬间在用"), + ("Active（Σ段 active_size）", f"{stats['active_mib']:.1f}", + "含 pending-free；与 Allocated 相等表示无待释放块"), + ("碎片（Reserved − Allocated）", f"{stats['fragmentation_mib']:.1f}", + "保留但未在用"), + ("峰值 max_memory_allocated", peak, "与 Allocated 的差即快照未覆盖的部分"), + ]) + + +def _category_rows(report: dict) -> str: + """Build the per-category table.""" + categories = report["stats"]["categories"] + if not categories: + return '[待解析] 快照内无活跃块' + total = sum(categories.values()) or 1.0 + ordered = sorted(categories.items(), key=lambda kv: kv[1], reverse=True) + return _rows([(name, f"{mib:.1f}", f"{mib / total * 100:.1f}%") for name, mib in ordered]) + + +def _block_rows(report: dict) -> str: + """Build the largest-blocks table.""" + rows = report["stats"]["top_blocks"] + if not rows: + return '[待解析]' + return _rows([(f"{mib:.2f}", category, site) for mib, category, site in rows]) + + +def _check_rows(report: dict) -> str: + """Build the invariant-check table.""" + return _rows([(name, "通过" if ok else '不通过', detail) + for name, ok, detail in report["checks"]]) + + +def _verdict(report: dict) -> str: + """Name the largest category and the largest single block.""" + stats = report["stats"] + if not stats["categories"]: + return '[待解析]' + name, mib = max(stats["categories"].items(), key=lambda kv: kv[1]) + biggest = stats["top_blocks"][0] + frag = stats["fragmentation_mib"] + parts = [f"占比最大的类别是 {name}（{mib:.1f} MiB）", + f"最大单块 {biggest[0]:.2f} MiB 来自 {biggest[2]}"] + if frag > stats["reserved_mib"] * 0.1: + parts.append(f"碎片 {frag:.1f} MiB 超过 Reserved 的 10%，值得先看分配器配置") + return "；".join(parts) + + +def _compare_rows(report: dict) -> str: + """Build the cross-snapshot growth table, or say none was requested.""" + diff = report.get("compare") + if not diff: + return ("未比对（传 --compare <另一份快照> " + "可查跨 step 增长，用于区分泄露与一次性峰值）") + cells = [("Allocated 增量", f"{diff['allocated_delta_mib']:+.2f}", "两份快照之差"), + ("Reserved 增量", f"{diff['reserved_delta_mib']:+.2f}", "分配器保留之差")] + for name, delta in sorted(diff["categories"].items(), key=lambda kv: -abs(kv[1])): + cells.append((f"类别 {name}", f"{delta:+.2f}", + "持续增长即疑似泄露" if delta > 0 else "")) + return _rows(cells) + + +def render(report: dict, template: Path) -> str: + """Fill the HTML template from a parsed snapshot record.""" + html = template.read_text(encoding="utf-8") + stats = report["stats"] + for marker, value in ( + ("", _compare_rows(report)), + ("", str(report["path"])), + ("", f"{stats['segments']} 段 / {stats['block_count']} 活跃块"), + ("", _total_rows(report)), + ("", _category_rows(report)), + ("", _block_rows(report)), + ("", _check_rows(report)), + ("", _verdict(report))): + html = html.replace(marker, value) + return html + + +def _summarize_line(report: dict) -> str: + """One console line: the totals plus any failed invariant.""" + stats = report["stats"] + failed = [name for name, ok, _ in report["checks"] if not ok] + tail = f"; 校验未通过: {', '.join(failed)}" if failed else "; 校验全通过" + return (f"reserved={stats['reserved_mib']:.0f}MiB " + f"allocated={stats['allocated_mib']:.0f}MiB " + f"frag={stats['fragmentation_mib']:.0f}MiB{tail}") + + +def main(argv: list[str] | None = None) -> int: + """Parse the snapshot and write the memory-breakdown HTML report.""" + parser = argparse.ArgumentParser( + prog="python .agent/skills/memory-analysis/scripts/breakdown.py", + description="Break an NPU memory snapshot down into an HTML report.") + parser.add_argument("--snapshot", required=True) + parser.add_argument("--out", default="memory-breakdown.html") + parser.add_argument("--json", dest="json_out", default=None, + help="also write the parsed record as JSON") + parser.add_argument("--compare", default=None, + help="an earlier snapshot to diff against, to expose growth " + "across steps (a leak) rather than a one-off peak") + args = parser.parse_args(argv) + + path = Path(args.snapshot).resolve() + if not path.is_file(): + print(f"[breakdown] snapshot not found: {path}", flush=True) + return 1 + report = analyse(path) + if args.compare: + earlier = Path(args.compare).resolve() + if not earlier.is_file(): + print(f"[breakdown] --compare snapshot not found: {earlier}", flush=True) + return 1 + report["compare"] = compare(summarize(load_snapshot(earlier)), report["stats"]) + print(f"[breakdown] vs {earlier.name}: " + f"allocated {report['compare']['allocated_delta_mib']:+.2f}MiB", flush=True) + template = Path(__file__).resolve().parent.parent / "templates" / "memory-breakdown.html" + Path(args.out).write_text(render(report, template), encoding="utf-8") + if args.json_out: + payload = {"path": str(path), "stats": report["stats"], "sidecar": report["sidecar"], + "compare": report.get("compare"), + "checks": [{"check": n, "passed": ok, "detail": d} + for n, ok, d in report["checks"]]} + Path(args.json_out).write_text(json.dumps(payload, indent=2, ensure_ascii=False), + encoding="utf-8") + print(f"[breakdown] {_summarize_line(report)}; report={args.out}", flush=True) + return 0 if all(ok for _, ok, _ in report["checks"]) else 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.agent/skills/memory-analysis/scripts/collect_snapshot.py b/.agent/skills/memory-analysis/scripts/collect_snapshot.py new file mode 100644 index 000000000..ef9cbf45c --- /dev/null +++ b/.agent/skills/memory-analysis/scripts/collect_snapshot.py @@ -0,0 +1,113 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Drive a run with NPU memory-snapshot recording and verify what it produced. + +The record and dump calls belong in the training process, which owns the +device, so this script passes the window and output path through the +environment and `mem_hook` is what reads them. Collection is judged by the +snapshot files actually being there with their counter sidecars, not by the +run's exit code alone. + +Pure stdlib; the device-side snapshot API is never imported here. +""" +from __future__ import annotations + +import argparse +import os +import shlex +import subprocess +import sys +from pathlib import Path + +SIDECAR_SUFFIX = ".stats.json" + + +def _snapshot_env(args: argparse.Namespace, snapshot: Path) -> dict: + """Environment carrying the window and output path that `mem_hook` reads.""" + env = dict(os.environ) + env["MEM_SNAPSHOT_PATH"] = str(snapshot) + env["MEM_SNAPSHOT_STEPS"] = args.steps + env["MEM_SNAPSHOT_STACKS"] = args.stacks + return env + + +def check_output(snapshot: Path) -> tuple[list[Path], list[str]]: + """Find the dumps `mem_hook` wrote for this base path. + + `mem_hook.dump` tags each snapshot, so the files are siblings named + ``.`` with a ``.stats.json`` sidecar beside each. + + Args: + snapshot: The base path handed to the run. + + Returns: + ``(complete, problems)`` — dumps that have their counter sidecar, and a + message per dump that is missing one. + """ + found = sorted(snapshot.parent.glob(f"{snapshot.stem}.*{snapshot.suffix}")) + if not found: + return [], [f"no snapshot matching {snapshot.stem}.*{snapshot.suffix} " + f"in {snapshot.parent}; the entry must call mem_hook.start()/step()"] + complete, problems = [], [] + for path in found: + sidecar = path.with_suffix(path.suffix + SIDECAR_SUFFIX) + if sidecar.is_file(): + complete.append(path) + else: + problems.append(f"{path.name}: missing {sidecar.name}; the peak and the " + f"allocator totals cannot be cross-checked without it") + return complete, problems + + +def main(argv: list[str] | None = None) -> int: + """Launch the run with snapshot recording; verify and report the dumps.""" + parser = argparse.ArgumentParser( + prog="python .agent/skills/memory-analysis/scripts/collect_snapshot.py", + description="Drive a run with NPU memory-snapshot recording.") + parser.add_argument("--run-dir", required=True) + parser.add_argument("--launch", required=True, + help="run command (argv string, shlex-split)") + parser.add_argument("--steps", default="1-3", + help="inclusive step window to dump (read by mem_hook)") + parser.add_argument("--stacks", default="all", + help="stack detail to record; 'all' keeps the Python frames " + "that attribution needs") + args = parser.parse_args(argv) + + run_dir = Path(args.run_dir).resolve() + run_dir.mkdir(parents=True, exist_ok=True) + snapshot = run_dir / "memory_snapshot.pickle" + launch = shlex.split(args.launch) + + print(f"[collect_snapshot] steps={args.steps} stacks={args.stacks} " + f"base={snapshot}", flush=True) + print(f"[collect_snapshot] {' '.join(launch)}", flush=True) + done = subprocess.run(launch, cwd=str(run_dir), + env=_snapshot_env(args, snapshot), check=False) + + complete, problems = check_output(snapshot) + for problem in problems: + print(f"[collect_snapshot] incomplete: {problem}", flush=True) + status = "ok" if done.returncode == 0 and complete else "incomplete" + print(f"[collect_snapshot] exit={done.returncode} dumps={len(complete)} " + f"status={status}", flush=True) + for path in complete: + print(f"[collect_snapshot] next: python " + f"{Path(__file__).with_name('breakdown.py')} --snapshot {path}", flush=True) + return 0 if status == "ok" else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.agent/skills/memory-analysis/scripts/mem_hook.py b/.agent/skills/memory-analysis/scripts/mem_hook.py new file mode 100644 index 000000000..05770b690 --- /dev/null +++ b/.agent/skills/memory-analysis/scripts/mem_hook.py @@ -0,0 +1,152 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""In-run memory-snapshot hook: record history and dump at chosen moments. + +Imported by the training script, inside the training process, so ``torch_npu`` +is imported at module top as elsewhere in this repository. It is never imported +by ``collect_snapshot.py``, which runs on the launching host. + +``collect_snapshot.py`` exports ``MEM_SNAPSHOT_PATH`` / ``MEM_SNAPSHOT_STEPS`` +/ ``MEM_SNAPSHOT_STACKS``; this module is what reads them, so the launch side +and the run side agree on one window:: + + import sys + sys.path.insert(0, ".agent/skills/memory-analysis/scripts") + import mem_hook + + mem_hook.start() # enable history recording + for step, batch in enumerate(loader): + train_step(batch) + mem_hook.step() # dumps inside the configured window + +``dump(tag)`` can be called from anywhere at any time — including from an +autograd hook mid-backward — so a run that is going badly can be inspected +while it is still running instead of after it finishes. + +Each dump writes two files: the snapshot pickle, and a ``.stats.json`` sidecar +holding ``memory_reserved`` / ``memory_allocated`` / ``max_memory_allocated``. +Those counters live only in the process, and ``breakdown.py`` needs them to +cross-check its totals and to say whether the snapshot is the peak — a +post-step snapshot holds no activations and is normally well below the peak. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +import torch +import torch_npu + +DEFAULT_PATH = "memory_snapshot.pickle" +DEFAULT_STEPS = "1-3" + +_state: dict = {"step": 0, "window": None, "base": None, "dumps": []} + + +def _parse_window(spec: str) -> tuple[int, int]: + """Parse ``"a-b"`` or ``"n"`` into an inclusive step window.""" + text = (spec or "").strip() + if "-" in text: + first, _, last = text.partition("-") + return int(first), int(last) + value = int(text) + return value, value + + +def start(path: str | None = None, steps: str | None = None, + stacks: str | None = None) -> None: + """Begin recording allocation history. + + Args: + path: Snapshot path; defaults to ``$MEM_SNAPSHOT_PATH`` or + ``./memory_snapshot.pickle``. Each dump adds its tag to the stem. + steps: Inclusive step window for automatic dumps, ``"a-b"`` or ``"n"``; + defaults to ``$MEM_SNAPSHOT_STEPS`` or ``1-3``. + stacks: Passed through to the recorder; defaults to + ``$MEM_SNAPSHOT_STACKS`` or ``all``. Python frames are what makes + attribution possible, so lowering this makes the report coarser. + """ + base = Path(path or os.environ.get("MEM_SNAPSHOT_PATH") or DEFAULT_PATH).resolve() + base.parent.mkdir(parents=True, exist_ok=True) + _state["base"] = base + _state["step"] = 0 + _state["dumps"] = [] + _state["window"] = _parse_window(steps or os.environ.get("MEM_SNAPSHOT_STEPS") + or DEFAULT_STEPS) + torch_npu.npu.memory._record_memory_history( # pylint: disable=protected-access + enabled="all", context="all", + stacks=stacks or os.environ.get("MEM_SNAPSHOT_STACKS") or "all") + + +def stop() -> None: + """Stop recording allocation history.""" + torch_npu.npu.memory._record_memory_history(enabled=None) # pylint: disable=protected-access + + +def _counters(device: int) -> dict: + """The allocator counters that the snapshot itself does not carry.""" + return { + "device": device, + "memory_reserved": int(torch.npu.memory_reserved(device)), + "memory_allocated": int(torch.npu.memory_allocated(device)), + "max_memory_allocated": int(torch.npu.max_memory_allocated(device)), + "max_memory_reserved": int(torch.npu.max_memory_reserved(device)), + } + + +def dump(tag: str = "snapshot", device: int | None = None) -> Path: + """Write a snapshot and its counter sidecar, and return the snapshot path. + + Safe to call at any moment, including from an autograd hook while the graph + is alive, which is the only way to see activations in the composition. + + Args: + tag: Appended to the snapshot stem to keep dumps apart. + device: Device whose counters to record; defaults to the current one. + + Returns: + The snapshot file path. + """ + base = _state["base"] or Path( + os.environ.get("MEM_SNAPSHOT_PATH") or DEFAULT_PATH).resolve() + target = base.with_name(f"{base.stem}.{tag}{base.suffix}") + target.parent.mkdir(parents=True, exist_ok=True) + torch_npu.npu.memory._dump_snapshot(str(target)) # pylint: disable=protected-access + # `_dump_snapshot` creates the file write-only (mode 0200), so not even its + # owner can read it back. Restore owner read/write; the snapshot records + # source paths, so it stays private to the owner. + target.chmod(0o600) + index = torch.npu.current_device() if device is None else device + sidecar = target.with_suffix(target.suffix + ".stats.json") + sidecar.write_text(json.dumps({"tag": tag, **_counters(index)}, indent=2), + encoding="utf-8") + _state["dumps"].append(target) + print(f"[mem_hook] dumped {target} (+{sidecar.name})", flush=True) + return target + + +def step() -> None: + """Advance the step counter and dump when inside the configured window.""" + _state["step"] += 1 + current = _state["step"] + window = _state["window"] or _parse_window(DEFAULT_STEPS) + if window[0] <= current <= window[1]: + dump(f"step{current}") + + +def dumps() -> list[Path]: + """Snapshot paths written so far in this process.""" + return list(_state["dumps"]) diff --git a/.agent/skills/memory-analysis/templates/memory-breakdown.html b/.agent/skills/memory-analysis/templates/memory-breakdown.html new file mode 100644 index 000000000..a26d681f4 --- /dev/null +++ b/.agent/skills/memory-analysis/templates/memory-breakdown.html @@ -0,0 +1,72 @@ + + + + +显存拆解报告 + + + +

显存拆解报告

+

快照：（）

+

规则版本：memory breakdown-dimensions v2（按 torch_npu 2.9.0 实采字段解析）。 +单位 MiB。

+ +

一、总量与碎片

+ + + +
项MiB说明
+

快照是取样瞬间，不等于峰值。步后取的快照里没有任何激活； +实测即便在 backward 中途取样，活跃量仍低于 max_memory_allocated。 +峰值构成需要在峰值时刻取样，或用 device_traces 回放。

+ +

二、按类别

+ + + +
类别MiB占活跃量
+

类别由记录栈按规则判定：torch/optim 帧→优化器状态、 +模块 .to() 转换→参数、native 帧含 AccumulateGrad→梯度 +（梯度由 C++ autograd 引擎分配，没有任何 Python 帧，这是唯一可用信号）、 +最内层自有代码帧在 forward 内→激活。

+ +

三、最大的活跃块

+ + + +
MiB类别调用点
+

调用点取最内层不在 site-packages 内的 Python 帧——记录栈是 +最内层优先，而其最内层几帧是快照机制自身（unwind / +CapturedTraceback / ??），取 frames[0] 会让 +每个块都归到同一处。

+ +

四、跨快照增长（泄露 vs 一次性峰值）

+ + + +
项Δ MiB说明
+

单份快照分不清「泄露」与「一次性大分配」：某类别在连续 step 间持续 +增长才是泄露。采集默认 dump step 1–3，正好可两两比对。

+ +

五、一致性校验

+ + + +
校验结果明细
+

校验不通过时对应数字不可直接采信。退出码：0=全通过，2=有不通过项。

+ +

六、结论与下一步

+

+

拆解只定位占用来源；省显存的手段见 perf-playbook skill +（重计算/offload/分配器配置等），性能影响须同时量化。

+ + diff --git a/AGENTS.md b/AGENTS.md index 12cc4d5dd..cde0816d0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -148,6 +148,7 @@ Configured in `.agent/settings.json` (Claude Code–style `PostToolUse` matchers | **add-unit-test** | How-to for `tests/ut` (procedures) | when adding UT / coverage | | **hf-model-integration** | Integrate HF-native or HF-component custom models through model-owned adapters | model onboarding / functional bring-up | | **hf-model-precision-validation** | Validate module parity and distributed numerical self-consistency | model numerical acceptance / diagnosis | +| **memory-analysis** | NPU memory snapshot breakdown: static/dynamic, peak tensor composition, recompute saving, to HTML | 显存分析 / 显存拆解 / OOM 定位 | ### Commands diff --git a/tests/ut/tools/test_memory_breakdown.py b/tests/ut/tools/test_memory_breakdown.py new file mode 100644 index 000000000..b833e179e --- /dev/null +++ b/tests/ut/tools/test_memory_breakdown.py @@ -0,0 +1,339 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Unit tests for the memory breakdown parser (CPU, synthetic fixture). + +The fixture reproduces the frame shapes of a real +``torch_npu.npu.memory._snapshot()`` taken on torch_npu 2.9.0.post2: the +recorded stack is innermost-first and starts with the capture machinery, the +optimizer and parameter blocks carry Python frames, and the gradient blocks +carry none at all because the autograd engine allocates them from C++. +""" +from __future__ import annotations + +import json +import pickle +import shutil +import sys +import tempfile +import unittest +from pathlib import Path + + +def _find_scripts() -> Path | None: + """Locate the memory-analysis scripts directory, or None when absent. + + The scripts live under ``.agent/skills`` (agent tooling, not a product + module), so they are reached by path rather than by import. The directory + is searched for among this file's ancestors instead of at a fixed depth: a + checkout has it beside ``tests/``, while a test tree copied on its own - as + the gate does - has no ``.agent`` sibling at all. + """ + for ancestor in Path(__file__).resolve().parents: + candidate = ancestor / ".agent" / "skills" / "memory-analysis" / "scripts" + if (candidate / "breakdown.py").is_file(): + return candidate + return None + + +_SCRIPTS = _find_scripts() +if _SCRIPTS is None: + raise unittest.SkipTest( + "memory-analysis scripts (.agent/skills/memory-analysis/scripts) are " + "not reachable from this test tree; they are agent tooling and are not " + "copied into a standalone test checkout" + ) +# Ahead of everything else on the path so the script's own sibling imports +# resolve here and not to an unrelated module. +sys.path.insert(0, str(_SCRIPTS)) + +import breakdown # noqa: E402 # pylint: disable=wrong-import-position +import collect_snapshot # noqa: E402 # pylint: disable=wrong-import-position + +MIB = 1024 * 1024 +SITE = "/opt/py/lib/python3.10/site-packages" + +# Every recorded stack starts with the snapshot machinery; `frames[0]` is the +# same for every block and says nothing about the allocation. +_NOISE = [ + {"name": "torch_npu::unwind::unwind()", "filename": "libtorch_npu.so", "line": 0}, + {"name": "torch_npu::CapturedTraceback::gather(bool, bool, bool)", + "filename": "libtorch_npu.so", "line": 0}, + {"name": "??", "filename": "libtorch_npu.so", "line": 0}, +] +# Allocator plumbing every device allocation goes through. +_PLUMBING = [ + {"name": "at::_ops::empty_strided::redispatch(c10::DispatchKeySet)", + "filename": "libtorch_cpu.so", "line": 0}, + {"name": "c10::Allocator::allocate(unsigned long)", "filename": "libc10.so", "line": 0}, +] + + +def _param_frames() -> list[dict]: + """A parameter block: module `.to()` conversion, then the caller.""" + return _NOISE + _PLUMBING + [ + {"name": "convert", "filename": f"{SITE}/torch_npu/utils/_module.py", "line": 76}, + {"name": "_apply", "filename": f"{SITE}/torch/nn/modules/module.py", "line": 957}, + {"name": "main", "filename": "/home/u/train.py", "line": 31}, + ] + + +def _optimizer_frames() -> list[dict]: + """An optimizer-state block: AdamW allocates through torch/optim/adam.py.""" + return _NOISE + _PLUMBING + [ + {"name": "_init_group", "filename": f"{SITE}/torch/optim/adam.py", "line": 177}, + {"name": "step", "filename": f"{SITE}/torch/optim/optimizer.py", "line": 485}, + {"name": "main", "filename": "/home/u/train.py", "line": 41}, + ] + + +def _gradient_frames() -> list[dict]: + """A gradient block: no Python frame at all, only the autograd engine.""" + return _NOISE + _PLUMBING + [ + {"name": "torch::autograd::AccumulateGrad::apply(std::vector&&)", + "filename": "libtorch_cpu.so", "line": 0}, + {"name": "torch::autograd::Engine::thread_main(std::shared_ptr const&)", + "filename": "libtorch_cpu.so", "line": 0}, + ] + + +def _activation_frames() -> list[dict]: + """An activation block: the innermost own-code frame is a forward.""" + return _NOISE + _PLUMBING + [ + {"name": "gelu", "filename": f"{SITE}/torch/nn/functional.py", "line": 2060}, + {"name": "forward", "filename": "/home/u/train.py", "line": 22}, + {"name": "main", "filename": "/home/u/train.py", "line": 44}, + ] + + +def _segment(blocks: list[tuple[int, str, list]], pad: int = 0) -> dict: + """Build one segment whose fields agree with its blocks.""" + entries = [{"address": 0, "size": size, "requested_size": size, + "state": state, "frames": frames} for size, state, frames in blocks] + if pad: + entries.append({"address": 0, "size": pad, "requested_size": 0, + "state": "inactive", "frames": []}) + live = sum(e["size"] for e in entries if e["state"].startswith("active")) + return {"device": 0, "segment_type": "large", "blocks": entries, + "total_size": sum(e["size"] for e in entries), + "allocated_size": live, "active_size": live, "frames": []} + + +def _snapshot() -> dict: + """Two segments covering every category the rules recognise.""" + return {"segments": [ + _segment([(8 * MIB, "active_allocated", _param_frames()), + (8 * MIB, "active_allocated", _optimizer_frames())], pad=2 * MIB), + _segment([(8 * MIB, "active_allocated", _gradient_frames()), + (4 * MIB, "active_allocated", _activation_frames())]), + ], "device_traces": []} + + +class MemoryBreakdownTest(unittest.TestCase): + """Parse a synthetic snapshot shaped like the real one.""" + + def setUp(self) -> None: + """Write the snapshot and its counter sidecar to a temp dir.""" + self.root = Path(tempfile.mkdtemp()) + self.snapshot = self.root / "memory_snapshot.step2.pickle" + with self.snapshot.open("wb") as handle: + pickle.dump(_snapshot(), handle) + self._write_sidecar(28 * MIB) + + def tearDown(self) -> None: + """Remove the temp tree.""" + shutil.rmtree(self.root, ignore_errors=True) + + def _write_sidecar(self, peak: int) -> None: + """(Re)write the allocator counters beside the snapshot.""" + sidecar = self.snapshot.with_suffix(self.snapshot.suffix + ".stats.json") + sidecar.write_text(json.dumps({ + "device": 0, "memory_reserved": 30 * MIB, "memory_allocated": 28 * MIB, + "max_memory_allocated": peak, "max_memory_reserved": 30 * MIB, + }), encoding="utf-8") + + def test_call_site_skips_the_capture_machinery(self): + """`frames[0]` is the snapshot machinery and must never be the answer.""" + site = breakdown.call_site(_param_frames()) + self.assertNotIn("unwind", site) + self.assertEqual(site, "train.py:31 main") + + def test_call_site_prefers_own_code_over_framework(self): + """The innermost frame outside site-packages wins.""" + self.assertEqual(breakdown.call_site(_optimizer_frames()), "train.py:41 main") + + def test_call_site_skips_allocator_plumbing_for_native_blocks(self): + """With no Python frame, the first non-`at::`/`c10::` symbol is used.""" + self.assertEqual(breakdown.call_site(_gradient_frames()), + "torch::autograd::AccumulateGrad::apply") + + def test_categories_match_each_rule(self): + """Each frame shape lands in its own category.""" + self.assertEqual(breakdown.categorize(_param_frames()), "参数") + self.assertEqual(breakdown.categorize(_optimizer_frames()), "优化器状态") + self.assertEqual(breakdown.categorize(_gradient_frames()), "梯度") + self.assertEqual(breakdown.categorize(_activation_frames()), "激活") + + def test_gradient_is_recognised_without_any_python_frame(self): + """Gradients carry no Python frame, so the native symbol is the signal.""" + self.assertEqual(breakdown._python_frames(_gradient_frames()), []) + self.assertEqual(breakdown.categorize(_gradient_frames()), "梯度") + + def test_totals_come_from_the_segment_fields(self): + """Reserved / allocated / fragmentation use the recorded segment sizes.""" + stats = breakdown.summarize(_snapshot()) + self.assertAlmostEqual(stats["reserved_mib"], 30.0) + self.assertAlmostEqual(stats["allocated_mib"], 28.0) + self.assertAlmostEqual(stats["fragmentation_mib"], 2.0) + self.assertEqual(stats["block_count"], 4) + + def test_category_totals_add_up_to_the_allocated_bytes(self): + """Every active block lands in exactly one category.""" + stats = breakdown.summarize(_snapshot()) + self.assertAlmostEqual(sum(stats["categories"].values()), stats["allocated_mib"]) + self.assertAlmostEqual(stats["categories"]["优化器状态"], 8.0) + self.assertAlmostEqual(stats["categories"]["激活"], 4.0) + + def test_checks_pass_when_the_snapshot_is_at_the_peak(self): + """A snapshot whose total equals the recorded peak satisfies every check.""" + report = breakdown.analyse(self.snapshot) + self.assertTrue(all(ok for _, ok, _ in report["checks"]), + [c for c in report["checks"] if not c[1]]) + + def test_peak_check_fails_when_the_snapshot_is_below_the_peak(self): + """A snapshot taken off-peak must say so rather than pass quietly.""" + self._write_sidecar(40 * MIB) + failed = [name for name, ok, _ in breakdown.analyse(self.snapshot)["checks"] if not ok] + self.assertTrue(any("峰值" in name for name in failed)) + + def test_peak_is_unknown_without_the_sidecar(self): + """With no counters recorded, the peak question is reported unanswered.""" + self.snapshot.with_suffix(self.snapshot.suffix + ".stats.json").unlink() + checks = breakdown.analyse(self.snapshot)["checks"] + peak = [c for c in checks if "峰值" in c[0]][0] + self.assertFalse(peak[1]) + self.assertIn("无 max_memory_allocated", peak[2]) + + def test_segment_identity_catches_a_corrupt_segment(self): + """A segment whose blocks do not tile it is reported, not averaged over.""" + snap = _snapshot() + snap["segments"][0]["total_size"] += 1234 + name, ok, detail = breakdown._segment_identity(snap) + self.assertFalse(ok) + self.assertIn("seg0", detail) + self.assertIn("total_size", name) + + def test_unreadable_snapshot_explains_the_write_only_mode(self): + """torch_npu writes snapshots mode 0200; the error must say how to fix it.""" + if sys.platform.startswith("win"): + self.skipTest("POSIX file modes are not enforced the same way on Windows") + self.snapshot.chmod(0o200) + with self.assertRaises(PermissionError) as caught: + breakdown.load_snapshot(self.snapshot) + self.assertIn("chmod 600", str(caught.exception)) + + def test_main_writes_the_report_and_the_json_record(self): + """The CLI writes both outputs and exits 0 when the checks pass.""" + out_html, out_json = self.root / "r.html", self.root / "r.json" + code = breakdown.main(["--snapshot", str(self.snapshot), "--out", str(out_html), + "--json", str(out_json)]) + self.assertEqual(code, 0) + record = json.loads(out_json.read_text(encoding="utf-8")) + self.assertAlmostEqual(record["stats"]["allocated_mib"], 28.0) + + def test_main_exits_nonzero_when_an_invariant_fails(self): + """An off-peak snapshot must not be reported as a clean parse.""" + self._write_sidecar(40 * MIB) + self.assertEqual(breakdown.main( + ["--snapshot", str(self.snapshot), "--out", str(self.root / "r.html")]), 2) + + def test_main_reports_a_missing_snapshot(self): + """A path that does not exist exits 1.""" + self.assertEqual(breakdown.main( + ["--snapshot", str(self.root / "nope.pickle"), + "--out", str(self.root / "r.html")]), 1) + + def test_render_leaves_no_unfilled_markers(self): + """Every template marker is replaced, with and without a comparison.""" + template = _SCRIPTS.parent / "templates" / "memory-breakdown.html" + report = breakdown.analyse(self.snapshot) + self.assertNotIn("