In un post pubblicato sulla piattaforma Hugging Face, Amazon apre una serie dedicata all’analisi delle prestazioni di inferenza, cioè dell’esecuzione di un modello già addestrato. Questa prima parte spiega come identificare dove viene impiegato il tempo, utilizzando strumenti di profiling applicati al chip Trainium di AWS, progettato per l’addestramento e l’inferenza dei modelli di deep learning.
- Il post introduce una serie sul profiling dell’inferenza con il chip AWS Trainium.
- TorchNeuron, backend open source di PyTorch, permette di catturare tracce visualizzabili in Neuron Explorer.
- La versione naive di RMSNorm impiega 2 355 µs, mentre le versioni fusa e compilata scendono a circa 400 µs.
- La fusione elimina le interruzioni tra le esecuzioni hardware e riduce il tempo a circa un sesto.
La guida si rivolge a lettori senza esperienza pregressa di profiling. Si basa su TorchNeuron, il backend open source nativo di PyTorch per Trainium, che supporta i flussi di lavoro standard di PyTorch, inclusa l’esecuzione in modalità eager e torch.compile, con poche modifiche al codice. TorchNeuron si integra anche con la Neuron Kernel Interface (NKI), che permette di creare operatori personalizzati.
Configurare il profiler
Il profiling si basa su due oggetti. NeuronConfig definisce i flussi di dati da catturare, i NeuronCores da tracciare e la cartella di output, poi viene passato a torch.profiler.profile. NeuronProfiler fornisce una funzione di esportazione che scrive la traccia del framework, in formato trace.json, insieme alle tracce di sistema di Neuron. Il parametro correlate_device_time=True permette di visualizzare il tempo trascorso sul dispositivo nella tabella key_averages(), dove può essere confrontato con il tempo CPU.
Le tracce possono poi essere aperte in Neuron Explorer, uno strumento di visualizzazione che funziona nel browser. Permette di seguire un’operazione dalla distribuzione nel framework fino alla sua esecuzione sull’hardware, passando per la schedulazione nel runtime. Per un NeuronCore, cioè un nucleo di calcolo del chip, il file NEFF contiene l’eseguibile compilato, mentre il file NTFF corrispondente contiene la traccia delle istruzioni hardware.
Tre versioni di una stessa operazione
Il post studia la funzione RMSNorm, un’operazione semplice composta da una serie di calcoli aritmetici e riduzioni. Vengono confrontate tre implementazioni: una versione scritta operazione per operazione in PyTorch, la chiamata alla funzione integrata torch.rms_norm(), e la versione scritta operazione per operazione ma compilata con torch.compile.
Nella versione naive, la tabella mostra che tre chiamate a aten::mul rappresentano la parte maggiore del tempo trascorso sul dispositivo. Il tempo Neuron totale raggiunge 2 355 µs per sei chiamate di operatori. Secondo l’autore, la funzione integrata torch.rms_norm(), che viene trasformata in un unico kernel di calcolo fuso, occupa solo un’operazione, aten::_fused_rms_norm, con un tempo Neuron di 403,3 µs. La versione compilata con torch.compile appare come una singola regione compilata, della durata di 375,5 µs.
Il post precisa che, in questi test, entrambi gli approcci richiedono circa un sesto del tempo di calcolo della versione naive. La fusione elimina le interruzioni tra le esecuzioni hardware ed evita di materializzare valori intermedi, come x_sq e mean_sq, tra operazioni inviate separatamente.
L’autore indica che il metodo può essere applicato a modelli più grandi: individuare un punto caldo nella tabella, verificarlo nella traccia, poi determinare se un’operazione fusa o torch.compile può ridurne il costo. La seconda parte della serie riguarderà il profiling a livello del dispositivo, con l’analisi del file NEFF.
