Długoterminowym celem jest uczynienie Shardy całkowicie samodzielnym komponentem, który będzie mógł działać z dowolnym dialektem MLIR. Obecnie Shardy zależy bezpośrednio od StableHLO, ale robimy postępy w eliminowaniu tej zależności dzięki różnym abstrakcjom i interfejsom, które zwiększają elastyczność Shardy.
Reguły podziału na fragmenty
Reguła podziału na fragmenty koduje sposób propagowania przez operację. Ponieważ Shardy zależy teraz od StableHLO, definiuje on reguły podziału na części dla każdej operacji stablehlo. Ponadto Shardy udostępnia ShardingRuleOpInterface, którego właściciele dialektów mogą używać w swoich operacjach do definiowania reguł podziału na części dla własnych operacji. Dopóki operacja wdraża ten interfejs, fragmenty będą mogły się przez niego propagować.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Operacje przepływu danych
Niektóre operacje, np. operacje oparte na regionach, wymagają innego podejścia, w którym reguły dzielenia, które opisują tylko zgodność między wymiarami we wszystkich operandach i wynikach, nie wystarczają. W takich przypadkach Shardy definiuje ShardableDataFlowOpInterface, aby właściciele dialektów mogli opisać propagowanie dzielenia za pomocą swoich operacji. Ten interfejs udostępnia metody umożliwiające pobieranie źródeł i celów każdej krawędzi przepływu danych za pomocą ich właściciela, a także pobieranie i ustawianie podziałów właścicieli krawędzi.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Ogólny opis sposobu, w jaki zarządzamy przepływem danych, znajdziesz w sekcji Data flow ops.
Interfejsy nie zostały jeszcze zaimplementowane
W przyszłości dodamy więcej interfejsów i właściwości, aby Shardy były bardziej elastyczne i niezależne od dialektu. Wymieniliśmy je poniżej.
Stałe dzielenie
Większość programów tensorowych w MLIR ma jedną instancję stałej, która jest używana wielokrotnie przez dowolną operację, która potrzebuje tej wartości. Ma to sens, gdy potrzebna jest stała o tej samej wartości. Jednak aby zapewnić optymalne dzielenie programu, chcemy, aby każde użycie stałej miało własny podział i nie było zależne od sposobu, w jaki inne operacje używają tej stałej.
Na przykład na rysunku poniżej, jeśli add jest podzielone na fragmenty, nie powinno to wpływać na sposób dzielenia na fragmenty zmiennych divide i subtract (w różnych częściach obliczeń).
Nazywamy to fałszywą zależnością: ponieważ stałe są tanie, nie ma prawdziwej zależności między operacjami, które używają tej samej stałej. Dzięki temu użytkownicy mogą decydować o dziedziczeniu stałych (i podobnych do nich) operacji. Każde użycie tej stałej może mieć inny podział, który może być propagowany oddzielnie do własnej kopii obliczeń podrzędnych stałej.
Aby to zrobić, użytkownicy Shardy muszą zdefiniować: – your_dialect.constant ->
sdy.constant pass; – sdy::ConstantLike trait, taki jak
iota; – mlir::Elementwise trait
do operacji na elementach, takich jak add i
multiply; –
sdy::ConstantFoldable do operacji takich jak
slice/broadcast.
Te operacje można teoretycznie obliczyć w czasie kompilacji, jeśli wszystkie ich operandy/wyniki są stałymi wartościami.
Priorytety op
W GSPMD operacje elementowe są propagowane najpierw, a następnie operacje takie jak matmul.
W Shardy chcemy umożliwić użytkownikom ustawianie własnych priorytetów op, ponieważ nie znamy ich dialektów a priori. Dlatego poproś o przekazanie listy operacji w kolejności, w jakiej Shardy mają je propagować.
Rysunek poniżej pokazuje, jak priorytety są używane w GSPMD do propagowania operacji w odpowiedniej kolejności.
Aby dowiedzieć się, dlaczego priorytety są ważne, zapoznaj się z artykułem GSPMD.
Niezależność od dialektu
Dopóki implementujesz poprzednie interfejsy, cechy i passy, Shardy będzie działać na Twoim dialekcie. Pracujemy nad tym, aby Shardy były bardziej elastyczne i niezależne od dialektu. Wkrótce podamy więcej informacji na ten temat.