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

Aristeidis Tsaris * Oak Ridge National Laboratory, Oak Ridge, United States
Chengming Zhang  Indiana University Bloomington, Bloomington, United States
Xiao Wang  Oak Ridge National Laboratory, Oak Ridge, United States  Junqi Yin
Oak Ridge National Laboratory, Oak Ridge, United States  Siyan Liu  Oak Ridge
National Laboratory, Oak Ridge, United States  Moetasim Ashfaq  Oak Ridge
National Laboratory, Oak Ridge, United States  Ming Fan  Oak Ridge National
Laboratory, Oak Ridge, United States
Jong Youl Choi  Oak Ridge National Laboratory, Oak Ridge, United States
Mohamed Wahib  RIKEN Center for Computational Science, Kobe, Japan  Dan Lu
Oak Ridge National Laboratory, Oak Ridge, United States  Prasanna Balaprakash
Oak Ridge National Laboratory, Oak Ridge, United States  Feiyi Wang  Oak Ridge
National Laboratory, Oak Ridge, United States

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

###### Index Terms:

Computing methodologies; Machine learning algorithms; Parallel algorithms;
Distributed deep learning

**footnotetext: Corresponding author<EMAIL_ADDRESS>
##  I Introduction

In the rapidly evolving landscape of artificial intelligence and machine
learning (AI and ML), Vision Transformers (ViTs) have become notable for their
ability to model long-range dependencies between tokens in time-series and
imaging data [1]. ViTs are proving effective in building foundational models
for scientific imagery applications, including research in Earth system
science and geospatial analysis [2, 3, 4]. Additionally, ViTs are commonly
used as the architecture of choice for image encoders in multimodal foundation
models, highlighting their importance in processing complex image inputs [5].
Despite their crucial role in foundation model (FM) development, ViTs have not
scaled to the same extent as Large Language Models (LLMs). To date, LLMs have
expanded to encompass over one trillion parameters [6], whereas the ViTs have
only reached a pinnacle with 22 billion parameter models, specifically the
ViT-22B [7].

ViTs face distinct challenges compared to the transformers used in LLMs, such
as GPTs. The architectural differences between these models significantly
influence their operation, training, and application. Transformers in LLMs are
primarily decoder-only models focused on processing and generating textual
data. Their input consists of embeddings of tokenized text, which operates
within a well-defined predictive space limited by the vocabulary size. This
bounded space simplifies the task of predicting the next word in a sequence,
allowing for efficient scaling primarily through increases in model size.
Conversely, ViTs embody an encoder-decoder architecture, adding complexity
layers to their functionality. Unlike LLMs, ViTs are trained on image and
video data converted into a sequence of tokens derived from patches. This
method of data representation is inherently more complex than the tokenized
text embeddings used by LLMs. The visual domain does not benefit from a
constrained predictive space like language. Instead, the space of possible
future images or visual outputs is virtually limitless, presenting a
significant challenge for ViTs when forecasting or generating new visual
content. This fundamental distinction between the textual versus visual data
necessitates different scaling strategies. For LLMs, enhancing model size has
been a focal point for improving performance. However, for ViTs, sequence
length distribution emerges as a crucial aspect for further scaling. The
infinite possibilities in visual representation demand that ViTs process
larger sequences of visual data. Moreover, the encoder-decoder framework of
ViTs must efficiently translate raw visual inputs into a sequence of tokens
and then interpret these tokens across a broad range of tasks. The challenge
is compounded by the need for ViTs to handle these tasks with a high degree of
accuracy and versatility, given the unbounded nature of visual data.

For LLMs and ViTs, extending sequence lengths has been identified as a method
to enhance model performance. While conventional training for ViTs on natural
images typically employs a sequence length of 256, even for expansive models
[7], scientific foundational models (FMs) such as those used in ClimaX [8], a
foundation model for weather and climate forecasting, leverage significantly
longer sequences, up to 8K or more. Scientific imagery, including microscopy
or remote sensing data, often surpasses the resolution of standard images,
reaching dimensions of 100,000 x 100,000 pixels [9, 10, 11]. This high
resolution, combined with the unique characteristics of scientific data, such
as integrating observational data with simulation data, necessitates much
longer sequence lengths for effective model training. For instance, training
on the native resolution of ERA5 reanalysis data [12] with the ClimaX
foundation model would require a sequence length of 277K, vastly exceeding the
lengths used for images from consumer devices. Moreover, the complexity of
scientific images is compounded by their channel diversity. Unlike the
correlated RGB channels of consumer device images, each channel in scientific
images typically represents a distinct physical quantity, leading to a linear
increase in the number of sequence lengths as the number of channels grows.
This diversity and the need for handling spatiotemporal images or unstructured
time-series data, such as video or scientific simulations, underscore the
demand for longer sequence lengths in scientific applications. However,
scaling sequence length is not without its challenges. The memory footprint of
transformers scales quadratically with sequence length, posing significant
constraints on training capacity with current hardware. As a result, training
on longer sequences often requires trade-offs, utilizing sparse attention
mechanisms or approximations of the full-attention matrix to manage
computational demands.

Dense attention within transformer models facilitates a comprehensive
interaction among all tokens in the input sequence, enabling the model to
capture long-range dependencies and intricate patterns [13]. Contrastingly,
sparse attention mechanisms [14], designed to mitigate the computational
demands of dense attention, limit each token’s interaction to a subset of
other tokens. While this approach reduces computational complexity—often from
quadratic to linear or sub-quadratic—it may not capture the full context as
effectively as dense attention, particularly in tasks where the relationship
between distant parts of the input is critical.

Justification for the use of dense attention despite its higher computational
cost comes down to its unmatched ability to model complex interactions and
dependencies across entire sequences. In scientific imagery, where the precise
arrangement and relationship of pixels or data points can convey essential
information, dense attention’s comprehensive feature learning capability is
invaluable. It ensures that no potential relationship or pattern is
overlooked, a critical factor in accurately modeling and interpreting
scientific imagery data. Moreover, dense attention’s robustness in handling
intricate data patterns makes it critical for advanced applications, including
high-resolution image analysis, complex sequence prediction, and detailed
textual interpretation, where the length of the context directly impacts model
performance.

