SageMaker分布式训练引擎技术揭秘

SageMaker 分布式训练引擎背后的科学

某机构机器学习副总裁在年度 re:Invent 大会上宣布了两项新功能，旨在让用户通过 SageMaker 训练大规模、数据密集型神经网络变得更便宜、更简单。SageMaker 是一种完全托管的服务，使开发者能够轻松地在云端和边缘构建、训练和部署机器学习模型。

SageMaker 的数据并行库即使在大规模计算实例参与训练的情况下，也能实现近乎线性的扩展效率，使得在大型数据集上训练模型更快、更具成本效益。SageMaker 的模型并行库则能自动协调那些规模过大、无法单机训练的神经网络的训练过程。

数据并行训练

数据并行训练将同一神经网络的副本发送到不同的分布式计算节点，每个节点用不同的数据批次训练其副本。然后将分别训练的结果聚合和分发，使所有节点以相同方式更新其模型。

数据并行训练通常依赖 All-Reduce 算法来聚合不同 GPU 计算出的梯度。SDP 则利用了某机构网络拓扑的优势。例如，一台特定的计算实例包含 8 个 GPU 和 96 个虚拟 CPU，全部通过高速连接互联。

SDP 将聚合梯度的主要责任卸载给 CPU，由 CPU 将梯度更新传输到其他计算节点的 CPU。当 CPU 在聚合和传输一批梯度时，GPU 可以开始处理下一批数据，这使得分布式训练能够更高效地扩展。

为了在 CPU 之间通信梯度更新，SDP 使用了 All-Reduce 操作。每个虚拟 CPU 会等待收到一定数量的梯度后，再将它们传递出去，确保每个虚拟 CPU 平等地参与跨节点的梯度平均，从而有效利用带宽。

在一篇发表于超级计算会议的论文中，某机构研究人员描述了对比实验：在 512 个 GPU 上训练 BERT 语言模型时，该方案将训练时间减少了 44%。在另一项实验中，使用 SDP 训练拥有约 4400 万个参数的 Mask-RCNN 神经网络，训练时间在 PyTorch 上为 6 分 45 秒，在 TensorFlow 上为 6 分 12 秒，比之前的记录提高了约 24%。

模型并行训练

模型并行面临的首要问题是如何在计算节点间划分神经网络。答案需要平衡两个目标：一是计算负载的均衡分布，即每个节点处理同一批训练数据的时间应大致相同；二是最小化节点间通信。

为了充分了解网络并合理划分，SMP 会执行一次初始跟踪运行，以确定模型拓扑结构和关键元数据，如可训练参数的大小、交换张量的大小以及执行模型每个部分所需的时间。

在模型并行中，模型操作具有顺序依赖性：第一个节点的输出传递到第二个节点，依此类推。实现并行的唯一途径是通过流水线：节点 1 处理一批输入并将输出发送给节点 2；当节点 2 开始工作时，节点 1 开始处理下一批输入，以此类推。

SMP 会针对给定的分区创建优化的流水线调度，使前向和后向计算可以协同流水线作业。例如，当一个 GPU 处理一批数据的前向传递时，另一个 GPU 可能正在处理另一批数据的后向传递。SMP 在后台协调每个训练步骤，管理跨 GPU 的所有工作，并根据需要传输必要的张量。

以前，在 256 个实例上训练一个 30 亿参数的模型需要数周的人工努力来将模型拆分到 GPU 上。而通过 SageMaker 自动化和优化模型分区，这一过程仅需 6 天。

快速上手

用户现在可以通过 GitHub 上的示例，或文档来了解 SageMaker 上的分布式训练，并开始使用新的分布式库。FINISHED
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）