查看PyTorch的GPU使用情况的工具

文章目录

torch.cuda API
PyTorch Snapshot
PyTorch Profiler
NVIDIA Nsight Systems
torchinfo

torch.cuda API

torch.cuda.memory_stats：返回给定设备的 CUDA 内存分配器统计信息字典。该函数的返回值是一个统计字典，每个字典都是一个非负整数。
torch.cuda.memory_summary：返回给定设备当前内存分配器统计信息的人类可读的打印输出。这对于在训练期间或处理内存不足异常时定期显示非常有用。

参考用法：Pytorch 如何获取Pytorch在CPU/主内存上的内存统计信息

torch.cuda.memory_snapshot：返回所有设备上的 CUDA 内存分配器状态的快照。解释该函数的输出需要熟悉内存分配器的内部结构。
torch.cuda.reset_peak_memory_stats：重置 CUDA 内存分配器跟踪的“峰值”统计数据。峰值统计数据对应于每个单独统计字典中的“峰值”键。
torch.cuda.memory_allocated：返回给定设备的张量当前占用的 GPU 内存（以字节为单位）。这可能小于 nvidia-smi 中显示的数量，因为缓存分配器可以保留一些未使用的内存，并且需要在 GPU 上创建一些上下文。
torch.cuda.max_memory_allocated：返回给定设备的张量占用的最大 GPU 内存（以字节为单位）。默认情况下，这将返回自该程序开始以来分配的内存峰值。reset_peak_memory_stats()可用于重置跟踪此指标的起点。例如，这两个函数可以测量训练循环中每次迭代的峰值分配内存使用情况。
torch.cuda.memory_reserved：返回给定设备的缓存分配器管理的当前 GPU 内存（以字节为单位）。
torch.cuda.max_memory_reserved：返回给定设备的缓存分配器管理的最大 GPU 内存（以字节为单位）。默认情况下，这将返回自该程序开始以来的峰值缓存内存。reset_peak_memory_stats() 可用于重置跟踪此指标的起点。例如，这两个函数可以测量训练循环中每次迭代的峰值缓存内存量。
torch.cuda.empty_cache：释放缓存分配器当前持有的所有未占用的缓存内存，以便这些内存可以在其他 GPU 应用程序中使用并在 nvidia-smi 中可见。empty_cache()不会增加 PyTorch 可用的 GPU 内存量。但是，在某些情况下，它可能有助于减少 GPU 内存碎片。

以下内容翻译自 Memory management 文档

PyTorch使用缓存内存分配器 (caching memory allocator) 来加快内存分配速度。这允许快速释放内存而无需设备同步。然而，由分配器管理的未使用内存在nvidia-smi中仍会显示为已使用。

你可以使用memory_allocated()和max_memory_allocated()来监控由张量占用的内存，
并使用memory_reserved()和max_memory_reserved()来监控缓存分配器管理的总内存量。
调用empty_cache()可以释放PyTorch中所有未使用的缓存内存，以便其他GPU应用程序可以使用这些内存。然而，由张量占用的GPU内存不会被释放，因此它不能增加PyTorch可用的GPU内存量。

对于更高级的用户，我们提供了通过memory_stats()进行更全面的内存基准测试的功能。我们还提供了通过memory_snapshot()捕获内存分配器状态的完整快照的功能，这可以帮助你理解代码产生的底层分配模式。

在环境中设置PYTORCH_NO_CUDA_MEMORY_CACHING=1以禁用缓存。

缓存分配器的行为可以通过环境变量PYTORCH_CUDA_ALLOC_CONF来控制。格式为PYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>...可用的选项有：

