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

在近期某机构网络服务年度re:Invent大会上，机器学习副总裁宣布了两项新功能，使客户能通过Amazon SageMaker更经济高效地训练大型数据密集型神经网络。SageMaker是完全托管服务，帮助开发者在云端和边缘端轻松构建、训练和部署机器学习模型。

数据并行库(SDP) 使神经网络训练能够以接近线性的效率扩展，即使大量EC2实例参与训练。这使得客户在大数据集上训练模型更快、更具成本效益。

模型并行库(SMP) 自动协调无法单机容纳的大型神经网络训练。以往分布式训练需要手动分区网络和手工调优代码，而SMP可自动完成这些工作。

数据并行技术

数据并行分布式训练将训练数据（橙色）分配到多个GPU。

如其名称所示，SDP采用数据并行技术：将相同神经网络副本发送到不同计算节点，每个节点在不同数据批次上训练其副本。然后将单独训练结果聚合分发，使所有节点以相同方式更新模型。

数据并行训练通常依赖all-reduce算法聚合不同GPU计算的梯度。SDP则利用某机构网络拓扑优势：例如p3dn.24xlarge机器包含8个Nvidia V100 GPU和96个虚拟CPU，均具备高速连接。

SDP将梯度聚合责任卸载给CPU，CPU同时向其他计算节点传输梯度更新。当CPU处理一批梯度时，GPU可处理下一批，从而提升分布式训练扩展效率。

在SC20会议上发表的论文中，研究人员通过实验比较了数据并行方案与集群内all-reduce方案。在512个GPU上训练BERT语言模型时，该方案将训练时间减少44%。

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

为获得足够网络信息进行合理分区，SMP执行初始追踪运行以确定模型拓扑和关键元数据，包括可训练参数大小、交换张量大小和模型组件执行时间。

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

SMP为给定分区创建优化流水线调度，支持前向传播和后向传播计算联合流水线处理。例如当一个GPU处理某数据批次的前向传播时，另一个GPU可处理另一批次的后向传播。根据流水线调度，SMP在底层协调每个训练步骤，管理GPU间所有工作并按需传输必要张量，使用为某机构基础设施优化的通信后端。

以往在256个实例上训练30亿参数模型需要数周手动分区工作，而通过SageMaker自动优化模型分区，仅需6天即可完成。

可通过GitHub示例立即开始使用SageMaker新分布式库，或阅读更多资料了解SageMaker分布式训练详情。