Scaling Vision Transformers: Evaluating DeepSpeed for Image-Centric Workloads

1. Problem Statement

Vision Transformers (ViTs) have emerged as a powerful architecture for image processing, utilizing self-attention mechanisms to capture global relationships within data, unlike traditional Convolutional Neural Networks (CNNs). However, ViTs face significant scalability challenges:

• Computational and Memory Demands: As model parameters scale from millions to billions and image resolutions increase, the memory and compute requirements become prohibitive for single-GPU training.
• Lack of Optimized Frameworks: While DeepSpeed is a highly efficient distributed training framework widely adopted for Large Language Models (LLMs) (e.g., BERT, GPT), its application to image-centric tasks and ViTs remains under-explored.
• Scalability Uncertainties: It is unclear how DeepSpeed's data parallelism performs across different hardware configurations (intra-node vs. inter-node) and how software parameters (batch size, gradient accumulation) impact training efficiency and communication overhead in ViT workloads.

2. Methodology

The study systematically evaluates DeepSpeed's performance for ViTs across three distinct remote clusters with varying hardware configurations:

• Clusters Used:
  • Nebula: Used for intra-node testing with homogeneous GPUs (RTX 3070).
  • Tesla: Used for inter-node testing with heterogeneous GPUs (mix of RTX 3070, GTX 1070, and Tesla P4).
  • Vector: Used for large-scale intra- and inter-node testing with 54 nodes equipped with Tesla T4 GPUs.
• Architecture & Datasets:
  • Model: ViT-Base-16 (ViT_b_16).
  • Datasets: CIFAR-10, CIFAR-100, and a planned (but uncompleted due to time) ImageNet-100.
• Technical Stack:
  • Framework: DeepSpeed integrated with PyTorch's torch.distributed and NCCL for GPU communication.
  • Orchestration: MPI (OpenMPI) for inter-node process management and mpirun for launching.
  • Parallelism Strategy: Focused on Data Parallelism (DP). The study explicitly excluded ZeRO (Zero Redundancy Optimizer) and Model/Pipeline Parallelism due to resource constraints, though ZeRO is noted for future work.
• Evaluation Metrics:
  • Strong Scaling: Fixed total workload (dataset size) while increasing GPU count.
  • Weak Scaling: Workload per GPU kept constant while increasing total dataset size and GPU count.
  • Key Metrics: Training time per epoch, communication overhead (synchronization costs), and model accuracy.
  • Variables: Batch size (16, 64, 128, 256) and gradient accumulation steps.

3. Key Contributions

1. Adaptation of DeepSpeed for ViTs: Successfully implemented and validated a DeepSpeed-based distributed training pipeline specifically for Vision Transformers, bridging the gap between NLP-focused frameworks and computer vision.
2. Hardware Homogeneity Analysis: Demonstrated that GPU heterogeneity (mixing different GPU models in a cluster) severely degrades scaling performance due to synchronization bottlenecks, whereas homogeneous clusters yield efficient scaling.
3. Parameter Optimization: Identified the optimal trade-off between batch size and synchronization overhead. The study found that while larger batches reduce synchronization frequency, they introduce CPU-to-GPU data loading bottlenecks.
4. Scalability Benchmarks: Provided empirical data on both intra-node and inter-node scaling trends for ViTs, establishing baselines for future research in distributed image training.

4. Key Results

• Impact of Hardware Heterogeneity:
  • On the Tesla cluster (heterogeneous GPUs), scaling failed to follow ideal trends. Adding weaker GPUs (GTX 1070, Tesla P4) to a cluster of stronger GPUs (RTX 3070) created computational bottlenecks, increasing training time due to synchronization waits.
• Batch Size and Synchronization:
  • On the Nebula cluster, small batch sizes (e.g., 16) resulted in disproportionately high synchronization costs.
  • Increasing batch size to 64 or 128 significantly reduced synchronization overhead and improved speedup.
  • However, scaling batch size beyond 128 (e.g., to 256) yielded diminishing returns due to the overhead of loading large batches from CPU to GPU memory.
  • Accuracy vs. Batch Size: Accuracy improved with batch size up to a point, then declined (likely due to overfitting), suggesting an optimal balance exists.
• Scaling Performance:
  • Vector Cluster (Homogeneous T4s): Achieved successful strong and weak scaling.
    • Strong Scaling: Training time decreased consistently as GPUs increased (e.g., 1 to 2 GPUs reduced time by nearly 50%).
    • Weak Scaling: Training time remained relatively constant as the dataset and GPU count scaled proportionally.
  • Intra-Node vs. Inter-Node: On the Vector cluster, there was no significant difference in performance between intra-node (multi-GPU on one node) and inter-node (single-GPU across multiple nodes) training when using a batch size of 64.
• Multi-Node Scaling: The system successfully scaled to 32 nodes (32 GPUs) on the Vector cluster for CIFAR-100, demonstrating robust inter-node communication capabilities.

5. Significance and Future Work

• Significance: This study provides a foundational basis for applying DeepSpeed to image-related tasks. It confirms that while DeepSpeed is viable for ViTs, success relies heavily on hardware homogeneity and careful tuning of batch sizes to balance communication overhead and memory constraints. It highlights that the "memory wall" and data loading bottlenecks are critical factors in image training, distinct from token-based LLM training.
• Future Directions:
  • ZeRO Optimization: Investigating ZeRO stages (1, 2, 3) to further reduce memory usage and enable larger models.
  • Advanced Parallelism: Adapting DeepSpeed-Ulysses (Sequence Parallelism) for ViTs. Instead of token sequences, the authors propose partitioning along the "image patch dimension" to handle high-resolution images and long sequences.
  • Broader Applications: Extending the framework to Vision Language Models (VLMs) and scientific imaging (medical MRI, robotics, genomics) where high-resolution and multi-channel data are common.
  • Benchmarking: Comparing DeepSpeed against other frameworks like Megatron-LM and HuggingFace Accelerate specifically for ViT workloads.