用DeepSpeed提升视觉Transformer的训练效率与可扩展性
Scaling Vision Transformers: Evaluating DeepSpeed for Image-Centric Workloads
- 将专为语言模型设计的DeepSpeed用于视觉Transformer分布式训练
- 发现批量大小和梯度累积是影响训练性能的关键因素
- 适合需要大规模图像模型训练的研究者参考
视觉变换器（ViTs）通过自注意力机制捕捉数据中的全局关系，在图像处理任务中展现出巨大潜力。然而，其大规模模型在参数量增加时面临显著的计算与内存挑战。本研究利用常用于语言模型的高效分布式训练框架DeepSpeed，提升ViTs的可扩展性与性能。在CIFAR-10、CIFAR-100等数据集上，评估了多种GPU配置下的节点内与跨节点训练效率，分析了分布式数据并行对训练速度、通信开销及整体可扩展性（强弱缩放）的影响。通过系统调整批大小与梯度累积等软件参数，识别出影响分布式训练性能的关键因素。实验结果为将DeepSpeed应用于图像任务提供了基础支撑。未来工作将进一步探索DeepSpeed的局限性，并优化ViT的分布式训练流程。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。