Yayılma

Genel Bakış

Bölme yayma, kullanıcı tarafından belirtilen bölmelere (veya tenzorların belirli boyutlarına) göre belirtilmeyen bölmelerin tahmin edilmesi için kullanıcı tarafından belirtilen bölmelerden yararlanır. Sabit bir noktaya ulaşılana kadar hesaplama grafiğinin veri akışını (kullanım-tanım zincirleri) her iki yönde de tarar. Yani, önceki bölme kararları geri alınmadan bölme artık değiştirilemez.

Yayma işlemi adımlara ayrılabilir. Her adımda, belirli bir işleme bakılır ve bu işlemin özelliklerine göre tensörler (operandlar ve sonuçlar) arasında yayılma işlemi gerçekleştirilir. Bir matmul'ü örnek olarak alırsak sol veya sağ tarafın daralmayan boyutu ile sonucun ilgili boyutu arasında ya da sol ve sağ tarafın daralan boyutu arasında yayılma işlemi yaparız.

Bir işlemin özellikleri, giriş ve çıkışlarındaki karşılık gelen boyutlar arasındaki bağlantıyı belirler ve işlem başına bölümleme kuralı olarak soyutlanabilir.

Çakışma çözümü olmadan bir dağıtım adımı, çakışan eksenleri yoksayarken mümkün olduğunca fazla dağıtım yapar. Buna (en uzun) uyumlu ana bölme eksenleri denir.

Ayrıntılı Tasarım

Çatışma çözümü hiyerarşisi

Hiyerarşide birden fazla çakışma çözümü stratejisi oluşturuyoruz:

1. Kullanıcı tanımlı öncelikler. Bölünme Temsili bölümünde, programın artımlı bölümlenmesine (ör. toplu paralellik -> megatron -> ZeRO bölme) izin vermek için boyut bölmelerine önceliklerin nasıl eklenebileceğini açıklamıştık. Bu, iterasyonlarda dağıtım uygulanarak elde edilir. i iterasyonunda, <=i önceliğine sahip tüm boyut bölmelerini dağıtır ve diğer tümünü yoksayarız. Ayrıca, önceki iterasyonlarda yoksayılmış olsalar bile kullanıcı tanımlı parçalandırmaların daha düşük öncelikli (>i) parçalara bölünmesini geçersiz kılmayacağından emin oluruz.
2. İşlem tabanlı öncelikler. Bölme işlemlerini işlem türüne göre dağıtırız. "Geçiş" işlemleri (ör. öğe bazında işlemler ve yeniden şekillendirme) en yüksek önceliğe sahipken şekil dönüşümü içeren işlemler (ör. nokta ve azaltma) daha düşük önceliğe sahiptir.
3. Agresif yayılma. Bölme işlemlerini agresif bir stratejiyle dağıtın. Temel strateji yalnızca çakışması olmayan parçaları dağıtırken agresif strateji çakışmaları çözer. Daha yüksek agresiflik, olası iletişim pahasına bellek kullanımını azaltabilir.
4. Temel Yayma Hiyerarşideki en düşük dağıtım stratejisidir. Çakışma çözümü yapmaz, bunun yerine tüm işlevler ve sonuçlar arasında uyumlu olan eksenleri dağıtır.

Bu hiyerarşi, iç içe yerleştirilmiş for döngüleri olarak yorumlanabilir. Örneğin, her kullanıcı önceliği için tam işlem önceliği yayımı uygulanır.

İşlem bölme kuralı

Bölme kuralıyla, belirli işlem türleri ve özellikleri hakkında düşünmek zorunda kalmadan bölme işlemlerini operatörlerden sonuçlara veya operatörler arasında dağıtmak için gereken bilgileri gerçek dağıtım algoritmasına sağlayan her işlemin soyut bir versiyonu sunulur. Bu, temel olarak işleme özgü mantığı ortadan kaldırır ve yalnızca dağıtım amacıyla tüm işlemler için paylaşılan bir temsil (veri yapısı) sağlar. En basit haliyle, yalnızca şu işlevi sağlar:

GetOpShardingRule(Operation *) -> OpShardingRuleAttr

