PyTorch显存可视化与Snapshot数据分析作者：吃果冻不吐果冻皮来源：吃果冻不吐果冻皮 ####**【点击】加入大模型技术交流群** 原文：https://zhuanlan.zhihu.com/p/677203832 显存优化和显存溢出(OOM)分析是调参过程中常见的两个问题，解决显存不

PyTorch显存可视化与Snapshot数据分析

By AiBard123
January 31, 2024 - 2 min read

作者：吃果冻不吐果冻皮来源：吃果冻不吐果冻皮

原文：https://zhuanlan.zhihu.com/p/677203832

显存优化和显存溢出(OOM)分析是调参过程中常见的两个问题，解决显存不足的问题一般而言需要分析显存消耗占比，对显存消耗较大的操作进行参数调优。显存分析的途径包括静态计算和运行时分析，静态计算（仿真）通过公式来预估显存消耗，这种方式有一定的误差 ；对于运行时分析一般是日志打印，整体理解性较弱。

在PyTorch2.1中推出了一个显存的snapshot功能，可以将显存消耗可视化，特点是查看简单、更易理解。本文简单讲解该特性的使用方式，并例举了几个数据的分析。一个transformer模型显存可视化后如下所示：

本文主要内容：

SnapshotAPI的使用
Snapshot数据分析
Profiler中显存可视化使用
Profiler数据分析（可读性较好）

文中示例代码：https://github.com/CalvinXKY/BasicCUDA/tree/master/pytorch/torch_mem_snapshot

SnapshotAPI的使用

Snapshot的工作原理：开启API后，torch会自动记录c10代码中CUDA allocator的显存消耗，显存的Python/C++跟踪调用堆栈、记录调用过程中的timeline。最后将这些数据保持下来生成pickle文件，用于可视化。

1.1 API调用方式：

调用步骤：

开始（训练/推理前）: torch.cuda.memory._record_memory_history(max_entries=80000)
保存（迭代结束后）: torch.cuda.memory._dump_snapshot(file_name)
停止（分析完成）: torch.cuda.memory._record_memory_history(enabled=None)

_record_memory_history(enabled=‘all’,context=‘all’,stacks=‘all’,max_entries=1,device=None)的参数解释：

max_entries：最多使用多少个alloc/free events来记录内存开销，内存的操作需要用events来记录。当记录数据溢出时，系统只保留最后的max_entries个events量的数据。参数设置过小保存数据会不足，设置过大可能会影响运行。
context：选择需要跟踪的数据类型[None,“state”,“alloc”,“all”]。“state"是指记录当前使用内存情况，“alloc"是指通过alloc调用过的内存跟踪（如果不会设置按默认值即可），all缺省。
stacks: [“python”,“all”], 记录python层的调用堆栈，或者增加c++的。默认"all” 表示C++调用堆栈也记录。
enabled：开关值。[None,“state”,“all”] “state”,“all” 含义见context。

相关Torch API帮助文档：https://pytorch.org/docs/main/torch_cuda_memory.html#understanding-cuda-memory-usage

函数编写的关键步骤示例：

def train(*args):
 # 建立模型...
  for data,label in range(iteration_data):
      y = model(data)
      loss(y, label).backward()
      op.step() 
  # 其它操作...
  

# 开启记录，并设置最多记录100000个数据点
torch.cuda.memory._record_memory_history(max_entries=100000)
  

# 训练调用
train(args)
  

# 保存数据
torch.cuda.memory._dump_snapshot(your_file_name.pickle)
  

# 停掉记录，关闭snapshot
torch.cuda.memory._record_memory_history(enabled=None)

可视化操作：

保存成功后会生成一个文件，如：“your_file_name.pickle” ，

方式一：将pickle文件拖拽到浏览器（如chrome）网址：https://pytorch.org/memory_viz 中。

方式二：将"your_file_name.pickle"转化为html，直接点击打开。转化方式：找到pytorch的安装包位置，里面有个torch/cuda/_memory_viz.py ，用它进行转换。如下给了一个示例：

