KT’s viaggio per ridurre il tempo di addestramento di un modello di visione transformers utilizzando Amazon SageMaker

KT Riduzione del Tempo di Addestramento di un Modello di Visione Transformers con Amazon SageMaker

KT Corporation è uno dei principali fornitori di telecomunicazioni in Corea del Sud, offre una vasta gamma di servizi tra cui telefonia fissa, comunicazione mobile, internet e servizi di intelligenza artificiale. KT’s AI Food Tag è una soluzione di gestione dietetica basata sull’intelligenza artificiale che identifica il tipo e il contenuto nutritivo degli alimenti nelle foto utilizzando un modello di visione artificiale. Questo modello di visione sviluppato da KT si basa su un modello pre-addestrato utilizzando una vasta quantità di dati di immagini non etichettate per analizzare il contenuto nutritivo e le informazioni sulle calorie di vari alimenti. L’AI Food Tag può aiutare i pazienti affetti da malattie croniche come il diabete a gestire la loro dieta. KT ha utilizzato AWS e Amazon SageMaker per addestrare questo modello AI Food Tag 29 volte più velocemente di prima e ottimizzarlo per la distribuzione in produzione con una tecnica di distillazione del modello. In questo post, descriviamo il percorso di sviluppo del modello di KT e il successo ottenuto utilizzando SageMaker.

Presentazione del progetto KT e definizione del problema

Il modello AI Food Tag pre-addestrato da KT si basa sull’architettura dei vision transformers (ViT) e ha più parametri di modello rispetto al loro modello di visione precedente per migliorare l’accuratezza. Per ridurre le dimensioni del modello per la produzione, KT sta utilizzando una tecnica di distillazione delle conoscenze (KD) per ridurre il numero di parametri del modello senza un impatto significativo sull’accuratezza. Con la distillazione delle conoscenze, il modello pre-addestrato viene chiamato modello insegnante e un modello di output leggero viene addestrato come modello studente, come illustrato nella figura seguente. Il modello studente leggero ha meno parametri di modello dell’insegnante, il che riduce i requisiti di memoria e consente la distribuzione su istanze più piccole e meno costose. Lo studente mantiene un’accuratezza accettabile nonostante sia più piccolo imparando dagli output del modello insegnante.

Il processo generale di addestramento per la distillazione delle conoscenze

Il modello insegnante rimane invariato durante la KD, ma il modello studente viene addestrato utilizzando le log-verosimiglianze di output del modello insegnante come etichette per calcolare la perdita. Con questo paradigma di KD, sia il modello insegnante che quello studente devono essere su una singola memoria GPU per l’addestramento. KT ha inizialmente utilizzato due GPU (A100 80 GB) nel loro ambiente interno in sede per addestrare il modello studente, ma il processo ha richiesto circa 40 giorni per coprire 300 epoche. Per accelerare l’addestramento e generare un modello studente in meno tempo, KT ha collaborato con AWS. Insieme, i team hanno notevolmente ridotto il tempo di addestramento del modello. Questo post descrive come il team ha utilizzato Amazon SageMaker Training, la SageMaker Data Parallelism Library, Amazon SageMaker Debugger e Amazon SageMaker Profiler per sviluppare con successo un modello AI Food Tag leggero.

Creazione di un ambiente di addestramento distribuito con SageMaker

SageMaker Training è un ambiente di addestramento di machine learning (ML) gestito su AWS che fornisce una serie di funzionalità e strumenti per semplificare l’esperienza di addestramento ed è utile nel calcolo distribuito, come illustrato nel seguente diagramma.

L'ambiente di addestramento distribuito del modello con SageMaker Training

I clienti di SageMaker possono anche accedere a immagini Docker integrate con vari framework di deep learning e i necessari pacchetti Linux, NCCL e Python per l’addestramento del modello. Gli scienziati dei dati o gli ingegneri di ML che desiderano eseguire l’addestramento del modello possono farlo senza il peso della configurazione dell’infrastruttura di addestramento o della gestione di Docker e della compatibilità di diverse librerie.

Durante un workshop di un giorno, siamo stati in grado di configurare una configurazione di training distribuito basata su SageMaker all’interno dell’account AWS di KT, accelerare gli script di training di KT utilizzando la libreria SageMaker Distributed Data Parallel (DDP) e persino testare un job di training utilizzando due istanze ml.p4d.24xlarge. In questa sezione, descriviamo l’esperienza di KT nel lavorare con il team di AWS e nell’utilizzare SageMaker per sviluppare il loro modello.