We developed an ultra-long vision transformer, with dense attention, that uses
distributed sequence parallel methods for scientific images. Our approach is
based on DeepSpeed-Ulysses [15] and on Long Sequence Segmentation (LSS) [16]
for distributing long sequences across GPUs as contiguous segments. DeepSpeed-
Ulysses aims to optimize the distribution of computational load for self-
attention across multiple GPUs, thereby enhancing the model’s scalability and
efficiency. LSS on the other hand addresses the direct dependency on self-
attention. It employs fused communication and a double gradient averaging
technique, reducing the complexity of aggregating partial self-attention
scores. This approach aims to simplify the computational process and improve
scalability by minimizing the communication overhead typically associated with
dense attention in long sequences.

These two methods exhibit different communication patterns per attention
layer. While DeepSpeed-Ulysses is expected to see quadratic growth in
communication demand with an increase in parallel GPUs due to its distribution
of self-attention, LSS aims for a more streamlined communication strategy,
with only two calls per attention layer. We analyzed the scaling efficiency of
each method separately, scaling the sequence length up to 1M sequence length,
for model sizes up to 10B of parameters.

Our analysis showed limitations and performance bottlenecks of each method,
for example, DeepSpeed-Ulysses scaling size is limited by the number of
attention heads of the model. And while flash-attention-v2 [17] can help to
fit very large sequences for small model sizes, for models that can’t fit on a
single GPU a hybrid parallelism is necessary. To that extent, we used Pipeline
Parallelism (PP) [18] with DeepSpeed-Ulysses and Tensor Parallelism (TP) [19]
with LSS so that we can train more efficient on longer sequence lengths
compared to single parallel mode.

There is a noted absence in the literature for ViT performance baselines, and
with this work, we provide an empirical study for the practical applicability
of each method in distributed environments at different scales. Also as top
HPC centers in America and Europe are using AMD GPUs, with this work we are
testing the AMD AI software stack on four state-of-the-art distributed
methods, while showing how to best use the Frontier Supercomputer for large-
scale AI on scientific images. We run at scale a 188K sequence length ViT
model, while scaling the batch size up to 2,048 AMD MI250X GPUs measuring a
94% scaling efficiency.

Additionally, we established the advantage of longer sequence training on a
climate application by showing that we can get better accuracy by increasing
the sequence length of the model. Specifically, we trained on the ERA5 climate
dataset, and by increasing either the image resolution or the number of
channels embedding to the model per time-step, we see a 20% higher accuracy in
temperature prediction. This is the first time a transformer model was trained
with a full attention matrix on 188K sequence length.

In summary, our contributions are the following:

  * •

We have developed a vision transformer capable of handling ultra-long sequence
lengths through sequence parallelism. Our method leverages DeepSpeed-Ulysses
and LSS parallel strategies, achieving sequence lengths of up to 1 million.
This marks the first successful application of sequence parallelism in
training Vision Transformers (ViTs), and by integrating model sharding, we
attained a 94% batch scaling efficiency across 2,048 AMD MI250X GPUs.

  * •

Our evaluation of sequence parallelism on ViTs covered models up to 10 billion
parameters. The analysis identified significant performance bottlenecks,
particularly in models with 5 and 10 billion parameters. To address these, we
employed distributed sequence strategies in combination with Pipeline
Parallelism (PP) and Tensor Parallelism (TP), as well as flash-attention-v2.
This approach enabled further scaling of sequence lengths for models that
exceed the memory capacity of a single GPU and included an empirical study
across various scales.

  * •

We have showcased the advantages of utilizing extended sequence lengths in
climate modeling, observing up to a 20% increase in accuracy for temperature
predictions. This achievement represents the first instance of successfully
training a transformer model to process a full-attention matrix across such
extensive sequence lengths, specifically 188K tokens.

The rest of the paper is organized as follows. In Sections II and III we
defined the concepts and the terms used in this paper, and discuss related
work. Section IV, shows the advantage of training on longer sequence lengths
for a climate application, and we calculate the computation training need for
scientific images on Frontier. Section V shows scaling solutions and best
approaches for different model sizes and sequence lengths. Finally,
convergence runs and weak scaling performance are shown in Section VI with
real climate data.

##  II Background

This section describes the ViT model and the patch embedding methods used in
this work, along with the various serial and distributed methods applied.

###  II-A Patch Embedding and Sequence Length

In most common cases, ViTs are applied to images captured by consumer devices,
where the typical resolution is 100s pixels. To make a sequence of tokens for
the encoder a patch size of 16x16 is typically selected. It is common practice
for larger models to have smaller patch sizes, leading to a larger sequence of
tokens [20, 7]. General RGB images have three overlapping channels on the same
object, so 3D patches are usually chosen. For a 224x224x3 image using 3D
patches of 16x16x3, we will end up with a sequence length of 192 tokens.

On the other hand, scientific simulations, such as in Earth system science, or
multi-sensor data acquisition systems, such as medical imaging, produce multi-
dimensional data, and the channels are independent. In this case, 2D patches
are more appropriate since the model needs to understand different physical
quantities from sensors or variables. For example, for the same 224x224x3
image, if we apply 2D 16x16 patches we will end up with a three times larger
sequence length. When we explicitly embed the full number of channels in the
transformer, we will call this Multi-Ch-ViT.

Some approaches aggregate all the channels into one, as proposed by [8], by
including a multi-head attention layer in the encoder, i.e., the data query is
a learnable vector. That way, the number of channels won’t increase the
sequence length. We call this Agg-Ch-ViT.

###  II-B Data Parallel (DP)

In data parallelism a copy of the entire model is replicated across each
worker. Zero Redundancy Optimizer (ZeRO) [21] removes the memory redundancies
across data-parallel processes by partitioning the three model states
(optimizer states, gradients, and parameters) across data-parallel processes
instead of replicating them.

###  II-C Pipeline Parallel (PP)

Pipeline parallelism [18] can improve the efficiency and scalability of
training large deep learning models, particularly those that are too large to
fit into the memory of a single device or that require substantial
computational power. Unlike data parallelism, which replicates the entire
model across multiple devices and divides the dataset among them, pipeline
parallelism divides the model itself into several smaller, sequential stages
or segments. Each stage of the model is then placed on a different device.

###  II-D Tensor Parallel (TP)

Tensor parallelism involves partitioning the model’s tensors across different
processors. Megatron-LM [19] exploits the inherent structure of Transformer
networks to implement tensor parallelism efficiently. It parallelizes both the
Multi-Layer Perceptron (MLP) blocks and the self-attention blocks of a
Transformer layer, using only a few synchronization primitives.