python /home/user/anaconda3/envs/py3.9/lib/python3.9/site-packages/torch/cuda/_memory_viz.py trace_plot your_file_name.pickle -o mem_snapshot.html

1.2 调用示例：

用torch原生API构造一个Transformer训练，完整的示例代码如下（git位置：https://github.com/CalvinXKY/BasicCUDA/blob/master/pytorch/torch_mem_snapshot/transformer_snapshot.py）：

# Author: kevin.xie  zhihu@kaiyuan
  

import torch
from torch import nn
from datetime import datetime
  

  

def train(num_iter=5, device="cuda:0"):
    model = nn.Transformer(d_model=512, nhead=2, num_encoder_layers=2, num_decoder_layers=2).to(device=device)
    x = torch.randn(size=(1, 1024, 512), device=device)
    tgt = torch.rand(size=(1, 1024, 512), device=device)
    model.train()
    labels = torch.rand_like(model(x, tgt))
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters())
    for _ in range(num_iter):
        y = model(x, tgt)
        loss = criterion(y, labels)
        loss.backward()
        print(loss.item())
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
  

  

def run():
    # Start recording memory snapshot history
    torch.cuda.memory._record_memory_history(max_entries=100000)
    # training running:
    train()
  

    timestamp = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
    file_name = f"visual_mem_{timestamp}.pickle"
    # save record:
    torch.cuda.memory._dump_snapshot(file_name)
  

    # Stop recording memory snapshot history:
    torch.cuda.memory._record_memory_history(enabled=None)
  

  

if __name__ == "__main__":
    run()

注意：

代码运行环境中PyTorch>=2.1 。
可视化API开启的系统开销比较高，正常运行时建议关闭。

运行后输出可视化后的结果：

Snapshot数据分析

2.1 激活内存数据

网页中选择“Active Memory Timeline”下拉框可看到激活内存数，这部分数据主要是记录tensor在计算过程中占用的内存数值以及其存活的周期；同时，能够查看每个tensor的内存消耗调用堆栈（Python/C++）。

整体数据分析：

如图是1.2运行中内存消耗情况的数据，一共有5个迭代。

第一个“三角”图形是“ labels = torch.rand_like(model(x, tgt))”操作产生的，后面的5个山峰形状图形是正反向运算时产生的内存变化。

第一个迭代的内存消耗小于后续迭代时的显存消耗，原因是：优化器的显存占用发生在第一个迭代结束后。

图中前向计算的内存消耗比后向计算的内存消耗更多；

局部数据分析：

波纹变化：图表中能够看到“数据条”有显示变化。首先需要明确单个tensor数据消耗显存在生命周期内是保持不变的。通过放大图片可以看到显存条纹有上/下行“波纹”，这种波纹形状的产生原因是因为其它的tensor被创建/释放了导致的显示绘图的变化，并非tensor占用的显存大小发生了改变。

尖峰值：后向计算时autograd会产生额外的内存消耗，能够形成一些尖峰值。可以点击一个尖峰值，其内容显示如下，可以看到触发尖峰的操作API源自autograd：

autograd_not_implemented_fallback.cpp:0:torch::autograd::autogradNotImplementedFallbackImpl(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::ivalue, std::allocator>*)
:0:c10::impl::BoxedKernelWrapper(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, double, std::array, bool, c10::optional), void>::call(c10::BoxedKernel const&, c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, double, std::array, bool, c10::optional)
??:0:at::_ops::_scaled_dot_product_efficient_attention_backward::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, double, std::array, bool, c10::optional)
??:0:torch::autograd::generated::ScaledDotProductEfficientAttentionBackward0::apply(std::vector<at::tensor, std::allocator>&&)
:0:torch::autograd::Node::operator()(std::vector<at::tensor, std::allocator>&&)
??:0:torch::autograd::Engine::evaluate_function(std::shared_ptr&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::shared_ptrconst&)
??:0:torch::autograd::Engine::thread_main(std::shared_ptrconst&)
??:0:torch::autograd::Engine::thread_init(int, std::shared_ptrconst&, bool)
??:0:torch::autograd::python::PythonEngine::thread_init(int, std::shared_ptrconst&, bool)
thread.cc:0:execute_native_thread_routine

