Kernel tam olarak nedir?

GPU'da İki Sayı

Python'da c = a + b yapmayı zaten biliyorsunuz. Bunu binlerce kez yaptınız. Şimdi bunu PyTorch tensor'ları ile konuşacağız. Bir tensor, sadece bir sayı dizisidir. Onu GPU'ya koymak, bu dizinin normal CPU belleği yerine GPU'nun belleğinde yaşadığı anlamına gelir. a ve b GPU'daki iki tensor olduğunda, bu tek satır o kadar hızlı tamamlanır ki üzerinde hiç düşünmezsiniz.

Şimdi küçültün. Diyelim ki a ve b, ikisi de GPU'da bulunan iki tek float. Aynı satır. c = a + b. Çipte aslında ne çalışır?

Cevap bir kernel'dir. Bu dünyada kernel, GPU'nun bir veri parçası üzerinde çalıştırdığı küçük bir programdır. Dizüstü bilgisayarınızın başlattığı işletim sistemi çekirdeği değil. Lineer cebir ders kitabındaki matematik çekirdeği değil. Bu kelime çok sık kullanılır ve bu sizin hatanız değil. GPU dünyasında kernel, sadece şu anlama gelir: GPU'ya hemen gitmesi, paralel olarak, ona verdiğiniz veri üzerinde çalışması söylenen küçük bir fonksiyon.

Bunun sonunda, bir PyTorch kodu parçasına bakıp GPU'nun kaç kernel çalıştıracağını sayabileceksiniz. Kulağa küçük bir numara gibi geliyor ve öyle de, ama aynı zamanda "GPU kara kutusu" hissini sarsmanın ilk adımı, modelinizin yavaş olması ve nedenini bilmemeniz durumu. GPU performansı hakkında sorabileceğiniz her soru sonunda "hangi kernel'ler çalıştı ve ne yapıyorlardı?" sorusuna gelir. İşte bu yüzden buradan başlıyoruz.

İlk Kernel'iniz

a ve b'yi biraz daha büyütelim. Bu sefer uzunluk-8 tensor'lar. Hala tek bir PyTorch satırı: c = a + b.

