Scaling Vision Transformers: Evaluating DeepSpeed for Image-Centric Workloads

1. Il Problema

I Vision Transformers (ViT) hanno dimostrato un potenziale straordinario nelle attività di elaborazione delle immagini grazie al meccanismo di self-attention, che permette di catturare relazioni globali all'interno dei dati, a differenza delle tradizionali CNN che si basano su convoluzioni gerarchiche. Tuttavia, la scalabilità dei ViT è fortemente limitata da elevate richieste computazionali e di memoria, specialmente quando si tratta di modelli su larga scala con miliardi di parametri.

Sebbene framework distribuiti come DeepSpeed siano ampiamente utilizzati e ottimizzati per i Large Language Models (LLM) per gestire la parallelizzazione dei dati e ridurre l'uso della memoria, la loro applicazione ai task basati su immagini (ViT) è stata poco esplorata. Esiste quindi la necessità di valutare se le tecniche di parallelizzazione dei dati di DeepSpeed possano essere efficacemente adattate per migliorare la scalabilità e le prestazioni dei ViT, affrontando le sfide specifiche legate alla comunicazione tra nodi e all'eterogeneità dell'hardware.

2. Metodologia

Gli autori hanno condotto una valutazione sistematica dell'efficienza dell'addestramento distribuito utilizzando DeepSpeed su diverse configurazioni hardware e dataset.

• Framework e Configurazione: L'implementazione si basa su DeepSpeed integrato con PyTorch (torch.distributed), NCCL per la comunicazione GPU-GPU e MPI (OpenMPI) per la gestione dei processi tra i nodi. Lo studio si concentra sul Data Parallelism (DP), replicando il modello su ogni dispositivo e distribuendo i mini-batch, senza inizialmente utilizzare le funzionalità ZeRO (Zero Redundancy Optimizer) a causa dei vincoli di risorse.
• Ambienti di Sperimentazione: Sono stati utilizzati tre cluster remoti con caratteristiche hardware diverse:
  • Nebula: Cluster locale con GPU di fascia media (partizioni single e dual-card).
  • Tesla: Cluster con GPU eterogenee (RTX 3070 e GTX 1070/Tesla P4), utilizzato per testare l'impatto dell'eterogeneità.
  • Vector: Cluster ad alte prestazioni con 54 nodi e GPU Tesla T4, utilizzato per test di scalabilità su larga scala (fino a 32 nodi).
