Vision Transformers (ViTs) are pivotal for foundational models in scientific imagery, including Earth science applications, due to their capability to process large sequence lengths. While transformers for text have inspired scaling sequence lengths in ViTs, adapting these for ViTs introduces unique challenges. We develop distributed sequence parallelism for ViTs, enabling them to handle up to 1M tokens. Our approach, leveraging DeepSpeed-Ulysses and Long-Sequence-Segmentation with model sharding, is the first to apply sequence parallelism in ViT training, achieving a 94% batch scaling efficiency on 2,048 AMD-MI250X GPUs. Evaluating sequence parallelism in ViTs, particularly in models up to 10B parameters, highlighted substantial bottlenecks. We countered these with hybrid sequence, pipeline, and flash attention strategies, to scale beyond single GPU memory limits. Our method significantly enhances climate modeling accuracy by 20% in temperature predictions, marking the first training of a vision transformer model to convergence with a sequence length of 188K tokens, using full self-attention.
BiKXieLZhangH, et al. (2023) Accurate medium-range global weather forecasting with 3d neural networks. Nature619(7970): 533–538.
3.
bodzásAKodytekPŽídekJ (2023) A high-resolution large-scale dataset of pathological and normal white blood cells. Scientific Data10: 466.
4.
ChaKSeoJLeeT (2023) A billion-scale foundation model for remote sensing images.
5.
ChenRJChenCLiY, et al. (2022) Scaling vision transformers to gigapixel images via hierarchical self-supervised learning. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. IEEE, pp. 16144–16155.
6.
ChenLZhongXZhangF, et al. (2023) Fuxi: a Cascade machine learning forecasting system for 15-day global weather forecast. npj Climate and Atmospheric Science6(1): 190.
7.
DaoT (2023) Flashattention-2: faster attention with better parallelism and work partitioning.
DehghaniMDjolongaJMustafaB, et al. (2023) Scaling vision transformers to 22 billion parameters.
11.
DengJDongWSocherR, et al. (2009) Imagenet: a large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition. IEEE, pp. 248–255.
12.
DingKZhouMWangH, et al. (2023) A large-scale synthetic pathological dataset for deep learning-enabled segmentation of breast cancer. Scientific Data10: 231. DOI: 10.1038/s41597-023-02125-y.
13.
DosovitskiyABeyerLKolesnikovA, et al. (2021) An image is worth 16x16 words: transformers for image recognition at scale.
14.
FedusWZophBShazeerN (2022) Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research23(1): 1532–4435.
HuangYChengYBapnaA, et al. (2019) Gpipe: efficient training of giant neural networks using pipeline parallelism.
19.
JacobsSATanakaMZhangC, et al. (2023) Deepspeed ulysses: system optimizations for enabling training of extreme long sequence transformer models.
20.
KhanSNaseerMHayatM, et al. (2022) Transformers in vision: a survey. ACM Computing Surveys54(10s): 1–41.
21.
KimYJJangHLeeK, et al. (2021) Paip 2019: liver cancer segmentation challenge. Medical Image Analysis67: 101854.
22.
KorthikantiVCasperJLymS, et al. (2022) Reducing activation recomputation in large transformer models.
23.
LiSXueFBaranwalC, et al. (2022) Sequence parallelism: long sequence training from system perspective.
24.
LiCGanZYangZ, et al. (2023) Multimodal foundation models: from specialists to general-purpose assistants.
25.
LiuZLinYCaoY, et al. (2021) Swin transformer: hierarchical vision transformer using shifted windows. In: Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV). IEEE.
26.
LiuHZahariaMAbbeelP (2023) Ring attention with blockwise transformers for near-infinite context.
NguyenTBrandstetterJKapoorA, et al. (2023) Climax: a foundation model for weather and climate. arXiv preprint arXiv:2301.10343.
29.
PathakJSubramanianSHarringtonP, et al. (2022) Fourcastnet: a global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214.
30.
RajbhandariSRuwaseORasleyJ, et al. (2021) Zero-infinity: breaking the gpu memory wall for extreme scale deep learning.
31.
RenJRajbhandariSAminabadiRY, et al. (2021) Zero-offload: democratizing billion-scale model training.
32.
ShoeybiMPatwaryMPuriR, et al. (2020) Megatron-lm: training multi-billion parameter language models using model parallelism.
33.
SimmonsASociCNicolasJ, et al. (2020) Global stratospheric temperature bias and other stratospheric aspects of era5 and era5. https://www.ecmwf.int/node/19362
34.
VaswaniAShazeerNParmarN, et al. (2023) Attention is all you need.
35.
WangSLiBZKhabsaM, et al. (2020) Linformer: self-Attention with linear complexity.
36.
WangXLyngaasITsarisA, et al. (2023) Ultra-long sequence distributed transformer.
XiongZWangYZhangF, et al. (2024) One for all: toward unified foundation models for Earth vision.
39.
YinJDashSWangF, et al. (2023) Forge: pre-Training open foundation models for science. In: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’23. New York, NY, USA: Association for Computing Machinery.
40.
YinJBoseACongG, et al. (2024) Comparative study of large language model architectures on frontier.
41.
ZhaoYGuAVarmaR, et al. (2023) Pytorch fsdp: experiences on scaling fully sharded data parallel.
42.
ZhaoXChengSZhengZ, et al. (2024) Dsp: dynamic sequence parallelism for multi-dimensional transformers.
