[论文解读] Out-of-Core GPU Gradient Boosting

Rong Ou|arXiv (Cornell University)|May 19, 2020

Domain Adaptation and Few-Shot Learning参考文献 9被引用 4

一句话总结

本文提出XGBoost中首个基于GPU的离外存储梯度提升算法，通过结合基于梯度的采样与高效的PCIe数据流技术，使模型可在远超GPU内存容量的数据集上进行训练。该方法在保持核心模型准确率与训练速度的同时，使单张16 GiB GPU的可训练数据集规模扩展至10倍。

ABSTRACT

GPU-based algorithms have greatly accelerated many machine learning methods; however, GPU memory is typically smaller than main memory, limiting the size of training data. In this paper, we describe an out-of-core GPU gradient boosting algorithm implemented in the XGBoost library. We show that much larger datasets can fit on a given GPU, without degrading model accuracy or training time. To the best of our knowledge, this is the first out-of-core GPU implementation of gradient boosting. Similar approaches can be applied to other machine learning algorithms

研究动机与目标

• 解决GPU内存容量限制在训练大规模机器学习模型时的瓶颈。
• 克服PCIe带宽限制对高效离外存储GPU计算的阻碍。
• 将XGBoost现有的基于CPU的外部内存支持扩展至GPU，且性能损耗最小化。
• 在不牺牲模型准确率或训练时间的前提下，实现对超过GPU内存容量的数据集的训练。
• 设计可扩展、可生产部署的实现，用于XGBoost开源生态。

提出的方法

• 利用基于梯度的采样技术，降低树构建过程中的工作内存占用。
• 通过流水线化、异步数据加载策略，从主内存或磁盘向GPU流式传输数据。
• 采用ELLPACK格式紧凑表示分箱特征，以减少内存访问流量。
• 采用两步GPU树构建方法：预处理（分箱与压缩）后接基于直方图的分裂评估。
• 通过确保采样不改变底层优化目标，保持模型行为的一致性。
• 将离外存储GPU流水线集成至XGBoost库中，支持向后兼容与模块化代码结构。

实验结果

研究问题

• RQ1 尽管存在PCIe带宽限制，是否可以使离外存储GPU训练在梯度提升中变得实用？
• RQ2 基于梯度的采样在多大程度上可减少GPU内存使用，同时不降低模型准确率？
• RQ3 在大规模数据集上，离外存储GPU训练的性能与核心GPU及CPU训练相比如何？
• RQ4 所提方法是否可无缝集成至广泛使用的开源库XGBoost中，且不破坏现有工作流？
• RQ5 与核心训练相比，该方法在单张GPU上可训练的最大数据集规模是多少？

主要发现

• 离外存储GPU实现支持最大达GPU内存10倍的数据集（例如，在16 GiB GPU上训练903 GiB数据集），采样比例f=0.1。
• 模型准确率与核心训练几乎完全一致：在Higgs数据集上，离外存储（f=1.0）AUC为0.8396，核心训练AUC为0.8398。
• 当f=0.5时，离外存储GPU训练耗时427.41秒，显著快于CPU离外存储（1228.53秒）与核心GPU（241.52秒）。
• 即使采用采样，离外存储GPU版本仍快于基于CPU的训练，证明GPU加速能力得以保留。
• 该方法实现接近核心性能，仅带来极小开销，具备生产环境部署可行性。
• 该实现已合并至开源XGBoost库中，推动更广泛应用与进一步研究。