Scaling Vision Transformers: Evaluating DeepSpeed for Image-Centric Workloads

Probleemstelling

Vision Transformers (ViTs) hebben zich bewezen als krachtige modellen voor beeldverwerking door gebruik te maken van self-attention-mechanismen om globale relaties in data te vangen. In tegenstelling tot traditionele Convolutional Neural Networks (CNN's), die features hiërarchisch extraheren, analyseren ViTs alle patches van een afbeelding gelijktijdig. Dit biedt echter een significant nadeel: de rekenkosten en geheugeneisen groeien exponentieel naarmate de modelgrootte en datasetgrootte toenemen (van miljoenen naar miljarden parameters).

Hoewel frameworks zoals DeepSpeed al succesvol worden ingezet voor de schaalbaarheid van taalmodellen (LLMs) via gedistribueerde training, is er weinig onderzoek gedaan naar de toepasbaarheid en efficiëntie van DeepSpeed voor beeldgerelateerde taken. De uitdaging ligt in het bepalen of de data-parallelisme-strategieën die voor tekst werken, ook effectief schalen voor ViTs, en welke softwareparameters (zoals batchgrootte) hierbij cruciaal zijn.

Methodologie

De auteurs hebben een experimenteel kader opgezet om de schaalbaarheid van ViTs met DeepSpeed te evalueren, met name gericht op data parallelism (zonder gebruik van ZeRO-optimisaties vanwege resource-beperkingen).

• Architectuur en Framework: Het onderzoek gebruikt de ViT-b16-architectuur, geïmplementeerd met PyTorch, DeepSpeed, NCCL (voor GPU-communicatie) en MPI (voor inter-node communicatie).
• Hardware-configuraties: Experimenten zijn uitgevoerd op drie verschillende clusters met variërende hardware:
  • Nebula: Voor intra-node training (deel van de ECE Linux faciliteit).
  • Tesla: Voor inter-node training, maar met een heterogene GPU-configuratie (RTX 3070, GTX 1070, Tesla P4).
  • Vector Institute: Voor uitgebreide intra- en inter-node tests met homogene Tesla T4 GPUs.
• Datasets: CIFAR-10, CIFAR-100 en ImageNet-100 (hoewel ImageNet vanwege tijdsbeperkingen niet volledig succesvol getraind kon worden).
• Schaalbaarheidstests:
  • Strong Scaling: Het werkload (datasetgrootte) blijft constant terwijl het aantal GPUs toeneemt.
  • Weak Scaling: De datasetgrootte per GPU blijft constant terwijl het totale aantal GPUs en de totale dataset toenemen.
• Variabele Parameters: De auteurs varieerden systematisch de batchgrootte en het aantal accumulatiestappen om de impact op trainingssnelheid, communicatie- overhead en nauwkeurigheid te meten.

Belangrijkste Bijdragen

1. Adaptatie van DeepSpeed voor ViT: Het eerste systematische onderzoek dat DeepSpeed toepast op Vision Transformers om intra- en inter-node schaalbaarheidstrends te observeren.
2. Evaluatie van Data Parallelism: Een gedetailleerde meting van trainingssnelheid en communicatie- overhead bij het verhogen van het aantal GPU's, zowel binnen één node als over meerdere nodes.
3. Analyse van Softwareparameters: Het identificeren van de kritieke invloed van batchgrootte en gradient accumulation op de prestaties van gedistribueerde training voor beeldmodellen.

Resultaten

De experimenten leverden de volgende inzichten op:

• Hardware Homogeniteit is Cruciaal: Training op de Tesla-cluster (met heterogene GPUs) presteerde slecht. De langzamere GPUs fungeerden als bottleneck, waardoor snellere GPU's moesten wachten tijdens de synchronisatie (AllReduce), wat leidde tot een toename in totale trainingstijd bij het toevoegen van meer GPUs.
• Batchgrootte en Communicatie- overhead:
  • Op de Nebula-cluster (krachtigere GPUs) bleek dat kleine batchgroottes (bijv. 16) leiden tot disproportioneel hoge synchronisatiekosten.
  • Grotere batchgroottes (64 of 128) verminderden de synchronisatiekosten aanzienlijk.
  • Er is een verzadigingspunt: bij het verhogen van de batchgrootte van 128 naar 256 nam de prestatieverbetering af, waarschijnlijk door de overhead van het laden van grote batches van CPU naar GPU-geheugen.
  • Optimale Batchgrootte: Een batchgrootte van 64 of 128 bleek de beste balans te bieden tussen synchronisatiekosten en geheugengebruik.
• Schaalbaarheidstrends:
  • Intra-node (Vector): Strong scaling toonde een consistente afname van de trainingstijd, met de grootste winst bij het gaan van 1 naar 2 GPU's (bijna halvering van de tijd). Weak scaling bleef stabiel.
  • Inter-node (Vector): Succesvolle strong scaling tot 32 nodes. Er was geen significante prestatieverschil tussen intra-node en inter-node training bij gebruik van homogene hardware.
• Nauwkeurigheid: De nauwkeurigheid verbeterde aanvankelijk met grotere batchgroottes, maar daalde bij te grote batches (mogelijk door overfitting). Inter-node training met meerdere GPU's toonde in latere epochs soms betere nauwkeurigheid dan single-node training, hoewel dit niet direct causaal is.

Betekenis en Toekomstperspectief

De studie biedt een fundamentele basis voor het toepassen van DeepSpeed op beeldgerelateerde taken. Het benadrukt dat:

1. DeepSpeed effectief kan schalen voor ViTs, mits de hardware homogeen is en de softwareparameters (zoals batchgrootte) correct worden geoptimaliseerd.
2. Communicatie- overhead een belangrijke factor blijft bij het schalen naar veel GPU's, wat aandacht vereist bij het ontwerp van pipelines voor grote modellen.

Toekomstig werk omvat het evalueren van ZeRO-stadia (voor geheugenefficiëntie), het testen van verschillende optimalizers (SGD, LAMB), en het uitbreiden naar Vision Language Models en wetenschappelijke beeldanalyse. De auteurs plannen ook om sequence parallelism (geïnspireerd op DeepSpeed-Ulysses) toe te passen op de "image patches-dimensie" om zeer lange sequenties van afbeeldingen te verwerken, wat relevant is voor medische beeldvorming, robotica en genomics.