ytk-mp4j 分布式机器学习通信库 - 关于猿辅导机器学习项目ytk-learn和ytk-mp4j分布式机器学习库

ytk-mp4j 分布式机器学习通信库
 

项目背景

目前可以用于分布式机器学习的通信主要基于MPI和RPC，其中MPI是分布式高性能计算的标配，虽然效率非常高，但是对于开发分布式机器学习任务来说有很多缺点： 开发难度大、数据支持太底层、只能用C/C++， Fortran编写等等；RPC 方式来实现类似 allreduce 这种操作，在特征维度特别高的场景，通信效率太低。所以我们开发了一套易用且高效的机器学习分布式通信库。

图2 ytk-mp4j 特性概略

项目简介

ytk-mp4j 是基于Java的高效分布式机器学习通信库，实现了类似 MPI Collective 通信中的大部分操作，包含gather， scatter， allgather， reduce-scatter， broadcast， reduce， allreduce，使用 ytk-mp4j 可以快速地把串行机器学习程序改造成支持多线程和多进程，ytk-learn 中所有涉及到分布式通信操作都是基于 ytk-mp4j 实现（表1中给出了部分例子）。

相比于MPI， ytk-mp4j 扩展实现了一些非常实用的特性：

1. 所有的通信操作都是基于最优算法实现［1，2］，性能非常高，同时支持多线程，多进程。同样的功能，在C/C++ 环境中，可能需要结合 MPI 和 OpenMP 才能实现

2. 不仅支持基本的数据类型（double， float， long， int， short， byte），而且还支持Java String及任意普通Java对象（Java 对象只需要实现 Kryo的 Serializer 接口）

3. 不仅支持传统数组类型的 Collective 通信，而且还支持Java Map数据类型，使用Map数据类型，用户可以实现非常复杂的通信操作（例如：集合求交、求并，链表的连接等操作）

4. 支持数据压缩传输，在网络资源很紧张的情况下，可以节约大量的带宽

5. 纯Java代码实现，可以无缝集成到 Hadoop， Spark 等分布式计算平台，构建自己的分布式机器学习系统

6. 使用 Java的SDP（Sockets Direct Protocol）可以实现高效的RDMA（Remote Direct Memory Access）

表1 ytk-mp4j在ytk-learn中的使用

ytk-mp4j 操作ytk-learn 中使用场景

allreduceLoss 求和，梯度求和，Hessian 求和，计算分位数，计算平均值，

计算评估指标（AUC， Confusion Matrix…），统计样本数量、

特征出现次数等等

reduce-scatterGBDT 高效梯度求和

allgatherL-BFGS 中计算Hv，GBDT 同步梯度

allreduce 操作是分布式机器学习中使用最多的通信操作，它对机群中所有的节点对应的数据进行归约操作，然后再分发给各个节点。下面给出了 ytk-mp4j 在多进程、多进程、数组，Map 下的 allreduce （归约操作为求和）示意图：

性能测试

表2给出了 ytk-mp4j 实现的Collective操作时间复杂度，其中  是网络连接延迟， 是传输1个字节需要的时间， 是需要传输的字节数量， 是进行1字节数据归约（reduction）操作需要的时间。可以看出，随着机器数量的增加，所有操作数据传输的时间是几乎不会增加的，只有连接和归约操作的时间会随机器数量增加，但在大数据通信时，连接和归约的时间占比很小。这个时间复杂度特性非常重要，它使得在特征维度、样本数量超过一定阈值的分布式机器学习训练任务中，训练加速比与机器数量接近线性关系。

表2 ytk-mp4j 实现的 Collective 操作时间复杂度

下图是测试在 1Gigabit Ethernet 网络下，10亿维 double 数组，各种 Collective 通信操作在不同的机器数量下的通信性能（时间单位： ms），从图中可以看出 ytk-mp4j 中的7种 Collective 操作的通信时间与机器数量的关系与理论值完全符合。