###  II-E Sequence Parallel (SP)

Sequence parallelism, which distributes the self-attention computation in
transformer models across the sequence dimension, has been an active area of
research, with many developed methods employing various approaches. DeepSpeed-
Ulysses [15]introduces a suite of system optimizations specifically designed
for the efficient training of extreme long sequence transformer models.
DeepSpeed-Ulysses partitions the input data along the sequence dimension,
which allows the system to handle much longer sequences. For attention
computation, it employs an all-to-all collective communication to ensure that
each GPU receives a complete sequence, but only for a non-overlapping subset
of the attention heads, allowing for parallel computation of attention across
GPUs. LSS [16] divides a long sequence into segments distributed among GPUs,
with each GPU computing a partial self-attention for its segment. It also
introduces a fused communication strategy to minimize the communication
overhead. Dynamic Sequence Parallelism (DSP) [22] enables efficient sequence
parallelism for multi-dimensional transformers by dynamically switching the
parallelism dimension according to the current computation stage. Ring
attention [23, 24] is another variant that organizes tokens in a ring
structure, allowing each token to attend to a fixed number of adjacent tokens
in the sequence. This method offers a compromise between dense and sparse
attention, balancing computational efficiency with the ability to capture
local and semi-distant relationships. However, ring attention may still fall
short in tasks where long-range dependencies are crucial.

###  II-F Flash Attention v2

FlashAttention [25] significantly speeds up attention computation by tiling
the attention matrix to reduce the number of memory reads and writes between
the GPU’s high bandwidth memory (HBM) and on-chip SRAM. FlashAttention-2 [17]
further accelerates attention computation by reducing the number of non-
MatMult FLOPs and parallelizing computation in the sequence length dimension.

##  III Related Work

As mentioned in Section I, literature reviews reveal an absence of references
utilizing ViTs with sequence parallelism for baseline comparisons. The
majority of existing studies concentrate on extending the sequence length
within a single GPU to examine its impact on convergence rates. Machine
learning datasets comprising natural images typically feature low resolutions,
such as 224x224 pixels. A common approach to enlarging the receptive field
involves reducing the patch size, thereby increasing the sequence length.
Conversely, high-resolution datasets, notably those in remote sensing imagery,
often resort to cropping images to 224x224 pixels [3, 4]. This practice
significantly hampers the transformer’s ability to apprehend the global
context, thereby undermining the potential advantages of utilizing high-
resolution images.

In the climate and weather domain, Vision Transformers (ViTs) and their
variants have been increasingly adopted for rapid and accurate weather
forecasting. The main objective of the ViT used on weather data is to predict
from a weather snapshot $X_{t+}$ of a specific time $t$ the future weather
forecast $X_{t+\Delta t}$ at lead-time $\Delta t$. The majority of the climate
ViTs handle the long sequence complexity by replacing the dense attention
matrix with some approximation. FourCastNet [26] employs a Fourier transform-
based token mixing scheme alongside a ViT backbone. This approach enables
flexible and precise forecasts of global weather conditions ranging from 6
hours to 7 days ahead, covering high-resolution, rapid-timescale variables
such as surface wind speed, precipitation, and atmospheric water vapor. Pangu-
weather has introduced a 3D Earth-specific Transformer model that also
leverages the ViT architecture. This model effectively captures the 3D
atmospheric dynamics, delivering more accurate 7-day weather forecasts [27]
compared to traditional numerical weather prediction methods. Additionally,
the FuXi model [28], based on ViT, demonstrates performance on par with
established forecasting methods, extending its forecasting capability to an
impressive 15 days.

The above models, even though are trained on the ERA5 high-resolution dataset
and use ViT as a backbone, they don’t use the full attention matrix for the
training. They mostly rely on approximations, or on sparse attention, such as
Swin Transformer [29], in order to fit that long sequence on a single GPU. On
the other hand, Climax [8] does use self-attention, but it is trained on lower
resolution compared to the previous models. Climax is a foundation model
leveraging extensive datasets from the 6th phase of the Coupled Models
Intercomparison Project (CMIP6) global climate models. It can be fine-tuned to
address various climate and weather tasks, including global/regional weather
forecasting and future climate projection. ClimaX and the original ViT paper
[20] showed a better and more robust solutions when increasing the sequence
length using the full attention matrix on ERA5 and ImageNet [30] dataset
respectively. In fact, [20] suggested that scaling the compute is be a better
predictor of the model quality rather than scaling the model size alone for
ViTs.

##  IV Sequence Length Impact and Bottlenecks

In this section, we describe how the image resolution and the multiple
channels can affect the sequence length and bottlenecks that are encountered
with the current approaches.

###  IV-A Compute Cost to Train ViTs on Frontier

To better understand the computation cost for training on large sequence
lengths and estimate the difference by scaling the model size, we can
calculate the ideal training time as a function of the number of FLOPs for the
ERA5 dataset, i.e. 350K time-steps. For the following calculations we used the
formulas as described by [31]. The total number of floating point operations
(FLOPs) needed to train a ViT, depends on the parameter of the model (P) and
the number of tokens (D), where for each element of a feed-forward weight
matrix, there is a total of 6 FLOPs (1 multiply-add operation during the
forward pass and 2 multiply-add operations during the backward pass) per input
token:

$T_{FLOPS}\sim 6*P*D$ (1)

Also, the training time ($t$) depends on the measured percentage of the GPU’s
theoretical peak ($R_{peak}$), the ratio of tokens over the parameters
($\gamma=T/P$) and the number of devices N:

$t\sim\frac{6*\gamma P^{2}}{R_{peak}N}$ (2)

The training time from the above equation assumes that the model will converge
on such a large batch size, the kernels libraries used by the ViT architecture
can fully utilize the compute of the device, and the scaling efficiency is
100% across Frontier (so we don’t account at all internode and internode
communication bandwidth, IO, algorithmic and libraries scaling
inefficiencies).

Figure 1: Computation cost in FLOPs by varying the number of tokens (left
plot), and the parameter size of the model (right plots). For both plots
equations (1), and (2) were used with a total number of tokens for the ERA5
datasets, i.e. around 350K images

Figure 1 was made using equations (1), and (2) by varying the number of tokens
(left plot) and the model size (right plot). As we can see from the plots
scaling the sequence length is as big of a problem as scaling the model size,
and even though the two are orthogonal to each other, they are both going to
add to the total computational cost and detailed study, as this, is necessary
to strike a balance of the two and to find the optimal ratio between sequence-
length and model size.

