梯度提升决策树 (GBDT) 驱动着从实时欺诈过滤到 PB 级需求预测的各种功能。由于其先进的准确性、SHAP-ready 解释能力以及在笔记本电脑、多 GPU 节点或 Spark 集群上运行的灵活性，XGBoost 开源库一直是首选工具。XGBoost 3.0 版本的开发以可扩展性为北极星。单个 NVIDIA GH200 Grace Hopper 超级芯片现在可以处理从千兆字节到 1 太字节 (TB) 规模的数据集。
借助一致性内存架构，新的外部内存引擎可以通过 900 GB/s NVIDIA NVLink-C2C 流式传输数据，因此 1 TB 模型可以在几分钟内完成训练，速度比 112 核（双插槽）CPU 机箱快 8 倍。这减少了对复杂多节点 GPU 集群的需求，并使可扩展性变得更简单。
本文介绍了里程碑式的 XGBoost 3.0 版本中的新功能和增强功能，包括深入了解外部内存以及如何利用 Grace Hopper 超级芯片达到 1 TB 的规模。
金融系统用例
XGBoost 驱动欺诈检测、信用风险预测和算法交易等关键金融系统。RBC 是全球市值最大的银行之一，其运行的潜在客户评分系统要求速度、准确性和可解释性。为了实现 ML 工作流的现代化，并处理对数十万条记录的持续模型调整，RBC 选择了 XGBoost。
RBC 基金管理公司 Gen AI 规划和估值总监 Christopher Ortiz 表示：“我们相信，在 NVIDIA GPU 的支持下，XGBoost 将使我们的预测潜在客户评分模型能够处理我们预计的数据量。通过利用 GPU，我们看到端到端速度提升了 16 倍，在我们的流水线测试中，我们看到模型训练的 TCO 显著降低了 94%。这代表了效率和成本效益方面的革命性飞跃。随着数据规模的扩大，我们可以更快地优化功能，从而提供更好的结果。
在 TB 级数据集上运行 XGBoost
为了使 XGBoost 能够实现 GPU 加速，我们引入了 GPU 直方图方法。与 CPU 上的 XGBoost 相比，这种方法实现了显著的加速。但是，仍然存在两个限制：
- 即使压缩了数据，GPU 显存仍然是可加速数据集的主要限制因素。
- 在 3.0 版本之前，实验外部内存会将流式批次串联在 GPU 内存中。
可以采用两种机制来解决这些限制：
- Quantile DMatrix 将每个特征预先分成固定数量的分位数桶。这可以压缩数据集，在不显著影响模型准确性的情况下大幅减少内存使用。
- 使用 Dask 和 Spark 等分布式框架对计算进行分片。
新的外部内存分位数 DMatrix
XGBoost 3.0 在外部内存分位数 DMatrix 中引入了第三种机制，可在单个 GH200 Grace Hopper 超级芯片上扩展到 TB 级数据集。这克服了设置分布式框架的复杂性，并利用了 Grace Hopper 超级芯片的超快 C2C 带宽。
外部内存分位数 DMatrix 是基于现有数据迭代器构建的，这些迭代器定义了读取数据集文件的方法。它负责处理数据集内存的所有管理，并用作 XGBoost 增强器对象的参数（图 1）。您可以将外部内存分位数 DMatrix 视为熟悉的 QuantileDMatrix，它预先对每个特征进行分箱，但使用相同的分位数逻辑。这意味着您可以保留所有现有超参数以获得相同的准确性。同时，数据本身位于主机 RAM 中，并在每次迭代时流式传输到 GPU。有关更多详细信息，请参阅 XGBoost 文档。
NVIDIA Grace Hopper 超级芯片如何使流式传输成为现实
这种设置非常适合基于 NVIDIA Grace 的超级芯片。GH200 超级芯片集成了 72 核 Grace CPU 和 Hopper GPU，并通过 NVLink C2C 连接。这意味着双向带宽为 900 GB/s，大约是 PCIe 5.0 x16 的 7 倍，延迟也远低于 PCIe 5.0 x16。
处理 1 TB 的训练作业通常需要一个大约 2 TB DRAM 的 CPU 机箱，或一个带有 8 到 16 个 NVIDIA H100 GPU 的小型 GPU 集群。这种设置可以提高速度，但增加了管理分布式框架的复杂性。借助 XGBoost 3.0 外部内存流，单个 GH200 超级芯片（80 GB HBM3 加上 480 GB LPDDR5X，由 900 GB/s NVLink-C2C 提供支持）现在可以自行处理相同的数据集，取代了 RAM 庞大的服务器和多 GPU Pod。
在 1 TB 数据集上对 Grace Hopper 进行基准测试
GPU 在处理密集（或近乎密集）的表时表现出色，因为 XGBoost 会大幅压缩这些表，从而减少总线流量并使直方图能够位于快速共享内存中。这意味着，对于 GPU 来说，形状几乎无关紧要：宽短和窄高的表在完成时间上几乎没有差别。
然而，ExtMemQuantileDMatrix 对形状很敏感。在对特征矩阵 (x) 和标签 (y) 进行 XGBoost 训练时，只有特征矩阵会根据行数进行分页，而标签或其他数据则不会。这意味着，对于较小的数据集，标签的大小（行数）是决定单个 GH200 超级芯片中可以容纳多少数据集的限制因素。
图 2 演示了这一点，其中总数据大小保持不变，但数据集中的行数和列数各不相同。
外部显存的最佳实践
要利用 GH200 超级芯片系统中的外部显存，请先阅读以下提示。有关更多详细信息，请参阅 XGBoost 文档。设置 grow_policy=’depthwise’ 以逐层构建树。这对于外部显存来说更好：
xgb_model = xgb.XGBRegressor(tree_method='hist',
device='cuda',
seed=42,
grow_policy='depthwise')
在与 RAPIDS 结合使用时，请始终从新的 RAPIDS Memory Manager (RMM) 池开始：
import cupy as cp, rmm, xgboost as xgb
from rmm.allocators.cupy
import rmm_cupy_allocator
mr = rmm.mr.ArenaMemoryResource(rmm.mr.CudaAsyncMemoryResource())
rmm.mr.set_current_device_resource(mr)
cp.cuda.set_allocator(rmm_cupy_allocator)
with xgb.config_context(use_rmm=True):
dtrain = xgb.ExtMemQuantileDMatrix(it_train, max_bin=256)
bst = xgb.train({“device”:“cuda”, “tree_method”:“hist”}, dtrain)
在 CUDA 12.8 或更高版本上运行，并使用支持 HMM 的驱动程序（Grace Hopper）。
XGBoost 3.0 中还有哪些新功能？
除了外部内存的改进之外，XGBoost 3.0 还提供了许多性能升级和 API 清理，包括：
- 对分布式外部内存的实验性支持。现在，您可以在集群中训练核心外模型，并在 GPU RAM 紧张时回退到主机内存。
- 在 DMatrix 构造期间减少 GPU 显存，并加快批量输入的初始化速度，从而使 XGBoost 3.0 速度更快。
- 对于“大部分密集”的数据，GPU hist 和 approx 方法的速度大约提高了 2 倍，并且降低了内存使用量。
- 外部内存现在支持分类特征、每个目标（包括分位数回归）和所有预测类型（包括 SHAP）。
开始使用 XGBoost 3.0
XGBoost 3.0 中的外部内存改进：GPU 外部内存的工作方式是，当数据超过 GPU 内存时，将其作为默认工作方式。亮点包括：
- 精度相同，占用空间更小：ExtMemQuantileDMatrix 在每个提升轮次中，都会及时将预先分箱的压缩页面从主机 RAM 传输到 GPU。
- Grace Hopper 扩展 GPU 显存：NVLink 将普通主机显存转变为 GPU 显存的快速扩展，因此一个 Grace Hopper 超级芯片可以无缝完成以前需要多个 GPU（跨一个或多个节点）才能完成的 TB 级工作。
- 直接升级：已经在使用 RAPIDS Memory Manager？在 GPU 工作流中，切换外部显存路径是一个简单的更改。现有工作流将继续不受影响地运行。
XGBoost 3.0 使您能够使用您一直使用的 XGBoost 调用在单个 Grace Hopper 芯片上处理 TB 级 GBDT 训练。
首先，下载 XGBoost 3.0 并查看安装指南。有关外部内存的更多信息，请参阅 XGBoost 文档。
加入我们的社区 Slack 频道，发布您对 XGBoost GPU 加速的反馈或问题。如果您是加速数据科学的新手，请查看加速数据科学学习路径，了解动手操作的研讨会。