Nel proof of concept, volevamo velocizzare un job di training utilizzando la libreria SageMaker DDP, ottimizzata per l’infrastruttura AWS durante il training distribuito. Per passare da PyTorch DDP a SageMaker DDP, è sufficiente dichiarare il pacchetto torch_smddp e cambiare il backend in smddp, come mostrato nel seguente codice:

import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend='smddp',rank=args.rank,world_size=args.world_size)

Per saperne di più sulla libreria SageMaker DDP, fare riferimento a SageMaker’s Data Parallelism Library.

Analisi delle cause della lentezza del training con il SageMaker Debugger e Profiler

Il primo passo per ottimizzare e accelerare un carico di lavoro di training consiste nel comprendere e diagnosticare dove si verificano i colli di bottiglia. Per il job di training di KT, abbiamo misurato il tempo di training per iterazione del data loader, del passaggio in avanti e del passaggio all’indietro:

1 iter time – dataloader : 0,00053 sec, forward : 7,77474 sec, backward: 1,58002 sec
2 iter time – dataloader : 0,00063 sec, forward : 0,67429 sec, backward: 24,74539 sec
3 iter time – dataloader : 0,00061 sec, forward : 0,90976 sec, backward: 8,31253 sec
4 iter time – dataloader : 0,00060 sec, forward : 0,60958 sec, backward: 30,93830 sec
5 iter time – dataloader : 0,00080 sec, forward : 0,83237 sec, backward: 8,41030 sec
6 iter time – dataloader : 0,00067 sec, forward : 0,75715 sec, backward: 29,88415 sec

Osservando il tempo nell’output standard per ogni iterazione, abbiamo notato che il tempo di esecuzione del passaggio all’indietro variava significativamente da un’iterazione all’altra. Questa variazione è insolita e può influire sul tempo totale di training. Per trovare la causa di questa velocità di training inconsistente, abbiamo cercato di identificare i colli di bottiglia delle risorse utilizzando il System Monitor (SageMaker Debugger UI), che consente di eseguire il debug di job di training su SageMaker Training e visualizzare lo stato delle risorse come CPU, GPU, rete e I/O della piattaforma di training gestita entro un determinato numero di secondi.

La SageMaker Debugger UI fornisce dati dettagliati ed essenziali che possono aiutare a identificare e diagnosticare i colli di bottiglia di un job di training. In particolare, il grafico a linea di utilizzo della CPU e le tabelle per le utilizzazioni della CPU/GPU per istanza ci hanno colpito.

Nel grafico a linea di utilizzo della CPU, abbiamo notato che alcune CPU venivano utilizzate al 100%.

Grafico a linea di utilizzo della CPU con un collo di bottiglia della CPU

Nella heat map (dove i colori più scuri indicano una maggiore utilizzazione), abbiamo notato che alcune core della CPU erano fortemente utilizzate durante tutto il training, mentre l’utilizzo della GPU non era costantemente alto nel tempo.

La mappa termica di utilizzo della CPU con un collo di bottiglia della CPU

Da qui abbiamo cominciato a sospettare che una delle ragioni per la lentezza della velocità di allenamento fosse un collo di bottiglia della CPU. Abbiamo revisionato il codice degli script di allenamento per vedere se ciò causava un collo di bottiglia della CPU. La parte più sospetta era il grande valore di num_workers nel caricatore dati, quindi abbiamo cambiato questo valore a 0 o 1 per ridurre l’utilizzo della CPU. Abbiamo quindi eseguito nuovamente il lavoro di allenamento e verificato i risultati.

Le seguenti schermate mostrano il grafico della linea di utilizzo della CPU, l’utilizzo della GPU e la mappa termica dopo aver mitigato il collo di bottiglia della CPU.

Il grafico della linea di utilizzo della CPU dopo aver mitigato un collo di bottiglia della CPU

L'utilizzo della CPU, l'utilizzo della GPU dopo aver mitigato un collo di bottiglia della CPULa mappa termica dell'utilizzo della CPU dopo aver mitigato un collo di bottiglia della CPU

Semplicemente modificando num_workers, abbiamo notato una significativa diminuzione dell’utilizzo della CPU e un aumento complessivo dell’utilizzo della GPU. Questo è stato un cambiamento importante che ha migliorato notevolmente la velocità di allenamento. Tuttavia, volevamo vedere dove potevamo ottimizzare ulteriormente l’utilizzo della GPU. Per questo abbiamo utilizzato SageMaker Profiler.

