Colossal-Auto: 大规模模型并行化与激活检查点的统一自动化
机器学习
2023-02-23 v2 人工智能
分布式、并行与集群计算
摘要
近年来,大规模模型在各个领域展现了最先进的性能。然而,训练此类模型需要各种技术来解决 GPU 等设备上计算能力和内存有限的问题。一些常用的技术包括流水线并行、张量并行和激活检查点。虽然现有工作集中于寻找高效的分布式执行计划(Zheng et al. 2022)和激活检查点调度(Herrmann et al. 2019, Beaumont et al. 2021),但尚未有方法被提出以联合优化这两个计划。此外,提前编译严重依赖于准确的内存和计算开销估计,这通常既耗时又具有误导性。现有的训练系统和机器学习流水线要么物理执行每个操作数,要么使用缩放的输入张量来估计内存使用量。为了应对这些挑战,我们引入了一个可以联合优化分布式执行和梯度检查点计划的系统。此外,我们提供了一个易于使用的符号分析器,能以最小的时间成本为任何 PyTorch 模型生成内存和计算统计数据。我们的方法允许用户在给定的硬件上以最少的代码改动并行化其模型训练。源代码已在 Colossal-AI GitHub 或 https://github.com/hpcaitech/ColossalAI 公开提供
引用
@article{arxiv.2302.02599,
title = {Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models},
author = {Yuliang Liu and Shenggui Li and Jiarui Fang and Yanjun Shao and Boyuan Yao and Yang You},
journal= {arXiv preprint arXiv:2302.02599},
year = {2023}
}