Profiling von TorchNeuron (Teil 1): Systemweites Profiling auf AWS Trainium
huggingface ·
Dieser Artikel stellt das System-Level-Profiling für TorchNeuron auf AWS Trainium vor, um Machine-Learning-Praktikern bei der Analyse von Leistungsengpässen beim Inferenzieren zu helfen. Es werden die Einrichtung von NeuronConfig und NeuronProfiler zur Erfassung von Geräte- und Laufzeit-Traces mithilfe der PrivateUse1-Erweiterung von PyTorch detailliert beschrieben. Die Anleitung zeigt, wie Neuron Explorer genutzt wird, um den Lebenszyklus von Operationen vom Framework-Dispatch bis zur Hardware-Ausführung zu visualisieren. Drei Implementierungen von RMSNorm werden verglichen: naive PyTorch-Operationen, die integrierte Funktion torch.rms_norm und eine naive Funktion, die mit torch.compile kompiliert wurde. Die Ergebnisse zeigen, dass sowohl die Fusion über torch.rms_norm als auch das Kompilieren die Neuron-Gerätezeit erheblich reduzieren, indem die Materialisierung intermediärer Tensoren und Lücken zwischen Operationen vermieden werden. Die fusionierten Ansätze benötigen etwa ein Sechstel der Gerätezeit, die die naive Implementierung erfordert. Der Autor schlussfolgert, dass das Identifizieren von Hotspots in key_averages-Tabellen und deren Überprüfung in Traces ein effektiver Workflow für Optimierungen ist.
huggingface