# Sequence Length Scaling in Vision Transformers for Scientific Images on Frontier

arXiv: https://arxiv.org/abs/2405.15780 · HTML: https://mlsanity.com/paper/2405.15780/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 42 references that are in the map.
- First posted on arXiv 17 Apr 2024.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Scaling Vision Transformers to Gigapixel Images via Hierarchical Self-Supervised Learning](https://mlsanity.com/paper/2206.02647/) (2022)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Multimodal Foundation Models: From Specialists to General-Purpose Assistants](https://mlsanity.com/paper/2309.10020/) (2023)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [FuXi: A cascade machine learning forecasting system for 15-day global weather forecast](https://mlsanity.com/paper/2306.12873/) (2023)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [A Billion-scale Foundation Model for Remote Sensing Images](https://mlsanity.com/paper/2304.05215/) (2023)
- [One for All: Toward Unified Foundation Models for Earth Vision](https://mlsanity.com/paper/2401.07527/) (2024)
- [xT: Nested Tokenization for Larger Context in Large Images](https://mlsanity.com/paper/2403.01915/) (2024)
- [Comparative Study of Large Language Model Architectures on Frontier](https://mlsanity.com/paper/2402.00691/) (2024)

## Built on by (most cited first, 1 of 1)

- [ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling](https://mlsanity.com/paper/2505.04802/) (2025)

## Abstract (arXiv)

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.