max_split_size_mb 阻止分配器拆分大于此大小（以MB为单位）的块。这可以帮助防止碎片化，并可能允许一些边缘工作负载在不耗尽内存的情况下完成。性能成本可能从‘零’到‘显著’不等，具体取决于分配模式。默认值是无限制，即所有块都可以被拆分。此选项应作为工作负载因‘内存不足’而中止并显示大量非活动拆分块时的最后手段。
roundup_power2_divisions 帮助将请求的分配大小四舍五入到最接近的2的幂除法，以更好地利用块。在当前的CUDACachingAllocator中，大小被四舍五入为512的块大小的倍数，因此这对于较小的大小来说很好。然而，对于较大的近邻分配，这可能效率低下，因为每个分配都会去不同大小的块，减少了这些块的重用。这可能会创建大量未使用的块，浪费GPU内存容量。此选项使分配大小四舍五入到最接近的2的幂除法。例如，如果我们需要四舍五入大小为1200，并且分割数为4，1200大小介于1024和2048之间，如果我们在它们之间进行4次分割，值分别为1024、1280、1536和1792。因此，分配大小为1200将被四舍五入到1280作为最接近的2的幂分割上限。
roundup_bypass_threshold_mb 对于超过阈值大小（以MB为单位）的分配请求，绕过四舍五入分配大小。这可以在进行预计会持久或寿命较长的大型分配时帮助减少内存占用。
garbage_collection_threshold 帮助主动回收未使用的GPU内存，以避免触发代价高昂的同步和回收所有操作（release_cached_blocks），这对延迟敏感的GPU应用程序（如服务器）可能不利。设置此阈值后（例如0.8），如果GPU内存容量使用超过阈值（即分配给GPU应用程序的总内存的80%），分配器将开始回收GPU内存块。算法倾向于首先释放旧的和未使用的块，以避免释放正在积极重用的块。阈值值应大于0.0且小于1.0。

PyTorch Snapshot

以下内容摘自博客：PyTorch显存可视化与Snapshot数据分析和[LLM]大模型显存计算公式与优化

模型显存内容分析：

在模型训练/推理时，显存（显卡的全局内存）分配一部分是给AI框架，另一部分给了系统（底层驱动）。其中系统层的显存消耗一般由驱动控制，用户不可控；框架侧的显存消耗用户可控。

如下图是一个模型训练过程中已用显存的数值随时间的变化：

在这里插入图片描述

内存中的许多微小峰值，通过将鼠标悬停在它们上，我们看到它们是卷积运算符临时使用的缓冲区。参考自：Understanding GPU Memory 1: Visualizing All Allocations over Time

注意：数据是具体的消耗值，不等于cudaMalloc创建的显存值。

显存消耗的内容包括：模型参数（parameter）、优化器状态值（optimizer_state）、激活值（activation）、梯度值（gradient）、输入数据（input）、临时变量（temporary）、自动梯度（autograd_detail）、未知数据（unknown)。从用户侧可以将这些数据进行一个分类：

可估算值：模型参数（parameter）、优化器状态值（optimizer_state）、激活值（activation）、梯度值（gradient）、输入数据（input）
未命名数据：临时变量（temporary）、未知数据（unknown)
其他（框架）：自动梯度（autograd_detail）

其中

“未命名数据”来源可能是用户创建的一些临时变量，这些变量未参与图的计算过程，所以未被统计；或者是一些未被框架跟踪（tracing）到的数据。
“自动梯度数据"是在反向传播求解梯度时产生的一些变量。

我们在显存计算时会发现“为什么有时显存估算值和实际测量值相差较大？”其中一个可能的原因是：未知的数据太大。即显存中可估算值占比相对较小，其它不可估算值的数据占比较大，导致计算值和实际值差距较大（误差可超过30%），比如估算得到的显存消耗为50GB，而实际测试达到了75GB。

如下图是运行一个LLM模型采集的一些过程数据，可以看到unknown占比有时能达到30%。

在这里插入图片描述
训练显存消耗（可估算部分）主要包括：模型参数（Model）+ 优化器状态（Optimizer status）+梯度值（Gradient）+激活值（Activation）。根据数值的变化，可将显存消耗分为静态/动态值。

训练过程中，模型参数、优化器状态一般不会变化，这两部分归属于静态值；
激活值、梯度值会随着计算过程发生变化，将它们归类到动态值。

在这里插入图片描述

Snapshot的使用：

在PyTorch2.1中的显存snapshot功能被增强（在1.X里面也有snapshot的记录操作），可以将显存消耗可视化，特点是查看简单、更易理解。

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

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)的参数解释：

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

可视化操作：

方式1：https://pytorch.org/memory_viz
方式2：

python .../lib/python3.8/site-packages/torch/cuda/_memory_viz.py trace_plot visual_mem_2024_06_02_10_19_14.pickle -o visual_mem_2024_06_02_10_19_14.html

Snapshot数据分析：

