Colossal-Auto：大规模模型的自动化和激活检查点的统一自动化
Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models
摘要：
近年来，大规模模型在各个领域都表现出了最先进的性能。然而，训练这样的模型需要各种技术来解决诸如GPU之类的设备上的有限计算能力和存储器的问题。一些常用的技术包括流水线并行、张量并行和激活检查点。虽然现有的工作集中在寻找有效的分布式执行计划（Zheng et al. 2022）和激活检查点调度（Herrmann et al. 2019，Beaumont et al. 2021），但还没有提出联合优化这两个计划的方法。此外，提前编译在很大程度上依赖于准确的内存和计算开销估计，这通常是耗时和误导的。现有的训练系统和机器学习流水线或者物理地执行每个操作数，或者利用缩放的输入张量来估计存储器使用。为了解决这些挑战，我们引入了一个系统，可以联合优化分布式执行和梯度检查点计划。此外，我们还提供了一个易于使用的符号分析器，可以以最小的时间成本为任何PyTorch模型生成内存和计算统计数据。我们的方法允许用户在给定的硬件上并行化他们的模型训练，并以最小的代码更改为基础。源代码可在Colossal-AI GitHub或https://github.com/hpcaitech/ColossalAI上公开获取
英文摘要：
In recent years, large-scale models have demonstrated state-of-the-art performance across various domains. However, training such models requires various techniques to address the problem of limited computing power and memory on devices such as GPUs. Some commonly used techniques include pipeline parallelism, tensor parallelism, and activation checkpointing. While existing works have focused on finding efficient distributed execution plans (Zheng et al. 2022) and activation checkpoint scheduling (Herrmann et al. 2019, Beaumont et al. 2021}, there has been no method proposed to optimize these two plans jointly. Moreover, ahead-of-time compilation relies heavily on accurate memory and computing overhead estimation, which is often time-consuming and misleading. Existing training systems and machine learning pipelines either physically execute each operand or estimate memory usage with a scaled input tensor. To address these challenges, we introduce a system that can jointly optimize distributed execution and gradient checkpointing plans. Additionally, we provide an easy-to-use symbolic profiler that generates memory and computing statistics for any PyTorch model with a minimal time cost. Our approach allows users to parallelize their model training on the given hardware with minimum code change based. The source code is publicly available at Colossal-AI GitHub or https://github.com/hpcaitech/ColossalAI
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
