ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

XGBoost 与 RAPIDS Memory Manager(RMM)插件集成指南:从编译启用、内存池配置到 CUDA Async Pool 迁移

XGBoost 与 RAPIDS Memory Manager(RMM)插件集成指南:从编译启用、内存池配置到 CUDA Async Pool 迁移 XGBoost 与 RAPIDS Memory ManagerRMM插件集成指南从编译启用、内存池配置到 CUDA Async Pool 迁移【免费下载链接】xgboostScalable, Portable and Distributed Gradient Boosting (GBDT, GBRT or GBM) Library, for Python, R, Java, Scala, C and more. Runs on single machine, Hadoop, Spark, Dask, Flink and DataFlow项目地址: https://gitcode.com/gh_mirrors/xg/xgboostRAPIDS Memory ManagerRMM为 NVIDIA GPU 提供了一套高效的内存分配器实现其核心卖点是用预先分配的池化内存规避cudaMalloc()的系统调用开销。本指南以 demo/rmm_plugin/README.rst 为骨架结合仓库内 CMake 构建脚本、全局配置源码与两个可直接运行的演示脚本完整讲解如何在 XGBoost 中启用 RMM 插件、通过全局配置use_rmm接管 GPU 内存分配、在多 GPU 场景下正确选择设备以及该插件自 3.5.0 起被弃用后如何平滑迁移到 CUDA async pool。一、RMM 插件是什么用池化分配解决cudaMalloc()的性能瓶颈在 GPU 计算中频繁调用cudaMalloc()/cudaFree()会产生显著的驱动开销尤其是在训练过程中需要反复分配临时缓冲区的场景。RMM 提供的pool sub-allocator池化子分配器正是针对这一问题它在初始化时一次性向驱动申请一大块显存后续的所有分配都从这块池子中划拨从而避免每次都直接调用cudaMalloc()。⚠️弃用声明自 3.5.0 起RMM 插件已被官方弃用新项目应改用 CUDA async pool见本文最后一节。本文介绍的 RMM 集成方式在旧版本中依然可用但读者应知晓其演进方向。从 CMakeLists.txt 可以看出PLUGIN_RMM选项默认关闭OFF需要显式开启才会编译 RMM 支持option(PLUGIN_RMM Build with RAPIDS Memory Manager (RMM) OFF)二、编译启用CMake 构建选项与前置条件2.1 基础构建命令要使用 RMM 插件必须同时开启 CUDA 支持。README 给出的标准构建命令为cmake -B build -S . -DUSE_CUDAON -DUSE_NCCLON -DPLUGIN_RMMON cmake --build build -j$(nproc)2.2 指定 RMM 库位置CMake 会在构建环境中自动查找 RMM 库你可以选择从源码构建 RMM或通过 Conda 安装。如果 CMake 找不到 RMM需要显式指定其前缀路径# 若使用 Conda 安装的 RMM cmake -B build -S . -DUSE_CUDAON -DUSE_NCCLON -DPLUGIN_RMMON -DCMAKE_PREFIX_PATH$CONDA_PREFIX # 若 RMM 安装在自定义路径 cmake -B build -S . -DUSE_CUDAON -DUSE_NCCLON -DPLUGIN_RMMON -DCMAKE_PREFIX_PATH/path/to/rmm2.3 编译前提与约束源码级校验CMakeLists.txt 在配置阶段对PLUGIN_RMM做了多项硬性校验任何一项不满足都会直接报错终止必须与USE_CUDAON同时启用PLUGIN_RMM不能脱离 CUDA 构建编译器必须是 GCC 或 ClangRMM 插件不支持其他编译器如 MSVC操作系统必须是 Linux该插件仅面向 Linux 平台同时 CMake 会打印弃用提示deprecation message提醒自 v3.5.0 起应改用 CUDA async pool。编译成功后预处理器宏XGBOOST_USE_RMM1会被注入到编译目标中见 cmake/Utils.cmake后续所有源码都通过该宏判断是否启用 RMM 代码路径。三、两个开箱即用的演示脚本demo/rmm_plugin目录下提供了两个演示分别覆盖单 GPU 与多 GPUDask场景脚本场景关键点rmm_singlegpu.py单机单卡通过rmm.reinitialize(pool_allocatorTrue)初始化 RMM 池rmm_mgpu_with_dask.py多卡Dask通过LocalCUDACluster(rmm_pool_size2GB)为每个 worker 配置池3.1 单 GPU 演示自动接管 RMM 池rmm_singlegpu.py 的核心流程如下import rmm from sklearn.datasets import make_classification import xgboost as xgb # 初始化 RMM 池分配器 rmm.reinitialize(pool_allocatorTrue) # 可选强制 XGBoost 的所有 GPU 内存分配都走 RMM见 README # xgb.set_config(use_rmmTrue) X, y make_classification(n_samples10000, n_informative5, n_classes3) dtrain xgb.DMatrix(X, labely) params { max_depth: 8, eta: 0.01, objective: multi:softprob, num_class: 3, tree_method: hist, device: cuda, } # XGBoost 会自动使用 RMM 池分配器 bst xgb.train(params, dtrain, num_boost_round100, evals[(dtrain, train)])只要 XGBoost 以 RMM 插件编译且进程内已初始化 RMM 池rmm.reinitialize(pool_allocatorTrue)XGBoost 在histdevicecuda下就会自动使用该池分配器。3.2 多 GPUDask演示pool size 即配置rmm_mgpu_with_dask.py 展示了分布式场景的接入方式——RMM 池的初始化被集成到LocalCUDACluster的构造参数中from dask.distributed import Client from dask_cuda import LocalCUDACluster import xgboost as xgb import dask.array def main(client): X, y make_classification(n_samples10000, n_informative5, n_classes3) # 实际生产中应优先用 dask 集合加载数据而不是 from_array X dask.array.from_array(X) y dask.array.from_array(y) dtrain xgb.dask.DaskDMatrix(client, X, labely) params { max_depth: 8, eta: 0.01, objective: multi:softprob, num_class: 3, tree_method: hist, eval_metric: merror, device: cuda, } output xgb.dask.train( client, params, dtrain, num_boost_round100, evals[(dtrain, train)] ) bst output[booster] history output[history] for i, e in enumerate(history[train][merror]): print(f[{i}] train-merror: {e}) if __name__ __main__: # 要在 GPU Dask 集群中使用 RMM 池分配器只需给 LocalCUDACluster 加上 rmm_pool_size 选项 with LocalCUDACluster(rmm_pool_size2GB) as cluster: with Client(cluster) as client: main(client)多卡场景下RMM 池由dask-cuda在每个 worker 上统一初始化XGBoost 侧无需额外配置。四、全局配置use_rmm强制所有分配走 RMM4.1 为什么需要这个开关XGBoost 内部对内存分配做了分级处理大多数大尺寸分配会走 RMM 分配器但性能关键路径上的一些小分配使用了另一套缓存分配器caching allocator以便更精细地控制内存分配行为。如果你希望所有GPU 分配都强制走 RMM可以通过全局配置use_rmm覆盖默认行为。4.2 Python 使用方式with xgb.config_context(use_rmmTrue): clf xgb.XGBClassifier(tree_methodhist, devicecuda)注意权衡是否强制use_rmmTrue取决于内存池大小与分配器类型的选择它可能让内存使用更稳定一致性更好但也可能带来轻微的性能退化。4.3 全局配置的源码定义include/xgboost/global_config.h 定义了GlobalConfiguration结构其中与 GPU 内存分配相关的三个全局字段为std::int32_t verbosity{1}; bool use_rmm{false}; bool use_cuda_async_pool{false};参数声明部分DMLC 参数系统对use_rmm的描述是Whether to use RAPIDS Memory Manager to allocate GPU memory in XGBoost默认值为false。4.4 各语言/框架的接入方式Pythonxgb.config_context(use_rmmTrue)或xgb.set_config(use_rmmTrue)Rxgb.set.config(use_rmm TRUE)对应测试见 R-package/tests/testthat/test_config.RSparkJVMxgboost4j-spark-gpu的GpuXGBoostPlugin在 GPU 场景默认写入Map(use_rmm - true)见 GpuXGBoostPlugin.scala而 Python Spark 接口则从全局配置读取后透传见 python-package/xgboost/spark/core.py。官方参数文档见 doc/parameter.rst其中同样标注了该配置自 3.5.0 起弃用并指明use_cuda_async_pool默认falsev3.2.0 引入是替代方案——注意当 XGBoost 以 RMM 支持编译时use_cuda_async_pool不可用因为它等价于 RMM 的CudaAsyncMemoryResource池。五、源码原理RMM 分配器如何接入 XGBoost 的底层内存管理5.1 分配器适配层XGBoost 的 GPU 内存管理集中在 src/common/device_vector.cuh。当XGBOOST_USE_RMM 1时设备分配器被适配为 RMM 的 Thrust 分配器#if defined(XGBOOST_USE_RMM) XGBOOST_USE_RMM 1 template typename T class ThrustAllocMrAdapter : public rmm::mr::thrust_allocatorT { ... }; template typename T using XGBBaseDeviceAllocator ThrustAllocMrAdapterT; #else // 非 RMM 编译路径 template typename T class XGBAsyncPoolAllocator : public thrust::device_malloc_allocatorT { ... };也就是说device_vectorT、caching_device_vectorT等容器见 device_vector.cuh在 RMM 编译路径下其底层内存最终由 RMM 的内存资源memory resource提供。5.2 缓存分配器与use_rmm的互斥逻辑在 src/common/device_vector.cuh 的缓存分配器实现中是否使用内部全局缓存分配器由以下条件决定use_cub_allocator_{!(GlobalConfigThreadLocalStore::Get()-use_rmm || GlobalConfigThreadLocalStore::Get()-use_cuda_async_pool)}当use_rmm或use_cuda_async_pool任一为真时XGBoost 会放弃内部缓存分配器把分配请求交给对应的外部内存池——这正是强制所有分配走 RMM这一行为的实现位置。5.3 外部内存场景的告警对于外部内存external memory训练src/data/ellpack_page_source.h 会检查全局配置若use_rmm与use_cuda_async_pool均未开启则打印警告若开启了use_rmm但 XGBoost 并未以 RMM 支持编译也会打印 XGBoost is not built with RMM support 警告。这提醒用户use_rmm配置项本身在任何构建下都可用但只有 RMM 编译路径才会真正生效。5.4 配置校验与运行时告警在 src/c_api/c_api.cc 的全局配置校验逻辑中以 RMM 编译时use_cuda_async_pool会被直接CHECK拒绝二者互斥而开启use_rmm会打印弃用告警DeprecatedFunc(RMM plugin, 3.5.0, CUDA async pool.)。此外src/c_api/c_api.cu 会把USE_RMM标记写入构建信息Python 端可通过xgb.build_info()[USE_RMM]查询当前构建是否包含 RMM 支持——tests/python-gpu/test_gpu_multi_target.py正是用它做条件测试的。六、多 GPU 注意事项不要用devicecuda:1选卡使用 RMM 时有一个关键陷阱RMM 内存池是预先在特定设备上分配的。如果在 XGBoost 中通过devicecuda:1这类参数切换 CUDA device ordinal可能触发cudaErrorIllegalAddress内存错误。正确做法是使用环境变量CUDA_VISIBLE_DEVICES选择设备CUDA_VISIBLE_DEVICES1 python your_training_script.py场景推荐做法说明单机多卡CUDA_VISIBLE_DEVICES设备号通过环境变量隐藏其他 GPU避免池与设备错位分布式训练由dask-cuda等分布式框架负责设备管理每个 worker 绑定独立设备并各自初始化池Scala-Spark GPU参考 Scala-Spark GPU 教程Spark 的 GPU 调度机制会自动分配设备七、内存超额订阅Over-Subscription利用主机内存扩展显存⚠️ 该特性仍处于实验阶段正在积极开发中。新一代 NVIDIA 平台如 Grace-Hopper 超级芯片通过NVLink-C2C实现了 CPU 与 GPU 的内存一致性模型coherent memory model。在此基础上用户可以借助最新版 RMM 提供的SamHeadroomMemoryResource把系统内存纳入 GPU 数据存储。这对 XGBoost 的意义在于可以让训练过程利用主机内存承载 GPU 计算所需的数据从而突破显存容量限制。代价是性能可能下降——因为 CPU 内存带宽更低且存在页面迁移page migration开销。从 src/data/ellpack_page_source.h 可以看到XGBoost 会检测SupportsPageableMem()HMM与SupportsAts()ATS能力并给出相应的警告信息这与超额订阅/异构内存管理的运行时支持是配套的。八、迁移到 CUDA Async Pool弃用后的替代方案自 v3.5.0 起RMM 插件正式弃用官方推荐使用CUDA async pool作为替代。迁移要点如下构建层面去掉-DPLUGIN_RMMON改用默认 CUDA 构建即可PLUGIN_RMM与USE_CUDA的联动约束、GCC/Clang Linux 的限制也随之消失运行层面将全局配置从use_rmmTrue切换为use_cuda_async_poolTruewith xgb.config_context(use_cuda_async_poolTrue): clf xgb.XGBClassifier(tree_methodhist, devicecuda)互斥关系use_rmm与use_cuda_async_pool不能同时启用以 RMM 编译的旧二进制中开启 async pool 会被直接拒绝见 src/c_api/c_api.cc功能对齐两者在 XGBoost 内部共享同一套跳过内部缓存分配器、将分配交给外部内存池的机制见 device_vector.cuh因此迁移后内存池化分配的行为保持一致。结语RMM 插件是 XGBoost GPU 内存管理演进中的重要一环它通过池化分配规避了cudaMalloc()的高昂开销并通过use_rmm全局配置让用户对分配行为拥有细粒度控制。理解其编译选项-DPLUGIN_RMMON、设备选择约束CUDA_VISIBLE_DEVICES与底层分配器适配逻辑device_vector.cuh无论你是在旧版本上继续使用 RMM还是向 CUDA async pool 迁移都能做到心中有数。仓库中的两个演示脚本 rmm_singlegpu.py 与 rmm_mgpu_with_dask.py 可直接作为起步模板。【免费下载链接】xgboostScalable, Portable and Distributed Gradient Boosting (GBDT, GBRT or GBM) Library, for Python, R, Java, Scala, C and more. Runs on single machine, Hadoop, Spark, Dask, Flink and DataFlow项目地址: https://gitcode.com/gh_mirrors/xg/xgboost创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表