2.2 内存块的使用与释放

当把选项切换到“Allocator State Hisotry”时，可以看到如下的可视化数据如下。该图展示从CUDA里面申请的内存块是如何被alloc分配成小的block，以及这些小的block什么时候被释放掉了。如果把光标放置到这些彩色的block上面，能够跟踪到是什么操作在使用该显存地址。

右侧框框表示segment（分片），在segment块中的彩色条表示block。

torch显存管理创建的顺序是先创建segment，然后创建block

block创建的堆栈查看：光标放到一个allc位置，显示内容如下:

(279.8MiB (293443588 bytes) allocated / 554.0MiB (580911104 bytes) reserved)
alloc b7f6a235dd800_2 6.0KiB (6144 bytes)
CUDACachingAllocator.cpp:0:c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::malloc(int, unsigned long, CUstream_st*)
:0:c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::malloc(void**, int, unsigned long, CUstream_st*)
:0:c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::allocate(unsigned long) const

主要看"alloc b7f6a235dd800_2 6.0KiB (6144 bytes)" 这个部分，表示使用了一个6KB大小的block块。其中b7f6a235dd800表示地址，_2表示该地址被使用的次数，整个数字标记是唯一的，目的是使其能够与tensor对应上。

block释放的堆栈查看：光标放到放到一个free位置，显示内容如下：

(304.9MiB (319707140 bytes) allocated / 554.0MiB (580911104 bytes) reserved)
free b7f6a235eb000_3 2.0KiB (2048 bytes)
CUDACachingAllocator.cpp:0:c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::free(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::Block*)
:0:c10::cuda::CUDACachingAllocator::Native::local_raw_delete(void*)
:0:c10::StorageImpl::~StorageImpl()

其中free表示“释放”掉了多少内存。

注意，block的释放仅从torch显存管理里面释放，而非cuda free操作，要发出cuda free得是segment的释放。

当一个segment的大小不能满足block需求时，系统会新申请一个segment块。这样的操作可能会产生显存碎片，下面给了一个显存碎片产生的示例。当创建一个约2M大小的tensor时，系统申请了一个20M的segment(s7faef4000000_0)，接着需要创建一个大小约24M的tensor时，系统重新申请了一个segment(s7faef2000000_0)。这样segment(s7faef4000000_0)有18M空间暂时时未被使用的，形成了碎片(fragment)。

复现代码：https://github.com/CalvinXKY/BasicCUDA/blob/master/pytorch/torch_mem_snapshot/block_fragment.py

2.3 Cache分片数据

把下拉框放到“Active Cached Segment Timeline”显示数据cache的segment分片的创建过程，这个数据比较直观能够看到cache是什么时候创建的，如下所示：

同样点击每个分片块，能够看到分片块创建的调用堆栈。通过这个数据能够直观看到哪些操作触发了大的segment的创建，比如图中蓝色部分就是上面提到的autograd产生的尖峰显存所对应的segment。

segment一般不会释放，上图中所有的segment创建后均未释放。tensor释放后可以通过empty_cache来释放segment（示例见附录）。

3 Profiler中显存可视化使用

PyTorch Profiler已支持把显存的snapshot数据记录到profiling中，通过prof.export_memory_timeline函数可导出snapshot数据。

由于Profiler能够给一些数据打标签，所以可更加详细的记录具体是由哪个过程消耗了显存。通过可视化图表能直观的看到激活值、优化器、输入等操作的显存消耗，相比snapshot图表更加容易读懂，但没有snapshot数据那么精细。

API参考文档：https://pytorch.org/docs/main/profiler.html

3.1 开启方式

torch.profiler中打开record功能：

with torch.profiler.profile(profile_memory=True,with_stack=True,on_trace_ready=trace_handler,)：

主要添加参数如下：

profile_memory (bool) – 是否记录 tensor memory allocation/deallocation.
with_stack (bool) – 追踪操作的调用堆栈.
on_trace_ready(Callable) – 记录完成后的后处理回调函数 。