SageMaker Profiler aiuta a identificare suggerimenti di ottimizzazione fornendo visibilità sull’utilizzo delle operazioni, inclusi il monitoraggio delle metriche di utilizzo della GPU e della CPU e il consumo del kernel della GPU/CPU all’interno degli script di allenamento. Aiuta gli utenti a capire quali operazioni stanno consumando risorse. Per utilizzare SageMaker Profiler, è necessario aggiungere ProfilerConfig alla funzione che invoca il lavoro di allenamento utilizzando il kit di sviluppo software (SDK) di SageMaker, come mostrato nel seguente codice:

from sagemaker import ProfilerConfig, Profilerfrom sagemaker.debugger import (ProfilerRule, rule_configs)rules=[ProfilerRule.sagemaker(rule_configs.ProfilerReport())]profiler_config = ProfilerConfig(profile_params = Profiler(cpu_profiling_duration=3600))from sagemaker.pytorch import PyTorchregion_name = 'us-west-2'image_uri=f'763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker'estimator = PyTorch(entry_point='train.py',source_dir='src',role=role,image_uri=image_uri,instance_count=4,instance_type='ml.p4d.24xlarge',distribution={'smdistributed': {'dataparallel': {'enabled': True}}},profiler_config=profiler_config,hyperparameters=hyperparameters,sagemaker_session=sagemaker_session,)

Nell’SDK di Python di SageMaker, hai la flessibilità di aggiungere le funzioni annotate per SageMaker Profiler per selezionare il codice o i passaggi nello script di allenamento che necessitano di un profilo. Di seguito è riportato un esempio del codice che devi dichiarare per SageMaker Profiler negli script di allenamento:

import smppySMProf = smppy.SMProfiler.instance()config = smppy.Config()config.profiler = {"EnableCuda": "1",}SMProf.configure(config)SMProf.start_profiling()…with smppy.annotate("Forward"):student_out = student_model(inp)with smppy.annotate("Backward"):loss.backward()…SMProf.stop_profiling()

Dopo aver aggiunto il codice precedente, se si esegue un lavoro di allenamento utilizzando gli script di allenamento, è possibile ottenere informazioni sulle operazioni consumate dal kernel della GPU (come mostrato nella figura seguente) dopo che l’allenamento viene eseguito per un certo periodo di tempo. Nel caso degli script di allenamento di KT, l’abbiamo eseguito per un’epoca e abbiamo ottenuto i seguenti risultati.

Time Spent By All GPU Kernels(1)

Quando abbiamo controllato i primi cinque tempi di consumo delle operazioni del kernel GPU tra i risultati di SageMaker Profiler, ci siamo accorti che per lo script di addestramento KT, la maggior parte del tempo viene consumata dall’operazione di prodotto matriciale, che è una moltiplicazione matriciale generale (GEMM) su GPU. Con questa importante conoscenza da SageMaker Profiler, abbiamo iniziato a investigare modi per accelerare queste operazioni e migliorare l’utilizzo della GPU.

Accelerare il tempo di addestramento

Abbiamo esaminato vari modi per ridurre il tempo di calcolo della moltiplicazione matriciale e applicato due funzioni PyTorch.

Shard dello stato dell’ottimizzatore con ZeroRedundancyOptimizer

Se si guarda all’ottimizzatore Zero Redundancy Optimizer (ZeRO), la tecnica DeepSpeed/ZeRO consente di addestrare un modello di grandi dimensioni in modo efficiente, con una velocità di addestramento migliore eliminando le ridondanze nella memoria utilizzata dal modello. ZeroRedundancyOptimizer in PyTorch utilizza la tecnica di frammentazione dello stato dell’ottimizzatore per ridurre l’utilizzo di memoria per un processo in Distributed Data Parallel (DDP). DDP utilizza gradienti sincronizzati nel passaggio di retropropagazione in modo che tutte le repliche dell’ottimizzatore iterino sugli stessi parametri e valori del gradiente, ma invece di avere tutti i parametri del modello, ciascuno stato dell’ottimizzatore è mantenuto mediante frammentazione solo per diversi processi DDP per ridurre l’utilizzo di memoria.

Per utilizzarlo, è possibile lasciare il proprio ottimizzatore esistente in optimizer_class e dichiarare un ZeroRedundancyOptimizer con il resto dei parametri del modello e del tasso di apprendimento come parametri.

student_optimizer = ZeroRedundancyOptimizer(student_model.parameters(),optimizer_class=torch.optim.AdamW,lr=initial_lr)

