TorchNeuron Profiling: AWS Trainium’da RMSNorm Performans Analizi (Bölüm 1)

AWS Trainium, derin öğrenme eğitim ve çıkarım iş yükleri için özel olarak tasarlanmış bir AI çipi iken, TorchNeuron bu çipi PyTorch ile entegre eden açık kaynaklı bir arka uç sunuyor. Makale, Neuron'un PyTorch'un PrivateUse1 uzantısı üzerinden nasıl kaydedildiğini ve torch_neuronx paketinin bu entegrasyonu nasıl sağladığını açıklıyor. Profil oluşturmak için iki temel nesne – NeuronConfig ve NeuronProfiler – kullanılıyor; NeuronConfig cihaz ve runtime izlerini, izlenecek NeuronCore sayısını ve zaman korelasyonunu tanımlarken, NeuronProfiler bu ayarları torch.profiler.profile içinde experimental_config olarak geçiriyor ve iz çıktısını JSON ve .neff/.ntff dosyalarıyla birlikte kaydediyor.

Profil aşamasında, RMSNorm işlemi üç farklı şekilde ele alınıyor. İlk örnek, ayrı ayrı PyTorch operasyonları (mul, mean, add, rsqrt) kullanarak naive bir implementasyon sunuyor; key_averages() tablosu, üç adet aten::mul çağrısının cihaz süresinin %49,5'ini oluşturduğunu gösteriyor ve toplam Neuron süresi 2,355 µs olarak ölçülüyor. İkinci örnek, PyTorch'un yerleşik torch.rms_norm fonksiyonunu kullanarak tek bir fused kernel'e indirgeniyor; bu durumda tek bir aten::_fused_rms_norm çağrısı %100 Neuron süresi (403,3 µs) alıyor. Üçüncü örnek ise aynı naive fonksiyonu @torch.compile(backend="neuron") ile derleyerek otomatik füzyon sağlıyor; sonuçta tek bir Torch-Compiled Region görülüyor ve Neuron süresi 375,5 µs olarak kaydediliyor.

Bu üç senaryo, fused kernel'lerin ve torch.compile'ın naive implementasyona kıyasla cihaz süresini yaklaşık altı kat azalttığını ortaya koyuyor. Neuron Explorer aracı, sistem izlerini görsel olarak inceleyerek her bir ATen operasyonunun framework, runtime ve donanım katmanlarındaki yaşam döngüsünü gösteriyor; fused ve derlenmiş sürümlerde tek bir dispatch ve tek bir donanım olayı görülürken, naive sürümde her bir aritmetik adım ayrı ayrı izleniyor. Makale, bu metodolojinin daha büyük modellerde hotspot tespiti ve optimizasyonu için de uygulanabileceğini vurguluyor ve bir sonraki bölümde cihaz‑seviye izleme ve roofline analizi yapılacağını belirtiyor.