一句话：提出时空张量划分原语，将子算子沿空间和时间维度分布，并设计搜索算法，提升大模型训练吞吐并降低显存占用。
大语言模型训练对并行技术需求迫切。张量划分是常用手段，但现有系统仅考虑空间维度上的子算子分布，限制了通信与显存优化。本文提出时空张量划分原语，允许子算子同时沿空间和时间维度分布，扩大并行化搜索空间，避免跨设备集合通信，重叠通信与计算，并降低峰值显存。作者还设计优化搜索算法，在可接受时间内从时空划分空间中寻找最优并行方案。实验表明，训练 LLM 时相比先进系统最高达 1.68 倍吞吐、69% 峰值显存；32 GPU 下几何平均加速 1.30 倍；3D 并行下吞吐提升 1.46 倍。
With the rapid up-scaling of transformer-based large language models (LLM), training these models is becoming increasingly demanding on novel parallel training techniques. Tensor partitioning is an extensively researched parallel technique, encompassing data and model parallelism, and has a significant influence on LLM training performance. However, existing state-of-the-art parallel training systems are based on incomplete tensor partitioning space, where the distribution of partitioned sub-operators is limited to the spatial dimension. We discover that introducing the temporal dimension into tensor partitioning of LLM training instance provides extra opportunities to avoid collective communication across devices, saving memory space and also overlapping device-to-device communication with computation. In this paper, we propose a new tensor partition primitive that distributes sub-operators along both the spatial and temporal dimensions to further explore communication and memory overhead reduction over current solutions. This new primitive creates a broader parallelization space and leads to parallel solutions that achieve better training throughput with lower peak memory occupancy compared to state-of-the-art techniques. To efficiently deploy optimized parallel transformer model training to multiple devices, we further present an optimization algorithm that can find optimal parallel solutions from our spatial-temporal tensor partition space with acceptable search time. Our evaluation shows that our optimized tensor partitioning achieves up to 1.68 × training throughput with 69% peak memory occupancy compared to state-of-the-art distributed training systems when training LLMs. Upon scaling to 32 GPUs, the geo-mean speedup across benchmarks is 1.30 ×. When applied in 3D parallelism, up to 1.46 × training throughput can be achieved.
@inproceedings{WangWXWL024,
title = {PrimePar: Efficient Spatial-temporal Tensor Partitioning for Large Transformer Model Training},
author = {Haoran Wang and Lei Wang and Haobo Xu and Ying Wang and Yuming Li and Yinhe Han},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620666.3651357}
}