analyze-memory-snapshot

分析内存快照

通过检查分配器状态和事件历史，分析 PyTorch CUDA 内存快照。始终优先依据快照中的证据，而非通用内存建议。

工作流

1. 如果用户未提供快照路径，请询问路径。接受 .pickle、.pkl、.json 文件，或包含各 rank 本地快照的目录。如果用户只有来自 memory_viz 的截图，请索要底层快照文件。

2. 在本地加载快照。PyTorch _dump_snapshot() 文件是 pickle 文件，因此只能对用户提供的文件或已存在于工作区中的文件进行反序列化。JSON 快照可以直接加载。

3. 如果可用，运行此技能的辅助工具进行初步汇总：

   python scripts/analyze_snapshot.py /path/to/snapshot.pickle
   

   如果有多个 rank，请对每个 rank 分别运行，并比较峰值、活动字节数、非活动字节数、OOM 事件和主要堆栈跟踪。

   如果用户有（或可以获取）同一运行中不同时间点的两份快照，请对它们进行差异比较——这样可以直接查出哪些分配位置有所增长，从而无需仅凭解读来排查泄漏：

   python scripts/analyze_snapshot.py step100.pickle --diff step200.pickle
   

   对于非常大的快照（包含数亿字节的跟踪历史），辅助工具会将跟踪分析限制在 --max-trace-events（默认值为 1M），并报告丢弃了多少条旧事件，以便分析无需加载全部内容也能顺利进行。

4. 如果辅助工具的输出不够充分，请读取原始快照。检查：

   • segments：由 cudaMalloc 返回的分配器段
   • blocks：每个段中存活、等待释放和非活动的内存块
   • device_traces：启用内存历史记录时按时间顺序排列的分配事件
   • frames：分配位置的堆栈跟踪

5. 根据分配器语义报告内存情况：

   • 已保留内存：各段 total_size 之和
   • 存活张量内存：处于 active_allocated 状态的内存块
   • 延迟释放：处于 active_awaiting_free 状态的内存块，通常与流有关
   • 非活动缓存内存：PyTorch 可以复用的 inactive 内存块
   • 请求大小与取整后大小：requested_size 与内存块 size 的对比

需要诊断的内容

内存的去向

根据最有用的用户代码帧，将存活的 active_allocated 内存块分组。首先报告占用最大的堆栈跟踪，并提供字节数、占存活内存的百分比、内存块数量，以及有代表性的文件、函数和行号。如果堆栈名称能够明确区分，请分别列出参数、优化器状态、激活值、临时缓冲区、通信缓冲区和用户容器。

