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 belgede, 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ölme yayma, bir programdaki her tenör için bir bölme işlemine karar verme işlemidir. Bu işlemde, tenörlerin bir alt kümesi için bölme kısıtlamaları kullanılır. Shardy'nin derleyici API'si, bölme yayılımını etkilemenin/kontrol etmenin çeşitli yollarını sunar. Ayrıca kullanıcıların manuel olarak parçalara ayrılmış 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 HİÇBİR ş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) bir parçalama ekleyerek söz konusu tenörün veya kullanımlarının bir alt kümesinin nasıl parçalanacağını belirtmek için kullanılır.

  • 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ümlendirilen bir alt hesaplamayı kapsar. Bu manuel eksenler boyunca yapılan 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 tenörü alan ve kendisine eklenmiş bir bölme özelliğine sahip bir MLIR işlemidir. İşlem şu şekilde olabilir:

  • Kullanımı yok (boş). Bu, eklenmiş bölme işleminin, 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ümlendirilmesi 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şkili herhangi bir 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 bölmemizi istiyoruz. 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 tenler 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ün parçalamasını senkronize tutmak amacıyla 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 açık bir kontrol sahibi olmak isteyebilir. Örneğin, bazı kullanıcılar toplu matmul'ü derleyiciye ertelemek yerine manuel olarak (ön uç API'den) 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 (diğer adıyla manuel eksenler) açısından yerel/manuel ve belirtilmeyen eksenler (diğer adıyla 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.