Sequence length scaling in vision transformers for scientific images on frontier

Aristeidis Tsaris; Chengming Zhang; Xiao Wang; Junqi Yin; Siyan Liu; Moetasim Ashfaq; Ming Fan; Jong Youl Choi; Mohamed Wahib; Dan Lu; Prasanna Balaprakash; Feiyi Wang

doi:10.1177/10943420251394758

Sequence length scaling in vision transformers for scientific images on frontier

Research output: Contribution to journal › Article › peer-review

Abstract

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.

Original language	English
Journal	International Journal of High Performance Computing Applications
DOIs	https://doi.org/10.1177/10943420251394758
State	Accepted/In press - 2025

Funding

This manuscript has been authored by UT-Battelle, LLC, under contract DE-AC05-00OR22725 with the US Department of Energy (DOE). The US government retains and the publisher, by accepting the article for publication, acknowledges that the US government retains a nonexclusive, paid-up, irrevocable, worldwide license to publish or reproduce the published form of this manuscript, or allow others to do so, for US government purposes. DOE will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan ( https://energy.gov/downloads/doe-public-access-plan ). This research used resources of the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725. The authors disclosed receipt of the following financial support for the research, authorship, and/or publication of this article: This study is supported by UT-Battelle; DE-AC05-00OR22725.

Keywords

computer vision
computing methodologies
distributed deep learning
machine learning algorithms
parallel algorithms

Access to Document

10.1177/10943420251394758

Cite this

@article{07b063c5da224cedae17149203e2330c,

title = "Sequence length scaling in vision transformers for scientific images on frontier",

abstract = "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.",

keywords = "computer vision, computing methodologies, distributed deep learning, machine learning algorithms, parallel algorithms",

author = "Aristeidis Tsaris and Chengming Zhang and Xiao Wang and Junqi Yin and Siyan Liu and Moetasim Ashfaq and Ming Fan and Choi, \{Jong Youl\} and Mohamed Wahib and Dan Lu and Prasanna Balaprakash and Feiyi Wang",

note = "Publisher Copyright: {\textcopyright} The Author(s) 2025",

year = "2025",

doi = "10.1177/10943420251394758",

language = "English",

journal = "International Journal of High Performance Computing Applications",

issn = "1094-3420",

publisher = "SAGE Publications",

}

TY - JOUR

T1 - Sequence length scaling in vision transformers for scientific images on frontier

AU - Tsaris, Aristeidis

AU - Zhang, Chengming

AU - Wang, Xiao

AU - Yin, Junqi

AU - Liu, Siyan

AU - Ashfaq, Moetasim

AU - Fan, Ming

AU - Choi, Jong Youl

AU - Wahib, Mohamed

AU - Lu, Dan

AU - Balaprakash, Prasanna

AU - Wang, Feiyi

PY - 2025

Y1 - 2025

N2 - 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.

AB - 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.

KW - computer vision

KW - computing methodologies

KW - distributed deep learning

KW - machine learning algorithms

KW - parallel algorithms

UR - https://www.scopus.com/pages/publications/105022209810

U2 - 10.1177/10943420251394758

DO - 10.1177/10943420251394758

M3 - Article

AN - SCOPUS:105022209810

SN - 1094-3420

JO - International Journal of High Performance Computing Applications

JF - International Journal of High Performance Computing Applications

ER -

Sequence length scaling in vision transformers for scientific images on frontier

Abstract

Funding

Keywords

Access to Document

Other files and links

Fingerprint

Cite this