MXNet源码分析 | Gluon接口分布式训练流程

本文主要基于MXNet1.6.0版本，对Gluon接口的分布式训练过程进行简要分析。

众所周知，KVStore负责MXNet分布式训练过程中参数的同步，那么它究竟是如何应用在训练中的呢？下面我们将从Gluon.Trainer这个接口入手，逐步分析分布式训练的梯度交换以及参数同步过程。下面这段代码摘自python/mxnet/gluno/trainer.py文件，相较于源代码删除了一些多余的信息（如某些判断、注释等），以便让我们更好地专注于通信过程。

代码中的step函数是进行梯度交换以及参数更新的主体，它首先调用_init_kvstore去初始化kvstore，然后调用_allreduce_grads进行梯度传输，最后调用_update实现参数更新。

class Trainer(object):
    def step(self, batch_size, ignore_stale_grad=False):
        if not self._kv_initialized:
            self._init_kvstore()
        if self._params_to_init:
            self._init_params()

        self._allreduce_grads()
        self._update(ignore_stale_grad)

首先，_init_kvstore这个函数会通过用户指定的参数来调用model.py中的_create_kvstore来初始化kvstore以及update_kv_store这两个变量。其中kvstore是KVStore类的一个实例化对象，而update_on_kvstore是一个布尔型变量，用来判断是否在ps端更新参数。换句话说，如果该变量为True，那么模型参数的更新发生在ps端；否则，模型参数的更新发生在worker端，ps端只做梯度的聚合操作（这种情况下，paramerter server是不是就变成了gradient server？😛）。然而，只有在同步训练模式下，我们才能设置update_on_kvstore=False，异步训练并不支持在worker端更新参数。在update_kv_store=True的情况下，我们需要告诉ps端训练过程中使用的优化器是什么，因此要调用kvstore.set_optimizer把优化器从worker端传给ps端。

from ..model import _create_kvstore
class Trainer(object):
    def _init_kvstore(self):
        """Create kvstore."""
        config = self._kvstore_params
        arg_arrays = {param.name: param.data(self._contexts[0]) for param in self._params}
        kvstore, update_on_kvstore = _create_kvstore(config['kvstore'], len(self._contexts),
                                                     arg_arrays)
        self._distributed = 'dist' in kvstore.type if kvstore else False
        if self._distributed and 'async' in kvstore.type:
            update_on_kvstore = True
            # raise err if user provides unsupported configs
            if config['update_on_kvstore'] is False:
                raise ValueError("Please set update_on_kvstore=True "
                                 "when training in async mode.")
        if config['update_on_kvstore'] is not None:
            update_on_kvstore = config['update_on_kvstore'

        if kvstore:
            if update_on_kvstore:
                # optimizer preferably needs to be set before init for multiprecision
                kvstore.set_optimizer(self._optimizer)
            self._kvstore = kvstore
            self._update_on_kvstore = update_on_kvstore
        else:
            self._kvstore = None
            self._update_on_kvstore = None
        self._kv_initialized = True

完成kvstore的初始化后，gluon.Trainer会调用_allreduce_grads来实现梯度的交换。欸，前面不是说MXNet是参数服务器架构吗，为啥为扯到Allreduce上呢？考虑update_on_kvstore=False的情况，最开始每个worker上都只有自己的本地梯度，把梯度push到ps并进行聚合后，每个worker从ps上pull回来的都是相同的、聚合后的梯度。整个过程中的push和pull操作，是不是就很像Reduce和Broadcast（worker上的梯度“Reduce”到ps上，然后ps端“Broadcast”聚合结果给各个worker）？观察_allreduce_grads的实现，可以发现，无论update_on_kvstore的值是什么，gluno.Trainer都会把梯度从worker端push到ps端，只不过当update_on_kvstore=True时，gluon.Trainer把梯度从worker上push到ps后就完事儿了；而当updata_on_kvstore=False时，gluon.Trainer还会从ps端把梯度的聚合结果pull回来，以便进行本地的参数更新。

class Trainer(object):
    def _allreduce_grads(self):
        if self._kvstore:
            for i, param in enumerate(self._params):
                if param.grad_req != 'null':
                    self._kvstore.push(i, param.list_grad(), priority=-i)
                    if not self._update_on_kvstore:
                        self._kvstore.pull(i, param.list_grad(), priority=-i,
                                           ignore_sparse=self._distributed)

gluon.Trainer._update函数会根据update_on_kvstore的值进行相应的参数拉取以及更新操作。在单机训练（kvstore is None）或者分布式训练的本地更新模式（update_on_kvstore=True）下，gluon.Trainer会使用用户设定的优化器在本地更新参数，以进行下一步的训练。在分布式训练的情况下，当我们设置update_on_kvstore=True时，模型参数会在ps端进行更新，所以在该函数只需要将模型参数从ps端pull到本地即可。

class Trainer(object):
    def _update(self, ignore_stale_grad=False):
        updates = [[] for _ in self._updaters]

        for i, param in enumerate(self._params):
            if self._kvstore and self._update_on_kvstore:
                if param._stype == 'default':
                    # 'row_sparse' parameters are not pulled immediately - they're pulled
                    # in `Block.forward`
                    self._kvstore.pull(i, param.list_data(), priority=-i)
                continue

            for upd, arr, grad in zip(updates, param.list_data(), param.list_grad()):
                if not ignore_stale_grad or arr._fresh_grad:
                    upd.append((i, grad, arr))
                    arr._fresh_grad = False

        if not (self._kvstore and self._update_on_kvstore):
            for updater, upd in zip(self._updaters, updates):
                if upd:
                    i, w, g = zip(*upd)
                    updater(i, w, g)

到这里，我们基本上就把python端的kvstore调用流程讲完了。