MXNet 分布式训练实战指南基于 Gluon 的多机多卡数据并行训练完整方案【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mx/mxnet本篇技术指南以 MXNet 仓库中的分布式训练官方示例example/distributed_training/README.md为核心骨架完整讲解如何利用 Gluon 接口在多个主机上启动分布式训练用多机多卡加速 ResNet 等大型模型。读者将掌握 MXNet 分布式训练的三类进程worker、parameter server、scheduler的协作原理、数据并行下的数据集切分方法、分布式 KVStore 的接入方式以及如何使用launch.py一键在多台机器上拉起训练任务。文中所有代码与配置均来自仓库内可运行的真实示例可直接复制落地。为什么需要分布式训练从单机多卡到多机多卡深度学习模型通常依赖 GPU 训练因为 GPU 的并行计算能力远超 CPU。即便使用现代 GPU训练大模型仍可能耗时数天。一个直接的加速思路是使用多块 GPU 并行训练但这受限于单台主机可挂载的 GPU 数量通常为 8 或 16 块。要进一步提速就必须把训练扩展到多台主机、多块 GPU的集群环境。MXNet 仓库中的官方教程示例 example/distributed_training/ 展示的正是这样一条路径在两台主机、每台 4 块 GPU共 8 块 GPU的环境下用数据并行data parallelism方式训练 ResNet18 网络与 CIFAR-10 数据集。数据并行的核心思想是把训练数据按 worker 数量均分到多台主机各主机在各自的数据子集上并行计算梯度再通过通信同步模型参数。由于多台主机同时在处理不同子集的数据整体训练时间会显著缩短。需要强调的是数据按 worker主机数量切分而不是按 GPU 数量切分——这是理解后续所有代码的关键前提。分布式训练架构三类进程的分工协作多主机分布式训练涉及三种不同类型的进程它们共同构成一个完整的训练集群进程类型数量职责Worker工作者每台主机一个每个迭代取一批数据在主机内所有 GPU 上执行前向与反向传播计算参数更新并发送给各主机的 parameter serverParameter Server参数服务器每台主机一个存储模型参数的一个子集接收各 worker 的梯度更新并维护参数的一致性Scheduler调度器整个集群唯一一个负责调度 worker 与 parameter server 的启动与连接Parameter Server参数的分布式存储由于多台主机协同训练同一个模型模型参数必须在所有主机间共享。为了让共享足够高效参数被切分成多份分布在不同主机上每台主机的 parameter server 只存储参数的一个子集。以教程中的两台主机为例参数被均匀切分为两份每台主机各存一半。每个迭代结束时每台主机都需要与其他所有主机通信以完成全部模型参数的更新。这种设计从源码层面也能得到印证Python 层的 KVStore 接口python/mxnet/kvstore/kvstore.py提供了push把梯度推入 store、pull从 store 拉取最新参数、pushpull合并推送与拉取等原子操作分布式模式下这些操作会经由 C 层 API 转发到远端 parameter server 进程执行。Worker前向反向与梯度上报每个 worker 在每次迭代中负责取一个 batch 的数据、在主机内所有 GPU 上并行执行前向与反向传播、计算参数更新、将更新发送到每台主机的 parameter server。因为有多个 worker 协同训练每个 worker 只需处理全部训练数据的 1/NN 为 worker 总数这是数据并行能线性加速的根本原因。Scheduler集群的大脑scheduler 负责 worker 与 parameter server 之间的调度和连接协调整个集群中只有一个 scheduler 实例通常由launch.py自动拉起。实战第一步接入分布式 Key-Value Store把单机训练改造成分布式训练最核心的改动是让Trainer使用一个**分布式 key-value storeKVStore**来存取和更新参数。在分布式训练中参数会被自动切分成 N 份并分布到 N 台主机上——这个切分过程由 KVStore 自动完成用户只需创建一个dist类型的 KVStore并让Trainer使用它即可import mxnet as mx from mxnet import gluon, kv # 创建分布式 key-value store store kv.create(dist)随后创建Trainer时把 store 通过kvstore参数传入。这样Trainer就不再在 GPU 或 CPU 内存中直接更新参数而是把梯度交给分布式 KVStore由各主机的 parameter server 完成参数存储与更新trainer gluon.Trainer(net.collect_params(), adam, {learning_rate: .001}, kvstorestore)在完整示例 example/distributed_training/cifar10_dist.py 中这两个关键步骤位于脚本开头先store kv.create(dist)再在创建Trainer时显式传入kvstorestore。除此之外网络定义、参数初始化、loss 与 optimizer 的选择与单机训练完全一致这也是 Gluon 高层 API 易用性的体现——分布式改造的侵入面被压缩到最小。从源码实现看kv.create(dist)返回的KVStore对象定义于 python/mxnet/kvstore/kvstore.py通过MXKVStoreInit、MXKVStorePush、MXKVStorePull等 C 层调用与远端服务器通信。分布式模式下若在 worker 节点调用set_optimizer优化器对象会被 pickle 序列化后发送给所有 parameter server见 kvstore.py从而在服务器端完成参数更新避免把原始梯度全部回传。实战第二步按 worker 切分训练数据数据并行要求训练数据在所有 worker 之间均分每个 worker 使用自己的数据子集进行训练。以两台机器为例每台机器运行一个 worker、每个 worker 管理 4 块 GPU数据按上图所示方式切分——注意切分依据是 worker 数量而非 GPU 数量。每个 worker 可以通过 KVStore 获知集群中的 worker 总数以及自己的 rank一个介于 0 到 N-1 的整数N 为 worker 总数store kv.create(dist) print(Total number of workers: %d % store.num_workers) print(This workers rank: %d % store.rank)输出示例Total number of workers: 2 This workers rank: 0这两个属性在 python/mxnet/kvstore/kvstore.py 中有明确定义rank返回当前 worker 节点的编号范围[0, num_workers)num_workers返回 worker 节点总数。知道 worker 总数和自己的 rank 后切分数据集就很简单了把数据集均匀分成 N 份根据 rank 选取其中一份。教程提供了一个自定义的SplitSampler完成此任务同样完整出现在 cifar10_dist.py 中class SplitSampler(gluon.data.sampler.Sampler): Split the dataset into num_parts parts and sample from the part with index part_index Parameters ---------- length: int Number of examples in the dataset num_parts: int Partition the data into multiple parts part_index: int The index of the part to read from def __init__(self, length, num_parts1, part_index0): # Compute the length of each partition self.part_len length // num_parts # Compute the start index for this partition self.start self.part_len * part_index # Compute the end index for this partition self.end self.start self.part_len def __iter__(self): # Extract examples between start and end, shuffle and return them. indices list(range(self.start, self.end)) random.shuffle(indices) return iter(indices) def __len__(self): return self.part_lenSplitSampler的工作原理一目了然part_len length // num_parts计算每个分区的大小整除向下取整start part_len * part_index当前 worker 分区的起始下标end start part_len当前 worker 分区的结束下标__iter__返回[start, end)区间内经过随机打乱的索引序列保证每个 epoch 内该分区的数据顺序被 shuffle__len__返回分区长度供DataLoader计算迭代步数。随后用SplitSampler构造DataLoader即可——以 CIFAR-10 的 50000 张训练图片、2 个 worker 为例每个 worker 恰好分到 25000 张# Load the training data train_data gluon.data.DataLoader(gluon.data.vision.CIFAR10(trainTrue).transform(transform), batch_size, samplerSplitSampler(50000, store.num_workers, store.rank))配套的测试数据则保持完整加载、不打乱顺序用于每个 epoch 结束后的精度评估见 cifar10_dist.py。实战第三步单 worker 内多 GPU 并行训练前文强调过数据集按 worker 数量切分而不是按 GPU 数量切分。这是因为把分区数据进一步摊到多块 GPU 上并行训练是每个 worker 自己的职责。首先需要指定该 worker 要使用的 GPU 列表# How many GPUs per machine gpus_per_machine 4 # Create the context (a list of all GPUs to be used for training) ctx [mx.gpu(i) for i in range(gpus_per_machine)]随后定义一个train_batch函数完成数据分发到各 GPU → 多 GPU 前向反向 → 参数更新的完整流程# Train a batch using multiple GPUs def train_batch(batch, ctx, net, trainer): # Split and load data into multiple GPUs data batch[0] data gluon.utils.split_and_load(data, ctx) # Split and load label into multiple GPUs label batch[1] label gluon.utils.split_and_load(label, ctx) # Run the forward and backward pass forward_backward(net, data, label) # Update the parameters this_batch_size batch[0].shape[0] trainer.step(this_batch_size)其中gluon.utils.split_and_load负责把一个 batch 的数据和标签均匀切分并加载到ctx列表中的每块 GPU 上。多 GPU 上的前向计算 loss与反向计算梯度由forward_backward实现——通过autograd.record()记录前向计算再对每块 GPU 上的 loss 依次调用backward()# Well use cross entropy loss since we are doing multiclass classification loss gluon.loss.SoftmaxCrossEntropyLoss() # Run one forward and backward pass on multiple GPUs def forward_backward(net, data, label): # Ask autograd to remember the forward pass with autograd.record(): # Compute the loss on all GPUs losses [loss(net(X), Y) for X, Y in zip(data, label)] # Run the backward pass (calculate gradients) on all GPUs for l in losses: l.backward()有了train_batch训练一个 epoch 的循环非常简洁for batch in train_data: # Train the batch using multiple GPUs train_batch(batch, ctx, net, trainer)需要留意 batch size 的语义示例中定义了batch_size_per_gpu 64而传给DataLoader的batch_size batch_size_per_gpu * gpus_per_machine即每块 GPU 独立处理 64 张图片一个全局 batch 共 256 张4 块 GPU。trainer.step(this_batch_size)传入的也是全局 batch 大小用于正确地归一化学习率的更新。整个 epoch 结束后调用evaluate_accuracy在测试集上评估精度并打印cifar10_dist.py。最后一步用 launch.py 一键拉起分布式训练分布式训练需要同时在多台机器上启动多个进程每台主机上要启动一个 worker 和一个 parameter serverscheduler 则只需在某一台主机上启动。手动管理这些进程繁琐且易错MXNet 提供了 tools/launch.py 工具来简化这一过程。启动命令与参数详解例如下面的命令在两台机器上启动分布式训练python ~/mxnet/tools/launch.py -n 2 -s 2 -H hosts \ --sync-dst-dir /home/ubuntu/cifar10_dist \ --launcher ssh \ python /home/ubuntu/cifar10_dist/cifar10_dist.py各参数含义与 tools/launch.py 中的 argparse 定义一一对应参数含义源码中的定义-n 2需要启动的 worker 进程数量--num-workers必填-s 2需要启动的 parameter server 数量--num-servers默认为 worker 数量--sync-dst-dir把当前目录内容 rsync 到各主机的目标目录--sync-dst-dir--launcher ssh使用 ssh 登录集群中每台机器并启动进程--launcher可选值local/ssh/mpi/sge/yarn默认sshpython .../cifar10_dist.py每个被启动的进程都要执行的命令command位置参数-H hosts集群中用于分布式训练的主机列表文件--hostfilelaunch.py内部会把上述参数转换为 dmlc-core 的 tracker 参数见 launch.py并依据--launcher选择的集群类型分发任务无 hostfile 时走local/sge/yarn提交有 hostfile 时走ssh/mpi提交。此外它还支持通过--env-server/--env-worker为服务器或 worker 进程单独注入环境变量例如OMP_NUM_THREADS:3、通过--env透传当前系统环境变量以及--p3开启 P3 分布式训练协议底层对应DMLC_PS_VAN_TYPEp3环境变量可以按需查阅 tools/launch.py。编写 hosts 文件与免密登录配置-H hosts指定的主机列表文件内容形如~/dist$ cat hosts d1 d2d1、d2是要运行分布式训练的主机名。launch.py应能仅凭主机名通过 ssh 登录这些主机例如~/dist$ ssh d1 Welcome to Ubuntu 16.04.3 LTS (GNU/Linux 4.4.0-1049-aws x86_64) ... Last login: Wed Jan 31 18:06:45 2018 from 72.21.198.67注意登录过程无需交互式输入认证信息。实现免密登录有多种方式一种简单做法是在~/.ssh/config中为每台主机指定 ssh 证书~$ cat ~/.ssh/config Host d1 HostName ec2-34-201-108-233.compute-1.amazonaws.com port 22 user ubuntu IdentityFile /home/ubuntu/my_key.pem IdentitiesOnly yes Host d2 HostName ec2-34-238-232-97.compute-1.amazonaws.com port 22 user ubuntu IdentityFile /home/ubuntu/my_key.pem IdentitiesOnly yes其中HostName是主机的实际可达地址user是登录用户IdentityFile指向私钥文件IdentitiesOnly yes限制只使用显式指定的密钥。更安全的方式是使用 ssh agent forwarding将本地持有的密钥通过 agent 转发给集群主机使用避免把私钥文件复制到每台机器上。运行效果与日志合并以下是分布式训练的一次实际运行输出$ python ~/mxnet/tools/launch.py -n 2 -s 2 -H hosts --sync-dst-dir /home/ubuntu/cifar10_dist --launcher ssh python /home/ubuntu/cifar10_dist/cifar10_dist.py 2018-06-03 05:30:05,609 INFO rsync /home/ubuntu/cifar10_dist/ - a1:/home/ubuntu/cifar10_dist 2018-06-03 05:30:05,879 INFO rsync /home/ubuntu/cifar10_dist/ - a2:/home/ubuntu/cifar10_dist Epoch 0: Test_acc 0.467400 Epoch 0: Test_acc 0.466800 Epoch 1: Test_acc 0.568500 Epoch 1: Test_acc 0.571300 Epoch 2: Test_acc 0.586300 Epoch 2: Test_acc 0.594000 Epoch 3: Test_acc 0.659200 Epoch 3: Test_acc 0.653300 Epoch 4: Test_acc 0.681200 Epoch 4: Test_acc 0.687900需要注意两点一是--sync-dst-dir会在启动前把本地代码目录 rsync 到集群各主机日志中的INFO rsync ...即此过程保证各主机运行的是同一份代码二是所有主机的输出会被合并打印到启动端控制台因此每个 epoch 会看到两行Test_acc——分别来自两个 worker且两个 worker 的测试精度略有差异是正常现象每个 worker 持有不同的数据分区与初始化随机性。进阶用 Horovod 替代参数服务器做梯度同步仓库中 example/distributed_training/cifar10_kvstore_hvd.py 提供了同一训练任务的 Horovod 版本展示了另一种分布式训练范式不依赖参数服务器而是通过环状 allreduce在所有进程中直接聚合梯度。与 KVStore 版本的核心差异在于# Use Horovod as the KVStore store kv.create(horovod) # Get the number of workers num_workers store.num_workers # Create the context based on the local rank of the current process ctx mx.cpu(store.local_rank) if args.no_cuda else mx.gpu(store.local_rank)kv.create(horovod)返回的 Horovod KVStore实现于 python/mxnet/kvstore/horovod.py在构造时调用hvd.init()初始化 Horovod 环境其pushpull操作直接映射为hvd.allreduce完成张量的跨进程求和rank/local_rank/num_workers属性则分别对应hvd.rank()/hvd.local_rank()/hvd.size()。由于 Horovod 是单进程多设备模型该版本使用store.local_rank选择当前进程绑定的 GPU每个进程管理一块 GPU并保留了同样的SplitSampler按num_workers与store.rank切分数据。该版本还额外支持--batch-size、--epochs、--lr、--no-cuda命令行参数并按 rank 0 过滤日志避免重复打印。小结本教程完整走通了 MXNet Gluon 分布式训练的五个关键环节理解架构worker、parameter server、scheduler 三类进程各司其职参数切分存储、数据按 worker 均分接入分布式 KVStorekv.create(dist)Trainer(..., kvstorestore)两行代码完成参数同步改造切分数据利用store.num_workers与store.rank配合自定义SplitSampler实现数据集分区多 GPU 训练gluon.utils.split_and_load把分区数据摊到主机内多块 GPU配合autograd.record()完成多卡前向反向一键启动用 tools/launch.py 配合 hosts 文件与 ssh 免密配置-n/-s/-H/--launcher ssh/--sync-dst-dir一条命令拉起全集群训练。完整的可运行代码见 example/distributed_training/cifar10_dist.pyHorovod 变体见 example/distributed_training/cifar10_kvstore_hvd.py。在动手前请确保所有主机的 MXNet 版本一致、GPU 驱动与 CUDA 环境可用并已完成 ssh 免密配置——这三个前提是分布式训练能否顺利启动的关键。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mx/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考