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("