使用NVIDIA多实例GPU与NUMA节点本地化加速数据处理
NVIDIA Ampere、NVIDIA Hopper和NVIDIA Blackwell系列旗舰级数据中心GPU都具有非一致性内存访问特性，但对外呈现单一内存空间。因此，大多数程序并不会受到内存非一致性问题的影响。然而，随着新一代GPU带宽的增加，考虑计算和数据的局部性可以带来显著的性能和功耗收益。
本文首先分析了NVIDIA GPU的内存层级结构，探讨了通过芯片间链路进行数据传输对功耗和性能的影响。接着，介绍了如何使用NVIDIA多实例GPU模式实现数据本地化。最后，以Wilson-Dslash模板算子为例，展示了运行MIG模式与未进行本地化处理的结果对比。
NVIDIA GPU中的内存层级结构
考虑图1所示包含两个NUMA节点的抽象内存层级结构视图。当节点0上的流多处理器需要访问节点1的DRAM中的内存位置时，必须通过L2结构传输数据。以NVIDIA Blackwell GPU为例，每个NUMA节点都是一个独立的物理芯片，这会增加延迟并提高数据传输所需的功耗。尽管增加了复杂性，但对NUMA无感知的代码仍然可以达到峰值DRAM带宽。
图1. 跨两个NUMA节点的GPU内存层级结构抽象视图
为了解决这些缺点，最小化NUMA节点间的数据传输是有益的。当向用户呈现单一内存空间时，NVIDIA架构在L2中采用一致性缓存来减少NUMA节点间的数据传输。这种机制有助于防止重复访问同一内存地址时，需要从L2结构接口重新获取数据。理想情况下，一旦地址被获取到本地L2缓存中，后续对该地址的所有访问都将命中缓存。
在引入一致性缓存之前，统一的L2缓存允许所有SM达到峰值带宽（如NVIDIA Volta），尽管延迟会根据SM与不同L2分段的距离而变化。从NVIDIA Ampere代开始，更大的芯片引入了NUMA节点层级，每个节点都有自己的L2缓存，并与其他节点通过一致性连接。
自NVIDIA Ampere架构以来，大型数据中心GPU都采用了这种设计（与小型游戏GPU不同），正如NVIDIA Blackwell Ultra架构中所述，L2结构连接维持了峰值带宽。随着GPU的持续发展，出现了两个挑战：增加的延迟和功耗限制。
- 增加的延迟：访问L2缓存的远端部分导致延迟增加，这会显著影响性能，尤其是对于同步操作。
- 功耗限制：在最大的GPU上，当张量核心处于活动状态时，功耗成为一个限制因素。通过本地化L2访问来降低功耗，可以实现降低L2结构时钟，并通过与GPU Boost相关的动态电压频率调整机制提高计算时钟。这样，可以显著提高张量核心的性能。
MIG减少NUMA节点间的数据传输
MIG功能随NVIDIA Ampere架构一同推出，支持将单个GPU划分为多个实例。通过使用MIG，开发人员可以为每个NUMA节点创建一个GPU实例，从而消除了通过L2结构接口的访问。
这种方法也带来了一些开销，包括使用PCIe在不同GPU实例之间进行通信的开销。下一节将展示运行MIG模式和未本地化内存的工作负载结果，以证明这种方法的有效性。
使用MIG实现数据本地化
MIG支持将兼容的NVIDIA GPU划分为多个隔离的实例，每个实例拥有专用的高带宽内存、缓存和计算核心。这可以实现跨多个用户或工作负载的高效、高性能的GPU利用。MIG可以在单个GPU上实现高达7倍的GPU资源利用率。它允许多个虚拟GPU以及虚拟机在单个GPU上并行运行，同时提供虚拟GPU所具备的隔离保证。
可以利用MIG提供的功能来实现NUMA节点本地化。通过为每个NUMA节点创建一个MIG实例，可以确保不同GPU实例之间的隔离。这种方法有助于消除NUMA节点间的通信流量。
MIG允许将实际GPU分割成GPU实例，在每个GPU实例中可以定义一个或多个计算实例。一个计算实例包含属于某个GPU实例的全部（当每个GPU实例只有一个计算实例时）或部分SM。为了实现GPU实例内的本地化，思路是创建两个映射到每个NUMA节点上的GPU实例。在Blackwell GPU上，可以启用MIG模式并列出可用的GPU实例配置文件，如图2中的代码所示。
由于Blackwell有两个NUMA节点（每个芯片一个），需要查找具有最多SM且存在两个实例的配置文件。如图2所示，这是ID为9的配置文件，可以有两个实例。每个实例将拥有89 GB内存和70个SM。使用两个这样的实例将总共只提供70×2=140个SM，而不是设备上全部的148个SM。
此时，需要在每个GPU实例中创建一个计算实例。可以使用图3中所示的命令完成。主GPU和GPU实例现在都有自己的标识符哈希码。为两个NUMA节点使用这些标识符：
MIG 3g.90gb 设备 0: (UUID: MIG-ee2ec0e5-0dda-5591-9ee7-4ae51028b6fa)
MIG 3g.90gb 设备 1: (UUID: MIG-2bbb368b-7cb0-53da-b1a4-7ace0652a197)
要使用这些设备，请将它们添加到CUDA_VISIBLE_DEVICES环境变量中。例如，要运行一个双进程MPI作业，可以创建一个包装脚本（wrapper.sh）：
#!/bin/bash
case $SLURM_PROCID in
0)
CUDA_VISIBLE_DEVICES="MIG-ee2ec0e5-0dda-5591-9ee7-4ae51028b6fa"
;;
1)
CUDA_VISIBLE_DEVICES="MIG-2bbb368b-7cb0-53da-b1a4-7ace0652a197"
;;
esac
$*
然后启动MPI作业：
$ mpirun -n 2 ./wrapper.sh my_executable
最后，当所有工作完成后，可以关闭MIG模式。
图2. 启用MIG模式并列出可用的GPU实例配置文件
图3. 为MIG创建计算实例
图4. 关闭MIG实例的命令
使用MIG进行本地化的优势是什么？
为了演示使用MIG进行本地化的优势，以Wilson-Dslash模板算子为例，这是格点量子色动力学的一个关键内核，来自QUDA库。该库用于加速多个大型LQCD代码，如Chroma和MILC。
Dslash内核是对四维环形晶格进行有限差分运算，其中每个晶格点的数据根据其八个正交相邻点的值进行更新。此处的四个维度是通常的空间维度（X, Y, Z）和时间维度（T）。该内核受限于内存带宽。
如果将晶格沿时间轴均匀分解到两个NUMA节点上，那么每个域将需要访问另一域在时间维度边界上的晶格点。如图5所示，子域边界上的绿色晶格点需要红色晶格点来完成其模板运算。理论上，晶格被布置在两个NUMA节点上。绿色点需要红色点来完成模板运算。可能的数据路径是未本地化时的常规内存访问（黑色箭头），或者在MIG本地化模式下通过主机进行MPI消息传递（黑色箭头）。
图5. MIG模式下Dslash内核的内存访问
访问相邻点最便捷的方式是通过共享的L2缓存和互连。然而，在MIG模式下运行时，此路径需要通过MPI使用PCIe或NVLink在MIG实例之间进行通信。因此，与访问连接到MIG实例的主内存相比，此路径将更慢。
对两个MIG实例之间几乎没有通信需求的工作负载，使用MIG模式往往获益更多。相反，需要将边界上的黑色点打包并通过MPI发送。此步骤会引入额外的延迟（缓冲区打包、发送和解包）。虽然通过不使用共享的L2缓存到缓存的互连来节省GPU功耗，但通过主机（例如，PCIe）进行传输时仍会消耗功耗。
需要在两个进程之间传输的数据量与消息中需要传输的表面点数量有关，具体来说，与分割方向正交的表面三维体积有关。在此示例中，分割始终沿时间方向，因此每个NUMA节点理论上最终获得(Ns * Nt)/2个点，其中Ns是空间体积中的点数，Nt是时间维度的长度。表面积与体积比为Ns / (Ns * Nt / 2) = 2/Nt。在问题中，考虑Nt=64，表面体积比保持恒定在1/32 ≈ 3.13%。
图6展示了未本地化的情况。全局内存由两个通过内存控制器连接到NUMA节点的内存组成。晶格上的彩色高亮表示数据可能来自本地DRAM，也可能通过共享的L2来自远程DRAM。
图6. 未本地化情况下的内存访问
这可以与未使用MIG的基线情况进行比较。在这种情况下，数据和计算都没有本地化，其场景更接近图6所示。每个NUMA节点都从其本地内存控制器以及另一个NUMA节点接收数据。事实上，只有一个全局晶格，图中为两个NUMA节点分开成两部分是人为的。
在这种情况下，处理一组点集的线程块完全由调度器随意分配给各个NUMA节点。由于数据均匀分布在两个NUMA内存上，通过共享L2传输的数据量远多于MIG本地化情况（后者仅传输最少需要的表面点）。这可能会导致显著的功耗成本。
另一方面，整个操作可以通过单个内核执行。可以避免为消息传递打包缓冲区和最后累积接收到的面所产生的延迟。
对于实验结果，我们关注在不同GPU功耗限制（瓦特）下工作负载执行的速度提升。速度提升是未本地化方法与MIG方法在相同功耗限制下（例如，均为700 W）运行所花费的实际时间的比值。
如图7所示，在400 W的GPU功耗限制下，MIG的性能优于未本地化数据，根据工作负载规模不同，速度提升最高可达2.25倍。其原因是，当GPU在低功耗限制下运行时，L2结构接口消耗的功率成为一个限制因素。在MIG模式下，由于没有消耗L2结构功率来在NUMA节点间传输数据，工作负载可以运行得更快。
然而，当GPU功耗限制提高时，在图7中由灰色、深绿色和黑色线条表示的实验中，以及部分绿色线条表示的情况中，MIG模式的表现略差。这是因为在更高的功耗限制下，消息传递所带来的额外延迟可能超过本地化的优势。
图7. 在不同工作负载规模下运行基于MIG的NUMA本地化
结果表明，较小规模的情况（尤其是图7中黑色和深绿色线条表示的情况）即使在未本地化的情况下，在更高的功耗限制下也从未耗尽可用功率。因此，它们从本地化节省的GPU功耗中获益甚微，而且在如此小的规模下，内核启动带来的延迟更加明显。较大规模的工作负载（例如，绿色线条）需要更多功耗，因此即使在更高的功耗限制下，也能比未本地化的设置获得优势。
开始使用基于MIG的NUMA节点本地化
NVIDIA数据中心GPU中的本地L2缓存会影响对NUMA无感知的工作负载的性能。我们在MIG模式下使用Wilson-Dslash算子进行的实验表明，当GPU在较低的功耗限制下运行，并且通过MPI（PCIe/NVLink）进行的数据传输相对于本地内存访问较少时，与相同功耗限制下的未本地化情况相比，基于MIG的NUMA节点本地化可以实现高达2.25倍的加速。
虽然系统在1000 W的更高功耗预算下可能比400 W配置实现更高的绝对性能，但在功耗受限条件下，基于MIG的本地化具有明显优势。在低功耗场景下，它能显著提高性能，使其成为在严格功耗限制内运行时特别有效的优化方法。
然而，总的来说，MIG并不能提供持续实现有效数据本地化所需的灵活性，尤其是在进程间通信开销在更高功耗限制下变得更加显著时。MIG仅支持那些规模太小而无法完全放入单个GPU的使用场景。因此，不建议将MIG用于本文中展示的案例。为了解决这些限制，目前正在研究其他替代方法。
要了解更多信息，请参阅《使用NVIDIA CUDA MPS无需更改代码即可提升GPU内存性能》。FINISHED
