TAPAS：面向大型神经网络张量并行策略的快速自动推导
机器学习
2025-08-06 v2 人工智能
分布式、并行与集群计算
摘要
张量并行是大型神经网络分布式训练的一项关键技术。然而，由于搜索空间随模型规模与张量维数呈指数增长，自动确定最优张量并行策略颇具挑战。这阻碍了自动并行系统在更大模型上的应用。我们观察到神经网络通常包含重复子结构，并构建了名为 TAPAS 的自动并行框架以消除冗余搜索开销。TAPAS 采用分治方法，通过识别那些唯一子结构高效折叠搜索空间。因此，其关于模型规模以亚线性复杂度运行，成为训练大规模网络的可扩展方案。我们的评估表明，在多种模型上，TAPAS 的搜索速度相较最先进自动并行框架最高快达 ，且所得策略性能相较专家设计的 Megatron-LM 库具竞争力甚至更优。
引用
@article{arxiv.2302.00247,
title = {TAPAS: Fast and Automatic Derivation of Tensor Parallel Strategies for Large Neural Networks},
author = {Ziji Shi and Le Jiang and Ang Wang and Jie Zhang and Chencan Wu and Yong Li and Xiaokui Xiao and Wei Lin and Jialin Li},
journal= {arXiv preprint arXiv:2302.00247},
year = {2025}
}
备注
Accepted by 54th International Conference on Parallel Processing (ICPP'25)