In een post gepubliceerd op het platform Hugging Face opent Amazon een reeks gewijd aan de analyse van inferentieprestaties, dat wil zeggen de uitvoering van een reeds getraind model. Dit eerste deel legt uit hoe je kunt achterhalen waar de tijd aan wordt besteed, met behulp van profilingtools toegepast op de Trainium-chip van AWS, ontworpen voor training en inferentie van deep learning-modellen.
- Amazon start een reeks over inferentieprestaties op AWS Trainium met profilingtools.
- TorchNeuron en de profiler tonen waar de tijd op het apparaat verloren gaat.
- De naïeve RMSNorm kost ongeveer zes keer meer rekentijd dan de gefuseerde versies.
- Het tweede deel behandelt de profilering op apparaatniveau via het NEFF-bestand.
De gids richt zich tot lezers zonder eerdere ervaring met profilering. Het maakt gebruik van TorchNeuron, de open source native backend van PyTorch voor Trainium, die standaard PyTorch-workflows ondersteunt, waaronder uitvoering in eager mode en torch.compile, met weinig codeaanpassingen. TorchNeuron integreert ook met de Neuron Kernel Interface (NKI), waarmee aangepaste operators kunnen worden gemaakt.
De profiler configureren
De profilering steunt op twee objecten. NeuronConfig definieert de te capteren datastromen, de te traceren NeuronCores en de uitvoermap, en wordt vervolgens doorgegeven aan torch.profiler.profile. NeuronProfiler biedt een exportfunctie die de framework-trace schrijft, in het formaat trace.json, naast de systeemtraceringen van Neuron. De parameter correlate_device_time=True maakt het mogelijk de tijd besteed aan het apparaat weer te geven in de tabel key_averages(), waar deze kan worden vergeleken met de CPU-tijd.
De traceringen kunnen vervolgens worden geopend in Neuron Explorer, een visualisatietool die in de browser werkt. Hiermee kan een bewerking worden gevolgd vanaf de verdeling in het framework tot aan de uitvoering op de hardware, via de planning in de runtime. Voor een NeuronCore, dat wil zeggen een rekenkern van de chip, bevat het NEFF-bestand de gecompileerde uitvoerbare code, terwijl het bijbehorende NTFF-bestand de trace van de hardware-instructies bevat.
Drie versies van dezelfde bewerking
De post onderzoekt de functie RMSNorm, een eenvoudige bewerking bestaande uit een reeks rekenkundige bewerkingen en reducties. Drie implementaties worden vergeleken: een versie geschreven bewerking voor bewerking in PyTorch, de aanroep van de ingebouwde functie torch.rms_norm(), en de versie geschreven bewerking voor bewerking maar gecompileerd met torch.compile.
In de naïeve versie laat de tabel zien dat drie aanroepen naar aten::mul het grootste deel van de tijd besteed aan het apparaat vertegenwoordigen. De totale Neuron-tijd bedraagt 2 355 µs voor zes operatoraanroepen. Volgens de auteur neemt de ingebouwde functie torch.rms_norm(), die wordt omgezet in een enkele gefuseerde rekenkernel, slechts één bewerking in beslag, aten::_fused_rms_norm, met een Neuron-tijd van 403,3 µs. De versie gecompileerd met torch.compile verschijnt als een enkele gecompileerde regio, met een duur van 375,5 µs.
De post preciseert dat, in deze tests, beide benaderingen ongeveer een zesde van de rekentijd van de naïeve versie in beslag nemen. De fusie elimineert onderbrekingen tussen hardware-uitvoeringen en voorkomt het materialiseren van tussenwaarden, zoals x_sq en mean_sq, tussen afzonderlijk verzonden bewerkingen.
De auteur geeft aan dat de methode kan worden toegepast op grotere modellen: identificeer een hotspot in de tabel, verifieer deze in de trace, en bepaal vervolgens of een gefuseerde bewerking of torch.compile de kosten ervan kan verminderen. Het tweede deel van de serie zal gaan over de profiel op apparaatniveau, met de analyse van het NEFF-bestand.
