Frontier上科学图像的视觉Transformer序列长度缩放
计算机视觉与模式识别
2024-05-28 v1 机器学习
摘要
视觉Transformer因其处理大序列长度的能力，在科学图像（包括地球科学应用）的基础模型中至关重要。虽然文本Transformer启发了ViT中序列长度的缩放，但将其适配到ViT引入了独特的挑战。我们开发了用于ViT的分布式序列并行技术，使其能够处理多达100万个token。我们的方法利用DeepSpeed-Ulysses和长序列分段与模型分片，是首次在ViT训练中应用序列并行，在2048块AMD-MI250X GPU上实现了94%的批处理缩放效率。评估ViT中的序列并行，特别是在高达100亿参数的模型中，揭示了显著的瓶颈。我们通过混合序列、流水线、张量并行和闪存注意力策略来应对这些瓶颈，以扩展到单个GPU内存限制之外。我们的方法在温度预测中将气候建模精度显著提高了20%，标志着首次在188K序列长度的全注意力矩阵上训练Transformer模型。
引用
@article{arxiv.2405.15780,
title = {Sequence Length Scaling in Vision Transformers for Scientific Images on Frontier},
author = {Aristeidis Tsaris and Chengming Zhang and Xiao Wang and Junqi Yin and Siyan Liu and Moetasim Ashfaq and Ming Fan and Jong Youl Choi and Mohamed Wahib and Dan Lu and Prasanna Balaprakash and Feiyi Wang},
journal= {arXiv preprint arXiv:2405.15780},
year = {2024}
}