一句话:提出 SwapAdvisor,用遗传算法联合优化算子调度、内存分配与 CPU-GPU 交换决策,突破 GPU 显存限制。
针对深层宽神经网络受限于 GPU 显存、难以继续扩大模型规模的问题,论文提出支持 GPU 与 CPU 内存交换的训练方案。已有交换工作仅面向特定模型且性能不佳。SwapAdvisor 将深度学习计算视为数据流图,在该图上联合优化三个维度:算子调度、内存分配以及交换决策。由于搜索空间巨大,作者设计了定制遗传算法进行探索。评估使用多种大模型,结果显示 SwapAdvisor 可训练超过 GPU 显存上限 12 倍的模型,并达到假设无限 GPU 显存基线吞吐的 53%–99%。
It is known that deeper and wider neural networks can achieve better accuracy. But it is difficult to continue the trend to increase model size due to limited GPU memory. One promising solution is to support swapping between GPU and CPU memory. However, existing work on swapping only handle certain models and do not achieve satisfactory performance. Deep learning computation is commonly expressed as a dataflow graph which can be analyzed to improve swapping. We propose SwapAdvisor, which performs joint optimization along 3 dimensions based on a given dataflow graph: operator scheduling, memory allocation, and swap decisions. SwapAdvisor explores the vast search space using a custom-designed genetic algorithm. Evaluations using a variety of large models show that SwapAdvisor can train models up to 12 times the GPU memory limit while achieving 53-99% of the throughput of a hypothetical baseline with infinite GPU memory.