[论文解读] Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models

Yuliang Liu, Shenggui Li+4 more|arXiv (Cornell University)|Feb 6, 2023

Parallel Computing and Optimization Techniques被引用 6

一句话总结

Colossal-Auto 引入了一套统一的自动化系统，可在多GPU系统上对大规模模型训练联合优化算子内并行和激活检查点技术。通过结合硬件感知性能分析与两阶段求解器，该系统以极少的代码修改生成接近最优的分布式执行计划，在A100 GPU上实现了优于现有方法的弱扩展性能。

ABSTRACT

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

研究动机与目标

• 为解决手动调优大规模模型分布式训练策略的挑战，该策略需要深厚的专业知识且易出错。
• 联合优化算子内并行和激活检查点技术，而现有系统通常将二者分开处理。
• 降低研究人员和无分布式训练经验的实践者使用大规模模型进行高效训练的门槛。
• 实现自动、硬件感知的执行计划生成，可适应复杂的GPU拓扑结构，包括NVLink和PCIe互连。
• 支持跨任意维度设备网格和新型硬件配置的灵活、可扩展的并行策略。

提出的方法

• 符号化性能分析器在计算图上执行元执行，以极低的运行时开销收集细粒度的内存和计算开销。
• 两阶段求解器首先识别最优的算子内并行策略，然后联合调度激活检查点以满足内存预算。
• 启发式张量布局转换算法可实现跨任意N维设备网格的高效分片，支持复杂拓扑结构。
• 系统通过将操作替换为带检查点的并行化函数，将PyTorch模型编译为分布式执行代码。
• 执行计划生成考虑了硬件拓扑结构，包括NVLink和PCIe带宽差异，以实现通信与计算的重叠。
• 系统与PyTorch集成，支持仅需极少用户代码修改的自动代码重新编译。

实验结果

研究问题

• RQ1 能否联合优化算子内并行和激活检查点技术，以在大规模模型训练中实现更好的内存和计算效率？
• RQ2 符号化性能分析器如何以低开销准确估算PyTorch模型的内存和计算成本？
• RQ3 统一求解器能否生成适应复杂GPU互连拓扑（如NVLink和PCIe）的硬件感知执行计划？
• RQ4 与顺序或孤立优化相比，联合优化并行和检查点在训练性能和内存使用方面表现如何？
• RQ5 该系统在无需手动调优的情况下，对不同模型架构和硬件配置的泛化能力如何？

主要发现

• Colossal-Auto 在所有评估方法中实现了最佳弱扩展性能，在8张A100 GPU上对145.5亿参数模型达到0.824 PFLOPS。
• 在所有实验设置下，系统均优于Megatron-LM、Optimus和3D张量并行，尤其在更大模型上表现更优。
• 执行计划自动利用相邻GPU之间的高带宽NVLink链路，并将低带宽PCIe通信与计算重叠。
• 即使在大规模场景下，系统也能避免内存溢出错误，而数据并行（DDP）在更大模型尺寸下会失败。
• 符号化性能分析器实现了对计算图中内存和计算使用情况的准确、低成本性能分析，从而支持有效优化。
• 并行与检查点的联合优化相比单独优化任一组件，能实现更优的内存利用率和性能表现。