3.2 调用示例

用transformer来进行调用示例，完整的代码如下所示（git位置https://github.com/CalvinXKY/BasicCUDA/blob/master/pytorch/torch_mem_snapshot/transformer_profile.py）：

# Author: kevin.xie  zhihu@kaiyuan
import torch
from torch import nn
from datetime import datetime
from torch.autograd.profiler import record_function
  

def trace_handler(prof: torch.profiler.profile):
   # 获取时间用于文件命名
   timestamp = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
   file_name = f"visual_mem_{timestamp}"
  

   # 导出tracing格式的profiling
   prof.export_chrome_trace(f"{file_name}.json")
  

   # 导出mem消耗可视化数据
   prof.export_memory_timeline(f"{file_name}.html", device="cuda:0")
  

  

def train(num_iter=5, device="cuda:0"):
    model = nn.Transformer(d_model=512, nhead=2, num_encoder_layers=2, num_decoder_layers=2).to(device=device)
    x = torch.randn(size=(1, 1024, 512), device=device)
    tgt = torch.rand(size=(1, 1024, 512), device=device)
    model.train()
    labels = torch.rand_like(model(x, tgt))
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters())
    
    with torch.profiler.profile(
            activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
            schedule=torch.profiler.schedule(wait=0, warmup=0, active=6, repeat=1),
            record_shapes=True,
            profile_memory=True,
            with_stack=True,
            on_trace_ready=trace_handler,
    ) as prof:
        for _ in range(num_iter):
            prof.step()
            with record_function("## forward ##"): 
                y = model(x, tgt)
  

            with record_function("## backward ##"):
                loss = criterion(y, labels)
                loss.backward()
                print(loss.item())
  

            with record_function("## optimizer ##"):
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
  

  

if __name__ == "__main__":
    # warm-up:
    train(1)
    # run:
    train(5)

运行后此代码生成了两个html文件，一个由warm up产生，另一个由run产生。打开后生成的html文件(直接用浏览器打开)，可以看到mem-timeline数据，如下所示，可从数据中看到各个部分的占用比例，以及产生与释放的时机。

上例的transformer模型完成了5个迭代步骤，每次迭代内存的消耗占比在图中踩了不同颜色进行区分，我们可看到显存占用内容主要如下：

parameter：模型参数
optimizer_state: 优化器状态参数
input：输入值
temporary：临时变量；
activation：激活值（前向运算产生）
gradient：梯度值
autograd_detail: 自动梯度产生的显存
unknown：无法分类的消耗

还有一个profiling的html文件，是warn-up动作产生的（这个数据暂未看出来有啥分析意义），打开后是这样：

注意，进行warn-up动作是必要的，不然产生的数据会不符合预期。

参考内容：

Understanding CUDA Memory Usage

https://pytorch.org/docs/main/profiler.html

https://pytorch.org/memory_viz

https://github.com/pytorch/tutorials/blob/main/beginner_source/transformer_tutorial.py

附1：segment释放示例

代码：https://github.com/CalvinXKY/BasicCUDA/blob/master/pytorch/torch_mem_snapshot/segment.py

cache数据如下：

segment和block创建的过程如下，能够与代码的操作进行一一对应。比如，第一个segment_alloc动作是为第一个tensor的创建准备了内存，之后alloc创建block用于tensor1（tensor1 = torch.randn(size=(10,1024, 1024, 512), device=device)）。

附2：“文本预测”训练

一个关于文本预测的训练的示例（torch官方给出的示例），主要是通过torch.nn.transformer实现。通过添加snapshotAPI，保存显存使用数据。代码：https://github.com/CalvinXKY/BasicCUDA/blob/master/pytorch/torch_mem_snapshot/predict_text_with_snapshot_example.py

源码的改动变化：

添加snapshot的运行时间变化（不同机器时长会有差异！）：

snapshot的结果如下所示：

历史文章：2023年12月大模型文章集锦

更多AI工具，参考Github-AiBard123，国内AiBard123

可关注我们的公众号：每天AI新工具