Bu kural, benzer kod parçalarını birçok işlemde kopyalamak yerine, dağıtım algoritmasını yalnızca bir kez bu veri yapısına (OpShardingRule) dayalı genel bir şekilde yazmamıza olanak tanır. Böylece, işlemlerde hata veya tutarsız davranış olasılığını büyük ölçüde azaltır.

matmul örneğine dönelim.

Yayma sırasında gerekli olan bilgileri (ör. boyutlar arasındaki ilişkiler) kapsayan bir kodlama, einsum gösterimi biçiminde yazılabilir:

(i, k), (k, j) -> (i, j)

Bu kodlamada her boyut tek bir faktörle eşlenir.

Yayınlama bu eşlemeyi nasıl kullanır? Bir operandın/sonucun boyutu bir eksen boyunca parçalara ayrılırsa yayınlama, bu eşlemede söz konusu boyutun faktörünü arar ve diğer operandları/sonuçları kendi boyutları boyunca aynı faktörle parçalara ayırır. Ayrıca (yinelemeyle ilgili daha önceki tartışmaya tabi olarak) söz konusu eksende bu faktöre sahip olmayan diğer operandları/sonuçları da çoğaltabilir.

Bileşik faktörler: yeniden şekillendirme kurallarının genişletilmesi

Birçok işlemde (ör. matmul) her boyutu tek bir faktörle eşlememiz yeterlidir. Ancak yeniden şekillendirme için yeterli değildir.

Aşağıdaki yeniden şekillendirme işleminde iki boyut tek bir boyutta birleştirilir:

%out = stablehlo.reshape(%in) : (tensor<2x4x32xf32>) -> tensor<8x32xf32>

Burada, girişin hem 0 hem de 1 boyutu, çıktının 0 boyutuna karşılık gelir. Giriş için faktörler vererek başlayalım:

(i,j,k) : i=2, j=4, k=32

Çıkış için aynı faktörleri kullanmak istiyorsak birden fazla faktöre referans vermek üzere tek bir boyuta ihtiyacımız olduğunu görebilirsiniz:

(i,j,k) -> ((ij), k) : i=2, j=4, k=32

Yeniden şekillendirme işlemi bir boyutu bölseydi de aynı işlem yapılabilir:

%out = stablehlo.reshape(%in) : (tensor<8x32xf32>) -> tensor<2x4x32xf32>

Burada,

((ij), k) -> (i,j,k) : i=2, j=4, k=32

Buradaki 8 boyutu temel olarak 2 ve 4 faktörlerinden oluşur. Bu nedenle faktörleri (i,j,k) faktörleri olarak adlandırıyoruz.

Bu faktörler, faktörlerden birine karşılık gelen tam bir boyutun olmadığı durumlarda da kullanılabilir:

%out = stablehlo.reshape(%in) : (tensor<8x4xf32>) -> tensor<2x16xf32>
// ((ij), k) -> (i,(jk)) : i=2, j=4, k=4

Bu örnekte, faktör boyutlarını neden saklamamız gerektiği de vurgulanmıştır. Çünkü bunları ilgili boyutlardan kolayca çıkaramayız.

Temel Yayma Algoritması

Bölme işlemlerini faktörlere göre dağıtma

Shardy'de tensörler, boyutlar ve faktörler hiyerarşisi vardır. Bunlar, verileri farklı düzeylerde temsil eder. Faktör, bir alt boyuttur. Bölme yayma işleminde kullanılan dahili bir hiyerarşidir. Her boyut bir veya daha fazla faktöre karşılık gelebilir. Boyut ile faktör arasındaki eşleme OpShardingRule ile tanımlanır.

Shardy, bölme eksenlerini boyutlar yerine faktörler boyunca dağıtır. Bunu yapmak için aşağıdaki şekilde gösterilen üç adımdan yararlanırız:

1. DimSharding - FactorSharding projesi
2. FactorSharding alanında bölme eksenlerini dağıtma
3. Güncellenmiş DimSharding'ı almak için güncellenmiş FactorSharding'ü projelendirin

Faktörler Arası Bölme Yayılımının Görselleştirilmesi

Bölme yayma sorununu ve algoritmasını görselleştirmek için aşağıdaki tabloyu kullanacağız.

