Shardy, tüm lehçeler için MLIR tabanlı bir tensör bölme sistemidir. Hem GSPMD hem de PartIR ekiplerinin ortak çalışmasıyla oluşturulan bu platform, her iki sistemin de en iyi özelliklerini ve hem ekiplerin hem de kullanıcıların ortak deneyimini bir araya getirir.
Avantajları
- GSPMD'nin dağıtımını PartIR'ın artımlı bölümlendirmesiyle birleştirerek kullanıcılar için daha fazla kontrol ve öngörülebilirlik.
- Paylaşılan deneyimden yararlanan yeni özellikler (ör. kullanıcılar bu özellikleri nasıl kullanacağını bilmedikleri sürece fazladan iletişime neden olan yeniden şekillendirmeler için yeni destek).
- Son kullanıcı hızını artırmak için daha iyi kullanılabilirlik ve hata ayıklama özelliği (ör. eksene dayalı bir bölme temsili kullanılarak).
- Kullanıcıları desteklemek için daha geniş bir aktif katkıda bulunanlar grubuna (dahili, harici ve çeşitli saat dilimlerinde) sahip, MLIR kullanan basit ve açık kaynak kod tabanı.
Bileşenler
- Bölme temsili: Belirli bir mantıksal ağa (olası birden fazla ağdan) bağlı olan ve boyut bölmelerini ve eksenleri kısıtlamayı, yeniden şekillendirme gibi işlemler için eksenleri bölme, artımlı bölümleme için öncelikleri ve daha fazlasını destekleyen eksen tabanlı bir bölme temsili.
- Derleyici API'leri: Bölme yayılımını etkilemek için bölme temsiliyle birlikte kullanılabilen bir dizi derleyici bileşeni.
- Giriş/çıkış bölme işlemleri: Ana işlevin bir girişine veya çıkışına bir bölme işlemi ekleyerek, işleve verildiğinde/işlevden döndürüldüğünde giriş/çıkış tenzorunun bu şekilde bölünmesi gerektiğini belirtin.
- Bölme Kısıtlaması: Bir ara tenöre (ör. bir matmul sonucu) bir bölme ekleyerek bu tenörün veya kullanımlarının bir alt kümesinin bu şekilde bölünmesi gerektiğini belirtin.
- Parçalama Şekli/Benzer: Birden fazla tenzoru aynı şekilde parçalanmaları gerektiğini belirtmek için bir kimliğe göre gruplandırın.
- Manuel Hesaplama: Örgü eksenlerinin bir alt kümesi kullanılarak manuel olarak bölümlendirilmiş bir alt hesaplamayı kapsar. Bu manuel eksenler boyunca yapılan bölme işlemleri tüm girişler ve çıkışlar için belirtilir ve alt hesaplama içinde tensor türleri bu bölme işlemlerine göre yereldir.
- Bölme Yayma: Kullanıcı önceliklerini ve bölme kısıtlamalarını derleyici maliyet modelleri ve sezgilerle birleştiren bir yayılma algoritması:
- Kullanıcı tanımlı öncelikler (ör. önce toplu paralellik, ardından ZeRO)
- İşlem tabanlı öncelikler (ör. önce öğe bazında işlemler, ardından matris çarpımları vb.)
- Daha ayrıntılı sezgisel kurallar (ör. toplu boyutları tercih etme).
- SPMD Bölmeleyici: Programı bir SPMD programına bölen, gerekli veri hareketini/biçimlendirmeyi ve toplu işlemleri sürece ekleyerek bölme yayma kararlarını azaltan bir bileşendir.
- İlk uygulamada kısa vadede mevcut GSPMD SPMD bölümleyicisi kullanılacak.
- Uzun vadede, MLIR tabanlı yeni bir SPMD bölümleyici oluşturmayı planlıyoruz.
Kod deposu
Shardy projesi aktif olarak geliştirilmektedir ve açık kaynak topluluğundan geri bildirim almak istiyoruz. Shardy koduna https://github.com/openxla/shardy adresinden ulaşabilirsiniz.