JaxPP通过灵活流水线并行提升大模型训练效率，自动分配任务与通信。
Scaling Deep Learning Training with MPMD Pipeline Parallelism
- 采用用户自定义流水线调度实现梯度累积，支持灵活并行配置。
- 相比最优SPMD配置，硬件利用率最高提升1.11倍。
- 适合需要高效扩展大规模模型训练的研究者和工程团队。
我们提出JaxPP，一个用于高效扩展大型深度学习模型训练的系统，支持灵活的流水线并行。引入无缝编程模型，允许用户定义流水线调度以实现梯度累积。JaxPP自动将对应于流水线阶段的任务分发到集群节点，并自动推断节点间的通信。我们实现了MPMD运行时，支持SPMD任务的异步执行。JaxPP的流水线并行实现相较于最佳SPMD配置，硬件利用率最高提升1.11倍。
原文摘要 · Abstract (English)
We present JaxPP, a system for efficiently scaling the training of large deep learning models with flexible pipeline parallelism. We introduce a seamless programming model that allows implementing user-defined pipeline schedules for gradient accumulation. JaxPP automatically distributes tasks, corresponding to pipeline stages, over a cluster of nodes and automatically infers the communication among them. We implement a MPMD runtime for asynchronous execution of SPMD tasks. The pipeline parallelism implementation of JaxPP improves hardware utilization by up to $1.11\times$ with respect to the best performing SPMD configuration.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。