• Dataset: CIFAR-10, CIFAR-100 e ImageNet-100 (sebbene l'addestramento su ImageNet non sia stato completato a causa dei limiti temporali).
• Metriche di Valutazione:
  • Strong Scaling: Carico di lavoro fisso (dataset completo) con aumento del numero di GPU.
  • Weak Scaling: Carico di lavoro proporzionale al numero di GPU (ogni GPU elabora la stessa quantità di dati).
  • Parametri Variabili: Dimensione del batch (batch size) e accumulo dei gradienti.
  • Osservazioni: Tempo di addestramento, overhead di comunicazione, accuratezza del modello e velocità di convergenza.

3. Contributi Chiave

1. Adattamento di DeepSpeed per ViT: Prima valutazione sistematica dell'uso di DeepSpeed per l'addestramento distribuito di Vision Transformers, adattando pipeline sviluppate per i linguaggi naturali al dominio visivo.
2. Analisi della Scalabilità Intra- e Inter-nodo: Confronto dettagliato tra l'addestramento su GPU nello stesso nodo (condivisione memoria) e su nodi diversi, evidenziando le differenze nelle prestazioni e nell'overhead di comunicazione.
3. Ottimizzazione dei Parametri Software: Identificazione dell'impatto critico della dimensione del batch e dell'accumulo dei gradienti sulle prestazioni. Gli autori hanno dimostrato come la scelta del batch size possa bilanciare i costi di sincronizzazione e l'utilizzo della memoria.
4. Valutazione dell'Eterogeneità Hardware: Dimostrazione pratica di come l'eterogeneità delle GPU (mix di modelli potenti e deboli) possa degradare le prestazioni complessive a causa dei colli di bottiglia nella sincronizzazione.

4. Risultati Principali

• Impatto dell'Eterogeneità Hardware (Cluster Tesla): L'addestramento inter-nodo su cluster con GPU eterogenee (Tesla) ha mostrato risultati non ottimali. La presenza di GPU più lente ha creato colli di bottiglia durante la fase di sincronizzazione (AllReduce), aumentando i tempi di addestramento e riducendo l'efficienza di scalabilità. Questo sottolinea l'importanza di hardware omogeneo per un scaling efficiente.
• Ottimizzazione del Batch Size (Cluster Nebula):
  • Batch size piccoli (es. 16) hanno generato costi di sincronizzazione sproporzionatamente alti, portando a una scarsa scalabilità.
  • L'aumento del batch size ha ridotto significativamente l'overhead di comunicazione.
  • È stato identificato un punto di equilibrio ottimale tra batch size 64 e 128, che massimizza l'efficienza senza saturare la memoria o introdurre nuovi colli di bottiglia nel caricamento dei dati.
  • Batch size eccessivamente grandi (es. 256) hanno mostrato un plateau nelle prestazioni, probabilmente a causa dell'overhead nel trasferimento dati CPU-GPU.
• Scalabilità su Cluster Vector:
  • Strong Scaling: Su CIFAR-10, l'aggiunta di risorse ha ridotto il tempo di addestramento quasi della metà passando da 1 a 2 GPU, con una tendenza decrescente costante.
  • Weak Scaling: I tempi di addestramento sono rimasti costanti all'aumentare delle GPU, confermando la corretta distribuzione del carico.
  • Confronto Intra vs Inter-nodo: Non sono state rilevate differenze significative nelle prestazioni di strong scaling tra l'uso di più GPU su un singolo nodo e l'uso di una GPU per nodo su più nodi (fino a 32 nodi), suggerendo che l'overhead di rete non è stato un fattore limitante in questa configurazione specifica.
• Accuratezza: L'accuratezza tende a migliorare con l'aumento del batch size fino a un certo punto, per poi declinare (possibile overfitting). In generale, l'uso di più GPU ha mostrato prestazioni leggermente migliori nelle epoche successive rispetto all'addestramento su singola GPU.

5. Significato e Lavori Futuri

Questo studio fornisce una base fondamentale per l'applicazione di DeepSpeed ai task di visione artificiale, dimostrando che il data parallelism è efficace per i ViT se gestito correttamente. I risultati evidenziano che l'ottimizzazione dei parametri software (come il batch size) è cruciale quanto la potenza hardware per ottenere una scalabilità efficiente.

Lavori Futuri Proposti:

• Valutazione di ZeRO: Estendere l'analisi includendo le diverse fasi di ZeRO (ZeRO-1, 2, 3) per misurare i risparmi di memoria e l'overhead aggiuntivo, fondamentale per modelli con attivazioni intermedie molto grandi (tipici dei ViT ad alta risoluzione).
• Confronto con Altri Framework: Benchmarking contro soluzioni come Megatron-LM o HuggingFace Accelerate.
• Modelli Vision-Language e Sequenze Lunghe: Adattare tecniche di Sequence Parallelism (simili a DeepSpeed-Ulysses per i LLM) ai ViT. L'idea è partizionare la dimensione delle "patch dell'immagine" invece delle sequenze di token, permettendo l'addestramento su immagini ad altissima risoluzione o dataset scientifici complessi (es. imaging medico, robotica, genomica).

In sintesi, il paper conferma che DeepSpeed è uno strumento promettente per scalare i ViT, ma richiede un'attenta configurazione hardware e software per mitigare i costi di comunicazione e gestire le specificità dei dati visivi.