数据中心/云端

用 Dask XGBoost 解锁多 GPU 模型训练

 

作为数据科学家,我们经常面临在大型数据集上训练模型的挑战。一种常用的工具是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 等高级功能,您可以显著加快训练时间。

Tags