Profilazione di TorchNeuron (Parte 1): Profilazione a livello di sistema su AWS Trainium
huggingface ·
Questo articolo introduce il profiling a livello di sistema per TorchNeuron su AWS Trainium, con l'obiettivo di aiutare i professionisti del machine learning ad analizzare i colli di bottiglia delle prestazioni di inferenza. Vengono dettagliate le configurazioni di NeuronConfig e NeuronProfiler per acquisire tracce del dispositivo e del runtime utilizzando l'estensione PrivateUse1 di PyTorch. La guida illustra come utilizzare Neuron Explorer per visualizzare il ciclo di vita delle operazioni, dal dispatch del framework all'esecuzione hardware. Vengono confrontate tre implementazioni di RMSNorm: operazioni PyTorch naive, torch.rms_norm integrato e una funzione naive compilata con torch.compile. I risultati mostrano che sia la fusione tramite torch.rms_norm che la compilazione riducono significativamente il tempo del dispositivo Neuron eliminando la materializzazione dei tensori intermedi e i gap tra le operazioni. Gli approcci fusi richiedono circa un sesto del tempo di dispositivo necessario rispetto all'implementazione naive. L'autore conclude che identificare i hotspot nelle tabelle key_averages e verificarli nelle tracce costituisce un flusso di lavoro efficace per l'ottimizzazione.
huggingface