Jalapeño MLA kernel'inin HotChips'te sunulması ve ardından SemiAnalysis'in yorumlarının ardından çok fazla tartışma oldu. OpenAI'in donanım ekibi olarak, bu küçük altın külçesine zar zor değindik: AI'ın kernel'lerimizi yazıyor olması ve bunu yaparken kernel'in ne yaptığını satır satır anlamamıza gerçekten gerek olmaması. Apaçık atladığımız şeyler: böyle bir şey nasıl mümkün olabilir? Bunu, daha geleneksel bir kod üretme yöntemine kıyasla doğru şekilde düşünme biçimi nedir? Optimize edilmiş kernel, optimize edilmemiş olan kadar sağlam mı?
Geçmişimden bahsetmem gerekirse, on yılı aşkın süredir hızlandırıcılar için derleyiciler üzerinde çalışıyorum. Mükemmel bir derleyici altyapısı ve üzerinde çalışan harika bir şirketler arası ekip ve çaba olan XLA'yı başlattım. OpenAI'de son 2+ yıldır, AI çağında derleyicilerin nasıl çalışması gerektiğini yeniden kavramsallaştırmaya çalışıyorum. Yeni derleyici formülasyonları mevcut güçlü yönlerden yararlanacak, ancak araç setinde kullanılacak güçlü yeni bir araç olduğunu inkar etmek imkansız.
Bu biraz uzun bir yolculuk olacak, ancak AI'ın bilgisayar programı iyileştirmenin otomasyonu için nasıl kullanıldığını aydınlatmayı umuyorum; yani optimize edici derleme. AI sayesinde, "derleyiciler 2.0" olarak düşündüğümüz bir şeyi deneyimleyebileceğimize inanıyorum. AI'ın önerebilecekleri temelde daha az kısıtlıdır ve önerdikleri, modelin eğitiminin ve bağlamının bir sonucudur; bu, onu "stokastik bir optimize edici" olarak sınıflandırmama yol açıyor – bu zorluklar yaratabilir ancak göreceğimiz gibi aynı zamanda büyük güçlerin de kaynağıdır…
Pek çok akademik araştırma ve endüstriyel uygulama zaten bu yönde ilerliyor ve AI'ın optimize edici derleyici alanına dahil olma potansiyelini hızla ortaya çıkarıyor, ancak bunun geniş kapsamlı bir açıklamasını yapmanın tam zamanı.
Arka Plan
Derleyiciler programları alır ve bu programların çevrilmiş veya iyileştirilmiş sürümlerini üretir.
Programlar, hem girdi hem de çıktı tarafında, bize programların ne anlama geldiğini, ne yapabileceklerini ve yapabilecekleri bu şeyler hakkında nasıl akıl yürüteceğimizi söyleyen anlam bilgisine (semantik) sahiptir.
Derleyiciler üzerinde çalışan bizler, onları saf işlevler gibi düşünürüz – bir veri yapısını alırlar ve karşılık gelen anlam bilgisine sahip olması gereken bir veri yapısı üretirler.
Bazen derleyicilerimiz "düşürme" veya "çeviri"ye odaklanır. Örneğin, C'yi alıp genellikle "daha düşük seviye" olarak kabul ettiğimiz x86-64 assembly'si üretebilirler. Ancak çoğu zaman bu sürecin bir alt parçası olarak sadece çeviriden daha fazlasını yaparlar…
Derleyicilerimiz pratikte "optimize etmeye" odaklanır. Programı temsil eden bir veri yapısını – bizim dilimizde bir "Ara Temsil" (IR) – alabilir ve bu programın daha iyi bir sürümünü üretmeye çalışırlar. Bazen "daha iyi", çalışması için daha az döngü gerektiği anlamına gelir, bazen daha az gereksiz kod içerdiği, bazen de program hakkında "doğru olması gerektiğini" kanıtlayabildiğimiz şeyler için uzmanlaştığı (kısmi değerlendirme) anlamına gelir.
Şimdi, kısaca, LLM'lerin aslen insan metnini bir dilden diğerine çevirmek için yaratıldığını düşünün. Çevirinin açıkça onların uzmanlık alanı olduğu görülüyor. Ve LLM'leri günlük görevlerde kullanımımız sayesinde, onların yeni çözümler yazabildiklerini ve mevcut çözümleri iyileştirebildiklerini de görebiliyoruz. Birçoğumuz kod yazanlar olarak bir LLM'den "bu kod parçacığını optimize et" isteme deneyimine sahibiz ve bunu dikkate değer bir şekilde yapabiliyorlar. (Ancak, kodu doğru bir şekilde optimize ettiklerini bilmemiz gerekiyor, buna geleceğiz!) Bu, sadece LLM'lerin optimize edici bir derleyicide aradığımız yeteneklere sahip olduğunu vurgulamak içindir.
Optimizasyon ve Optimallik
Optimize edici derleyiciler, tahmin edilebileceği gibi, üzerinde çalıştıkları programın optimalliğini belirli bir amaca (genellikle yürütme süresi) göre artırmaya çalışırlar. Bunu genel durumda, keyfi bir program için yapmak o kadar zordur ki, derleyici mühendisleri için tam istihdam teoremi adında bir teorem vardır. (Bunu ancak derleyici mühendisi olmayı seçtikten sonra öğrendim, ama yine de beni rahatlattı!)
"Süper optimize ediciler", optimize edici derleyicilerin harika küçük bir alt alanıdır. Belirli bir program olduğunu ve anlam bilgisi yoluyla ne yaptığını söyleyebildiğimizi hayal edin. Aynı anlam bilgisine sahip en optimal program nedir? Süper optimize edicilerin üstesinden gelmeye çalıştığı şey budur ve bu etkili bir şekilde bir arama problemidir…
Aynı anlam bilgisine sahip en kısa programı bulmaya çalıştığımı ve aday bir programın aynı anlam bilgisine sahip olup olmadığını sormanın bir yolu olduğunu hayal edin. Varsayımsal olarak, her programı amaç sırasına göre sıralayabilir ve aynı anlam bilgisine sahip en küçük olanı seçebilirim.
Ancak, her programı amaç sırasına göre sıralamak oldukça zor görünüyor. 2013 tarihli "STOKE" (Stokastik Süper Optimizasyon) başlıklı en sevdiğim akademik makalelerden biri şunu sordu: "Peki ya programları tekrar tekrar rastgele değiştirirsek, sonunda en iyi programı gözlemler miyiz?" Rastgele bir yürüyüşle (ve orijinal makine öğrenimi arkadaşımız Markov Chain Monte Carlo / Metropolis-Hastings ile) sonunda o optimal programı göreceğinizi öne sürdüler.
Monte Carlo değişiklikleri genellikle aptalcadır (rastgele bir değişiklik seçersiniz), ancak aynı zamanda hızlıdır. LLM'ler ise çok akıllıdır (birçok akıl yürütme token'ı), ancak karşılaştırmalı olarak yavaştır.
Peki ya aptal/hızlı Monte Carlo değişiklikleri yerine, programları hangi yönlere götüreceğini LLM'lerin bulmasını sağlasaydık? Programımızı optimize edilmiş program alanında yürüten, çok zeki bir stokastik optimize edicimiz olurdu.
Optimizasyon için Sezgiler
Bir adım geri gidelim. Tanıdığınız, "kod parçacıklarını en iyi şekilde optimize eden" kişiyi düşünün. Kısaca "Optimize Edici Okan" diyelim. Okan'ın muhtemelen ne tür kod değişikliklerinin meyve verebileceğine dair bir içgüdüsü vardır. Okan muhtemelen işe yarayıp yaramadıklarını görmek için bazı şeyleri dener ve eğer işe yaramazsa geri alır ve başka bir şey dener. Ancak ne tür şeylerin mümkün olduğuna ve derleyiciyi nasıl yenebileceklerine dair bir sezgileri vardır.
Okan'ın sahip olduğu bu sezgiler genellikle derleyicilerin yaptıklarının ötesindedir. Modern optimize edici derleyiciler sonuçlarıyla oldukça etkileyici olsa da, oldukça basit kurallara ve buluşsal yöntemlere dayanırlar. Teknik jargonda, sabit noktaya kadar çalıştırılan yerel bir veri akışı dönüşümü fikrine dayanırlar. Ayrıca aşama sıralaması dikkate alırız; yani A'yı ve ardından B'yi dikkate almak için derleyici boru hatları oluştururuz, ancak bileşik AB problemini dikkate almayız. Zamanlayıcılar ve kayıt ayırıcılar bunun kötü şöhretli bir örneğidir, bileşik zamanlayıcı-kayıt ayırıcı üzerinde birçok doktora tezi denenmiştir (aşama sıralamasını çökertmenin faydalarını elde etmek için), ancak bunları pratikte çalıştırmak zor olmuştur.
Okan'ın uzmanlığının değerli olmasının nedeni budur. Okan genellikle birkaç NP-tam problemi, duruma özel buluşsal yöntemlerle nasıl dengeleyeceğini bilir. Dolayısıyla daha özel bağlam farkındalığı ve duyarlılığı vardır. Okan ayrıca, optimize edici derleyicilerin karlı bir şekilde uygulayamayabileceği teknikleri, özellikle kombinasyon halinde, kullanabilir; örneğin ana hatları belirleme, özel ABI'ler oluşturma veya vektörleştirmeyi etkinleştirmek için dönüşümler veya "keşke derleyicinin bunu yapmanın bir yolu olsaydı" diye söylendiğimiz diğer bir sürü şey.
Şimdi, AI'ın, sahip olduğu akıl yürütme yetenekleri aracılığıyla, küçük bir Okan gibi davranabileceğini düşünün. Neyin meyve vereceği konusunda eşleşmiş bir sezgiye sahip olmayabilir, ancak neyin karlı olabileceğine dair bir ipucuna sahiptir ve çok, çok sayıda deneme yapabilir.
Bu yaklaşımla, STOKE makalesinin aksine, zaman sonsuza gittikçe optimal programı görebileceğimizi garanti edemeyiz, ancak AI'ın "daha insani" akıl yürütme yetenekleri olduğu için, birim zaman başına önemli ölçüde insan benzeri ilerleme kaydedebilir.
Konuyu Bağlamak: MLA Kernel
Şunu söyleyerek başlayayım: AI'ın Jalapeño MLA kernel'imiz için hangi düşük seviyeli kodu ürettiğini bilmiyorum, ancak MLA için numpy'yi nasıl yazacağımı biliyorum.
Daha önce üzerinde çalıştığım XLA derleyicisinde, bu numpy işlemlerini kümeler halinde birleştirir ve ardından bunları döngülere, talimatlara ve daha düşük seviyeli temel öğelere indirgemek için "yayıcı" adı verilen bir meta program kullanırdık.
XLA derleyicisi/yayıcı programı bunu yaptığında, arkadan hangi assembly'nin çıktığını umursamama gerek yoktu. Stokastik optimize edicimiz için AI, kavramsal olarak yayıcı meta programın yerini alır – hem düşürür hem de optimize eder ve ondan giderek daha fazla tavana doğru optimize etmesini isteyebiliriz.
Umarım bu, AI'ın nereye yerleştiğini ve mevcut bir optimize edici derleyici sistemindeki bir bileşene nasıl benzediğini açıklığa kavuşturur. Şöyle düşünmek de faydalıdır: "assembly kodu" olarak kabul ettiğimiz katman şimdi yukarı doğru hareket ediyor. Normal C++ yazıp -O3'te (en yüksek tipik optimizasyon seviyesi) derlediğinizde, yazdığınız C++'ı anlasanız bile ortaya çıkan assembly'yi anlamayı beklemezsiniz. Burada da benzer bir şey yapıyoruz, ancak daha yüksek seviyeli ve daha matematiksel bir girdi tanımıyla.
Şimdi, kilit bir soru, AI'dan aldığımız programın gerçekten de üst düzey tanıma / numpy'ye eşdeğer olup olmadığını nasıl kontrol ettiğimizdir. Bu kontrol mekanizması, AI stokastik optimizasyon sürecinin sağlamlığını tesis eder. Gelecekteki bir blog yazısının bunun hakkında daha fazla ayrıntıya gireceğini tahmin ediyorum, ancak şimdilik, anlamsal eşdeğerlik için test etmenin mümkün olduğunu ve bunu yaptığımızı söylemek yeterli olacaktır. Hızlandırıcı programları, geniş bağlamlarında oldukça matematiksel ve veri akışı odaklı olduklarından, "AI tarafından optimize edilen programın tam olarak yaptığı şey budur" diye doğrulayabileceğimiz güçlü, eksiksiz sözleşmelere özellikle uygundur.
Bu alandaki birçok ilgili tekniğin, program sentezi alt alanındaki çabalar tarafından öncülük edildiğini unutmayın. Optimize edici derleyiciler "işte anlam bilgisine sahip bir program, onu daha iyi hale getir ancak eşdeğer anlam bilgisiyle!" derken, program sentezi "bu anlam bilgisine sahip bir program olduğuna inanıyorum, lütfen bulabileceğin en iyisini bulmaya çalış" der. Program sentezi, optimize edici derlemeden daha zor bir problemdir, ancak aynı zamanda temelde daha az kısıtlıdır. Okan gibi insanların optimize edici derleyiciden daha iyisini yaptıklarında yaptıkları şey budur ve AI'ın artık otomatikleştirmemize yardımcı olabileceği bir şeydir. AI, orijinal programdan "ilham" alabilir, ancak sadece üzerinde küçük yerel dönüşümler yapması gerekmez. Klasik optimize edici derleyiciler "ah, bir kabarcık sıralaması yazmışsın" görmez ve sözleşmeyi anlayarak bunu hızlı sıralamaya çevirmez, ancak hem Okan hem de AI bunu yapabilir. Bizi, klasik optimize edici derleyici rejiminden ziyade, stokastik optimizasyonla program sentezi rejimine daha çok sokan şey budur.
Tüm bunlar, "numpy'dan çok uzak olmayan" bir şeyle başlayabileceğiniz, 48 saat bekleyebileceğiniz ve HotChips konuşmamızda gösterdiğimiz gibi aynı anlam bilgisine sahip optimize edilmiş bir kernel elde edebileceğiniz gerçeğinde birleşiyor:
Slaytın da belirttiği gibi, makinemizde AI'ın, oldukça iyi ayarlandığını düşündüğümüz kernel'lerde bile bile insan uzmanlarımızı geçerek performansı artırdığını sıklıkla gözlemleyebiliyoruz. Keşfedilmesi gerekebilecek çok sayıda kombinasyon/permütasyon çeşidi nedeniyle genellikle hala elde edilebilir makul bir yüzde kalır. Bunlar genellikle bir insan performans mühendisi için dayanılmaz derecede sıkıcıdır.
Özet ve Sonuç
Bir derleyici, günün sonunda, sadece bir işlevdir. Programımızı bu işleve veririz ve programımızın daha iyi bir sürümünü geri alırız. Çıktı olarak aldığımız programın ve girdi olarak verdiğimiz programın aynı anlam bilgisine sahip olmasını bekleriz.
Geleneksel optimize edici derleyiciler, programları veri akışı kuralları ve buluşsal yöntemler aracılığıyla daha iyi hale getirir. Bunların kökeni tamamen anlaşılabilirdir, ancak yapabilecekleri hamleler açısından daha sınırlı olabilirler.
Buna karşılık AI, stokastik bir optimize edici olarak, sadece "çok düşünmek" ve bir şeyler üretmek zorundadır. Hamleleri temelde o kadar sınırlı değildir, bu da onları insan uzman optimize edicimize daha benzer kılar. Ürettiği programların sağlam olduğunu ve girdiğimiz anlam bilgisini uyguladığını kontrol etmenin yollarına ihtiyacımız var ve bunları yerine koyduk. Ve bu tür AI optimizasyonu, çok güçlü sözleşmelere sahip matematiksel işlemler için özellikle uygundur. Sözleşmeler, kernel'in ne yaptığını satır satır anlama ihtiyacını ortadan kaldırır.
AI tarafından üretilen MLA kernel'ini işte böyle elde ettik!
