Derleyici API'si

Arka plan

Okuyucuların, bir tensörün bölümlendirilmesinin Shardy'de nasıl ifade edilebileceğini açıklayan bölümlendirme temsilinin en azından temellerini bildiği varsayılır. Bu dokümanda, bölme temsillerinin bir programda nasıl kullanılabileceği (ör. programın belirli bir tenörüne bölme eklemek için) gösterilmektedir.

Bölümlendirme yayma, bir programdaki her bir tensör için bölümlendirmeye karar verme işlemidir. Bu işlemde, tensörlerin bir alt kümesi için bölümlendirme kısıtlamaları kullanılır. Shardy'nin derleyici API'si, parçalama yayılımını etkilemenin/kontrol etmenin birkaç yolunu sunar. Ayrıca kullanıcıların manuel olarak bölümlenmiş hesaplamaları programlarına eklemesine olanak tanır.

Hedef

Bu dokümanda, Shardy'deki bu tür API bileşenlerinin tasarımı, davranışları ve değişmezlikleri açıklanmaktadır. Bu API, bölme yayma işlemini kontrol etmek için kullanılır. Bu dokümanda, yayma işleminin davranışı veya nasıl tasarlandığı hakkında NOT şey tartışılmayacaktır.

Genel Bakış

  • Giriş/çıkış bölme işlemleri: Ana işlevin 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.

  • Parçalama Kısıtlaması: Bu, bir ara tenöre (ör. bir matmul sonucu) eklenerek söz konusu tenörün veya kullanımlarının bir alt kümesinin nasıl parçalanacağını belirtir.

  • Parçalama Grubu: 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 bölümlendirmeler tüm girişler ve çıkışlar için belirtilir ve alt hesaplama içinde tensor türleri bu bölümlendirmelere göre yereldir.

Ayrıntılı Tasarım

Giriş/çıkış bölme işlemleri

Kullanıcıların ana işlevin giriş ve çıkışları için bir bölme belirlemesine olanak tanır.

MLIR'de özellikler işlev bağımsız değişkenlerine ve sonuçlarına eklenebilir. Bu nedenle kullanıcılar, işleve bölme özelliklerini bu şekilde ekleyebilir.

Örneğin:

@mesh_xy = <["x"=2, "y"=2]>

// The 1st input has a sharding specified, but the 2nd input doesn't.
// The output has a sharding specified.
func @main(%arg0: tensor<8x8xf32>
            {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"}, {}]>},
            %arg1: tensor<8x16xf32>)
    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_xy, [{}, {"y"}]>}) {
  ...
}

Parçalama Kısıtlaması

Kullanıcıların programlarındaki bir ara tenöre bir bölme eklemelerine olanak tanır. Bu, bölme işlemini yapana söz konusu tenörün veya kullanımlarının bir alt kümesinin bu şekilde bölünmesi gerektiğini söyler.

Bu, giriş olarak tenzoru alan ve kendisine eklenmiş bir bölme özelliğine sahip bir MLIR işlemidir. İşlem şu şekilde olabilir:

  • Kullanımı yok (sarkık) demektir. Yani, eklenmiş bölme işlemi, tensörün kendisinin nasıl bölünmesi gerektiği anlamına gelir.
  • Kullanım alanı var: Bu, ekteki bölme işleminin, bölme kısıtlaması işleminin kullanım alanlarının nasıl bölünmesi gerektiği anlamına gelir. Giriş tenörünün diğer kullanım alanları farklı bir bölme işlemine sahip olabilir (Giriş tenörünün başka bir kullanım alanı yoksa davranış, kullanım alanı yok durumuyla aynıdır). Yayma, tenzorun bölümlendirilmesini belirler ve gerekirse yeniden bölümlendirir.

Açık boyut bölmelerine sahip olabilir. Bu, işlenenin mevcut eksenler boyunca daha da bölünebileceği anlamına gelir.

@mesh_xy = <["x"=2, "y"=2]>

%0 = ... : tensor<8x8xf32>
%1 = sdy.sharding_constraint %0 <@mesh_xy, [{"x"}, {?}]> : tensor<8x8xf32>

Parçalama Grubu

İki veya daha fazla tenör arasında veri bağımlılığı ya da güçlü veri bağımlılığı olmadığında, kullanıcılar bu tenlerin aynı veya benzer şekilde bölümlenmesi gerektiğini bilse de Shardy API bu ilişkiyi belirtmenin bir yolunu sunar. Bu sayede kullanıcılar, tensörlerin birbirine göre bölümlendirilmesi gerektiğini açıkça belirtebilir.

Bunu başarmak için her grubun aynı parça grubu kimliğiyle ilişkilendirilmiş herhangi sayıda talimat içerdiği bir parça grubu kavramı sunuyoruz. Bölme grupları, aynı gruptaki bölmelerin aynı olmasını zorunlu kılar.

