Come sfruttare al meglio la tua GPU moderna

fotografato da Deva Darshan SU Unsplash

Gli ultimi anni hanno visto progressi rivoluzionari nel campo dell’intelligenza artificiale, forse meglio esemplificati dalla recente popolarità e proliferazione di applicazioni basate su LLM come ChatGPT. Queste scoperte sono state alimentate da sviluppi altrettanto entusiasmanti nei macchinari utilizzati per addestrare i modelli di intelligenza artificiale. Architetture nuove e innovative, sofisticati core di elaborazione tensore e acceleratori HW dedicati hanno consentito la convergenza di modelli IA di dimensioni sempre crescenti, a velocità sempre più elevate. In questo post, ci concentreremo su un particolare progresso nell’HW specializzato nell’intelligenza artificiale: l’inclusione di core di elaborazione tensore dedicati a virgola mobile a 8 bit (FP8).. Apparendo nelle più moderne architetture HW AI (ad esempio, Nvidia Hopper, Nvidia Ada LovelaceE L’Avana Gaudi2) i tensor core FP8 consentono un aumento significativo delle operazioni in virgola mobile al secondo (FLOPS), nonché opportunità di ottimizzazione della memoria e risparmio energetico sia per l’addestramento dell’IA che per i carichi di lavoro di inferenza.

Per sfruttare le funzionalità dell’8° PQ a livello hardware è necessario un supporto adeguato nello stack SW e nel framework di sviluppo che utilizziamo per creare le nostre applicazioni di formazione e inferenza sull’intelligenza artificiale. In questo post descriveremo come fare modificare uno script di formazione PyTorch in modo da utilizzare il supporto integrato per il tipo di dati FP8 di un GPU Nvidia H100. Inizieremo fornendo alcune motivazioni per l’utilizzo del tipo di dati FP8. Esamineremo quindi il supporto API PyTorch specifico dell’FP8 esposto da Motore del trasformatore libreria e mostrare come integrarli in un semplice script di formazione. Anche se lo faremo non approfondiremo la teoria alla base dell’uso dell’8° PQ per la formazione sull’intelligenza artificiale, noteremo le potenziali sfide legate al suo utilizzo. Infine, dimostreremo le significative opportunità di ottimizzazione del tipo di dati FP8.

Disclaimer

Si prega di non interpretare la nostra menzione di componenti, metodologie o servizi SW come un’approvazione al loro utilizzo. La migliore progettazione per lo sviluppo ML varierà notevolmente in base ai dettagli specifici del tuo carico di lavoro AI. Tieni inoltre presente che le API e i comportamenti di alcuni dei pacchetti SW e dei componenti che menzioneremo potrebbero cambiare nel momento in cui leggerai questo post. Ti invitiamo vivamente a valutare eventuali decisioni di progettazione basate sull’HW e sul SW più aggiornati disponibili.

Man mano che i modelli di intelligenza artificiale diventano sempre più sofisticati, aumentano anche i macchinari necessari per addestrarli. IL GPU Nvidia H100che si dice supporti “prestazioni e scalabilità senza precedenti”, è (al momento in cui scriviamo) il più nuovo e potente acceleratore AI di Nvidia, appositamente progettato con l’obiettivo di consentire lo sviluppo dell’IA della prossima generazione. Con l’attuale campagna pubblicitaria sull’intelligenza artificiale in pieno svolgimento, la domanda per queste GPU è stata enorme (ad esempio, vedi Qui). Di conseguenza, e non sorprende, il costo di queste GPU è stato estremamente elevato, forse addirittura proibitivo per molti dei nostri lettori. Fortunatamente, i fornitori di servizi cloud come AWS, GCP e Microsoft Azure offrono l’accesso “pay as you go” (per ora/al secondo) alle macchine basate su H100, aprendo così l’opportunità del loro utilizzo a una comunità molto più ampia di sviluppatori di intelligenza artificiale. .

In AWS, le GPU H100 sono offerte come componente di recentemente annunciato Famiglia di istanze AWS EC2 p5. Si afferma che questi casi “accelera i tempi di implementazione della soluzione fino a 4 volte rispetto alle istanze EC2 basate su GPU della generazione precedente e riduci i costi per l’addestramento dei modelli ML fino al 40%”.

In un post recente abbiamo discusso alcune delle considerazioni che dovrebbero essere prese in considerazione nella scelta di un’istanza di formazione ML. Abbiamo evidenziato il fatto che il tipo di istanza ottimale dipenderà molto dal progetto in questione. Nello specifico, quando si tratta di istanze di formazione ML: più grande è non sempre meglio. Ciò è particolarmente vero per la famiglia di istanze p5. È vero: il p5 probabilmente supererà qualsiasi altro tipo di istanza; dopo tutto, l’H100 è una bestia indiscussa in termini di prestazioni. Ma una volta considerato il costo del p5 ($ 98,32 l’ora per l’istanza p5.48xlarge da 8 GPU – al momento in cui scriviamo), potresti trovare altri tipi di istanza più adatti.

