Uzun vadeli hedefimiz, Shardy'yi herhangi bir MLIR lehçesiyle çalışabilen tamamen bağımsız bir bileşen haline getirmektir. Şu anda Shardy doğrudan StableHLO'ya bağlıdır ancak Shardy'ı daha esnek hale getirmek için çeşitli soyutlamalar ve arayüzler aracılığıyla bu bağımlılığı kaldırma konusunda ilerleme kaydediyoruz.
Bölme Kuralları
Bölme kuralı, bir işlemde nasıl dağıtım yaptığımızı kodlar. Shardy artık StableHLO'ya bağlı olduğundan her stablehlo işlemi için bölme kurallarını tanımlar. Ayrıca, lehçe sahipleri kendi işlemleri için bölme kurallarını tanımlamak üzere işlemlerinde kullanabilecekleri ShardingRuleOpInterface öğesini sağlar. Bir işlem bu arayüzü uyguladığı sürece Shardy bu arayüz üzerinden dağıtım yapabilir.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Veri akışı işlemleri
Bölgeye dayalı işlemler gibi bazı işlemler, yalnızca tüm operatörler ve sonuçlardaki boyutlar arasındaki eşlemeyi açıklayan bölme kurallarının yeterli olmadığı farklı bir yaklaşım gerektirir. Bu durumlarda Shardy, lehçe sahiplerinin operasyonları aracılığıyla bölme işleminin yayılmasını tanımlayabilmesi için bir ShardableDataFlowOpInterface tanımlar. Bu arayüz, her veri akışı kenarının kaynaklarını ve hedeflerini sahibi aracılığıyla alma ve ayrıca kenar sahiplerinin parçalanmasını alma ve ayarlama yöntemleri sağlar.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Veri akışı işlemlerini nasıl ele aldığımızla ilgili üst düzey bir genel bakış için Veri akışı işlemleri bölümüne de bakın.
Henüz uygulanmamış arayüzler
Gelecekte Shardy'yi daha esnek ve lehçeden bağımsız hale getirmek için daha fazla arayüz ve özellik eklenecektir. Bunları aşağıda bulabilirsiniz.
Sabit bölme
MLIR'deki çoğu tensör programında, bu değere ihtiyaç duyan tüm işlemler tarafından yeniden kullanılan bir sabit örneği bulunur. Bu, gerekli sabit aynı olduğunda anlamlıdır. Ancak bir programın optimum şekilde bölümlendirilmesi için bir sabitin her kullanımının kendi bölümlendirmesine sahip olmasına ve diğer işlemlerin bu sabiti nasıl kullandığından etkilenmemesine izin vermek isteriz.
Örneğin, aşağıdaki şekilde add parçalara ayrılırsa bu, divide ve subtract'nin (hesaplamanın farklı bölümlerinde) nasıl parçalara ayrıldığını etkilemez.
Buna yanlış bağımlılık denir: Sabitler ucuz olduğundan aynı sabit değerin kullanıldığı işlemler arasında gerçek bir bağımlılık yoktur. Bu nedenle, kullanıcılar sabit (ve sabit benzeri) işlemlerinin bölümlendirilmesine karar verebilir. Bu sabitin her kullanımı, sabit alt hesaplamanın kendi kopyasına ayrı olarak yayılabilen farklı bir bölme işlemine sahip olabilir.
Bunu başarmak için Shardy kullanıcılarının şunları tanımlaması gerekir: - your_dialect.constant -> sdy.constant geçişi; - iota gibi bir sdy::ConstantLike özelliği; - add ve multiply gibi öğe bazında işlemler için bir mlir::Elementwise özelliği; - slice/broadcast gibi işlemler için bir sdy::ConstantFoldable.
Tüm operatörleri/sonuçları sabitse bu işlemler teknik olarak derleme zamanında hesaplanabilir.
Operasyon öncelikleri
GSPMD'de önce öğe bazında işlemler, ardından matmul gibi işlemler dağıtılır.
Shardy'de, lehçeleri hakkında önceden bilgi sahibi olmadığımız için kullanıcıların kendi işlem önceliklerini belirlemelerine izin vermek istiyoruz. Bu nedenle, Shardy'nin dağıtmasını istedikleri sırayla bir işlem listesi iletmelerini isteyeceğiz.
Aşağıdaki şekilde, GSPMD'de önceliklerin doğru sırada işlem yaymak için nasıl kullanıldığı gösterilmektedir.
Üst önceliklerin neden önemli olduğuyla ilgili bir tartışma için GSPMD makalesine bakın.
Lehçeden bağımsızdır.
Önceki arayüzleri, özellikleri ve geçişi uyguladığınız sürece Shardy, lehçeniz için çalışabilir. Shardy'yi daha esnek ve lehçeden bağımsız hale getirmek için çalışıyoruz. Gelişmeler için takipte kalın.