• Her sütun bir faktörü temsil eder. F0, dizini 0 olan faktör anlamına gelir. Bölme işlemlerini faktörlere (sütunlar) dağıtırız.
• Her satır bir tensörü temsil eder. T0, 0 dizine sahip tenzoru ifade eder. Tensörler, belirli bir işlemle ilgili tüm operatörler ve sonuçlardır. Bir satırdaki eksenler çakışamaz. Bir eksen (veya alt eksen), bir tensörü birden çok kez bölmek için kullanılamaz. Açıkça kopyalanan bir eksen, tensörü bölümlemek için kullanılamaz.

Bu nedenle, her hücre bir faktör bölme işlemini temsil eder. Kısmi tensörlerde bir faktör eksik olabilir. C = dot(A, B) için tablo aşağıda verilmiştir. N içeren hücreler, faktörün tensörde olmadığını gösterir. Örneğin, F2, T1 ve T2'de ancak T0'da değildir.

Bölme eksenlerini toplama ve dağıtma

Yayılımı görselleştirmek için aşağıda gösterilen basit bir örneği kullanırız.

1.adım: Her faktör boyunca dağıtılacak eksenleri (diğer adıyla (en uzun) uyumlu ana bölme eksenleri) bulun. Bu örnekte, ["a", "b"] değerini F0 boyunca, ["c"] değerini F1 boyunca ve F2 boyunca hiçbir değeri dağıtmayız.

2.adım: Aşağıdaki sonucu elde etmek için faktör bölmelerini genişletin.

Veri akışı işlemleri

Yukarıdaki dağıtım adımı açıklaması çoğu işlem için geçerlidir. Ancak bazı durumlarda bölümleme kuralının uygun olmadığı durumlar vardır. Bu durumlarda Shardy, veri akışı işlemlerini tanımlar.

Bazı X işleminin veri akışı kenarı, bir kaynak grubu ile bir hedef grubu arasında bir köprü tanımlar. Bu köprü, tüm kaynakların ve hedeflerin aynı şekilde bölümlenmesini sağlar. Bu tür işlemlere örnek olarak stablehlo::OptimizationBarrierOp, stablehlo::WhileOp, stablehlo::CaseOp ve sdy::ManualComputationOp verilebilir. Sonuç olarak, ShardableDataFlowOpInterface'ı uygulayan tüm işlemler veri akışı işlemi olarak kabul edilir.

Bir işlem, birbirine dik olan birden fazla veri akışı kenarı içerebilir. Örneğin:

    y_0, ..., y_n = while (x_0, ..., x_n)
                    ((pred_arg_0,... , pred_arg_n) { ... })
                    ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                    })

Bu işlemde n veri akışı kenarı vardır: i. veri akışı kenarı, x_i, return_value_i kaynakları ile y_i, pred_arg_i, body_arg_i hedefleri arasındadır.

Shardy, bir veri akışı kenarının tüm kaynakları ve hedefleri arasında, kaynaklar operatör, hedefler sonuç ve sdy.op_sharding_rule kimliği normal bir işlemmiş gibi parçalama işlemlerini dağıtır. Yani ileri yayılma, kaynaklardan hedeflere, geri yayılma ise hedeflerden kaynaklara olur.

Kullanıcı tarafından, her veri akışı kenarının kaynak ve hedeflerini sahibi aracılığıyla nasıl alacağını ve ayrıca kenar sahiplerinin bölümlemelerini nasıl alıp ayarlayacağını açıklayan çeşitli yöntemler uygulanmalıdır. Sahip, Shardy'nin dağıtımı tarafından kullanılan veri akışı kenarının kullanıcı tarafından belirtilen hedefidir. Kullanıcı bunu keyfi olarak seçebilir ancak statik olması gerekir.

Örneğin, aşağıda tanımlanan custom_op değeri için:

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

Bu custom_op, veri akışı kenarları için iki türe sahiptir: return_value_i (kaynaklar) ile y_i (hedefler) arasında her biri n kenar ve x_i (kaynaklar) ile body_arg_i (hedefler) arasında n kenar. Bu durumda kenar sahipleri, hedeflerle aynıdır.