Sequence Length Scaling in Vision Transformers for Scientific Images on Frontier
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2024
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866909210280722432 |
|---|---|
| author | Tsaris, Aristeidis Zhang, Chengming Wang, Xiao Yin, Junqi Liu, Siyan Ashfaq, Moetasim Fan, Ming Choi, Jong Youl Wahib, Mohamed Lu, Dan Balaprakash, Prasanna Wang, Feiyi |
| author_facet | Tsaris, Aristeidis Zhang, Chengming Wang, Xiao Yin, Junqi Liu, Siyan Ashfaq, Moetasim Fan, Ming Choi, Jong Youl Wahib, Mohamed Lu, Dan Balaprakash, Prasanna Wang, Feiyi |
| contents | 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 has inspired scaling sequence lengths in ViTs, yet 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, tensor parallelism, 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 transformer model on a full-attention matrix over 188K sequence length. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_15780 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Sequence Length Scaling in Vision Transformers for Scientific Images on Frontier Tsaris, Aristeidis Zhang, Chengming Wang, Xiao Yin, Junqi Liu, Siyan Ashfaq, Moetasim Fan, Ming Choi, Jong Youl Wahib, Mohamed Lu, Dan Balaprakash, Prasanna Wang, Feiyi Computer Vision and Pattern Recognition Machine Learning 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 has inspired scaling sequence lengths in ViTs, yet 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, tensor parallelism, 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 transformer model on a full-attention matrix over 188K sequence length. |
| title | Sequence Length Scaling in Vision Transformers for Scientific Images on Frontier |
| topic | Computer Vision and Pattern Recognition Machine Learning |
| url | https://arxiv.org/abs/2405.15780 |