cs.LGFeb 24, 2026

Scaling Vision Transformers: Evaluating DeepSpeed for Image-Centric Workloads

Authors: Huy Trinh, Rebecca Ma, Zeqi Yu, Tahsin Reza

Organizations: Electrical & Computer Eng University of Waterloo

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Vision Transformers (ViTs) have demonstrated remarkable potential in image processing tasks by utilizing self-attention mechanisms to capture global relationships within data. However, their scalability is hindered by significant computational and memory demands, especially for large-scale models with many parameters. This study aims to leverage DeepSpeed, a highly efficient distributed training framework that is commonly used for language models, to enhance the scalability and performance of ViTs. We evaluate intra- and inter-node training efficiency across multiple GPU configurations on various datasets like CIFAR-10 and CIFAR-100, exploring the impact of distributed data parallelism on training speed, communication overhead, and overall scalability (strong and weak scaling). By systematically varying software parameters, such as batch size and gradient accumulation, we identify key factors influencing performance of distributed training. The experiments in this study provide a foundational basis for applying DeepSpeed to image-related tasks. Future work will extend these investigations to deepen our understanding of DeepSpeed's limitations and explore strategies for optimizing distributed training pipelines for Vision Transformers.

Figures & tables

Fig. 1: ViT Model Overview [ 3 ]
Fig. 2: Data Parallelism Gradient Computation [ 6 ]
Fig. 3: Remote Cluster Setup
DatasetNo. of ClassesNo. of ImagesResolution
CIFAR-10 [ 12 ]1060,00032x32
CIFAR-100 [ 12 ]10060,00032x32
ImageNet-100 [ 13 ] *100100,000224x224
TABLE I: Datasets for evaluation
Fig. 4: Tesla Strong Scaling
Fig. 5: Tesla Weak Scaling
Fig. 6: Nebula Strong Scaling vs Batch Sizes
Fig. 7: Nebula Accuracy vs Batch Sizes
Fig. 8: CIFAR-10 Strong Scaling (Batch size 64)
Fig. 9: CIFAR-10 Weak Scaling (Batch size 64)
Fig. 10: CIFAR-10 and CIFAR-100 Accuracies
Fig. 11: Loss and Accuracy of Strong Scaling (Batch size 64)
Fig. 12: CIFAR-10 Strong Scaling Speedup (Batch size 16)
Fig. 13: CIFAR-10 Strong Scaling Speedup (Batch size 64)
Fig. 14: Multi-node Single GPU Strong Scaling Result On CIFAR-100 (Batch size 64)
Fig. 15: Multi-node Single-GPU vs Single-Node Multi-GPU Strong Scaling Result On CIFAR-100 (Batch size 64)
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Fig. 16: CIFAR-100 Strong Scaling (Batch size 64)
Fig. 17: CIFAR-100 Weak Scaling (Batch size 64)

Explore similar work

CardsList

  1. Quick ViTs: Speeding up Vision Transformers through Equivariance

    May 21, 2025David Nordström, Johan Edstedt, Fredrik Kahl +1Vision TransformerEquivariant Neural Networks

    Chalmers University of Technology · Linköping University · University of Amsterdam

  2. ScalePredictor: Instance-aware Scale Learning for Accurate Quantization of Vision Transformers

    Jun 20, 2026Changjun Li, Runqing Jiang, Lian Xu +3Vision TransformerActivation Quantization

    School of Electronics and Communication Engineering, Shenzhen Campus, Sun Yat-sen University, Shenzhen, China · Computer Science and Software Engineering, The University of Western Australia, Perth, Australia · University of Oxford, Oxford, United Kingdom

  3. Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs

    Apr 17, 2026Jun-Liang Lin, Kamesh Madduri, Mahmut Taylan KandemirDistributed TrainingGraph Transformer

    The Pennsylvania State University University Park, PA, USA