随着我们向更密集的计算基础设施迈进，拥有更多的计算、更多的 GPU、加速网络等，多 GPU 训练和分析变得越来越流行。随着开发者和从业者从 CPU 集群转向 GPU 集群，我们需要工具和最佳实践。RAPIDS 是一套开源的 GPU 加速数据科学和 AI 库。借助 Spark 和 Dask 等工具，这些库可以轻松地针对更大的工作负载进行横向扩展。本博文简要概述了 RAPIDS Dask，并重点介绍了多 GPU 数据分析的三种最佳实践。
在使用 GPU 实现最高性能时，用户通常会面临内存压力和稳定性问题。虽然 GPU 在计算方面比 CPU 更强大，但与系统内存相比，GPU 的内存通常更少。例如，GPU 工作负载通常在核外场景中执行，此时 GPU 内存小于一次处理工作负载所需的内存总量。此外，CUDA 生态系统提供了 多种类型的内存 ，用于不同的目的和应用。
Dask 是什么？
Dask 是一个高度灵活的开源分布式 Python 库。 Dask 可以帮助扩展复杂的自定义 Python 代码，但更重要的是，它可以使用 Dask-Array 和 Dask-Dataframe 扩展数组和 Dataframe 工作负载。 RAPIDS 使用 Dask-DataFrame 帮助扩展 GPU ETL 和 ML 工作负载，例如 Dask-cuDF 和 cuML/XGBoost，并采用类似 Pandas 的界面。
df = dd.read_parquet("/my/parquet/dataset/")
agg = df.groupby('B').sum()
agg.compute()
适用于 CPU 和 GPU 的 Dask 最佳实践
Dask 的众多优势之一是用户可以同时面向 CPU 和 GPU 后端 。然而，开发两个代码库（一个用于 CPU，另一个用于 GPU）非常耗时且难以维护。Dask 支持在 CPU 和 GPU 后端之间轻松切换。
与 PyTorch 配置设备的机制类似：
device = 'cuda' if torch.cuda.is_available() else 'cpu'
Dask 还可以配置后端目标：
# GPU Backends
dask.config.set({"array.backend": "cupy"})
dask.config.set({"dataframe.backend": "cudf"})
# CPU Backends
dask.config.set({"array.backend": "numpy"})
dask.config.set({"dataframe.backend": "pandas"})
# Configure with Environment Variables
DASK_DATAFRAME__BACKEND=cudf
现在，我们可以轻松开发代码，而无需调用 特定的 后端 I/O 指令，也无需以与硬件无关的方式编写分布式 Python 分析代码：
# Dataframe Example with
dask.config.set({“dataframe.backend”: “cudf”}):
ddf = dd.read_parquet('example.parquet')
res = ddf.groupby("col").agg({'stats': ['sum', 'min', 'mean']})
res.compute()
通过使用这种简单的 CPU/GPU 设置，Dask-RAPIDS 用户可以轻松地针对这两种设备进行开发，而不会产生任何开发开销。更完整的示例是， NeMo 中的数据管护框架通过此机制支持 CPU/GPU 部署 。
集群和内存配置
数据工作流既是计算密集型应用，也是内存密集型应用。配置和管理分布式应用程序使用内存的方式可能意味着完成作业与失败之间的区别。如果不使用“正确的”内存，工作流程很容易导致性能损失，甚至更糟糕的是，由于内存不足（OOM）错误而导致故障。挑选“正确”的内存并配置分布式系统具有挑战性，通常需要高水平的专业知识，并且可能需要进行耗时的基准测试。在对性能和内存分配/碎片进行多次实验后，我们发现以下配置是各种基于表格的工作负载（例如常见的 ETL（过滤器、连接、聚合等）和重复数据删除（NeMo Curator））的良好起点。
dask cuda worker <SCHEDULER:IP> --rmm-async --rmm-pool-size <POOL_SIZE> --enable-cudf-spill
使用 RMM 选项：rmm-async 和 rmm-pool-size 可以显著提高性能和稳定性。
rmm-async 使用底层 cudaMallocAsync 内存分配器，可大大减少内存碎片，但性能成本微乎其微。内存碎片很容易导致 OOM 错误。如需深入了解，请参阅 CUDA 流有序内存分配器 博客系列的 第 1 部分 和 第 2 部分 。
rmm-pool-size 参数预先分配 GPU 上的 GPU 内存池。与原始 CUDA 分配 (直接 cudaMalloc) 相比，通过预分配内存池，进行子分配的性能显著提升。这在很大程度上是由于 GPU 内存分配和销毁 (free) 的成本相当高，而且对于数据应用，在工作流中可能会进行成千上万 (如果不是数百万) 的 alloc/free 调用，并且总体而言会降低性能。
将数据从设备转移到主机，即“溢出”功能不仅仅是实现一次。溢出通常可以实施，但通常会牺牲性能。Dask-CUDA 和 cuDF 有几种 溢出机制 ：device-memory-limit、memory-limit、jit-unspill、enable-cudf–spill。enable-cudf-spill 支持 cuDF 内部 溢出到主机内存 的功能。cuDF 的内部溢出机制将在对象或中间产品所需的内存超过 GPU 可用内存时，将数据（cuDF 缓冲区）从设备移至主机。我们发现，对于基于表格的工作负载，与其他 Dask-CUDA 选项（包括 --device-memory-limit）相比，使用 enable-cudf-spill 通常更快、更稳定。
加速网络
如上所述，密集多 GPU 系统通过 NVLink 使用加速网络进行架构设计。借助 Grace Blackwell 和最新的 NVLink ，我们可以获得 900 GB/s 的双向带宽。常见的 ETL 例程，如 join、shuffle 等，需要在多台设备上移动数据——CPU 和 GPU 之间的溢出对于核外算法也至关重要。在这些密集系统上使用 NVLink 对于实现更高的性能至关重要。这些硬件系统可以在启用 UCX 的 Dask 中轻松启用：
# Single Node Cluster
cluster = LocalCUDACluster(protocol="ucx")
# CLI
dask-cuda-worker <SCHEDULER_IP> --protocol ucx
结束语
通过为 Dask CUDA 工作器配置最佳内存设置并启用 UCX 加速网络，用户可以实现稳定的核外 GPU 计算和最大性能。此外，通过使用 Dask 的数组和 dataframe 后端选项，可以轻松实现对 CPU 和 GPU 等多硬件后端的支持。
有关如何更好地使用 Dask RAPIDS 的文档和其他详细信息，我们建议您阅读 Dask-cuDF 和 Dask-CUDA 最佳实践 文档。