激活内存数据：网页中选择“Active Memory Timeline”下拉框可看到激活内存数，这部分数据主要是记录tensor在计算过程中占用的内存数值以及其存活的周期；同时，能够查看每个tensor的内存消耗调用堆栈（Python/C++）。
内存块的使用与释放：当把选项切换到“Allocator State Hisotry”时，可以看到从CUDA里面申请的内存块是如何被alloc分配成小的block，以及这些小的block什么时候被释放掉了。如果把光标放置到这些彩色的block上面，能够跟踪到是什么操作在使用该显存地址。框框表示segment（分片），在segment块中的彩色条表示block。torch显存管理创建的顺序是先创建segment，然后创建block。注意，block的释放仅从torch显存管理里面释放，而非cuda free操作，要发出cuda free的是segment的释放。
Cache分片数据：把下拉框放到“Active Cached Segment Timeline”显示数据cache的segment分片的创建过程，这个数据比较直观能够看到cache是什么时候创建的。通过这个数据能够直观看到哪些操作触发了大的segment的创建。segment一般不会释放。tensor释放后可以通过empty_cache来释放segment。

以下内容摘自博客：Pytorch 显存管理机制与显存占用分析方法

Snapshot 是 PyTorch 2.1 及以上版本提供的一种自动化显存分析工具。在代码的开始和结束处添加指定语句然后运行代码，PyTorch 会自动记录 CUDA allocator 的显存消耗、显存的 Python/C++ 调用堆栈和调用过程中的时间线，最后将这些数据保存并生成 .pickle 文件，将文件拖入网页即可查看显存占用。

值得注意的是，Snapshot 同样只关注当前进程，而且无法关注到 CUDA Context 部分的显存占用。

Active Memory Timeline

在这里插入图片描述横轴是程序执行的时间轴，纵轴是已分配的显存（torch.cuda.memory_allocated(device)），通过该视图可以查看 tensor 在程序运行过程中的显存占用和生命周期。

色块起点表示 tensor 或 Block 的分配，终点表示 tensor 或 Block 的释放，长度代表生命周期，色块的滑坡代表此前有其他 tensor 被释放（这里的释放并非真正意义上的空间释放）。

从上图中任选一个色块：

红框 1 表示该 tensor 或 Block 的编号（同一个 tensor 在三个视图中的编号一致）
红框 2 表示该 tensor 或 Block 的地址
红框 3 表示该 tensor 或 Block 的 size
红框 4 表示在色块起点时刻显存管理器已分配的显存总量（区别于 2.2.3 已缓存的显存总量）

Allocator State History

torch.cuda.empty_cache() 调用前右侧有 4 个空白 Segment
torch.cuda.empty_cache() 调用后，之前 4 个空白的 Segment 得以真正释放

torch.cuda.empty_cache() 调用前右侧有 4 个空白 Segment，torch.cuda.empty_cache() 调用后，之前 4 个空白的 Segment 得以真正释放。

上图右侧是某一时刻 Segment 和 Block 的分配情况，白框是 Segment，色块是已分配的 Block。

上图左侧记录着 Segment 和 Block 随时间的申请、分配、释放历史，左侧第一列表示动作，第二列表示 Segment 或 Block 的地址，第三列表示显存大小：

segment_alloc：显存管理器此时调用 cudaMalloc 从 GPU 申请一个新的 Segment 缓存块
alloc：显存管理器从 Segment 中划分一块空间给 Block
free：表示 tensor 或 Block 的释放（将 Block 所在空间归还给显存管理器）
segment_free：表明程序此时调用了torch.cuda.empty_cache()，显存管理器会将一些完全未分配的 Segment 释放

通过该视图可以查看程序运行过程中 Segment 和 Block 的申请、分配、释放历史。

在上图右侧的左上角，有一个 2MB 大小的 Segment 在torch.cuda.empty_cache()调用后看起来并没有得到释放，这是因为该 Segment 其实并非为空，而是分配了一个 8KB 大小的 Block。

Active Cached Segment Timeline

在这里插入图片描述
类似 2.2.1 的 Active Memory Timeline 图，横轴是程序执行的时间轴，纵轴是已缓存的显存（torch.cuda.memory_reserved(device)），色块是 Segment（2.2.1 中的色块是 Block）。
通过该视图可以直观地查看各 Segment 的生命周期，以及是由哪些操作触发了 Segment 的创建。如果不是用户主动调用torch.cuda.empty_cache()，Segment 一般不会释放。