Örneğin, aşağıda gösterilen gibi varsayımsal bir kullanıcı programında, programın çıkışını programın girişiyle tam olarak aynı şekilde parçalara ayırmak isteriz. Bu iki program arasında veri bağımlılığı yoktur.

Bu programı çalıştırırsak bölme yayma, %1 ve %2 tenörlerinin bölünmesiyle ilgili çıkarım yapamaz ve bu vektörler kopyalanır. Ancak, giriş %0 ve çıkış %2'nin aynı shard_group içinde olduğunu belirten bir shard_group özelliği ekleyerek @mesh_xy, [{"x"},{"y"}]> parçalamasının giriş %0'ten çıkış %2'ye ve ardından burada sabit %1 olarak yayınlanan grafiğin geri kalanına dağıtılmasına izin veririz. sdy.sharding_group işlemi ile bir gruba değer atayabiliriz.

@mesh_xy = <["x"=2, "y"=2]>

module @"jit_zeros_like" {
  func.func @main(%arg0: tensor<8x2xi64> {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"},{"y"}]>} }) -> (tensor<8x2xi64>) {
    %0 = sdy.sharding_group %arg0, id=0 : tensor<8x2xi64>
    %1 = stablehlo.constant dense<0> : tensor<8x2xi64>
    %2 = sdy.sharding_group %1, id=0 : tensor<8x2xi64>
    return %2 : tensor<8x2xi64>
  }
}

Yukarıdaki basit örnekte, alternatif olarak çıkışta girişle aynı parçalamayı açıkça belirtebilirdik. Bu da aynı etkiyi elde etmemizi sağlardı. Çünkü girişe hangi parçayı atamak istediğimizi önceden biliyorduk. Ancak daha gerçekçi durumlarda, Shardy'nin geri kalanını halletmesi ve onlara atanacak en iyi parçalamayı bulması için Shardy'yi kullanırken, birden fazla tenör için parçalamayı bilmek zorunda kalmadan birden fazla tenörün parçalamasını senkronize tutmak için parçalamayı kullanırız.

Manuel Hesaplama

Kullanıcılar, hesaplamalarının bölümlerinin nasıl bölündüğü ve hangi koleksiyonların kullanıldığı konusunda net bir kontrol sahibi olmak isteyebilir. Örneğin, bazı kullanıcılar derleyiciye ertelemek yerine toplu matmul'ü manuel olarak (ön uç API'sinden) uygulamak ister. Bunu yapmalarına olanak tanıyan Manuel Hesaplama API'si sunuyoruz.

Bu, manuel alt hesaplama için tek bir bölgeye sahip MLIR işlemidir. Kullanıcılar, ağ eksenlerinin bir alt kümesini (mümkünse tümü dahil) kullanarak bu alt hesaplama için giriş/çıkış bölmelerini belirtir. Alt hesaplama, belirtilen ağ eksenleri (manuel eksenler) açısından yerel/manuel ve belirtilmeyen eksenler (serbest eksenler) açısından küresel/bölümlenmemiş olur. Alt hesaplama, bu işlemin dışındaki hesaplama gibi dağıtım sırasında serbest eksenler boyunca daha da bölünebilir.

Örneğin:

@mesh_name = <["data"=2, "model"=2]>

%0 = ... : tensor<16x32xf32>
%1 = sdy.manual_computation(%0)
    in_shardings=[<@mesh_name, [{"data"}, {"model",?}]>]
    out_shardings=[<@mesh_name, [{"data"}, {?}]>]
    manual_axes={"data"}
    (%arg1: tensor<8x32xf32>) {
  // body
  return %42 : tensor<8x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

Değişmezler

  1. Tüm in_shardings, out_shardings ve manual_axes aynı ağa referans vermelidir. manual_axes, ağa göre sıralanır.

  2. manual_axes, tüm giriş/çıkış bölmelerinde açıkça kullanılmalıdır. Yani her bölme için tüm manuel eksenler bir boyutu bölmeli veya açıkça çoğaltılmalıdır.

  3. Giriş/çıkış bölmelerinden birinde serbest bir eksen (manual_axes içinde olmayan herhangi bir ızgara ekseni) varsa bu eksen, aynı boyut bölmesindeki herhangi bir manuel eksenden küçük olmalıdır (yukarıdaki örnekte, {"model", "data"} boyut bölmesindeki bir boyut bölme geçersiz olur).

  4. Hesaplamanın bölgesi/gövdesi yerel hesaplamadır (ör. kullanıcı tarafından belirtilen koleksiyonlar dahil). Manuel eksenler boyunca giriş/çıkış bölme işlemine göre yerel olmalıdır (yukarıdaki notu inceleyin).

Manuel hesaplamaları iç içe yerleştirme

Her biri kendi benzersiz manuel eksen grubunda çalıştığı sürece birden fazla manuel hesaplamayı birbirinin içine yerleştirebilirsiniz.