Nella sezione successiva addestreremo un modello di visione artificiale relativamente grande su un p5.48xlarge e confronteremo le sue prestazioni con un p4d.24xlarge contenente 8 GPU Nvidia A100.

Modello giocattolo

Nel blocco di codice seguente definiamo a Trasformatore di visione Modello di classificazione supportato da (ViT) (utilizzando il popolare timm Pacchetto Python versione 0.9.10) insieme a un set di dati generato casualmente. Le dorsali ViT sono disponibili in molte forme e dimensioni. Qui abbiamo scelto quella che viene spesso definita configurazione ViT-Huge — with 632 milioni di parametri, per sfruttare al meglio la capacità dell’H100 per i modelli di grandi dimensioni.

import torch, time
import torch.optim
import torch.utils.data
import torch.distributed as dist
from torch.nn.parallel.distributed import DistributedDataParallel as DDP
import torch.multiprocessing as mp

# modify batch size according to GPU memory
batch_size = 64

from timm.models.vision_transformer import VisionTransformer

from torch.utils.data import Dataset

# use random data
class FakeDataset(Dataset):
def __len__(self):
return 1000000

def __getitem__(self, index):
rand_image = torch.randn((3, 224, 224), dtype=torch.float32)
label = torch.tensor(data=(index % 1000), dtype=torch.int64)
return rand_image, label

def mp_fn(local_rank, *args):
# configure process
dist.init_process_group("nccl",
rank=local_rank,
world_size=torch.cuda.device_count())
torch.cuda.set_device(local_rank)
device = torch.cuda.current_device()

# create dataset and dataloader
train_set = FakeDataset()
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=batch_size,
num_workers=12, pin_memory=True)

# define ViT-Huge model
model = VisionTransformer(
embed_dim=1280,
depth=32,
num_heads=16,
).cuda(device)
model = DDP(model, device_ids=(local_rank))

# define loss and optimizer
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

model.train()

t0 = time.perf_counter()
summ = 0
count = 0

for step, data in enumerate(train_loader):
# copy data to GPU
inputs = data(0).to(device=device, non_blocking=True)
label = data(1).squeeze(-1).to(device=device, non_blocking=True)

# use mixed precision to take advantage of bfloat16 support
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, label)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

# capture step time
batch_time = time.perf_counter() - t0
if step > 10: # skip first steps
summ += batch_time
count += 1
t0 = time.perf_counter()
if step > 50:
break
print(f'average step time: {summ/count}')

if __name__ == '__main__':
mp.spawn(mp_fn,
args=(),
nprocs=torch.cuda.device_count(),
join=True)

Abbiamo addestrato questo modello sia su p5.48xgrande E p4d.24xlarge tipi di istanza utilizzando PyTorch 2.1 Contenitore di apprendimento profondo AWS (763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-ec2).

Non sorprende che le prestazioni di step-time p5 superino le prestazioni p4d: 0,199 secondi per step rispetto a 0,41 – più del doppio della velocità!! Ciò significherebbe dimezzare il tempo necessario per addestrare i tuoi modelli ML di grandi dimensioni. Tuttavia, se si tiene conto della differenza di costo (32,77 dollari l’ora per il p4d contro 98,32 dollari l’ora per il p5 – al momento in cui scriviamo) si sviluppa una storia completamente diversa. IL il rapporto prezzo-prestazioni del p5 è ~30% peggiore rispetto al p4d!! Questo è molto lontano dal miglioramento del 40% osservato nel Annuncio p5.

A questo punto potresti trarre una delle due possibili conclusioni. La prima possibilità è che, nonostante tutto il clamore pubblicitario, il p5 semplicemente non sia la macchina giusta per te. La seconda è che il p5 potrebbe essere ancora praticabile, ma sarebbero necessari degli adattamenti al modello per sfruttarne appieno le potenzialità. Nelle prossime sezioni adotteremo il secondo approccio e dimostreremo come l’utilizzo del tipo di dati FP8, unico per il tipo di istanza p5, può alterare completamente i risultati comparativi in ​​termini di prezzo-prestazioni.

La prima cosa che dovremmo sottolineare è che, al momento della stesura di questo articolo, PyTorch (versione 2.1) non include un tipo di dati mobile nativo a 8 bit. Per programmare il nostro script per utilizzare FP8 utilizzeremo Motore del trasformatore (TE) una libreria dedicata per accelerare i modelli Transformer su GPU NVIDIA. TE (versione 0.12) è preinstallato nel contenitore AWS PyTorch 2.1 DL.

Sebbene la teoria alla base dell’uso dell’8° PQ per la formazione vada oltre lo scopo di questo post (ad es Qui), è importante esserne consapevoli i meccanismi di utilizzo dell’8°PQ sono molto più complessi di quelli dell’8°PQ Alternative a 16 bit (float16 e bfloat16). Fortunatamente, l’implementazione di TE nasconde all’utente tutti i dettagli complicati. Si prega di consultare il funzionario documentazione così come questo semplice esempio per istruzioni su come utilizzare le API TE. Per saperne di più su cosa succede dietro le quinte, assicurati di vedere i due tutorial video seguenti.