PyTorch Profiler

待更新

NVIDIA Nsight Systems

待更新

torchinfo

torchinfo，可实现模型参数量计算、各层特征图形状计算和计算量计算等功能。

使用参考：6.5 模型参数打印 · PyTorch实用教程(第二版)

支持版本：PyTorch 1.4.0+

安装：pip install torchinfo

def summary(
    model: nn.Module,
    input_size: Optional[INPUT_SIZE_TYPE] = None,
    input_data: Optional[INPUT_DATA_TYPE] = None,
    batch_dim: Optional[int] = None,
    cache_forward_pass: Optional[bool] = None,
    col_names: Optional[Iterable[str]] = None,
    col_width: int = 25,
    depth: int = 3,
    device: Optional[torch.device] = None,
    dtypes: Optional[List[torch.dtype]] = None,
    mode: str | None = None,
    row_settings: Optional[Iterable[str]] = None,
    verbose: int = 1,
    **kwargs: Any,
) -> ModelStatistics:

Summarize the given PyTorch model. Summarized information includes: 总结给定的 PyTorch 模型。汇总信息包括：

Layer names,
input/output shapes,
kernel shape,
# of parameters,
# of operations (Mult-Adds),
whether layer is trainable

可选参数：

model (nn.Module)：要汇总的 PyTorch 模型。模型应该完全处于 train() 或 eval() 模式。如果层不全处于同一模式，运行 summary 可能会对 batchnorm 或 dropout 统计数据产生副作用。
input_size (Sequence of Sizes)：输入数据的形状，格式为 List/Tuple/torch.Size（数据类型必须与模型输入匹配，默认是 FloatTensors）。应在元组中包括批量大小。默认值：None
input_data (Sequence of Tensors)：模型前向传递的参数（数据类型推断）。如果 forward() 函数需要多个参数，请传入 args 列表或 kwargs 字典（如果 forward() 函数只接受一个字典参数，请将其包装在一个列表中）。默认值：None
batch_dim (int)：输入数据的批量维度。如果 batch_dim 为 None，则假设 input_data / input_size 包含批量维度，并在所有计算中使用它。否则，扩展所有张量以包含 batch_dim。指定 batch_dim 可以作为运行时优化，因为如果指定了 batch_dim，torchinfo 在前向传递时使用批量大小为 1。默认值：None
cache_forward_pass (bool)：如果为 True，使用模型类名称作为键缓存 forward() 函数的运行。如果前向传递是一个耗时的操作，这可以更容易地修改模型摘要的格式，例如更改深度或启用的列类型，特别是在 Jupyter Notebooks 中。警告：在启用此功能时修改模型架构或输入数据/输入大小不会使缓存失效或重新运行前向传递，可能会导致不正确的摘要。默认值：False
col_names (Iterable[str])：指定输出中要显示的列。目前支持的有：“input_size”, “output_size”, “num_params”, “params_percent”, “kernel_size”, “mult_adds”, “trainable”。默认值：“output_size”, “num_params”。如果未提供 input_data / input_size，则仅使用 “num_params”。
col_width (int)：每列的宽度。默认值：25
depth (int)：显示嵌套层的深度（例如，Sequentials）。低于此深度的嵌套层将不会在摘要中显示。默认值：3
device (torch.Device)：为模型和 input_data 使用此 torch 设备。如果未指定，则使用 input_data 的数据类型（如果已提供）或模型的参数。否则，使用 torch.cuda.is_available() 的结果。默认值：None
dtypes (List[torch.dtype])：如果使用 input_size，torchinfo 假设输入使用 FloatTensors。如果模型使用不同的数据类型，请指定该数据类型。对于多个输入，请指定每个输入的大小，并在此处指定每个参数的类型。默认值：None
mode (str)：“train” 或 “eval” 中的一个，决定在调用 summary() 之前调用 model.train() 还是 model.eval()。默认值：“eval”
row_settings (Iterable[str])：指定在一行中显示哪些特性。目前支持的有：“ascii_only”, “depth”, “var_names”。默认值：“depth”
verbose (int)：0（静默）：无输出；1（默认）：打印模型摘要；2（详细）：详细显示权重和偏置层。默认值：1。如果使用 Jupyter Notebook 或 Google Colab，默认值为 0。
**kwargs：模型 forward 函数中使用的其他参数。不再支持传递 *args。