分布式机器学习：同步并行SGD算法的实现与复杂度分析

1 分布式机器学习概述

大规模机器学习训练常面临计算量大、训练数据大（单机存不下）、模型规模大的问题，对此分布式机器学习是一个很好的解决方案。

1）对于计算量大的问题，分布式多机并行运算可以基本解决。不过需要与传统HPC中的共享内存式的多线程并行运算（如OpenMP）以及CPU-GPU计算架构做区分，这两种单机的计算模式我们一般称为 计算并行 ）。

2）对于训练数据大的问题，需要将数据进行划分并分配到多个工作节点(Worker)上进行训练，这种技巧一般被称为 数据并行 。每个工作节点会根据局部数据训练出一个本地模型，并且会按照一定的规律和其他工作节点进行通信（通信内容主要是本地模型参数或者参数更新），以保证最终可以有效整合来自各个工作节点的训练结果并得到全局的机器学习模型。

如果是训练数据的样本量比较大，则需要对数据按照样本进行划分，我们称之为“数据样本划分”，按实现方法可分为“随机采样法”和“置乱切分法”。样本划分的合理性在于机器学习中的经验风险函数关于样本是可分的，我们将每个子集上的局部梯度平均，仍可得到整个经验风险函数的梯度。

如果训练数据的维度较高，还可对数据按照维度进行划分，我们称之为“数据维度划分”。它相较数据样本划分而言，与模型性质和优化方法的耦合度更高。如神经网络中各维度高度耦合，就难以采用此方式。不过，决策树对维度的处理相对独立可分，将数据按维度划分后，各节点可独立计算局部维度子集中具有最优信息增益的维度，然后进行汇总。此外，在线性模型中，模型参数与数据维度是一一对应的，故数据维度划分常与下面提到的模型并行相互结合。

3）对于模型规模大的问题，则需要对模型进行划分，并且分配到不同的工作节点上进行训练，这种技巧一般被称为 模型并行 。与数据并行不同，模型并行的框架下各个子模型之间的依赖关系非常强，因为某个子模型的输出可能是另外一个子模型的输入，如果不进行中间计算结果的通信，则无法完成整个模型训练。因此，一般而言，模型并行相比数据并行对通信的要求更高。

这里提一下数据样本划分中的几种划分方式。给定 \(n\) 个 \(d\) 维样本和 \(K\) 个工作节点，数据样本划分需要完成的任务是将 \(n\) 个样本以某种形式分配到 \(K\) 个工作节点上。

随机采样法中我们独立同分布地从 \(n\) 个样本中有放回随机采样，每抽取一个样本将其分配到一个工作节点上。这个过程等价于先随机采 \(n\) 个样本，然后均等地划分为 \(K\) 份。

随机采样法便于理论分析，但基于实现难度的考虑，在工程中更多采用的是基于置乱切分的划分方法。即现将 \(n\) 个样本随机置乱，再把数据均等地切分为 \(K\) 份，再分配到 \(K\) 个节点上进行训练。置乱切分相当于无放回采样，每个样本都会出现在某个工作节点上，每个工作节点的本地数据没有重复，故训练数据的信息量一般会更大。我们下面的划分方式都默认采取置乱切分的方法。

我们在后面的博客中会依次介绍针对数据并行和模型并行设计的各种分布式算法。本篇文章我们先看数据并行中最常用的同步并行SGD算法（也称SSGD）是如何在Spark平台上实现的。

2 同步并行SGD算法描述与实现

SSGD [1] 对应的伪代码可以表述如下：

其中，SSGD算法每次依据来自 \(K\) 个不同的工作节点上的样本的梯度来更新模型，设每个工作节点上的小批量大小为 \(b\) ，则该算法等价于批量大小为 \(bK