使用 PyTorch 在 Windows 上通过 AMD 部署 LLM 入门指南
使用 AMD 消费级显卡，通过 PyTorch 在 Windows 上运行 LLM。
在之前的两篇关于拉普拉斯算子的博文中，我们开发了基于拉普拉斯算子的有限差分代码的 HIP 实现，并应用了两种可能的代码优化来优化 L2 缓存和全局内存之间的内存移动。第三部分将介绍一些额外的优化和通用技巧，以微调内核的性能。快速回顾一下，回想一下拉普拉斯算子的形式是标量场 u(x,y,z) 的梯度的散度
∇⋅∇u=∇2u=∂x2∂2u+∂y2∂2u+∂z2∂2u,我们在 第 1 部分 中开始的基线 HIP 实现的性能达到了理论峰值的约 50%[1]。然而，根据初步的 rocprof 分析，我们预测有限差分内核的性能应至少达到理论峰值的 71%[1]。为了实现这一目标，我们应用了两种优化：
-
引入循环分块以显式重用加载的模板
-
重新排序模板点的读取访问模式
有关完整的代码实现，请参阅 第 2 部分 中的重新排序读取访问模式部分。通过这些更改，我们达到了性能预测的 95%，但仍有一些悬而未决的问题：
-
之前的优化需要手动调整某些参数，这可能导致性能出现奇怪的特征，即性能会随着较大的分块因子而突然下降。修复此性能下降的根本原因是否能让我们更接近 第 1 部分 中定义的图例指标 (FOM)？
-
我们一直专注于优化内核的缓存和获取操作的数据重用。能否通过改进写入操作的相同方面来获得一些收益？
-
我们引入的优化需要对代码进行非平凡的更改。是否存在替代的优化方法，可以在不增加代码复杂性的情况下获得显著的性能提升？
这篇博文将探讨其中一些遗留问题。接下来的几节将介绍和讨论以下概念：
-
生成临时文件以理解寄存器使用情况和暂存器内存
-
应用启动边界以控制寄存器使用
-
应用非时序存储以释放更多缓存
寄存器压力与暂存器使用
在前一篇博文中描述的循环分块优化中，分块因子 m=16 导致内核的 FOM 恶化。rocprof 指标 FETCH_SIZE 上升到理论限制的 4 倍以上，而 L2CacheHit 指标下降到 50% 以下。我们怀疑高分块因子导致寄存器使用泛滥，造成溢出。为此，我们引入了一个新的编译标志 —save-temps，它告诉编译器为每个 GPU 内核生成有关寄存器使用、溢出、占用率等重要信息。它还包括主机和设备代码的指令集架构 (ISA) 转储。未来的 Lab Notes 博文将详细介绍 AMD GPU ISA。
我们检查了四个关键指标：
-
SGPR
-
VGPR
-
Scratch
-
Occupancy
SGPR 和 VGPR 分别指代标量寄存器和向量寄存器，Scratch 指代暂存器内存，这可能是寄存器溢出的一个指标，而 Occupancy 则代表执行单元 (EU) 上可以运行的波形 (wavefronts) 的最大数量。请注意，寄存器和暂存器使用情况的统计数据可以直接从 rocprof 输出文件中找到，而占用率和 ISA 汇编等其他详细信息只能从临时文件中找到。用户只需在提供的 makefile 中取消注释 #TEMPS=true 行，即可在名为 temps/laplacian_dp_kernel1-hip.s 的文件中生成包含此信息的临时文件。以下是基线 HIP 内核 1 的一些示例输出：
.section .AMDGPU.csdata; Kernel info:; codeLenInByte = 520; NumSgprs: 18; NumVgprs: 24; NumAgprs: 0; TotalNumVgprs: 24; ScratchSize: 0; MemoryBound: 0; FloatMode: 240; IeeeMode: 1; LDSByteSize: 0 bytes/workgroup (compile time only); SGPRBlocks: 2; VGPRBlocks: 2; NumSGPRsForWavesPerEU: 18; NumVGPRsForWavesPerEU: 24; AccumOffset: 24; Occupancy: 8; WaveLimiterHint : 1; COMPUTE_PGM_RSRC2:SCRATCH_EN: 0; COMPUTE_PGM_RSRC2:USER_SGPR: 8; COMPUTE_PGM_RSRC2:TRAP_HANDLER: 0; COMPUTE_PGM_RSRC2:TGID_X_EN: 1; COMPUTE_PGM_RSRC2:TGID_Y_EN: 1; COMPUTE_PGM_RSRC2:TGID_Z_EN: 1; COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: 2; COMPUTE_PGM_RSRC3_GFX90A:ACCUM_OFFSET: 5; COMPUTE_PGM_RSRC3_GFX90A:TG_SPLIT: 0除了 ISA 转储之外，这里还有很多信息需要解读。我们建议所有感兴趣的读者参考 此演示文稿 以及 寄存器压力 博文，了解有关寄存器、暂存器、占用率等更多详细信息。
下表包含基线和优化内核在不同分块因子 m 下的上述四个关键指标：
| SGPR | VGPR | Scratch | Occupancy | |
|---|---|---|---|---|
| 内核 1 – 基线 | 18 | 24 | 0 | 8 |
| 内核 3 – 重新排序加载 m=1 | 24 | 18 | 0 | 8 |
| 内核 3 – 重新排序加载 m=2 | 26 | 28 | 0 | 8 |
| 内核 3 – 重新排序加载 m=4 | 34 | 54 | 0 | 8 |
| 内核 3 – 重新排序加载 m=8 | 52 | 90 | 0 | 5 |
| 内核 3 – 重新排序加载 m=16 | 90 | 128 | 180 | 4 |
分块因子与寄存器/暂存器使用之间存在很强的相关性。在 m=16 时，寄存器使用量增加到“溢出”到暂存器空间，即不再适合寄存器空间的寄存器被卸载到全局内存中。占用率（定义为每个 EU 的波形数量）与寄存器使用之间也存在很强的反相关性——随着寄存器使用量的增加，占用率会下降。那么，为了防止溢出或提高占用率，可以做些什么？
启动边界
控制寄存器使用的一个快速方法是为内核应用启动边界。默认情况下，HIP 编译器会根据最大允许的线程块大小 1024 个线程来限制每个线程的寄存器数量。如果线程块大小在编译时已知，则为内核设置启动边界是一个好习惯。设置启动边界需要以下参数：
__launch_bounds__(MAX_THREADS_PER_BLOCK,MIN_WAVES_PER_EU)第一个参数 MAX_THREADS_PER_BLOCK 会告知编译器线程块的尺寸，以便它能够针对特定的块大小优化寄存器使用。第二个参数 MIN_WAVES_PER_EU 是一个可选参数，用于指定每个 EU 上必须激活的最小波形数量。默认情况下，第二个值设置为 1，无需修改，而默认的 MAX_THREADS_PER_BLOCK 值 1024 需要更改，因为我们没有使用全部 1024 个线程。
到目前为止，我们一直使用 256 × 1 × 1 的线程块大小。因此，这是如何为内核 3 设置启动边界，将 MAX_THREADS_PER_BLOCK 设置为 256：
template <typename T>__launch_bounds__(256)__global__ void laplacian_kernel(...) {
...我们将这一单行更改指定为“内核 4”，并检查其对寄存器和暂存器空间使用情况的影响。
| SGPR | VGPR | Scratch | Occupancy | |
|---|---|---|---|---|
| 内核 1 – 基线 | 18 | 24 | 0 | 8 |
| 内核 3/内核 4 – 重新排序加载 m=1 | 24/24 | 18/18 | 0/0 | 8/8 |
| 内核 3/内核 4 – 重新排序加载 m=2 | 26/26 | 28/28 | 0/0 | 8/8 |
| 内核 3/内核 4 – 重新排序加载 m=4 | 34/34 | 54/54 | 0/0 | 8/8 |
| 内核 3/内核 4 – 重新排序加载 m=8 | 52/52 | 90/94 | 0/0 | 5/5 |
| 内核 3/内核 4 – 重新排序加载 m=16 | 90/84 | 128/170 | 180/0 | 4/2 |
将启动边界应用于分块因子 m=4 及以下对寄存器使用没有影响。当 m=8 时，只有 VGPR 有轻微增加，而对于 m=16，VGPR 大幅增加，暂存器使用完全消除。请注意，占用率显著下降，这引发了关于这是否会对 m=16 的性能产生负面影响的问题。让我们看看 FOM 性能：
| 加速比 | 目标百分比 | |
|---|---|---|
| 内核 1 – 基线 | 1.00 | 69.4% |
| 内核 3 – 重新排序加载 m=1 | 1.20 | 82.9% |
| 内核 3 – 重新排序加载 m=2 | 1.28 | 88.9% |
| 内核 3 – 重新排序加载 m=4 | 1.34 | 93.1% |
| 内核 3 – 重新排序加载 m=8 | 1.37 | 94.8% |
| 内核 3 – 重新排序加载 m=16 | 0.42 | 29.4% |
| 内核 4 – 启动边界 m=1 | 1.20 | 82.9% |
| 内核 4 – 启动边界 m=2 | 1.28 | 88.9% |
| 内核 4 – 启动边界 m=4 | 1.34 | 93.1% |
| 内核 4 – 启动边界 m=8 | 1.39 | 96.1% |
| 内核 4 – 启动边界 m=16 | 1.34 | 93.2% |
毫不奇怪，启动边界对寄存器、暂存器或占用率没有影响的内核的性能与之前相同。启动边界对 SGPR、VGPR、Scratch 和 Occupancy 统计数据影响的程度与性能提升的程度之间存在明显的相关性。分块因子 m=8 和 m=16 的内核的性能有所提升。让我们检查相应的 rocprof 指标：
| FETCH_SIZE (GB) | 获取效率 (%) | L2CacheHit (%) | |
|---|---|---|---|
| 理论值 | 1.074 | – | – |
| 内核 1 – 基线 | 2.014 | 53.3 | 65.0 |
| 内核 3 – 重新排序加载 m=1 | 1.347 | 79.7 | 72.0 |
| 内核 3 – 重新排序加载 m=2 | 1.166 | 92.1 | 70.6 |
| 内核 3 – 重新排序加载 m=4 | 1.107 | 97.0 | 68.8 |
| 内核 3 – 重新排序加载 m=8 | 1.080 | 99.4 | 67.7 |
| 内核 3 – 重新排序加载 m=16 | 3.915 | 27.4 | 44.5 |
| 内核 4 – 启动边界 m=1 | 1.346 | 79.8 | 72.0 |
| 内核 4 – 启动边界 m=2 | 1.167 | 92.1 | 70.6 |
| 内核 4 – 启动边界 m=4 | 1.107 | 97.0 | 68.8 |
| 内核 4 – 启动边界 m=8 | 1.080 | 99.4 | 67.3 |
| 内核 4 – 启动边界 m=16 | 1.094 | 98.2 | 66.1 |
注意：虽然这些实验未显示
WRITE_SIZE或 Write efficiency(%)，但报告的Kernel 3 - Reordered loads m=16的WRITE_SIZE和 Write efficiency(%) 分别为 2.547 GB 和 41.7%。没有暂存器溢出的内核的写入效率接近 100%。
现在，当 m=16 结合启动边界时，不再发生寄存器溢出到暂存器，我们在加速比、获取效率和 L2CacheHit 方面看到了显著的提升。对于那些努力优化高寄存器使用率内核的用户，可以通过应用启动边界来快速获得性能提升。尽管 m=8 的内核与 m=16 相比，寄存器使用量少得多，但应用启动边界仍然对 VGPR 使用产生了影响，从而仅将整体性能提升到足以成为迄今为止性能最佳的内核。
非时序内存访问
我们的大部分优化工作都集中在提高空间局部性，但我们尚未考虑的是时间局部性——也就是说，如何在时间上优先缓存变量。加载 u 的元素以及存储 f 的元素都会占用缓存行。但不同之处在于，根据 第 1 部分 中描述的数据布局，u 的每个元素理论上最多可以重用六次，而 f 的每个元素只访问一次。因此，我们可以使用 clang 的内置非时序存储内建函数，让 f 绕过 L2 缓存，从而增加 u 条目可用缓存。还应注意，这些内建函数是 AMD GPU 特有的。
AMD clang 编译器提供了两个重载的内建函数，允许生成非时序加载和存储：
T __builtin_nontemporal_load(T *addr);void __builtin_nontemporal_store(T value, T *addr);在拉普拉斯算子的示例中，我们只需要非时序存储。让我们首先将此内建函数应用于初始基线内核：
| 内核 1（之前） | 内核 1（之后） |
| |
为了评估这个简单的代码修改的影响，我们将基线实现（带或不带非时序存储）的性能与 m=1 时内核 3 的性能进行比较：
| 加速比 | 目标百分比 | |
|---|---|---|
| 内核 1 – 基线 | 1.00 | 69.4% |
| 内核 1 – 非时序存储 | 1.19 | 82.5% |
| 内核 3 – 重新排序加载 m=1 | 1.20 | 82.9% |
一行代码的更改所带来的改进与重构整个基线内核以利用循环分块和重新排序的内存访问模式所带来的改进相当。检查 rocprof 统计数据：
| FETCH_SIZE (GB) | 获取效率 (%) | L2CacheHit (%) | |
|---|---|---|---|
| 理论值 | 1.074 | – | – |
| 内核 1 – 基线 | 2.014 | 53.3 | 65.0 |
| 内核 1 – 非时序存储 | 1.429 | 75.2 | 71.4 |
| 内核 3 – 重新排序加载 m=1 | 1.347 | 79.7 | 72.0 |
当 m=1 时，非时序存储和重新排序加载之间的统计数据也相当。这些发现表明，利用用于非时序内存访问的重载内建函数实际上可以是用户应用的第一项优化，因为它是一个“低垂的果实”，只需要修改一行代码。下一个显而易见的问题是，当我们将内建非时序存储与循环分块因子 m=8 和启动边界结合使用时会发生什么？让我们再次对内核 4 进行一项更改：
| 内核 4（之前） | 内核 5（之后） |
| |
这个新的内核 5 是涉及循环分块、重新排序加载、应用启动边界和利用非时序存储的优化的累积。当循环分块因子 m=8 时的性能如下所示：
| 加速比 | 目标百分比 | |
|---|---|---|
| 内核 1 – 基线 | 1.00 | 69.4% |
| 内核 3 – 重新排序加载 m=8 | 1.37 | 94.8% |
| 内核 4 – 启动边界 m=8 | 1.39 | 96.1% |
| 内核 5 – 非时序存储 m=8 | 1.44 | 100% |
通过结合所有优化，我们实现了 1.44 倍的加速比，并达到了目标性能的 100%！
让我们再次检查 rocprof 指标：
| FETCH_SIZE (GB) | 获取效率 (%) | L2CacheHit (%) | |
|---|---|---|---|
| 理论值 | 1.074 | – | – |
| 内核 1 – 基线 | 2.014 | 53.3 | 65.0 |
| 内核 3 – 重新排序加载 m=8 | 1.080 | 99.4 | 67.7 |
| 内核 4 – 启动边界 m=8 | 1.080 | 99.4 | 67.3 |
| 内核 5 – 非时序存储 m=8 | 1.074 | 100 | 67.4 |
在这些改进过程中，我们将从全局内存的加载次数减少了一半。测量到的获取和写入大小已达到理论限制，因此进一步的性能改进必须来自其他方面。由于报告的有效内存带宽已达到预测目标，因此进一步改进的空间可能很小。
总结与后续步骤
拉普拉斯有限差分系列文章的第三部分介绍了另外两项优化，这两项优化都只需要修改一行源代码。我们还向读者介绍了临时文件，这些文件可以进一步深入了解内核的暂存器使用情况、寄存器压力和占用率，而所有这些都可以通过为内核应用启动边界来简单地改变。与利用循环分块和重新排序内存加载相比，应用内建的非时序加载更容易实现，并且为初始 HIP 实现提供了显著的性能提升，因此应优先于内核重构。但我们再次强调，这些内建函数是不可移植的，并且是 AMD GPU 特有的。到目前为止提出的所有四项优化结合起来，使我们的有效内存带宽达到了预测目标。
然而，仍然存在一些悬而未决的问题。这个拉普拉斯系列文章的最后三篇文章侧重于针对问题规模 nx,ny,nz = 512, 512, 512 在 MI250X GPU 的单个 GCD 上运行的优化。如果我们对其他硬件和问题规模运行内核 5 会怎样？是否会出现其他性能问题？拉普拉斯系列文章的下一篇也是最后一篇文章将对此进行深入探讨。
如果您有任何问题或评论，请在 GitHub 讨论区 与我们联系
使用 ROCm 版本 5.3.0-63 和 MI250X 单 GCD 进行的测试。基准测试结果不是经过验证的性能数据，仅用于演示代码修改的相对性能改进。实际性能结果取决于多种因素，包括系统配置和环境设置，结果的重现性无法保证。