始终将内存归因于用户的模型代码，而非分配器或框架内部。快照堆栈通常主要由 torch/nn/modules/*、torch/autograd/* 和 torch/optim/* 中的帧构成；跳过这些帧，找到用户自己的文件中的第一个帧，并指出占用内存的模块或代码行（例如，“38% 的存活内存是 model.decoder.layers[*].mlp 的激活值，分配位置为 modeling.py:214”）。只有当堆栈中不存在用户代码帧时，才回退到 torch 内部帧。

内存何时被释放

使用 device_traces 区分 free_requested 和 free_completed。待处理或延迟的释放可能意味着其他流仍需要这些张量，或相关工作尚未执行到释放点。在多步跟踪中，将重复出现的分配位置与其对应的释放事件进行比较：每一步都出现、却没有对应释放事件的分配位置，比一次性的模型或优化器分配更能说明存在内存滞留。如果跟踪中包含 OOM 事件，请报告分配失败的大小、device_free、OOM 堆栈帧，以及 OOM 发生前哪些存活的分配位置占用了内存。

泄漏与滞留模式 查找随步骤增长的活跃字节数、来自同一调用栈且未释放的重复活跃分配，或大量张量不断累积、随后一起释放的锯齿形模式。最有力的证据是双快照差异（辅助工具中的 --diff）：如果间隔 k 个步骤拍摄的快照之间，某个分配位置的存活字节数和块数都在增长，那么它正在泄漏或被保留。

根据快照特征识别以下已命名模式，并给出相应修复方法：

• 保留的 autograd 图：第 N 步的活跃激活块在第 N+k 步时仍然存在，通常还会伴随每步累积一个损失或输出张量。典型原因是记录或累积仍持有计算图的张量——追加前使用 loss.detach()（或 .item()），或者查找导致计算图持续存活的 retain_graph=True。
• 不断增长的 Python 容器：同一用户代码分配位置的块数随步骤线性增长。修复方法是先分离张量并移至 CPU，再追加；或限制容器大小。
• 优化器状态在第一步激增：首次调用 optimizer.step() 时出现大量归因于 torch/optim/* 的分配，之后保持稳定。这属于预期保留，并非泄漏——务必明确说明；此处 OOM 的解决方法是使用更小的模型或批次、占用更少内存的优化器，或进行分片，而不是排查泄漏。
• 激活内存与检查点：活跃内存主要由前向传播激活占用，这些激活都会在反向传播时释放（device_traces 中呈锯齿形）。建议使用激活检查点（torch.utils.checkpoint），并根据快照量化可释放的激活字节数与重新计算成本之间的权衡。
• DDP 与 FSDP 下的梯度内存：每个 rank 上都有完整大小的梯度和参数块，说明 DDP 正在复制数据；如果用户受内存限制，且梯度、参数和优化器状态占据了大部分活跃内存，请指出 FSDP/ZeRO 式分片是结构性解决方案。
• Python 引用环：只有在垃圾回收运行时才会批量释放。建议使用 from torch.utils.viz._cycles import warn_tensor_cycles，并在初始化期间调用 warn_tensor_cycles()，以查找导致 CUDA 张量持续存活的对象引用环。

碎片化与分配器行为

不要把 nvidia-smi 显示的用量等同于活跃张量内存。PyTorch 使用缓存型 CUDA 分配器，因此保留内存可能远大于已分配的张量内存。大型非活跃块可供 PyTorch 重用；empty_cache() 可能会将完全非活跃的内存段归还给 CUDA，但不会释放仍被活跃张量占用的内存。

最常见的易混淆情况是预留空间充足但仍发生 OOM：失败的请求大小超过了最大的单个连续非活跃块，即使非活跃内存总量更大。请根据快照进行诊断：将 OOM 请求大小与最大非活跃块进行比较（辅助工具会报告 "largest single reusable inactive block"），并检查非活跃内存是否分散在许多内存段中。符合这种模式时，按以下顺序提出建议：

1. PYTORCH_ALLOC_CONF=expandable_segments:True——让内存段可以扩展，而不是分配新的固定大小内存段；这通常是解决碎片化 OOM 最有效的方法。
2. PYTORCH_ALLOC_CONF=max_split_size_mb:<N>——阻止分配器拆分大于 N MiB 的块，减少难以重用的碎片；建议 N 接近失败请求的大小。
3. 减少用户代码中的分配大小变化（批次或序列大小变化是常见原因；填充至固定形状会有所帮助）。

如果碎片化并非问题，请明确说明：如果仅活跃张量内存就已接近设备容量，调整分配器参数也无济于事，解决方法应着眼于模型层面（批次大小、检查点、分片、精度）。

非 PyTorch 内存 快照只能看到由 PyTorch 分配器管理的 CUDA 内存。NCCL、自定义 CUDA 扩展、cuDNN/cuBLAS 工作区或直接 CUDA 分配可能不会显示。如果怀疑快照之外的设备内存，请将 torch.cuda.device_memory_used(device) 或 nvidia-smi 的结果与快照的预留字节数进行比较。

如果用户需要捕获快照

给出最简捕获模式：

import torch

torch.cuda.memory._record_memory_history(max_entries=100000)
try:
    run_your_code()
    torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
finally:
    torch.cuda.memory._record_memory_history(enabled=None)

对于运行时间较长的任务，请使用较大的 max_entries，但要提醒用户，跟踪历史记录会带来内存开销。

要调试 OOM，请注册 OOM 观察器，使快照在故障发生时自动转储——用户不应该需要重现两次崩溃：

import torch


def dump_on_oom(device, alloc, device_alloc, device_free):
    torch.cuda.memory._dump_snapshot(f"oom_snapshot_dev{device}.pickle")


torch.cuda.memory._record_memory_history(max_entries=100000)
torch._C._cuda_attach_out_of_memory_observer(dump_on_oom)

如果怀疑训练步骤之间存在泄漏，请捕获一对快照——一份较早的快照（在第一个优化器步骤之后，此时优化器状态已分配），以及一份在许多步骤之后捕获的快照——然后使用 analyze_snapshot.py earlier.pickle --diff later.pickle 对它们进行差异比较。在步骤中的相同位置（例如 optimizer.step() 之后立即）转储，可避免差异受到步骤内噪声的影响。

输出格式

先用一句话说明主要发现。然后提供：

1. 总量：预留、实时已分配、活动但等待释放、非活动缓存，以及在跟踪历史记录允许的情况下提供峰值。
2. 按实时内存占用列出排名靠前的分配位置，并提供堆栈帧和百分比。
3. 释放时间线：已完成的释放、待处理的释放、延迟释放和 OOM 事件。
4. 诊断：指出匹配的模式（保留计算图、容器增长、优化器状态峰值、激活值压力、碎片化／缓存、引用循环或非 PyTorch 内存差额），并说明哪些快照证据支持该判断。明确区分预期保留与泄漏。
5. 具体后续步骤：需要修改的用户代码位置（快照堆栈中的文件:行号）、针对匹配模式的具体修复方案、怀疑存在碎片化时可尝试的分配器设置，以及如果需要通过差异比较确认诊断，应捕获的第二份快照。

明确说明不确定之处。如果快照缺少历史记录，请说明无法根据该快照判断哪些释放和峰值问题。