NVIDIA cuML 分布式机器学习模块 `cuml.dask` 全解析：多节点多 GPU 算法 API 指南

NVIDIA cuML 分布式机器学习模块cuml.dask全解析多节点多 GPU 算法 API 指南【免费下载链接】cumlNVIDIA cuML: GPU-Accelerated Machine Learning项目地址: https://gitcode.com/GitHub_Trending/cu/cumlcuml.dask是 NVIDIA cuML 面向多节点多 GPUMNMG场景的分布式算法子包基于 Dask 生态把 cuDF、CuPy 与 RAFT 通信层串联起来让 KMeans、PCA、随机森林、线性模型等经典算法可以跨 GPU 集群并行训练与推理。本文以仓库 API 文档页 docs/source/api/cuml.dask.rst 为骨架结合python/cuml/cuml/dask/下各模块源码完整梳理cuml.dask的模块结构、每个公开类的核心参数、基类设计原理与一套可落地的分布式训练工作流帮助读者快速上手 cuML 的多 GPU 机器学习开发。一、cuml.dask是什么cuml.dask是 cuML 中使用 Dask 实现的多节点、多 GPU 算法的统一入口其官方定位在 API 文档页中一句话概括为Multi-node, multi-GPU algorithms using Dask.它并非一套独立的算法实现而是由以下几层拼装而成数据层以dask_cudf.DataFrame或 CuPy 后端的dask.array承载分布式数据每个 Dask worker 持有若干分区调度层借助dask.distributed的Client与 Future 机制在集群上分发任务通信层通过 raft_dask 的Comms在 worker 之间建立 NCCL 通信组供 MNMG 算法做全局归约如质心同步、特征协方差聚合算法层每个 worker 上实际调用的是 cuML 的*_mgmulti-GPU单机实现例如cuml.cluster.kmeans_mg.KMeansMG、cuml.linear_model.linear_regression_mg.LinearRegressionMG。从仓库结构看python/cuml/cuml/dask/init.py 在包导入时即完成三件事检查 Dask 依赖是否齐全、注册全部 12 个子模块、并强制关闭 Dask 的 p2p shuffledask.config.set({dataframe.shuffle.method: tasks})以确保分布式 DataFrame 混洗行为在 cuML 场景下稳定可控。二、安装与依赖缺什么会怎样cuml.dask是 cuML 的可选扩展不是cuml本体。其依赖检查逻辑位于 python/cuml/cuml/dask/init.py导入时会尝试加载dask、dask.distributed、dask_cudf与raft_dask任缺其一即抛出ModuleNotFoundError并给出修复指引。# 包导入时的依赖检查源码节选 try: import dask import dask.distributed as _ # noqa import dask_cudf as _ # noqa import raft_dask as _ # noqa except ModuleNotFoundError as exc: raise ModuleNotFoundError( Not all requirements for using cuml.dask are installed.\n\n # Install with Conda:\n conda install rapids-dask-dependency dask-cudf raft-dask\n\n # Or install with pip:\n pip install cuml-{cu_version}[dask] )对应的两种安装方式为# Conda 方式 conda install rapids-dask-dependency dask-cudf raft-dask # pip 方式{ver} 为当前 CUDA 版本对应的标签如 cu12 pip install cuml-{ver}[dask]除了 Python 侧依赖部分 MNMG 算法如 KMeans、PCA还依赖 cuML C 层的多 GPU 编译产物。相关代码通过mnmg_import装饰器见 python/cuml/cuml/dask/common/base.py包裹*_mg模块的惰性导入若当前构建未启用多 GPU 支持会抛出带--multigpu构建提示的RuntimeError。三、API 全景模块与公开类总览API 文档页 docs/source/api/cuml.dask.rst 按功能域划分了 12 个命名空间与 python/cuml/cuml/dask/ 目录一一对应。完整映射如下功能域命名空间公开 APIClustercuml.dask.clusterDBSCAN、KMeansDecompositioncuml.dask.decompositionPCA、TruncatedSVDEnsemblecuml.dask.ensembleRandomForestClassifier、RandomForestRegressorLinear Modelscuml.dask.linear_modelLinearRegression、Ridge、Lasso、ElasticNetManifoldcuml.dask.manifoldUMAPNaive Bayescuml.dask.naive_bayesMultinomialNBNeighborscuml.dask.neighborsNearestNeighbors、KNeighborsClassifier、KNeighborsRegressorPreprocessingcuml.dask.preprocessingLabelBinarizer、OneHotEncoderFeature Extractioncuml.dask.feature_extraction.textTfidfTransformerDatasetscuml.dask.datasetsmake_blobs、make_classification、make_regressionSolverscuml.dask.solversCDBase Classes and Mixinscuml.dask.common.baseBaseEstimator、DelayedParallelFunc、DelayedPredictionMixin、DelayedTransformMixin、DelayedInverseTransformMixin注意cuml.dask下还有metrics如confusion_matrix、common分布式数据/工具类等目录但文档页列出的 12 个命名空间是公开 API 的主要入口。下面按域逐一展开。四、集群DBSCAN 与 KMeans4.1 KMeans多轮通信的最小化数据搬运cuml.dask.cluster.KMeans实现见 python/cuml/cuml/dask/cluster/kmeans.py是文档中最典型的 MNMG 算法。源码 docstring 明确指出其设计目标每个迭代只在 worker 之间共享质心从而把数据搬运量降到最低而 predict 阶段则退化为纯并行embarrassingly parallel直接调用单 GPU KMeans。核心参数与单 GPU 版本对齐默认值取自源码参数默认值说明n_clusters8质心簇数量必须是正整数max_iter300EM 迭代上限越大越准但越慢tol1e-4质心变化小于该阈值时提前收敛initscalable-k-means初始化策略scalable-k-means/k-means||/random或传入(n_clusters, n_features)的 ndarray 作为初始质心oversampling_factor2scalable k-means 采样倍数越大初始质心越好但内存开销越大总采样数为oversampling_factor * n_clusters * 8max_samples_per_batch32768分批计算两两距离时的样本批大小影响显存占用每批约max_samples_per_batch * n_clusters个元素n_clusters很大时可调小random_stateNone随机种子保证可复现verboseFalse日志级别fit的分布式流程kmeans.py可以拆成四步值得作为理解 MNMG 的样板数据预处理DistributedDataHandler.create把 Dask 集合按 worker 组织成分区支持sample_weight会自动归一化全局预检preflight_validate_n_clusters在客户端校验n_clusters_fetch_worker_sizes汇总各 worker 行数若n_samples n_clusters直接报错——把可预测的全局错误在提交分布式任务前暴露出来建通信组Comms(comms_p2pFalse)初始化 RAFT/NCCL 通信会话随后向每个 worker 提交_func_preflight_fit参数校验与_func_fit真实训练内部实例化cuml.cluster.kmeans_mg.KMeansMG结果汇总质心经 NCCL 同步后各 worker 副本一致因此只从第一个 worker 拉取完整模型省内存inertia_由各 worker 的(inertia, n_samples)汇总求和labels_则保留为分布式 Dask 集合dask_cudf 或 dask.array避免大数组回传客户端。训练完成后可直接访问cluster_centers_、labels_、inertia_等属性它们通过基类的属性代理机制从远端 Future 或本地模型中取回。4.2 DBSCANcuml.dask.cluster.DBSCAN实现见 python/cuml/cuml/dask/cluster/dbscan.py是密度聚类算法的分布式版本同样遵循每 worker 本地计算 全局通信合并标签的模式。与 KMeans 的连续迭代不同DBSCAN 的核心开销在于跨分区的邻域查询与标签传播适合样本量大、簇形状不规则的数据。五、分解PCA 与 TruncatedSVDcuml.dask.decomposition.PCA实现见 python/cuml/cuml/dask/decomposition/pca.py的源码 docstring 明确了两点输入约定MNMG PCA 期望 Dask cuDF 对象作为输入两种算法Full默认对数据做完整特征分解后取前 K 个特征向量Jacobi通过迭代修正前 K 个特征向量速度更快但精度可能略低。其公开参数与单 GPU 版本一致主要包括n_components默认1主成分个数、whiten是否白化以及svd_solver相关选项类继承自BaseDecomposition、DecompositionSyncFitMixin与DelayedTransformMixin/DelayedInverseTransformMixin意味着它属于同步拟合 延迟变换型算法fit时通过DecompositionSyncFitMixin在集群上做全局协方差聚合见 python/cuml/cuml/dask/decomposition/base.pytransform/inverse_transform则按分区并行执行。TruncatedSVD面向稀疏/截断场景适合降维后仍需保留数据局部结构的任务接口风格与 PCA 一致。六、集成学习随机森林分类与回归cuml.dask.ensemble.RandomForestClassifier实现见 python/cuml/cuml/dask/ensemble/randomforestclassifier.py采用纯并行embarrassingly-parallel策略设森林共N棵树、集群有w个 worker则每个 worker 只在本地数据上构建N/w棵树互不通信。源码还给出两条实用经验若每个 worker 只持有数据子集通常要求数据预先充分洗牌效果才接近全量训练若把全部数据复制到每个 workerfit收到w个内容相同的分区结果将近似单 GPU 拟合。核心参数默认值取自 randomforestclassifier.py参数默认值说明n_estimators100森林总树数不是每 worker 树数会被均分到各 workersplit_criterion0(gini)分裂准则0/gini、1/entropy、2/mse、4/poisson、5/gamma、6/inverse_gaussian分类任务仅 gini/entropy 有效bootstrapTrue是否对每棵树做有放回采样False则每棵树用全量数据max_samples1.0每棵树使用的样本行比例max_depthNone最大深度None表示不设限直到叶子纯净max_leaves-1每棵树最大叶子数-1表示不设限软约束max_featuresauto每个节点分裂时考虑的特征数策略分类器额外支持predict_proba通过DelayedPredictionProbaMixinRandomForestRegressor实现位于 randomforestregressor.py参数体系一致但准则以mse等回归型为主。七、线性模型家族LinearRegression、Ridge、Lasso、ElasticNetcuml.dask.linear_model提供四件套均继承自BaseEstimator、SyncFitMixinLinearModel与DelayedPredictionMixin。以 LinearRegression 为例拟合算法algorithmeig时基于协方差矩阵的特征分解速度快适合高瘦tall and skinny数据文档同时提示 SVD 更慢但数值稳定性有保证特征数增大时 eig 的精度可能下降截距fit_interceptTrue默认会额外拟合常数项c把模型表达为y x·β c置False要求数据已中心化属性拟合后提供coef_n_features长度的 cuDF Series与intercept_。Ridgeridge.py加入 L2 正则alpha参数控制强度Lassolasso.py与ElasticNetelastic_net.py则通过坐标下降求解稀疏解。这些线性模型的fit走SyncFitMixinLinearModel._fitbase.py先在集群上初始化Comms把每个 worker 的分区大小通过parts_to_ranks广播给各 rank再提交model_func创建LinearRegressionMG等 worker 端模型最后统一等待 Future 完成。八、流形学习UMAPcuml.dask.manifold.UMAP实现见 python/cuml/cuml/dask/manifold/umap.py是单 GPU UMAP 的分布式版本面向超大样本量的降维与可视化。其核心流程是先构建分布式 kNN 图复用cuml.dask.neighbors的索引能力再在图上做嵌入优化主要参数与单 GPU 版本一致包括n_neighbors、n_components、min_dist、spread、learning_rate、random_state等其中 kNN 图构建阶段的分布式查询由NearestNeighbors支撑。九、朴素贝叶斯MultinomialNBcuml.dask.naive_bayes.MultinomialNB实现见 python/cuml/cuml/dask/naive_bayes/naive_bayes.py提供多项式朴素贝叶斯分类器的分布式实现适合文本分类等计数特征场景。核心参数包括平滑系数alpha默认1.0即 Laplace 平滑与fit_prior是否学习类别先验。fit阶段在各 worker 上统计类别计数与特征计数再跨 worker 归约得到全局先验与条件概率predict/predict_proba走延迟并行通道。十、近邻NearestNeighbors 与 KNN 分类/回归cuml.dask.neighbors包含三个类NearestNeighborsnearest_neighbors.py在分布式数据上构建近邻索引。关键参数n_neighbors默认5与batch_size默认2_000_000——batch_size决定每次查询处理的最大行数直接影响吞吐且每个承载索引分区的 worker 需要约batch_size * n_features * 4字节的额外显存需按显存预算调整KNeighborsClassifierkneighbors_classifier.py在分布式索引上做 kNN 投票分类支持predict与predict_probaKNeighborsRegressorkneighbors_regressor.pykNN 回归预测值为邻居标签的加权/平均聚合。三者共享BaseEstimator与DistributedDataHandler基础设施fit建立分布式索引查询阶段按分区并行。十一、预处理与特征提取cuml.dask.preprocessing提供LabelBinarizerpreprocessing/_label.py把类别标签二值化为 one-hot 风格的分布式表示OneHotEncoderpreprocessing/encoders.py类别特征独热编码支持dtype、handle_unknown、sparse_output等参数transform/inverse_transform均走DelayedTransformMixin/DelayedInverseTransformMixin的延迟并行通道。cuml.dask.feature_extraction.text.TfidfTransformerfeature_extraction/text/把词频矩阵转换为 TF-IDF 权重用于文本向量化流水线接口与 scikit-learn 对应组件对齐。十二、分布式数据集生成make_blobs 等cuml.dask.datasets的三个生成器是快速验证 MNMG 工作流的最佳工具它们的工作方式一致在每个 Dask worker 上调用单 GPU 版本的生成器再把结果拼接成分布式集合。以make_blobsdatasets/blobs.py为例签名与关键参数from cuml.dask.datasets import make_blobs X, y make_blobs( n_samples100, # 总行数跨 worker 分配 n_features2, # 特征数 centersNone, # 簇数或固定质心数组None 时生成 3 个中心 cluster_std1.0, # 簇内点标准差 n_partsNone, # 分区数可大于 worker 数None 时等于 worker 数 center_box(-10, 10), # 质心所在边界框 shuffleTrue, # 是否在 worker 内打乱样本 random_stateNone, # 随机种子 return_centersFalse, # 是否同时返回质心 dtypefloat32, # 数据类型 clientNone, # 显式传入 Dask Client workersNone, # 限定使用的 worker 地址列表None 表示全部 )返回值为 CuPy 后端的dask.arrayX形状(n_samples, n_features)y形状(n_samples,)return_centersTrue时额外返回质心数组。源码把n_samples均分到n_parts个分区并分派到各 worker用独立的随机种子保证可复现。make_classification与make_regression位于 classification.py 和 regression.py接口风格一致。十三、求解器CD坐标下降cuml.dask.solvers.CDpython/cuml/cuml/dask/solvers/是坐标下降Coordinate Descent求解器的分布式版本对应单 GPU 的cuml.solvers.CD是 Lasso/ElasticNet 内部使用的优化器也可独立使用。核心参数包括alpha正则强度、max_iter、tol、selection坐标选择策略等。由于坐标下降天然适合行分区并行——各 worker 在本地分区上计算梯度/残差后做全局归约——它是理解分布式优化器的良好示例。十四、基类与 Mixin分布式估算器的骨架文档页的最后一部分列出了cuml.dask.common.base中的基础设施类它们支撑了上述所有算法值得单独解读源码见 python/cuml/cuml/dask/common/base.py。14.1BaseEstimator一切分布式估算器的根构造函数__init__(*, clientNone, verboseFalse, **kwargs)client显式传入dask.distributed.Client不传时通过get_client()从全局获取或自动创建verbose既保存在本地也会注入kwargs[verbose]一并传给 worker 端模型其余超参数收进self.kwargs在提交分布式任务时透传给*_mg模型。核心机制是internal_model它是一个dask.distributed.Future[cuml.Base]、本地cuml.Base实例或None三态对象统一通过_set_internal_model管理。_check_internal_model会做类型校验Future 的类型必须是cuml.Base子类get_combined_model()则把集群上训练的模型收敛为一个可序列化的单 GPU 模型——这对pickle保存/加载分布式模型至关重要__getstate__/__setstate__已实现该逻辑。此外BaseEstimator.__getattr__base.py实现了属性代理用户访问model.cluster_centers_时如果本地没有该属性会优先查找_cluster_centers_再递归代理到internal_model——若模型在远端则通过dask.delayed惰性取回。这就是为什么分布式模型用起来和单 GPU 模型一样直观。14.2DelayedParallelFunc与四个 MixinDelayedParallelFunc._run_parallel_funcbase.py把predict/transform/score这类纯并行操作抽象成统一通道把X切成X.to_delayed()分区列表把内部模型与目标函数分别包装为dask.delayed模型带pureTrue以便跨任务复用对每个分区提交func(model_delayed, part, **kwargs)按output_collection_typecupy或cudf缺省沿用 fit 时的输入类型把结果拼接为dask.array或dask_cudf.DataFrame并支持delayed惰性与output_futures直接返回 Future两种返回形态。基于它派生的四个 Mixin 分别对应DelayedPredictionMixin_predict封装model.predict、DelayedPredictionProbaMixin_predict_proba封装model.predict_proba、DelayedTransformMixin_transform封装model.transform、DelayedInverseTransformMixin_inverse_transform封装model.inverse_transform。几乎所有文档列出的算法类都是BaseEstimator 若干 Mixin 算法特有 fit 逻辑的组合。十五、从入门到实战一个完整的 MNMG 工作流综合上文一个典型的多 GPU KMeans 训练流程如下代码模式取自 make_blobs 与 KMeans 的 docstring 示例from dask_cuda import LocalCUDACluster from dask.distributed import Client from cuml.dask.cluster import KMeans from cuml.dask.datasets import make_blobs # 1. 启动本地 GPU 集群每个 worker 一个线程 cluster LocalCUDACluster(threads_per_worker1) client Client(cluster) # 2. 生成跨 worker 分布的示例数据1000 行、10 特征、42 个簇 workers list(client.scheduler_info()[workers].keys()) X, y make_blobs( n_samples1000, n_features10, centers42, cluster_std0.1, workersworkers, ) # 3. 分布式训练 model KMeans(n_clusters42, random_state42) model.fit(X) # 4. 读取结果cluster_centers_ 经属性代理取回labels_ 保持分布式 print(model.cluster_centers_) print(model.inertia_) labels model.labels_ # dask_cudf / dask.array labels.compute() # 需要时再物化 # 5. 清理 client.close() cluster.close()工作流要点回顾数据形态fit接受dask_cudf.DataFrame或 CuPy 后端的dask.array内部由DistributedDataHandlerpython/cuml/cuml/dask/common/input_utils.py统一处理因此生成器产出的分布式集合可直接喂给所有算法client 生命周期估算器构造时get_client()会解析client参数建议显式传入并统一管理惰性与物化predict/transform默认返回惰性 Dask 集合delayedTrue调用.compute()/.persist()才真正执行便于构建流水线可复现性各算法均支持random_stateKMeans 内部还会用check_random_seed做归一化确保集群上每次调用种子一致见 kmeans.py。十六、测试与验证如何确认分布式行为仓库为cuml.dask提供了完备的测试套件位于 python/cuml/tests/dask/如test_dask_kmeans.py、test_dask_linear_regression.py、test_dask_datasets.py、test_dask_dbscan.py、test_dask_kneighbors_classifier.py、test_dask_base.py等。这些测试覆盖了分布式输入处理、基类属性代理、各算法与单 GPU 结果的一致性对比、数据集生成器的形状与可复现性等。在 CI 中它们由 ci/run_cuml_dask_pytests.sh 驱动本地验证分布式功能可参考该脚本的 pytest 入口方式执行对应测试文件。对开发者而言这些测试既是回归保障也是分布式估算器应该如何被使用的权威示例——例如test_dask_base.py演示了BaseEstimator的内部模型管理test_dask_datasets.py验证了n_parts与 worker 的映射关系阅读它们能比文档更快理解每个类的契约。结语cuml.dask把 cuML 的单 GPU 算法库扩展到了多节点多 GPU 维度数据分布在 Dask 集合上通信由 RAFTComms承接算法主体是各 worker 上的*_mg实现预测/变换则统一走DelayedParallelFunc的纯并行通道。本文依据 API 文档页 docs/source/api/cuml.dask.rst 的分类框架逐一解析了聚类、分解、集成、线性模型、流形、朴素贝叶斯、近邻、预处理、特征提取、数据集生成、求解器以及基类 Mixin 的接口与实现要点。上手时只需三步装齐dask-cudf/raft-dask依赖、用LocalCUDACluster起集群、把单 GPU 的 fit/predict 换成cuml.dask的对应类——其余交给 Dask 与 NCCL 编排。【免费下载链接】cumlNVIDIA cuML: GPU-Accelerated Machine Learning项目地址: https://gitcode.com/GitHub_Trending/cu/cuml创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考