###  IV-B Sequence Length for Scientific Images

Scientific images are usually found in very high resolution on multiple
channels. Data resolution in Earth system science can vary depending on the
targeted phenomena and whether the data comes from simulations or satellites.
In our current work, the primary dataset utilized for convergence runs was the
ECMWF Reanalysis v5 (ERA5) [12]. ERA5 represents the fifth generation of
reanalysis by the European Centre for Medium-Range Weather Forecasts (ECMWF),
covering global climate and weather for the past eight decades. It provides
hourly estimates across a range of variables that represent the atmosphere,
ocean, and land surface. Typically, a ViT patch size of 2 is employed to embed
the climate data, which, when trained on low resolution, results in an 8K
sequence length [8]. The native resolution of ERA5 is 770x1440 grid points,
leading to a 277K sequence, which is impractical to accommodate on a single
GPU. Moreover, weather models, such as those referenced in [32], can have
resolutions up to 100 times greater than those of climate models in CMIP6.
Additionally, these datasets may feature more than 100 channels, culminating
in sequence lengths that extend into the millions.

###  IV-C Training Results for Different Sequence Length

Given the many ways the sequence length can be increased with climate data,
along with the growing interest in developing Earth system foundational models
based on ViTs, we think this application space is ideal for studying the
effect of sequence length.

With the ERA5 dataset, we choose 5 input atmospheric variables: geopotential,
temperature, u-component of wind, v-component of wind and specific humidity on
17 pressure levels: 10, 20, 30, 50, 70, 100, 150, 200, 250, 300, 400, 500,
600, 700, 850, 925, 1000, 3 surface level variables: 2m temperature, 10m
u-component of wind, 10m v-component of wind, and 2 statics: orography and
land-sea mask. Collectively, this combination leads to a total of 92 variables
(or channels). The performance is evaluated for 4 variables: 2m temperature
(T2m), 10m u-component of wind (U10), geopotential at 500hPa (Z500), and
temperature at 850hPa (T850). We also re-gridded the original dataset from
0.25° (770 × 1440) to 3 low resolutions for the experiments: 1.0° (180 × 360),
1.40625° (128 × 256), and 5.625° (32 × 64). The re-gridding is done by
utilizing the xESMF [33] based on the high-performance Earth System Modeling
Framework (ESMF) as the backend, and it performs re-gridding between the
general curvilinear grids with different re-gridding algorithms, e.g.,
bilinear, nearest neighbor and conservative, we applied the bilinear method to
all of our re-gridding tasks.

Figure 2: Validation accuracy on half-year of the ERA5 dataset of the Z500,
T850, T2m, and U10 variables. The effect in the accuracy is shown for three
different spatial resolutions: 5.625°, 1.40625°, and 1.0°. A patch size of 4
was used, and so the sequence length for each resolution is 128, 2048, and
4050 respectively. Figure 3: Validation accuracy on half-year of the ERA5
dataset of the Z500, T850, T2m, and U10 variables. The effect in the accuracy
is shown for including all 92 variables in the model, Multi-Ch-ViT, versus
Agg-Ch-ViT. The sequence length of the first approach is 11,776, while it
reduces to 128 for the second approach.

We used the first half of the 2018 year for training and the second half for
testing. All the hyper-parameters were kept constant between each experiment,
using a predicted range of 28 hours, with a lead-time of 6 hours. Figure 2
shows the effect of three different spatial resolutions: 5.625°, 1.40625°, and
1.0°. Since we use a constant patch size of 4, it will result in the following
sequence lengths: 128, 2048, and 4050, respectively. Figure 3 shows the effect
of Multi-Ch-ViT, by using all 92 channels for each time-step, against the Agg-
Ch-ViT as implemented in [8]. This will also increase the sequence length from
128 to 11,776 tokens per time step.

Even though experiments in Figures 2 and 3 are different, they both show that
increasing the sequence length results in a better overall accuracy. Also it
is worth noting, increasing the resolution leads to a better accuracy from the
very first epoch, while for the multi-ch-ViT, we see improvements for the
majority of the predicted variables after the 4th epoch. The reason for the
latter is in multi-ch-ViT the sequence length increases 92 times from the
baseline, compare to 32 time for the resolution increase, and so it takes
longer time for the attention layers to build a good set of weights, and learn
the relationships between the 92 different channels. Another observation is
that in both cases, the Z500 doesn’t see as big of an improvement as the other
variables. Insensitivity of Z500 is partly expected as it exhibits spatially
less varying synoptic-scale features in the middle troposphere. Therefore,
increasing the resolution does not add significantly distinct details. On the
other hand, the T850 in the lower troposphere expectedly has more spatial
variability due to the influence of the land-atmosphere flux exchanges. Hence,
it shows improvements in results for the Multi-Ch-ViT compared to Z500.

From the experiments conducted thus far, we have identified a need to increase
the sequence length for ViTs. Unstructured raw data, such as images and
videos, can significantly escalate the number of tokens on which transformers
are trained, presenting substantial computational challenges. The optimal
balance between image resolution and the number of channels likely varies
depending on the domain, necessitating a careful consideration of
computational costs. For small models dealing with very high-resolution
images, aggregating the channels may be the preferable strategy. Conversely,
when dealing with a large number of channels but with moderate resolution, it
is advisable not to aggregate the channels. Instead, increasing the patch size
to manage the sequence length more effectively in relation to the image
resolution is recommended.

##  V Scaling Solutions for Large Sequence Length and Evaluation

In this section, we evaluate the performance implications of extending the
sequence length using different distribution methods on Frontier’s AMD GPU.
Specifically, we incorporate two sequence parallelism strategies into ViTs:
DeepSpeed-Ulysses and LSS, each adopting a significantly different approach to
distributing the sequence length. We independently assess PP and TP, in
addition to integrating them with the aforementioned sequence parallelism
techniques. DeepSpeed-Ulysses and PP are implemented within the DeepSpeed
framework[34], whereas LSS and TP utilize the Fully Sharded Data Parallel
(FSDP) framework [35]. Given the variance in framework integration, a direct
comparison between LSS and DeepSpeed-Ulysses is not straightforward; however,
we can compare their scaling efficiency. Additionally, we offer a practical
guide to determine the most suitable approach under varying conditions.