Precisione mista automatica

Precisione mista automatica (AMP) utilizza il tipo di dati torch.float32 per alcune operazioni e torch.bfloat16 o torch.float16 per altre, per la comodità di rapido calcolo e ridotta utilizzo di memoria. In particolare, poiché i modelli di deep learning sono generalmente più sensibili ai bit di esponente che ai bit di frazione nei loro calcoli, torch.bfloat16 è equivalente ai bit di esponente di torch.float32, permettendo loro di apprendere rapidamente con una perdita minima. torch.bfloat16 funziona solo su istanze con architettura NVIDIA A100 (Ampere) o superiore, come ml.p4d.24xlarge, ml.p4de.24xlarge e ml.p5.48xlarge.

Per applicare AMP, è possibile dichiarare torch.cuda.amp.autocast nei file di script di addestramento come mostrato nel codice sopra e dichiarare dtype come torch.bfloat16.

with torch.cuda.amp.autocast(dtype="torch.bfloat16"):teacher = teacher_model(input_data)student = student_model(input_data)loss = loss(teacher, student, target)loss.requires_grad_(True)loss.backward()student_optimizer.step()student_optimizer.zero_grad(set_to_none=True)

Risultati in SageMaker Profiler

Dopo aver applicato le due funzioni agli script di addestramento e aver eseguito nuovamente un lavoro di addestramento per un’epoca, abbiamo controllato i tempi di consumo delle prime cinque operazioni per il kernel GPU in SageMaker Profiler. La figura seguente mostra i nostri risultati.

Time Spent By All GPU Kernels(2)

Possiamo vedere che l’operazione GEMM, che era al primo posto nella lista prima dell’applicazione delle due funzioni di Torch, è scomparsa dalle prime cinque operazioni, sostituita dall’operazione ReduceScatter, che di solito si verifica nell’addestramento distribuito.

Risultati di velocità di addestramento del modello KT distillato

Abbiamo aumentato la dimensione del batch di addestramento di altri 128 per tener conto dei risparmi di memoria derivanti dall’applicazione delle due funzioni di Torch, ottenendo una dimensione finale del batch di 1152 invece di 1024. L’addestramento del modello studente finale è stato in grado di eseguire 210 epoche al giorno; i tempi di addestramento e l’accelerazione tra l’ambiente di addestramento interno di KT e SageMaker sono riassunti nella seguente tabella.

Ambiente di addestramento Spec GPU di addestramento Numero di GPU Tempo di addestramento (ore) Epoche Ore per epoca Riduzione del rapporto
Ambiente di addestramento interno di KT A100 (80GB) 2 960 300 3.20 29
Amazon SageMaker A100 (40GB) 32 24 210 0.11 1

La scalabilità di AWS ci ha permesso di completare il lavoro di addestramento 29 volte più velocemente rispetto al passato utilizzando 32 GPU invece di 2 in locale. Di conseguenza, l’uso di più GPU su SageMaker avrebbe ridotto significativamente il tempo di addestramento senza differenze nei costi complessivi di addestramento.

Conclusione

Park Sang-min (Vision AI Serving Technology Team Leader) del AI2XL Lab nel Convergence Technology Center di KT ha commentato sulla collaborazione con AWS per lo sviluppo del modello AI Food Tag:

“Recentemente, con l’aumento dei modelli basati su trasformatori nel campo della visione, i parametri del modello e la memoria GPU richiesta sono in aumento. Stiamo utilizzando una tecnologia leggera per risolvere questo problema, e ci vuole molto tempo, circa un mese per imparare una volta. Attraverso questa PoC con AWS, siamo stati in grado di identificare i colli di bottiglia delle risorse con l’aiuto di SageMaker Profiler e Debugger, risolverli, e quindi utilizzare la libreria di parallelismo dei dati di SageMaker per completare l’addestramento in circa un giorno con il codice del modello ottimizzato su quattro istanze ml.p4d.24xlarge.”

SageMaker ha aiutato il team di Sang-min a risparmiare settimane di tempo nell’addestramento e nello sviluppo del modello.

Basandoci su questa collaborazione sul modello di visione, AWS e il team di SageMaker continueranno a collaborare con KT su vari progetti di ricerca AI/ML per migliorare lo sviluppo dei modelli e la produttività del servizio mediante l’applicazione delle capacità di SageMaker.

Per saperne di più sulle funzionalità correlate in SageMaker, controlla i seguenti link: