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.

利用分布式序列并行化技术以及混合序列、流水线和张量并行化策略，使得 Vision Transformers 在处理大序列长度和训练 10B 参数模型时能够取得更好的性能表现，从而提高气候模型的精度。