###  V-A Frontier Architecture

All the experiments were performed using the Frontier Supercomputer [36] at
the Oak Ridge Leadership Computing Facility. Each Frontier node has a single
64-core AMD EPYC CPU and four AMD Instinct MI250X GPU accelerators. The MI250X
GPU is comprised of two Graphics Compute Dies (GCDs), connected with Infinity
Fabric CPU-GPU, while the four MI250X GPUs are connected with Infinity Fabric
GPU-GPU of 50GB/s. The system identifies each GCD independently, so from the
application perspective, it can be considered that each node has 8 GPUs, each
with 64 GB of high-bandwidth memory. For simplicity, we will use the term GPU
when referring to a GCD. The nodes are connected via a Slingshot-11
interconnect with 100GB/s to a total of 9408 nodes, making it the first true
exascale machine. For the software stack, we used Pytorch 2.4 nightly build
03/16/2024. ROCm v5.7.0, MIOpen v2.19.0, RCCL v2.13.4 with libfabric v1.15.2
plugin.

###  V-B Flash Attention v2 on Frontier

Figure 4 shows the compute performance in FLOPs and the memory footprint as a
function of the sequence length, for Flash-Attention-v2 (FA2) [17] on
Frontier’s AMD MI250X GPU. For the base model, a hidden dimension of 1024 was
used with 16 attention heads, while for the large model, a 2048 hidden
dimension was used with 32 attention heads. We scale the sequence length from
512, to 16K, and set the batch size so that the total number of tokens is 16K
through each run.

The first main observation is that we can fit 4 times larger sequence lengths
using FA2 for the base model and 32 times larger for the large model, which is
a significant improvement when scaling the sequence length. For the base
model, we also see improvements in the number of computational efficiencies,
especially as we move to a larger sequence length. Also, the max reserved
memory is reduced by 50% or more in all cases.

Overall, we don’t observe big performance improvements using FA2 regarding
computational efficiency, but we see savings in memory footprints. It is worth
noting that the small performance gain doesn’t reflect the AMD hardware itself
but rather the readiness of the software implementation of FA2 on the ROCm
software stack. For example, on a GPT model, [37] has shown a performance gain
of 20% on the Frontier MI250X GPUs using the open-sourced development [38] of
flash attention v1. In this work, we directly used Pytorch fused attention
kernels, while we didn’t write from scratch rocm code for our specific model
architecture.

Given that FA2 performance is still in the experimental stage within Pytorch
on AMD GPUs, we will not use it when searching for the best-performing
distributed approach but only after finding the best parallel strategy to push
the highest sequence length per GPU.

Figure 4: Single GPU measurements with and without flash-attention-2, on
Frontier MI250X GPU (1 GCD). The top plot shows the GPU throughput for two
model sizes, ViT-Base and ViT-Large, while the bottom plot shows the max GPU
memory reserved throughout each run. For the flash-attention-2 measurements,
Pytorch 2.4 nightly build 03/16/2024 was used with rocm5.7.

Figure 5: Weak scale the sequence length for each distributed method
separately. Specifically, the top plot shows the DeepSpeed-Ulysses and
Pipeline (PP) methods within the DeepSpeed framework, while we also scale the
batch size linearly with the DeepSpeed Zero Redundancy Optimizer (DP). The
bottom plot shows the Long Short-Sequence method (LSS) and the tensor
parallelism (TP) methods within Pytorch’s FSDP framework, while we also show
scaling linearly the batch size with the fully sharded method (DP). The model
for the two implementations was set to a 2560 hidden dimension with 16 heads
total.

###  V-C Different Distributed Approaches

We use DeepSpeed-Ulysses, PP, LSS, and TP separately to scale the sequence
length beyond the single GPU limits. The top plot of Figure 5 shows DeepSpeed-
Ulysses and PP used separately by weak scaling the sequence length within the
DeepSpeed framework. The bottom plot of Figure 5 shows LSS and TP used
separately by weak scaling the sequence length within Pytorch’s FSDP
framework. In all four cases, we also measured the performance of scaling the
batch for each method separately using data-parallel approach. A ViT model
architecture of 2560 hidden dimensions was used, with 16 attention heads in
total.

As we can see DeepSpeed-Ulysses shows the best scaling efficiency from 1 to 16
GPUs. Next, we see that LSS scales better than TP and PP. That behavior was
expected since PP and TP are designed to distribute the model size rather than
the sequence length. We also observe that the DeepSpeed data-parallel scale is
better for DeepSpeed-Ulysses and PP, compared to FSDP model sharding for LSS
and TP. Probably, that’s an indication that some tuning is necessary between
FSDP and LSS. Finally, even though DeepSpeed-Ulysses scales better than LSS,
there is a fundamental limit in the number of sequence parallel ranks that we
can use, which is imposed by the number of heads. For example, the model
architecture we used had 16 attention heads, so we can’t use more than 16
sequence parallel ranks, limiting the sequence length we can run to 16K.

Figure 6: The overview of combining sequence parallelism and pipeline
parallelism.

Figure 7: Weak scale the sequence length by pairing the DeepSpeed-Ulysses with
Pipeline parallelism (PP) and the Long Short-Sequence method (LSS) with Tensor
Parallelism (TP). The top plots show the GPU throughput, while the bottom
plots show the max GPU memory reserved throughout each run. The number of PP
(TP), shows how many pipelines (tensor parallel) ranks were used, and the
DeepSpeed-Ulysses (LSS) ranks were adjusted to keep the number of GPUs
constant. For example, for 8 GPUs, the Ulysses-2PP means that we did sequence
parallelism within 4 GPUs, and between each GPU pair, we did pipeline
parallelism, while the Ulysses-4PP means that we did pipeline parallelism
within 4 GPUs, and between each GPU pair, we did sequence parallelism.

###  V-D Hybrid Distributed Approaches

Here, we present the results of combining DeepSpeed-Ulysses with Pipeline
Parallelism (PP) and Long-Short Sequence (LSS) with Tensor Parallelism (TP) to
scale the sequence length.

Often, it is beneficial to integrate distributed methods to optimize their
advantages and mitigate potential bottlenecks. This approach is justified
because: (1) Sequence parallelism, TP, and PP are orthogonal to each other.
(2) Communication overhead can significantly restrict scalability if solely
relying on sequence parallelism, TP, or PP. To address this, we combined LSS
with TP to reduce the communication overhead involved in distributing
activations from sequence groups to tensor parallel groups. Furthermore, we
integrated DeepSpeed-Ulysses with PP (see an example configuration in Figure
6). DeepSpeed-Ulysses primarily utilizes two all-to-all communication calls in
both the forward and backward passes, complemented by an all-reduce per layer,
and concludes with an all-gather call at the end of each step for
optimization. In contrast, PP mainly involves all-reduce operations between
partitioned layers.

Additionally, the top plot in Figure 7 illustrates the compute throughput in
TFLOPs for pairing two methods, whereas the bottom plot depicts the memory
footprint of the application. The labels 2PP, 4PP, and 8PP indicate the number
of pipeline ranks, with the rest of the configurations dedicated to
distributing the sequence length. For instance, ”Ulysses-2PP” on 8 GPUs
signifies that there were four groups employing sequence parallelism with
DeepSpeed-Ulysses, each consisting of two GPU pairs engaged in pipeline
parallelism. This rationale is similarly applied to LSS-2TP and other
configurations.

Our first observation is that we successfully overcame the limitation of
DeepSpeed-Ulysses, which was only capable of handling a number of attention
heads equal to the distributed ranks. We are now able to scale beyond a 16K
sequence length, thanks to the ability to utilize additional ranks with
pipeline parallelism, scaling up to 65K on 128 GPUs. This is achieved by
employing 8 GPUs per node for pipeline parallelism and distributing the
sequence length across nodes. Using DeepSpeed-Ulysses alone in the previous
subsection, we observed substantial scaling benefits as the ranks increased.
However, integrating it further, we now witness a significant performance
enhancement. We attribute the primary reason for this improvement to
communication dynamics; as indicated by the bottom left plot, our memory
footprint exceeds 95%, suggesting that memory consumption increases with more
ranks, contrary to our experience with solely DeepSpeed-Ulysses.

On the other hand, LSS is not confined by the number of sequence ranks
deployable, yet its scaling efficiency did not match that of DeepSpeed-
Ulysses. Thus, in the right plot of Figure 7, we experimented with combining
LSS with Tensor Parallelism (TP) to enhance compute performance. The memory
footprint data reveal that, despite employing additional ranks with TP, we
were unable to reduce the memory usage below 50% of the maximum, which would
allow for an increase in the local batch size and overall computation.
Interestingly, a comparison between the memory footprints of DeepSpeed-Ulysses
with Pipeline Parallelism (PP) versus LSS with TP shows an increase in memory
usage in the former case and a decrease in the latter. This outcome is
promising for LSS and TP configurations, as it suggests potential for further
optimization to accommodate larger local batch sizes, thereby improving
compute throughput.

###  V-E Combining DeepSpeed-Ulysses with FA2

In this subsection, we use the best performing distributed approach we found
so far, i.e., DeepSpeed-Ulysses with Flash-attention-v2 to push the limits of
the sequence length on Frontier with our current software stuck. For our
tests, we used models with four different sizes, as defined in Table I.

TABLE I: A Summary of the different model
architectures explored in this work Model | Width | Depth | Heads | Params [M]
---|---|---|---|---
ViT-Base | 1024 | 8 | 16 | 84
ViT-Large | 2048 | 16 | 32 | 671
ViT-5B | 4096 | 32 | 32 | 5,370
ViT-10B | 4608 | 48 | 64 | 10,194

Figure 8: Measuring the GPU throughput, the memory footprint, and the image-
per-second for four model sizes (see Table I), by scaling the sequence length.
Only the DeepSpeed-Ulysses distributed method was used, along with DeepSpeed
Zero Redundancy Optimizer and flash-attention-2. The numbers on top of the bar
charts show the number of GPUs for each measurement.

Figure 8 shows the compute throughput, the memory footprint, and the image-
per-second for the four different models for different sequence lengths. The
ViT-Base and the ViT-Large model can scale up to 1M sequence length on 16 and
32 GPUs respectively. The compute scaling efficiency looks very good in both
cases, which indicates that we are not compute bound for those models, even up
to 1M sequence length. On the other hand, we can see the memory increases
almost linearly for the base and large model, mostly an expected behavior
since we are using flash-attention on the compute, and communication doesn’t
seem to be affecting the performance. Finally, we can see for the base and
large model, 16 times less image-per-second performance from 64K to 1M
sequence length. This is also expected since we are increasing the sequence
length, and so the number of tokens per training step, almost by the same
amount, going from 64K to 1M.

Figure 8 also shows the performance of the ViT-5B and ViT-10B models. To
accommodate a 64K sequence length, a minimum of 8 GPUs is necessary for the
ViT-5B, and a minimum of 32 GPUs is required for the ViT-10B. However, for
both models, scaling up to a 1M sequence length using FA2 and DeepSpeed-
Ulysses alone is not feasible. This limitation arises because, within
DeepSpeed-Ulysses the number of sequence ranks cannot exceed the number of
attention heads. Consequently, the maximum number of sequence ranks for the
ViT-5B is 32, and for the ViT-10B, it is 64. We can overcome this limitation
by integrating pipeline parallelism as shown before.

Finally, we can see from Figure 8 that the ViT-5B and ViT-10B models show
poorer compute performance compared to base and large models. Both of those
models won’t be able to fit on a single GPU, so we are relying on DeepSpeed to
share the parameters and activations of the model. This will increase the
overall communication of the application since apart from the all-to-all
communication that DeepSpeed-Ulysses is primarily using, we are going to have
additional all-gather and scatter due to the larger model size. We can also
see that the memory footprint of the two larger models is constant as we scale
the sequence length, which is an indication that the application at that scale
doesn’t have ideal compute and communication overlap and some further fine-
tuning might be necessary.

###  V-F Scaling Sequence Length

Based on these experiments, we found that if we want to use a single
distributed strategy to scale the sequence length for smaller models the best
performing approach is DeepSpeed-Ulysses along with FA2. Also combining
DeepSpeed-Ulysses with DeepSpeed Zero parallel modes shows better out-of-the-
box scaling performance compared to comping FSDP with other distributed
strategies. For models that can’t fit on a single GPU DeepSpeed-Ulysses alone
is limited, and an overlap between PP and DeepSpeed-Ulysses is necessary. The
latter will probably have its upper bound limit since the memory footprint
keeps increasing, while we increase the number of ranks, even by keeping the
local batch size constant. On the other hand, even though the LSS method does
show less scaling efficiency, compared to DeepSpeed-Ulysses, it doesn’t suffer
from a limit in the number of sequence ranks to use. Furthermore, when
combined with TP, since the memory footprint does decrease with more ranks, it
shows more potential to fit larger local batch size, and so at a very large
scale of model and sequence length probably a combination of LSS and TP would
be the best approach.

##  VI Training Long Sequence ViTs for Climate Data

Here, we study the performance and training results on climate data, using FA2
and DeepSpeed-Ulysses on two model sizes, ViT-Base and ViT-Large (as defined
in Table I) for up to 188,416 sequence length. We first show the performance
of weak scaling the batch size of a model with 188,416 sequence length up to
2048 GPUs on Frontier.

###  VI-A Performance at Scale

Figure 9: The top plot shows a convergence comparison between a single GPU
sequence length and distributing the sequence to 16 GPUs using DeepSpeed-
Ulysses. Real ERA5 data was used for this test, along with scaling the batch
size to 4 with DeepSpeed Zero Optimizer in both cases. In the bottom plot we
weak scale the batch size of a 188,416 sequence length base model up to 2,048
GPUs. The real plot is with real ERA5 data, the synthetic plot is with random
data, and the ”no DistAttent” is with synthetic data but turning off
distributed attention. A minimum of 16 GPUs was necessary to fit the base
model with a 188,416 sequence length with a local batch size of 4 that was
used for the convergence runs.

First, we think it is essential when a distributed method, beyond data
parallel, is integrated into a model to showcase numerical correctness when
converging to a solution on real data. To that extent, the top plot of Figure
9 compares the same model and sequence length for 4 DeepSpeed Zero ranks
against 16 sequence parallel ranks with DeepSpeed-Ulysses along 4 DeepSpeed
Zero ranks in each group. We can see that for the first steps the loss is
identical, as expected, and we further test that in the first forward pass,
all matrices are the same. We then see some small differences as training
progresses, which primarily come from the non-deterministic parts of the
model, such as normalization layers and dropout. Additionally, we see that the
model and the data are seeded correctly as we are performing the comparison
experiments next.

Figure 10: The plot shows the validation accuracy on half-year of the ERA5
dataset of the Z500, T850, T2m, and the U10 variables. The effect in the
accuracy is shown between a base Agg-Ch-ViT model with 2,048 sequence length
(aggregate 92 variables to 1), a large Agg-Ch-ViT model with 2,048 sequence
length (aggregate 92 variables to 1), and a base Multi-Ch-ViT model with
188,416 sequence length (all 92 variables of the ERA5 dataset were used).

The bottom plot of Figure 9 shows a weak scaling of the batch size for the
188,416 sequence length model on Frontier. In the training, we found a better
convergence with a local batch size of 4, and for this, we need a minimum of
16 GPUs to fit the 188K sequence length for the ViT-Base model. With this as
the baseline, using DeepSpeed Zero, we scale up to 2,048 GPUs. The blue curve
runs are with real climate data, showing a scaling efficiency of 94%. The red
curve shows a 96% scaling efficiency with synthetic data, confirming that the
application isn’t IO-bound as expected. The green curve also represents a run
with synthetic data, but this time, we don’t wrap the attention layers with
DeepSpeed-Ulysses. This would be a way to estimate the communication overhead
from distributing the sequence length, which, as we mentioned before, is
mostly two all-to-all calls in the forward pass and two all-to-all calls + all
reduce in the backward pass per layer. Since we get the same scaling
efficiency as in the red curve, that probably means that the communication
part of DeepSpeed-Ulysses scales very well with the batch size, but the
constant 20% reduction in the number of FLOPS that we see probably relates to
the overlap between compute and communication within the sequence parallel
groups.

Second, we show the training results using the ERA5 dataset, with the same
setup as in Section IV. trained with a resolution of 1.40625°. All the hyper-
parameters were kept constant between each experiment, using a predicted range
of 28 hours, with a lead-time of 6 hours. Three experiments are shown: one
with Agg-Ch-ViT and the ViT-Base model, one with Agg-Ch-ViT and the ViT-Large
model, and one with Multi-Ch-ViT and the ViT-Base model. The sequence length
for the first two experiments is the same, 2,048 but the model size is
increased from 100M to 800M, while in the latter experiment, the sequence
length is 188,416.

Figure 10 shows results from all three experiments. With larger sequence
lengths, Z500 and U10 don’t show improvements, while T850 and T2m seem to show
better accuracy. The advantage of the Multi-Ch-ViT is that we explicitly pass
the model to all the channels per time-step, compared to the Agg-Ch-ViT, where
all channels are aggregated into one. As discussed in Section IV for the Z500,
it is expected to not see much improvement since Z500 features won’t have fast
spatiotemporal variations. In comparison, the temperature on the surface and
lower troposphere exhibits substantially more variations. Also, we observed
the same behavior as in Section IV, where the convergence starts at lower
performance, and with more epochs, the model can utilize better the longer
sequence length. In principle, if the model could understand the full climate
phenomena, this should reflect on the better performance of the Z500 as well,
but the number of epochs was not as large to see this effect. Finally, we see
that for those experiments, increasing the model size didn’t result in a
better solution. Probably 8 times larger model size isn’t enough to see any
major improvements, but because we only used 1 year of data for those
experiments, we can’t justify a larger model.

Finally, we show that Ulysses also scales well with the batch size, using
DeepSpeed Zero as the data-parallel mode, while IO can successfully hide
behind compute even for 188K sequence length. Even though we only trained for
10 epochs, we still observe almost 20% improvements in the accuracy for the
temperature channels, while we recommend as the sequence length becomes
larger, a longer number of epochs is necessary. Also for ViTs, a practitioner
should consider scaling the sequence length first, rather than the model size,
though a balance between the two would be the optimal approach.

##  VII Conclusion and Future Work

We developed a ViT capable of processing ultra-long sequence lengths of up to
1 million tokens, employing sequence parallelism with DeepSpeed-Ulysses and
LSS strategies. This approach, tested across models with up to 10 billion
parameters, addressed performance bottlenecks and utilized distributed
sequence strategies alongside Pipeline and Tensor Parallelism to scale beyond
single GPU memory capacities. Notably, applying these methodologies to climate
modeling increased temperature prediction accuracy by up to 20%, marking a
significant advancement in handling extensive sequence lengths with full
attention.

Future research will focus on refining parallelism strategies for enhanced
efficiency and exploring the application of ultra-long sequence processing
across various domains. Investigating alternative attention mechanisms to
optimize computational demands while preserving model performance will also be
a key area of interest, potentially broadening the utility and scalability of
ViTs in data-intensive fields.

## Acknowledgments

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 (http://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.

## References

  * [1] Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah.  Transformers in vision: A survey.  ACM Computing Surveys, 54(10s):1–41, January 2022.
  * [2] Richard J Chen, Chengkuan Chen, Yicong Li, Tiffany Y Chen, Andrew D Trister, Rahul G Krishnan, and Faisal Mahmood.  Scaling vision transformers to gigapixel images via hierarchical self-supervised learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16144–16155, 2022.
  * [3] Zhitong Xiong, Yi Wang, Fahong Zhang, and Xiao Xiang Zhu.  One for all: Toward unified foundation models for earth vision, 2024.
  * [4] Keumgang Cha, Junghoon Seo, and Taekyung Lee.  A billion-scale foundation model for remote sensing images, 2023.
  * [5] Chunyuan Li, Zhe Gan, Zhengyuan Yang, Jianwei Yang, Linjie Li, Lijuan Wang, and Jianfeng Gao.  Multimodal foundation models: From specialists to general-purpose assistants, 2023.
  * [6] William Fedus, Barret Zoph, and Noam Shazeer.  Switch transformers: scaling to trillion parameter models with simple and efficient sparsity.  J. Mach. Learn. Res., 23(1), jan 2022.
  * [7] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetić, Dustin Tran, Thomas Kipf, Mario Lučić, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby.  Scaling vision transformers to 22 billion parameters, 2023.
  * [8] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover.  Climax: A foundation model for weather and climate.  arXiv preprint arXiv:2301.10343, 2023.
  * [9] Yoo Jung Kim, Hyungjoon Jang, Kyoungbun Lee, Seongkeun Park, Sung-Gyu Min, Choyeon Hong, Jeong Hwan Park, Kanggeun Lee, Jisoo Kim, Wonjae Hong, Hyun Jung, Yanling Liu, Haran Rajkumar, Mahendra Khened, Ganapathy Krishnamurthi, Sen Yang, Xiyue Wang, Chang Hee Han, Jin Tae Kwak, Jianqiang Ma, Zhe Tang, Bahram Marami, Jack Zeineh, Zixu Zhao, Pheng-Ann Heng, Rüdiger Schmitz, Frederic Madesta, Thomas Rösch, Rene Werner, Jie Tian, Elodie Puybareau, Matteo Bovio, Xiufeng Zhang, Yifeng Zhu, Se Young Chun, Won-Ki Jeong, Peom Park, and Jinwook Choi.  Paip 2019: Liver cancer segmentation challenge.  Medical Image Analysis, 67:101854, 2021.
  * [10] Alexandra bodzás, Pavel Kodytek, and Jan Žídek.  A high-resolution large-scale dataset of pathological and normal white blood cells.  Scientific Data, 10, 07 2023.
  * [11] Kexin Ding, Mu Zhou, He Wang, Olivier Gevaert, Dimitris Metaxas, and Shaoting Zhang.  A large-scale synthetic pathological dataset for deep learning-enabled segmentation of breast cancer.  Scientific Data, 10, 04 2023.
  * [12] Adrian Simmons, Cornel Soci, Julien Nicolas, Bill Bell, P. Berrisford, Rossana Dragani, Johannes Flemming, Leopold Haimberger, Sean Healy, Hans Hersbach, Andras Horányi, Antje Inness, J. Munoz-Sabater, Raluca Radu, and Dinand Schepers.  Global stratospheric temperature bias and other stratospheric aspects of era5 and era5.1, 01/2020 2020.
  * [13] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need, 2023.
  * [14] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity, 2020.
  * [15] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He.  Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, 2023.
  * [16] Xiao Wang, Isaac Lyngaas, Aristeidis Tsaris, Peng Chen, Sajal Dash, Mayanka Chandra Shekar, Tao Luo, Hong-Jun Yoon, Mohamed Wahib, and John Gouley.  Ultra-long sequence distributed transformer, 2023.
  * [17] Tri Dao.  Flashattention-2: Faster attention with better parallelism and work partitioning, 2023.
  * [18] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen.  Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019.
  * [19] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.  Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020.
  * [20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
  * [21] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He.  Zero-offload: Democratizing billion-scale model training, 2021.
  * [22] Xuanlei Zhao, Shenggan Cheng, Zangwei Zheng, Zheming Yang, Ziming Liu, and Yang You.  Dsp: Dynamic sequence parallelism for multi-dimensional transformers, 2024.
  * [23] Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You.  Sequence parallelism: Long sequence training from system perspective, 2022.
  * [24] Hao Liu, Matei Zaharia, and Pieter Abbeel.  Ring attention with blockwise transformers for near-infinite context, 2023.
  * [25] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher R’e.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  ArXiv, abs/2205.14135, 2022.
  * [26] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al.  Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators.  arXiv preprint arXiv:2202.11214, 2022.
  * [27] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian.  Accurate medium-range global weather forecasting with 3d neural networks.  Nature, 619(7970):533–538, 2023.
  * [28] Lei Chen, Xiaohui Zhong, Feng Zhang, Yuan Cheng, Yinghui Xu, Yuan Qi, and Hao Li.  Fuxi: a cascade machine learning forecasting system for 15-day global weather forecast.  npj Climate and Atmospheric Science, 6(1):190, 2023.
  * [29] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
  * [30] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [31] Junqi Yin, Sajal Dash, Feiyi Wang, and Mallikarjun Shankar.  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, 2023. Association for Computing Machinery.
  * [32] The High-Resolution Rapid Refresh.  https://rapidrefresh.noaa.gov/hrrr/.
  * [33] xesmf: Universal regridder for geospatial data, 2020.
  * [34] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He.  Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning, 2021.
  * [35] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li.  Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023.
  * [36] The Frontier supercomputer.  https://www.olcf.ornl.gov/frontier/.
  * [37] Junqi Yin, Avishek Bose, Guojing Cong, Isaac Lyngaas, and Quentin Anthony.  Comparative study of large language model architectures on frontier, 2024.
  * [38] AMD. Flash Attention on ROCm.  https://github.com/ROCmSoftwarePlatform/flash-attentio, 2023.