这篇文章最初发表在 NVIDIA 技术博客上。
作为数据科学家，我们经常面临在大型数据集上训练模型的挑战。一种常用的工具是XGBoost，这是一种稳健且高效的梯度提升框架，因其在处理大型表格数据时的速度和性能而被广泛采用。
理论上，使用多个 GPU 可以显著提高计算能力，从而加快模型训练。然而，许多用户发现，当试图通过 Dask 和 XGBoost 进行训练时，Dask 是一个用于并行计算的灵活的开源 Python 库，而 XGBoost 则提供 Dask API 来训练 CPU 或 GPU 的 Dask DataFrames。
训练 Dask XGBoost 的一个常见障碍是处理不同阶段的内存不足（OOM）错误，包括
- 加载训练数据
- 将 DataFrame 转换为 XGBoost 的 DMatrix 格式
- 在实际模型培训期间
解决这些记忆问题可能很有挑战性，但非常有益，因为多 GPU 训练的潜在好处很诱人。
顶级外卖
这篇文章探讨了如何在多个 GPU 上优化 Dask XGBoost 并管理内存错误。在大型数据集上训练 XGBoost 带来了各种挑战。我使用Otto Group Product Classification Challenge 数据集来演示 OOM 问题以及如何解决它。该数据集有 1.8 亿行和 152 列，加载到内存中时总计 110 GB。
我们要解决的关键问题包括：
- 使用最新版本的 RAPIDS 以及正确版本的 XGBoost 进行安装。
- 设置环境变量。
- 处理 OOM 错误。
- 利用 UCX-py 实现更多加速。
请务必按照每个章节随附的笔记本进行操作。
先决条件
利用 RAPIDS 的力量进行多 GPU 训练的第一步是正确安装 RAPIDS 库。需要注意的是，有几种安装这些库的方法 – pip，conda，docker，和从源代码构建，每种方法都与 Linux 和 Windows Subsystem for Linux 兼容。
每种方法都有其独特的考虑因素。对于本指南，我建议在遵守 conda 安装说明的同时使用 Mamba。 Mamba 提供了与 conda 类似的功能，但速度要快得多，尤其是在依赖关系解析方面。具体来说，我选择了 全新安装 mamba。
安装最新的 RAPIDS 版本
作为最佳实践，我们建议您始终安装最新的 RAPIDS 库以使用最新的功能。您可以在 RAPIDS 安装指南 中找到最新的安装说明。
这篇文章使用 23.04 版本，可以通过以下命令进行安装：
mamba create -n rapids-23.04 -c rapidsai -c conda-forge -c nvidia \
rapids=23.04 python=3.10 cudatoolkit=11.8
本说明安装所有所需的库，包括 Dask、Dask-cuDF 、XGBoost 等。特别是，您需要检查使用以下命令安装的 XGBoost 库：
曼巴列表 xgboost
输出如表 1 所示：
| 名称 | 版本 | 建筑 | 频道 |
| XGBoost | 1.7.1 月 24 日星期四 | CUDA _11_py310_3 | 夜间急流 |
表 1。安装正确的 XGBoost，其通道应为 rapidsai 或 rapidsai 夜间
避免手动更新 XGBoost
一些用户可能会注意到 XGBoost 的版本不是最新的，它是 1.7.5。使用 pip 或 conda-forge 手动更新或安装 XGBoost 与 UCX 一起训练 XGBoost 时出现问题。
错误消息将显示如下内容：
异常：“XGBoostError（’ 14:14.27 /opt/conda/conda-bld/work/rabit/include/rabit/internal/utils.h:86:Allreduce 失败’）”
相反，请使用从 RAPIDS 安装的 XGBoost。验证 XGBoost 版本正确性的快速方法是曼巴列表 xgboost并检查 xgboost 的“通道”，该通道应为“rapidsai”或“Rapidsainightly”。
rapidsai 通道中的 XGBoost 是在启用 RMM 插件的情况下构建的，在多 GPU 训练方面提供了最佳性能。
多- GPU 训练演练
首先，我将浏览 Otto 数据集的 multi-GPU 训练笔记本，并介绍使其工作的步骤。稍后，我们将讨论一些高级优化，包括 UCX 和溢出。
您还可以在 XGB-186-CLICKS-DASK GitHub 上找到笔记本。或者，我们还提供了一个具有完全命令行可配置性的 python 脚本。
我们要使用的主要库是 xgboost、dask、dask_ CUDA 和 dask- cuDF 。
import os
import dask
import dask_cudf
import xgboost as xgb
from dask.distributed import Client
from dask_cuda import LocalCUDACluster
环境设置
首先，让我们设置我们的环境变量，以确保我们的 GPU 是可见的。此示例使用八个 GPU ，每个 GPU 32 GB 内存，这是在没有 OOM 复杂性的情况下运行此笔记本的最低要求。在下面的启用内存溢出一节中，我们将讨论将此要求降低到 4 GPU 的技术。
GPUs = ','.join([str(i) for i in range(0,8)])
os.environ['CUDA_VISIBLE_DEVICES'] = GPUs
接下来，定义一个助手函数，为多个 GPU 单个节点创建一个本地 GPU cluster。
def get_cluster():
cluster = LocalCUDACluster()
client = Client(cluster)
return client
然后，为您的计算创建一个 Dask 客户端。
client = get_cluster()
正在加载数据
现在，让我们加载 Otto 数据集。我们将使用 dask_cuDF 的 read_parquet 函数，该函数利用多个 GPU 将 parquet 文件读取到 dask_cuDF.DataFrame 中。
users = dask_cudf.read_parquet('/raid/otto/Otto-Comp/pqs/train_v152_*.pq').persist()
该数据集由 152 列组成，这些列代表工程特征，提供了关于用户查看或购买特定产品的频率信息。目标是根据用户的浏览历史来预测用户下一步会点击哪个产品。您可以在writeup中查看此数据集的详细信息。
即使在这个早期阶段，也可能出现内存不足的错误。这个问题通常是由于镶木地板锉刀。为了解决这个问题，我们建议使用较小的行组来重写镶木地板文件。如果您想了解更深入的解释，请参阅 Parquet Large Row Group Demo 笔记本。
加载数据后，我们可以检查其形状和内存使用情况。
users.shape[0].compute()
users.memory_usage().sum().compute()/2**30
“点击”列是我们的目标，这意味着如果用户点击了推荐的项目。我们忽略 ID 列，并使用其余列作为功能。
FEATURES = users.columns[2:]
TARS = ['clicks']
FEATURES = [f for f in FEATURES if f not in TARS]
接下来，我们创建了一个用于训练 xgboost 模型的输入数据格式，即DaskQuantileDMatrix。当使用直方图树方法时，DaskQuantileDMatrix是 DaskDMatrix 的一个替代品，它有助于减少总体内存使用量。
这一步骤对于避免 OOM 错误至关重要。如果我们使用 DaskDMatrix，即使使用 16 个 GPU 也会发生 OOM 错误。相反，DaskQuantileDMatrix 能够在没有 OOM 错误的情况下以八个或更少的 GPU 训练 xgboot。
dtrain = xgb.dask.DaskQuantileDMatrix(client, users[FEATURES], users['clicks'])
XGBoost 模型训练
然后，我们设置 XGBoost 模型参数并开始训练过程。给定目标列“点击”是二进制的，我们使用二进制分类目标。
xgb_parms = {
'max_depth':4,
'learning_rate':0.1,
'subsample':0.7,
'colsample_bytree':0.5,
'eval_metric':'map',
'objective':'binary:logistic',
'scale_pos_weight':8,
'tree_method':'gpu_hist',
'random_state':42
}
现在，您已经准备好使用全部八个 GPU 来训练 XGBoost 模型了。
输出：
[99] train-map:0.20168
CPU times: user 7.45 s, sys: 1.93 s, total: 9.38 s
Wall time: 1min 10s
就是这样！您已经完成了使用多个 GPU 训练 XGBoost 模型。
启用内存溢出
在上一个XGB-186-CLICKS-DASK笔记本电脑中，我们在 Otto 数据集上训练 XGBoost 模型，至少需要八个 GPU。假设该数据集占用 110GB 的内存，而每个 V100 GPU 提供 32GB，那么数据与 GPU 内存的比率仅为 43%（计算为 110/（32*8））。
最理想的情况是，我们只需要使用四个 GPU 就可以将其减半。然而，在我们之前的设置中直接减少 GPU 总是会导致 OOM 错误。此问题源于创建生成DaskQuantileDMatrix从 Dask cuDF 数据帧和在训练 XGBoost 的其他步骤中。这些变量本身消耗了 GPU 内存的相当大的份额。
优化相同的 GPU 资源以训练更大的数据集
在XGB-186-CLICKS-DASK-SPILL 笔记本电脑中，我介绍了之前设置的一些小调整。现在，通过启用溢出，您只需使用四个 GPU 就可以在同一数据集上进行训练。这项技术允许您使用相同的 GPU 资源来训练更大的数据。
溢出是一种技术，当 GPU 内存中所需的空间被其他数据帧或序列占用，导致原本会成功的操作耗尽内存时，该技术会自动移动数据。它能够在不适合内存的数据集上进行核心外计算。 RAPIDS cuDF 和 dask- cuDF 现在支持从 GPU 溢出到 CPU 内存。
实现溢出非常容易，我们只需要用两个新参数重新配置集群，设备内存限制和jit_unspill：
def get_cluster():
ip = get_ip()
cluster = LocalCUDACluster(ip=ip,
device_memory_limit='10GB',
jit_unspill=True)
client = Client(cluster)
return client
device_memory_limit=’10GB’设置在触发溢出之前每个 GPU 可以使用的 GPU memory 的量的限制。我们的配置有意为设备内存限制10GB，基本上小于 GPU 的总 32GB。这是一种深思熟虑的策略，旨在抢占 XGBoost 训练期间的 OOM 错误。
同样重要的是要理解 XGBoost 的内存使用不是由 Dask-CUDA 或 Dask-cuDF 直接管理的。因此，为了防止内存溢出，Dask- CUDA 和 Dask- cuDF 需要在 XGBoost 操作达到内存限制之前启动溢出过程。
Jit_unspill启用实时取消溢出，这意味着当 GPU 内存不足时，集群将自动将数据从 GPU’内存溢出到主内存，并及时取消溢出以进行计算。
就这样！笔记本的其余部分与上一个笔记本相同。现在，它只需四个 GPU 就可以进行训练，节省了 50%的计算资源。
想要了解更多详细信息，请参阅 XGB-186-CLICKS-DASK-SPILL 笔记本。
使用统一通信 X（UCX）实现最佳数据传输
UCX-py 是一种高性能通信协议，特别适用于 GPU – GPU 通信，能够提供优化的数据传输能力。
为了有效地使用 UCX，我们需要设置另一个环境变量RAPIDS _NO_INITIALIZE：
os.environ["RAPIDS_NO_INITIALIZE"] = "1"
它阻止 cuDF 在导入时运行各种诊断，这需要创建 NVIDIA CUDA 上下文。当运行分布式并使用 UCX 时，我们必须在创建 CUDA 上下文之前打开网络堆栈（由于各种原因）。通过设置该环境变量，导入 cuDF 的任何子进程都不会在 UCX 有机会创建 CUDA 上下文之前创建。
重新配置群集：
def get_cluster():
ip = get_ip()
cluster = LocalCUDACluster(ip=ip,
device_memory_limit='10GB',
jit_unspill=True,
protocol="ucx",
rmm_pool_size="29GB"
)
client = Client(cluster)
return client
protocol=’cx’参数将 ucx 指定为用于在集群中的工作程序之间传输数据的通信协议。
使用 prmm_pool_size=’29GB’参数为每个工作线程设置 RAPIDS 内存管理器（RMM）池的大小。RMM 允许有效地使用 GPU 存储器。在这种情况下，池大小被设置为 29GB，这小于 32GB 的总 GPU 内存大小。这种调整至关重要，因为它解释了 XGBoost 创建某些存在于 RMM 池控制之外的中间变量的事实。
通过简单地启用 UCX，我们的训练时间得到了显著的加速——在溢出时，速度提高了 20%，而在不需要溢出时，速度提高了 40.7%。请参阅XGB-186-CLICKS-DASK-UCX-SPILL笔记本了解详细信息。
配置本地目录
有时会出现警告消息，例如“UserWarning:创建临时目录花费了惊人的长时间。”磁盘性能正成为一个瓶颈。
为了规避这个问题，我们可以设置本地目录属于dask-CUDA，指定本地计算机上存储临时文件的路径。这些临时文件在 Dask 的磁盘溢出操作中使用。
建议的做法是设置本地目录到快速存储设备上的某个位置。例如，我们可以设置本地目录到/raid/dask_dir如果它在高速本地 SSD 上。进行这种简单的更改可以显著减少临时目录操作所需的时间，从而优化您的整体工作流程。
最终的集群配置如下：
def get_cluster():
ip = get_ip()
cluster = LocalCUDACluster(ip=ip,
local_directory=’/raid/dask_dir’
device_memory_limit='10GB',
jit_unspill=True,
protocol="ucx",
rmm_pool_size="29GB"
)
client = Client(cluster)
return client
结果
如表 2 所示，两种主要的优化技术是 UCX 和溢出。我们设法用四个 GPU 和 128GB 的内存来训练 XGBoost。我们还将很好地展示性能扩展到更多 GPU 。
| 溢出 | 溢出 | |
| UCX 关闭 | 135s/8 GPU /256 GB | 270s/4 GPU /128 GB |
| UCX 打开 | 80s/8 GPU /256 GB | 217s/4 GPU /128 GB |
表 2。四种优化组合概述
在每个单元中，这些数字表示端到端执行时间、所需的最小 GPU 数量以及可用的总 GPU memory。所有四个演示都完成了加载和训练 110 GB Otto 数据的相同任务。
总结
总之，利用 Dask 和 XGBoost 的多个 GPU 可能是一次激动人心的冒险，尽管偶尔会出现内存不足等问题。
您可以通过以下方式缓解这些记忆挑战，并挖掘多 GPU 模型训练的潜力：
- 在输入镶木地板文件中仔细配置行组大小等参数
- 确保 RAPIDS 和 XGBoost 的正确安装
- 利用 Dask Quantile DMatrix
- 启用溢出
此外，通过应用 UCX-Py 等高级功能，您可以显著加快训练时间。
阅读原文
