1 导引
我们在博客《分布式机器学习：同步并行SGD算法的实现与复杂度分析（PySpark）》和博客《分布式机器学习：模型平均MA与弹性平均EASGD（PySpark） 》中介绍的都是同步算法。同步算法的共性是所有的节点会以一定的频率进行全局同步。然而，当工作节点的计算性能存在差异，或者某些工作节点无法正常工作（比如死机）的时候，分布式系统的整体运行效率不好，甚至无法完成训练任务。为了解决此问题，人们提出了异步的并行算法。在异步的通信模式下，各个工作节点不需要互相等待，而是以一个或多个全局服务器做为中介，实现对全局模型的更新和读取。这样可以显著减少通信时间，从而获得更好的多机扩展性。
2 异步SGD2.1 算法描述与实现
异步SGD[9]是最基础的异步算法，其流畅如下图所示。粗略地讲，ASGD的参数更新发生在工作节点，而模型的更新发生在服务器端。当参数服务器接收到来自某个工作节点的参数梯度时，就直接将其加到全局模型上，而无需等待其它工作节点的梯度信息。
下面我们用Pytorch实现的训练代码（采用RPC进行进程间通信）。首先，我们设置初始化多个进程，其中0号进程做为参数服务器，其余进程做为worker来对模型进行训练，则总的通信域（world_size）大小为workers的数量+1。这里我们设置参数服务器IP地址为localhost，端口号29500。
def run(rank, world_size): os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '29500' options=rpc.TensorPipeRpcBackendOptions( num_worker_threads=16, rpc_timeout=0 # infinite timeout ) if rank == 0: rpc.init_rpc( "ps", rank=rank, world_size=world_size, rpc_backend_options=options ) run_ps([f"trainer{r}" for r in range(1, world_size)]) else: rpc.init_rpc( f"trainer{rank}", rank=rank, world_size=world_size, rpc_backend_options=options ) # trainer passively waiting for ps to kick off training iterations # block until all rpcs finish rpc.shutdown()if __name__=="__main__": world_size = n_workers + 1 mp.spawn(run, args=(world_size, ), nprocs=world_size, join=True)下面我们定义参数服务器的所要完成工作流程，包括将训练数据划分到各个worker，异步调用所有worker的训练流程，最后训练完毕后在参数服务器完成对模型的评估。
def run_trainer(ps_rref, train_dataset): trainer = Trainer(ps_rref) trainer.train(train_dataset)def run_ps(trainers): transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform) local_train_datasets = dataset_split(train_dataset, n_workers) print(f"{datetime.now().strftime('%H:%M:%S')} Start training") ps = ParameterServer() ps_rref = rpc.RRef(ps) futs = [] for idx, trainer in enumerate(trainers): futs.append( rpc.rpc_async(trainer, run_trainer, args=(ps_rref, local_train_datasets[idx])) ) torch.futures.wait_all(futs) print(f"{datetime.now().strftime('%H:%M:%S')} Finish training") ps.evaluation()这里数据集的划分代码采用我们在《Pytorch：单卡多进程并行训练》中所述的数据划分方式：
class CustomSubset(Subset): '''A custom subset class with customizable data transformation''' def __init__(self, dataset, indices, subset_transform=None): super().__init__(dataset, indices) self.subset_transform = subset_transform def __getitem__(self, idx): x, y = self.dataset[self.indices[idx]] if self.subset_transform: x = self.subset_transform(x) return x, y def __len__(self): return len(self.indices) def dataset_split(dataset, n_workers): n_samples = len(dataset) n_sample_per_workers = n_samples // n_workers local_datasets = [] for w_id in range(n_workers): if w_id < n_workers - 1: local_datasets.append(CustomSubset(dataset, range(w_id * n_sample_per_workers, (w_id + 1) * n_sample_per_workers))) else: local_datasets.append(CustomSubset(dataset, range(w_id * n_sample_per_workers, n_samples))) return local_datasets 以下是参数服务器类ParameterServer的定义：
class ParameterServer(object): def __init__(self, n_workers=n_workers): self.model = Net().to(device) self.lock = threading.Lock() self.future_model = torch.futures.Future() self.n_workers = n_workers self.curr_update_size = 0 self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9) for p in self.model.parameters(): p.grad = torch.zeros_like(p) self.test_loader = torch.utils.data.DataLoader( datasets.MNIST('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=32, shuffle=True) def get_model(self): # TensorPipe RPC backend only supports CPU tensors, # so we move your tensors to CPU before sending them over RPC return self.model.to("cpu") @staticmethod @rpc.functions.async_execution def update_and_fetch_model(ps_rref, grads): self = ps_rref.local_value() for p, g in zip(self.model.parameters(), grads): p.grad += g with self.lock: self.curr_update_size += 1 fut = self.future_model if self.curr_update_size >= self.n_workers: for p in self.model.parameters(): p.grad /= self.n_workers self.curr_update_size = 0 self.optimizer.step() self.optimizer.zero_grad() fut.set_result(self.model) self.future_model = torch.futures.Future() return fut def evaluation(self): self.model.eval() self.model = self.model.to(device) test_loss = 0 correct = 0 with torch.no_grad(): for data, target in self.test_loader: output = self.model(data.to(device)) test_loss += F.nll_loss(output, target.to(device), reduction='sum').item() # sum up batch loss pred = output.max(1)[1] # get the index of the max log-probability correct += pred.eq(target.to(device)).sum().item() test_loss /= len(self.test_loader.dataset) print('\nTest result - Accuracy: {}/{} ({:.0f}%)\n'.format( correct, len(self.test_loader.dataset), 100. * correct / len(self.test_loader.dataset))) 以下是Trainer类的定义：
class Trainer(object): def __init__(self, ps_rref): self.ps_rref = ps_rref self.model = Net().to(device) def train(self, train_dataset): train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) model = self.ps_rref.rpc_sync().get_model().cuda() pid = os.getpid() for epoch in range(epochs): for batch_idx, (data, target) in enumerate(train_loader): output = model(data.to(device)) loss = F.nll_loss(output, target.to(device)) loss.backward() model = rpc.rpc_sync( self.ps_rref.owner(), ParameterServer.update_and_fetch_model, args=(self.ps_rref, [p.grad for p in model.cpu().parameters()]), ).cuda() if batch_idx % log_interval == 0: print('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( pid, epoch + 1, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))完整代码我已经上传到了GitHub仓库 [Distributed-Algorithm-PySpark]，感兴趣的童鞋可以前往查看。
运行该代码得到的评估结果为：
Test result - Accuracy: 9696/10000 (97%)可见该训练算法是收敛的，但在10个epoch下在测试集上只能达到97%的精度，不如我们下面提到的在10个epoch就能在测试集上达到99%精度的Hogwild!算法。注意，ASGD和Hogwild都是异步算法，但ASGD是分布式算法（当然我们这里是单机多进程模拟），进程间采用RPC通信，不会出现同步错误的问题，根本不需要考虑加不加锁。而Hogwild!算法是单机算法，进程/线程间采用共享内存通信，需要考虑加不加锁的问题，不过Hogwild！算法为了提高训练过程中的数据吞吐量，直接采用了无锁的全局内存访问。
2.2 收敛性分析
ASGD避开了同步开销，但会给模型更新增加一些延迟。我们下面将ASGD的工作流程用下图加以剖析来解释这一点。用\(\text{worker}(k)\)来代表第\(k\)个工作节点，用\(w^t\)来代表第\(t\)轮迭代时服务端的全局模型。按照时间顺序，首先\(\text{worker}(k)\)先从参数服务器获取全局模型\(w^t\)，再根据本地数据计算模型梯度\(g(w_t)\)并将其发往参数服务器。一段时间后，\(\text{worker}(k’)\)也从参数服务器取回当时的全局模型\(w^{t+1}\)，并同样依据它的本地数据计算模型的梯度\(f(w^{t+1})\)。注意，在\(\text{worker}(k’)\)取回参数并进行计算的过程中，其它工作节点（比如\(\text{worker}(k)\)）可能已经将它的梯度提交给服务器并进行更新了。所以当\(\text{worker}(k’)\)将其梯度\(g(w^{t+1})\)发给服务器时，全局模型已经不再是\(w^{t+1}\)，而已经是被更新过的版本。
我们将上面这种现象称为梯度和模型的失配，也即我们用一个比较旧的参数计算了梯度，而将这个“延迟”的梯度更新到了模型参数上。这种延迟使得ASGD和SGD之间在参数更新规则上存在偏差，可能导致模型在某些特定的更新点上出现严重抖动，设置优化过程出错，无法收敛。后面我们会介绍克服延迟问题的手段。
3 Hogwild!算法3.1 算法描述与实现
异步并行算法既可以在多机集群上开展，也可以在多核系统下通过多线程开展。当我们把ASGD算法应用在多线程环境中时，因为不再有参数服务器这一角色，算法的细节会发生一些变化。特别地，因为全局模型存储在共享内存中，所以当异步的模型更新发生时，我们需要讨论是否将内存加锁，以保证模型写入的一致性。
Hogwild！算法[2]为了提高训练过程中的数据吞吐量，选择了无锁的全局模型访问，其工作逻辑如下所示：
这里使用我们在《Pytorch：单卡多进程并行训练》所提到的torch.multiprocessing来进行多进程并行训练。多进程原本内存空间是独立的，这里我们显式调用model.share_memory()来讲模型设置在共享内存中以进行进程间通信。不过注意，如果我们采用GPU训练，则GPU直接就做为了多进程的共享内存，此时model.share_memory()实际上为空操作（no-op）。
我们用Pytorch实现的训练代码如下：
from __future__ import print_functionimport torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.multiprocessing as mpfrom torchvision import datasets, transformsimport osimport torchimport torch.optim as optimimport torch.nn.functional as Fbatch_size = 64 # input batch size for trainingtest_batch_size = 1000 # input batch size for testingepochs = 10 # number of global epochs to trainlr = 0.01 # learning ratemomentum = 0.5 # SGD momentumseed = 1 # random seedlog_interval = 10 # how many batches to wait before logging training statusn_workers = 4 # how many training processes to usecuda = True # enables CUDA trainingmps = False # enables macOS GPU trainingdry_run = False # quickly check a single passdef train(rank, model, device, dataset, dataloader_kwargs): torch.manual_seed(seed + rank) train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs) optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum) for epoch in range(1, epochs + 1): model.train() pid = os.getpid() for batch_idx, (data, target) in enumerate(train_loader): optimizer.zero_grad() output = model(data.to(device)) loss = F.nll_loss(output, target.to(device)) loss.backward() optimizer.step() if batch_idx % log_interval == 0: print('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( pid, epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) if dry_run: breakdef test(model, device, dataset, dataloader_kwargs): torch.manual_seed(seed) test_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs) model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: output = model(data.to(device)) test_loss += F.nll_loss(output, target.to(device), reduction='sum').item() # sum up batch loss pred = output.max(1)[1] # get the index of the max log-probability correct += pred.eq(target.to(device)).sum().item() test_loss /= len(test_loader.dataset) print('\nTest set: Global loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))) class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) if __name__ == '__main__': use_cuda = cuda and torch.cuda.is_available() use_mps = mps and torch.backends.mps.is_available() if use_cuda: device = torch.device("cuda") elif use_mps: device = torch.device("mps") else: device = torch.device("cpu") print(device) transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('../data', train=False, transform=transform) kwargs = {'batch_size': batch_size, 'shuffle': True} if use_cuda: kwargs.update({'num_workers': 1, 'pin_memory': True, }) torch.manual_seed(seed) mp.set_start_method('spawn', force=True) model = Net().to(device) model.share_memory() # gradients are allocated lazily, so they are not shared here processes = [] for rank in range(n_workers): p = mp.Process(target=train, args=(rank, model, device, train_dataset, kwargs)) # We first train the model across `n_workers` processes p.start() processes.append(p) for p in processes: p.join() # Once training is complete, we can test the model test(model, device, test_dataset, kwargs)运行得到的评估结果为：
Test set: Global loss: 0.0325, Accuracy: 9898/10000 (99%)可见该训练算法是收敛的。
3.2 收敛性分析
当采用不带锁的多线程的写入（即在更新\(w_j\)的时候不用先获取对\(w_j\)的访问权限），而这可能会导致导致同步错误[10]的问题。比如在线程\(1\)加载全局参数\(w^t_j\)之后，线程\(2\)还没等线程\(1\)存储全局参数更新后的值，就也对全局参数\(w^t_j\)进行加载，这样导致每个线程都会存储值为\(w^t_j – \eta^t g(w^t_j)\)的更新后的全局参数值，这样就导致其中一个线程的更新实际上在做“无用功”。直观的感觉是这应该会对学习的过程产生负面影响。不过，当我们对模型访问的稀疏性（sparity）做一定的限定后，这种访问冲突实际上是非常有限的。这正是Hogwild！算法收敛性存在的理论依据。
假设我们要最小化的损失函数为\(l: \mathcal{W}\rightarrow \mathbb{R}\)，对于特定的训练样本集合，损失函数\(l\)是由一系列稀疏子函数组合而来的：
\[l(w) = \sum_{e\in E}f_e(w_e)\]
也就是说，实际的学习过程中，每个训练样本涉及的参数组合\(e\)只是全体参数集合中的一个很小的子集。我们可以用一个超图\(G=(V, E)\)来表述这个学习过程中参数和参数之间的关系，其中节点\(v\)表示参数，而超边\(e\)表示训练样本涉及的参数组合。那么，稀疏性可以用下面几个统计量加以表示：
\[\Omega:=\max_{e\in E}|e|\\\Delta:=\frac{\underset{1\leqslant v \leqslant n}{\max}|\{e\in E: v\in e\}|}{|E|}\\\rho:=\frac{\underset{e\in E}{\max}|\{\hat{e}\in E: \hat{e}\cap e \neq \emptyset \}|}{|E|}\]
其中，\(\Omega\)表达了最大超边的大小，也就是单个样本最多涉及的参数个数；\(\Delta\)反映的是一个参数最多可以涉及多少个不同的超边；而\(\rho\)则反映了给定任意一个超边，与其共享参数的超边个数。这三个值的取值越小，则优化问题越稀疏。在\(\Omega\)、\(\Delta\)、\(\rho\)都比较小的条件下，Hogwild！算法的收敛性保证还需要假设损失函数是凸函数，并且是Lipschitz连续的，详细的理论证明和定量关系请参考文献[2]。
参考
[1] Agarwal A, Duchi J C. Distributed delayed stochastic optimization[J]. Advances in neural information processing systems, 2011, 24.
[2] Recht B, Re C, Wright S, et al. Hogwild!: A lock-free approach to parallelizing stochastic gradient descent[J]. Advances in neural information processing systems, 2011, 24.
[3] 刘浩洋，户将等. 最优化：建模、算法与理论[M]. 高教出版社, 2020.
[4] 刘铁岩，陈薇等. 分布式机器学习：算法、理论与实践[M]. 机械工业出版社, 2018.
[5] Stanford CME 323: Distributed Algorithms and Optimization (Lecture 7)
[6] Bryant R E等.《深入理解计算机系统》[M]. 机械工业出版社, 2015.
数学是符号的艺术，音乐是上界的语言。