Per modificare il nostro modello per utilizzare TE, avvolgiamo il Transformer Layer specializzato di TE con una classe di blocco trasformatore personalizzata conforme a timm firma del livello di blocco.

import transformer_engine.pytorch as te
from transformer_engine.common import recipe

class TE_Block(te.transformer.TransformerLayer):
def __init__(
self,
dim,
num_heads,
mlp_ratio=4.,
qkv_bias=False,
qk_norm=False,
proj_drop=0.,
attn_drop=0.,
init_values=None,
drop_path=0.,
act_layer=None,
norm_layer=None,
mlp_layer=None
):
super().__init__(
hidden_size=dim,
ffn_hidden_size=int(dim * mlp_ratio),
num_attention_heads=num_heads,
hidden_dropout=proj_drop,
attention_dropout=attn_drop
)

Successivamente, modifichiamo l’inizializzazione di VisionTransformer per utilizzare il nostro livello di blocco personalizzato:

  model = VisionTransformer(
embed_dim=1280,
depth=32,
num_heads=16,
block_fn=TE_Block
).cuda(device)

Fino ad ora non abbiamo apportato alcuna modifica specifica per H100: lo stesso codice può essere eseguito sul nostro tipo di istanza p4d basato su A100. L’ultima modifica è avvolgere il modello di passaggio in avanti con a te.fp8_autocast gestore del contesto. Questa modifica richiede una GPU che supporti FP8:

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
with te.fp8_autocast(enabled=True):
outputs = model(inputs)
loss = criterion(outputs, label)

Alcune osservazioni cautelative sull’uso dell’8° PQ

L’utilizzo di una rappresentazione a virgola mobile a 8 bit (in contrapposizione a una rappresentazione a 16 o 32 bit) implica una precisione inferiore e un intervallo dinamico inferiore. Questi possono avere un impatto significativo sulla raggiungibilità e/o sulla velocità della convergenza del modello. Sebbene l’implementazione sottostante del TE FP8 sia progettata per affrontare questa sfida, non vi è alcuna garanzia che funzionerà per il tuo modello. Potrebbe essere necessario armeggiare con i meccanismi sottostanti dell’FP8 (ad esempio, utilizzando il TE ricetta API), ottimizzare alcuni iperparametri e/o limitare l’applicazione dell’8° PQ a sottosezioni del modello. Potresti scoprire che, nonostante tutti i tuoi tentativi, il tuo modello semplicemente non è compatibile con FP8.

Nella tabella seguente riassumiamo i risultati dei nostri esperimenti su entrambi i tipi di istanze EC2 p4d.24xlarge e p5.48xlarge, con e senza la libreria TE. Per gli esperimenti p5.48xlarge abbiamo raddoppiato la dimensione del batch per aumentare l’utilizzo della memoria GPU da 80 GB. L’utilizzo dell’FP8 riduce il consumo di memoria della GPU consentendo un ulteriore aumento della dimensione del batch.

Risultati dell’esperimento (per autore)

Possiamo vedere che l’uso del blocco trasformatore TE ha aumentato il rapporto prezzo-prestazioni sia sui tipi di istanza p4d (~19%) che su quelli p5 (~32%). L’utilizzo dell’FP8 aumenta le prestazioni del p5 di un ulteriore 20% circa. A seguito delle ottimizzazioni TE e FP8 il rapporto prezzo/prestazioni del p5.48large basato su H100 batte quello del p4d.24large basato su A100 — anche se non con un margine molto ampio (~2%). Tenendo conto dell’aumento di 3 volte della velocità di addestramento, possiamo tranquillamente concludere che p5 sarebbe il tipo di istanza migliore per l’addestramento del nostro modello ottimizzato.

Si noti che l’aumento relativamente piccolo del rapporto prezzo-prestazioni (molto inferiore al 40% menzionato nell’annuncio di p5) ci lascia desiderare ulteriori ottimizzazioni specifiche per H100… ma per quelle bisognerà aspettare un altro post :).

In questo post abbiamo dimostrato come programmare uno script di addestramento PyTorch per utilizzare tipi mobili a 8 bit. Abbiamo ulteriormente dimostrato come l’uso dell’FP8 possa essere un fattore chiave per ottenere le migliori prestazioni dalle moderne GPU come Nvidia H100. È importante sottolineare che la fattibilità dell’8° PQ e il suo impatto sui risultati della formazione possono variare notevolmente in base ai dettagli del modello.

Questo post continua una lunga serie di pubblicazioni sul tema dell’ottimizzazione dei carichi di lavoro di machine learning. Assicurati di vedere alcuni dei nostri altri post su questo importante argomento.

Fonte: towardsdatascience.com

Lascia un commento

Il tuo indirizzo email non sarà pubblicato. I campi obbligatori sono contrassegnati *