Repository navigation
Conversation
…apshots 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).
weich97
requested review from
liuchongming74,
suteng,
yangzhenzhang02 and
yaoyifanx
as code owners
October 11, 2026 15:13
Contributor
Author
|
/check-pr |
Contributor
Author
|
/retest |
|
Contributor
Author
|
/retest |
1 similar comment
Contributor
Author
|
/retest |
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
这次改了什么
本 PR 原先提交的是 memory-analysis skill 的骨架。在真机上实采快照后发现骨架的
归属逻辑是坏的：它取
frames[0]作为分配来源，而记录栈最内层优先、最内几帧恰是快照机制自身（
torch_npu::unwind::unwind/CapturedTraceback::gather/??），于是每一个块都会被归到同一个
unwind()，归属完全无效。骨架未合入，故原地重做并强推覆盖（原骨架提交
3f002da备份在本地）。现在字段、帧形态、恒等式全部对着实采数据写，分类规则对着一次已知构成的运行
逐项对账。
实采验证
torch_npu 2.9.0.post2 / CANN 9.0.0 / Ascend 910。全链路
collect_snapshot → mem_hook → breakdown真跑贯通，一次运行产出 3 份 dump（step2 / step3 / backward 中途）。
分类逐项对账（地面真值已知）
模型：3 ×
Linear(2048, 2048, bias=False)，bf16；AdamW；输入 512×2048 bf16。allocated_size112.01 MiB一致性校验（五项，前四项实测零偏差）
total_size、Σ活跃块 =allocated_sizeallocated_sizetotal_size=memory_reserved()allocated_size=memory_allocated()快照不是峰值（同一次运行实测）
optimizer.step()之后max_memory_allocated()步后快照里没有任何激活，比真实峰值低 30 MiB；即便在 backward 中途取样仍低
18 MiB。所以：
mem_hook.dump(tag)支持任意时刻调用（含从 autograd hook 内），每份 dump 旁写
max_memory_allocated，报告把"快照是否为峰值"列为校验项、差距超 2% 即标红并退出码 2。不标出来就会被当成峰值构成读。
跨快照比对（区分泄露与一次性分配）
单份快照分不清"泄露"与"一次性大分配"，
--compare正是为此。两个方向都验过。归属规则（实测得出，容易写错的地方）
frames[0]不是归属。 记录栈最内层优先，最内几帧是快照机制自身；第一个
.py帧实测在第 26/62 位，位置不固定，必须扫不能索引。归属取最内层不在 site-packages 内的
.py帧 → 调用方自己的代码行。torch::autograd::AccumulateGrad。按 Python 帧分类会把梯度整类漏掉。at::/c10::分配器管道，否则每个梯度块都报at::_ops::empty_strided::redispatch。跳过后拿到torch::autograd::AccumulateGrad::apply，才有意义。/optim/：AdamW 走torch/optim/adam.py，不是 adamw.py。total_size/allocated_size/active_size），重新累加的值只用于校验。附带修掉的三个问题
_dump_snapshot写出的文件权限是--w-------（0200，只写），连属主自己都读不回（真机上直接 PermissionError）。
mem_hook落盘后补chmod 0600，解析器遇到非 mem_hook 产出的 dump 会给出可操作提示而非裸 traceback。
collect_snapshot.py导出的MEM_SNAPSHOT_STEPS/MEM_SNAPSHOT_PATH仓库里没有任何消费者，是死接口；且仅凭退出码判成功。现改为逐份校验 dump
与计数器旁车，新增的
mem_hook就是读这些变量的那一端。AclrtGetMemUsageInfo failed, ret:207000的 ERROR 日志，但快照正常产出 —— 已记入 references，避免据此误判采集失败。
SKILL.md 据实收敛
骨架的 SKILL.md 宣称了 Module/Layer 粒度汇总、重计算收益估算、峰值时刻构成，
这些当前并未实现。已在 SKILL.md 里单列「尚未实现」一节如实标注，不当成已有
能力；
device_traces回放峰值构成的口径标[待对齐]。测试
tests/ut/tools/test_memory_breakdown.py，21 条 + 1 条按平台跳过，CPU only，fixture 复刻真实帧形态（含噪声帧、分配器管道帧、无 Python 帧的梯度块）。覆盖归属
优先级、管道跳过、四条类别规则、段恒等式、峰值校验三种状态、比对两个方向、
不可读 dump 的提示、三种退出码。
按 #1072 的教训，测试向上搜索定位
.agent/skills/memory-analysis/scripts，找不到则unittest.SkipTest——CI 会把tests/单独拷出，固定深度路径假设会导致 collectionerror 中断整个 UT 会话。
pylint 干净（仅剩
C9004 sys-path-mutation，与已过门禁的 #1072 同形），lizard 最大 CCN 10 / NLOC 30，
check_agents_catalog通过。