PatrickStar：基于块内存管理的预训练模型并行训练
机器学习
2022-11-11 v4 分布式、并行与集群计算
摘要
预训练模型（PTM）正在革新人工智能（AI）技术。然而，PTM 训练的硬件要求高得令人望而却步，使其成为少数人的游戏。因此，我们提出了 PatrickStar 系统以降低 PTM 的硬件要求，使其人人可用。PatrickStar 使用 CPU-GPU 异构内存空间来存储模型数据。与现有工作不同，我们将模型数据组织为内存块（chunk）并动态分布于异构内存中。在预热迭代中收集的运行时内存统计信息指导下，块在异构内存中被高效编排，产生更低的 CPU-GPU 数据传输量和更高的带宽利用率。与 Zero Redundancy Optimizer 共生，PatrickStar 可扩展到多节点多 GPU。该系统能以更大模型和更大批量训练任务，这是现有工作无法完成的。实验结果表明，PatrickStar 将模型规模扩展到 DeepSpeed 的 2.27 和 2.5 倍，并持续展现出显著更高的执行速度。PatrickStar 还在 32 GPU 集群上成功运行了 175B GPT3 训练任务。我们的代码公开于 https://github.com/Tencent/PatrickStar。
引用
@article{arxiv.2108.05818,
title = {PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management},
author = {Jiarui Fang and Zilin Zhu and Shenggui Li and Hui Su and Yang Yu and Jie Zhou and Yang You},
journal= {arXiv preprint arXiv:2108.05818},
year = {2022}
}
备注
13 pages