Bunu çalıştırdığınızda, CPU'nuz (asıl Python'u yapan makine) GPU'ya şunu söyler: hey, bu programı bu veri üzerinde çalıştır. Bu talimata bir başlatma denir. Başlatılan şey bir kernel'dir: çalışmaya hazır bir program. Başlatmaların kendileri ucuzdur, her biri mikrosaniyeler. Başlatmanın etrafındakiler (GPU'ya giden veri, geri gelen sonuçlar) asıl maliyetin olduğu yerdir ve ilerledikçe sayacağımız şey de budur.

Kernel'in içinde, asıl iş thread adı verilen küçük işçiler tarafından yapılır. Bir GPU'nun binlercesi mevcuttur. 8 elemanlı toplamamız için 8 thread işi alır: thread 0 eleman 0'ı, thread 1 eleman 1'i ele alır ve bu şekilde thread 7'ye kadar devam eder. Her thread aynı küçük programı çalıştırır: a'nın bir elemanını oku, b'nin karşılık gelen elemanını oku, bunları topla, sonucu c'ye yaz.

(Pratikte GPU, thread'leri warp adı verilen sabit boyutlu gruplar halinde başlatır, NVIDIA kartlarında her zaman 32 thread ve diziniz tam bölünemediğinde fazlalıkları maskeler. Şimdilik bunu görmezden gelebiliriz.)

Yani bir PyTorch satırımız, bir başlatmamız, bir kernel'imiz, 8 toplama yapan 8 thread'imiz var. Şimdi çipte gerçekte neyin geçtiğini sayalım. Toplamayı yapmak için her thread'in a'nın elemanına ve b'nin elemanına ihtiyacı vardır. Bu, a'dan 8 okuma ve b'den 8 okuma demektir. Ardından her thread sonucunu c'ye yazar. Bu da 8 yazma demektir.

Bu okuma ve yazmalar, GPU çipinin hemen yanında bulunan büyük belleğe gider. Veri merkezi kartlarında (A100, H100) bu belleğe HBM (yüksek bant genişlikli bellek) denir. Tüketici kartlarında (RTX 4090) ve insanların bunu denemesi muhtemel olan Colab tarzı T4'lerde ise GDDR olarak adlandırılır. Her iki durumda da, çipin yanında bulunan hızlı bir bellektir ve biz buna sadece GPU belleği diyeceğiz. Hızlıdır, ama bedava değildir ve ona her gidişin bir maliyeti vardır.

Bir kernel = bir başlatma = veri üzerinde bir geçiş. Kernel, gövdesi içinde ne yaparsa yapsın, kenarlarındaki okuma ve yazmalar (GPU belleğine gidip girdileri getirme, çıktıları yazmak için geri dönme yolculukları) maliyetli olan kısımdır. Bütün şekil budur.

Tensor'lar büyüdüğünde bunların hiçbiri değişmez. Aynı PyTorch satırı, aynı tek kernel, sadece daha fazla thread. Eğer a ve b birer milyon elemansa, GPU aynı kernel'i daha büyük bir thread sürüsüyle başlatır. Matematik ölçeklenir, taşınan baytlar ölçeklenir, zihinsel model değişmez. Tek satır, tek kernel.

İki İşlem Arasında Ne Olur

1c = (a + b).relu()

Python'u, bunun bir toplama ve ardından bir relu olmak üzere zincirleme iki işlem olduğunu bilecek kadar iyi biliyorsunuz. Bir yorumlayıcıda bu, iki fonksiyon çağrısıdır. Bir GPU'da, hevesli PyTorch'ta bu, iki kernel başlatmasıdır: biri toplama, biri relu için. Şimdiye kadar şaşırtıcı bir şey yok.

Asıl ilginç olan, iki kernel arasında ne olduğudur.

Toplama bittiğinde, sonucunu bir yere koymak zorundadır. O bir yer GPU belleğidir. Toplama, tam bir ara dizi (buna tmp diyelim) belleğe yazar. Ardından bir an sonra relu başlatılır ve ilk işi aynı tmp dizisini bellekten geri okumaktır. Her şeyi okur, her elemana relu uygular, sonucu c'ye yazar.

Bu iki kernel için bellek trafiğini sayın:

• Toplama: a'yı okur, b'yi okur, tmp'yi yazar. Üç dizi boyutunda transfer.
• Relu: tmp'yi okur, c'yi yazar. İki tane daha.

Toplamda beş dizi boyutunda transfer. Bunu, bir önceki bölümdeki tek başına 8 elemanlı toplamayla karşılaştırın, ki o üç taneydi. Zincire .relu() eklemek size sadece relu'nun hesaplama maliyetine mal olmadı. Size dizinin GPU belleği üzerinden tam bir gidiş-dönüş yolculuğuna mal oldu, çünkü tmp'nin, bir sonraki kernel onu geri okuyabilsin diye yazılması gerekiyordu.

Hiçbir şey önbelleğe alınmadı. tmp bir kayıt defterinde veya hızlı bir yerel önbellekte kalamadı. GPU belleğine (yavaş ve uzak tür) gitti ve hemen geri geldi. İki kernel birbirine yabancıdır. Verileri, ikisinin de konuşmayı bildiği tek ortam olan GPU belleği aracılığıyla aktarmak zorundadırlar.

PyTorch bunu neden böyle yapıyor? Çünkü hevesli modda, a + b yazdığınızda, PyTorch onu hemen çalıştırır. Bir sonraki adımda .relu() çağıracağınızı bilmez. Her işlem, Python satırı çalıştırıldığı anda gönderilir. Hiçbir plan, hiçbir ileriye bakış yoktur. Her işlem kendi başına durur, gerçek bir dizi üretir ve onu bellek aracılığıyla gelen her şeye aktarır.

Bu, tutunmanız gereken kalıptır. PyTorch kodunuzdaki her ara değer, fiziksel olarak GPU belleğine yazılır ve bir sonraki işlem tarafından geri okunur. Hemen hepsi. "Kernel sayısı"nın gerçekte ölçtüğü şey budur. Her ekstra kernel, verilerinizin GPU belleği üzerinden yapması gereken başka bir gidiş-dönüş yolculuğudur.

Füzyon: İki İşlem, Tek Kernel

Tüm bunları tek seferde yapan bir kernel hayal edin: a'nın bir elemanını okur, b'nin bir elemanını okur, bunları toplar, sonuca relu uygular (tümü kernel'in içinde, çipten asla ayrılmayan küçük thread başına geçici alanında) ve ancak o zaman nihai değeri c'ye yazar. Ara (a + b) hala vardır, ancak yalnızca kernel'in içinde, her thread'in özel geçici alanında. GPU belleğine asla yazılmaz. tmp, gerçek bir dizi olarak, hiç var olmaz.

Şimdi transferleri sayın. a'nın okumaları: eleman başına 1. b'nin okumaları: eleman başına 1. c'nin yazmaları: eleman başına 1. Üç dizi boyutunda transfer. İki kernel'li sürümle aynı matematik, ancak iki gidiş-dönüş daha az.

8 uzunluğunda bu bir yuvarlama hatasıdır. Kimse umursamaz. 1 milyon veya 100 milyon uzunluğunda, bu ekstra gidiş-dönüşler çalışma süresinin büyük bir kısmını oluşturur ve duvar saati bunu yansıtır. *Bellek trafiğinin neden bu şekilde baskın olduğu, bu serideki 2. Makalenin tüm konusudur, bu yüzden "neden"i burada bırakacağım. Şu anki nokta sadece: aynı matematik, daha az yolculuk, pratikte daha hızlı.

Ayrı kernel'ler olacak işlemleri, ara değerin GPU belleğini ziyaret etmesi gerekmesin diye tek bir kernel'de birleştirme numarasının bir adı vardır. Buna füzyon denir. Bütün kelime bu. Bütün fikir bu.

Şimdi garip kısım. Bu birleştirilmiş kernel'i elle yazmak, toplama + relu için kolay görünüyor. İki işlem. Ortada bir satır "hesaplama". Ancak gerçek PyTorch kodunun, her biri kendi şekilleri, veri türleri ve yayın kuralları olan düzinelerce zincirleme işlemi vardır. Tüm bunları doğru bir şekilde ele alan birleştirilmiş bir kernel yazmak gerçek bir mühendislik işidir. Genellikle bu rutin elemansal kernel'leri elle yazmazsınız.

İyi haber: PyTorch, tam bu tür durumlar için, sizin için otomatik olarak yeniden yazma işlemini yapan bir araçla birlikte gelir. Buna torch.compile denir.

Neredeyse kesinlikle torch.compile'in şu şekilde kullanıldığını görmüşsünüzdür:

1model = torch.compile(model)

Tek satır. İnternette birileri size bunun işleri hızlandırdığını söyledi. İşte gerçekte ne yaptığı, düz İngilizceyle: hevesli modun işlemlerinizi birer birer çalıştırması yerine, torch.compile, fonksiyonunuzun gerçekleştirdiği tensor işlemlerini yakalar, bunları birleştirme fırsatlarını arar ve optimize edilmiş kod üretir. Aynı varsayımlarla eşleşen sonraki çağrılar bu kodu yeniden kullanabilir.

Yukarıda kağıt üzerinde yaptığımız füzyon (toplama ve relu, tek bir kernel paylaşarak, tmp belleğe asla dokunmadan), işlemler yeterince basit olduğu sürece torch.compile'in kodunuza otomatik olarak yapacağı türden bir şeydir. İnsanlar torch.compile'in "PyTorch'u hızlandırdığını" söylediğinde, kastettikleri şeyin büyük bir kısmı budur.

torch.compile'in kendi başına birleştiremediği durumlar için (tanımadığı özel işlemler, olağandışı indirgemeler, garip bellek düzenleri), birisinin yine de elle bir kernel yazması gerekir. Triton ve CUDA gibi araçların kullanıldığı yer burasıdır. Ayrı bir makale.

Kendiniz Görün

Yukarıdaki her şey, kernel'leri kağıt üzerinde sayıyordu. Şimdi bunları gerçek bir GPU'da sayma zamanı. CUDA GPU'su olan herhangi bir makineniz varsa (bir iş istasyonu, bir Colab not defteri, bir bulut örneği), bunu birkaç dakika içinde kendiniz çalıştırabilirsiniz.

Araç torch.profiler'dır. PyTorch'un içine yerleşiktir. Tek yaptığı, kodunuz çalışırken GPU'nun gerçekte ne yaptığını kaydetmek ve size okuyabileceğiniz bir tablo vermektir.

Adım 1: hevesli sürüm.

İki işlemli satırı, çağırabileceğimiz bir şey olsun diye bir fonksiyonun içine sarın:

1import torch
2from torch.profiler import profile, ProfilerActivity
3
4def add_relu(a, b):
5    return (a + b).relu()
6
7a = torch.randn(1_000_000, device="cuda")
8b = torch.randn(1_000_000, device="cuda")
9
10with profile(activities=[ProfilerActivity.CUDA]) as prof:
11    add_relu(a, b)
12    torch.cuda.synchronize()
13
14print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

torch.cuda.synchronize() sadece, zamanlamaları okumadan önce GPU'nun işini bitirdiğinden emin olmak içindir. GPU işi eşzamansız olarak çalışır ve senkronizasyon olmadan bazen asıl kernel işi yerine başlatma yükünü ölçersiniz.

Adım 2: çıktıyı okuyun.

Gerçek profil oluşturucu çıktınız, beklediğinizden daha fazla satıra sahip olacaktır. Bir sürü bellek ayırma ve PyTorch defter tutma satırı karışmış olacaktır. Önemsediğimiz satırlar, CUDA kernel'leri, GPU'nun çalıştırdığı asıl fonksiyonlardır. İsminde kernel geçen satırları arayın. İkisi kabaca şöyle görünecektir:

1vectorized_elementwise_kernel<...BinaryFunctor_add...>    1    12us
2vectorized_elementwise_kernel<...threshold_kernel...>     1     9us

Tam şablon adları PyTorch sürümleri arasında değişir (relu genellikle threshold olarak görünür çünkü temel işlem budur ve toplama bazen CUDAFunctor_add olarak görünür). Her şeyi ayrıştırmaya çalışmayın. Sadece satırları sayın. İki satır. İki kernel. Bir bölüm önce olacağını söylediğimiz şeyin aynısı: biri toplama, biri relu için.

Adım 3: derlenmiş sürüm.

Bir satır değişiklik. Fonksiyonu torch.compile ile sarın:

1compiled = torch.compile(add_relu)

Şimdi, çalıştırmadan önce bilmenizde fayda olan bir püf nokta var. torch.compile ile derlenmiş bir fonksiyona ilk çağrı yavaştır, bazen garip bir şekilde yavaştır, çünkü torch.compile işini o zaman yapar: kodunuzu analiz eder, neyin birleştirileceğini bulur, birleştirilmiş kernel'i üretir. İlk çağrıyı profillerseniz, kernel'i değil, derleme adımını ölçüyorsunuzdur. Bu yüzden kalıp şudur: bir kez ısıtmak için çağırın, sonucu atın, ardından profilleme yapın.

1compiled(a, b)                # ısıtma, at
2torch.cuda.synchronize()
3
4with profile(activities=[ProfilerActivity.CUDA]) as prof:
5    compiled(a, b)
6    torch.cuda.synchronize()
7
8print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Adım 4: çıktıyı tekrar okuyun.

Bu sefer:

1triton_poi_fused_add_relu_0    1    14us

Bir satır. Bir kernel. İsim size ne yaptığını bile söylüyor: birleştirilmiş toplama ve relu. Öncekiyle aynı matematik, iki yerine bir başlatma.

Bu makalenin bahsettiği şeyi, tek bir cümlede yaptınız: PyTorch'tan iki işlemi tek bir kernel'de birleştirmesini istediniz, profil oluşturucuyu izlediniz ve sayının ikiden bire düştüğünü doğruladınız. Füzyon, vahşi doğada, sizin makinenizde.

Daha çarpıcı görmek isterseniz, birkaç farklı tensor boyutunda deneyin. 100 uzunluğunda, her iki sürüm de o kadar hızlı çalışır ki fark gürültüde kaybolur. 10 milyon veya 100 milyon uzunluğunda, derlenmiş sürüm açıkça öne geçmeye başlar, çünkü kestiğimiz gidiş-dönüş, bu ölçekteki işin gerçek bir parçasıdır.

Kernel saymak artık soyut bir tavsiye değil. Kontrol etmenin bir yolunuz var.

Toparlayalım

İşte her şey tek seferde.

PyTorch kodunuz, bir GPU'da çalıştığında, bir dizi kernel'e dönüşür. Her kernel bir başlatma, verileriniz üzerinde bir geçiş, girdileri almak ve çıktıları yazmak için GPU belleği üzerinden bir gidiş-dönüş yolculuğudur. Basit işlemler tek bir kernel olur. İşlem zincirleri, varsayılan olarak, aralarında ara değerlerin bellek üzerinden gidiş-dönüş yaptığı işlem başına bir kernel olur. torch.compile, bu ara değerlerin belleğe asla dokunmaması için basit zincirleri sizin için birleştirebilir. Daha az kernel genellikle daha az bellek trafiği anlamına gelir. Ve daha az bellek trafiği genellikle daha hızlı anlamına gelir.