L'objectif à long terme est de faire de Shardy un composant complètement autonome, capable de fonctionner avec n'importe quel dialecte MLIR. Actuellement, Shardy dépend directement de StableHLO, mais nous progressons vers la suppression de cette dépendance via diverses abstractions et interfaces pour rendre Shardy plus flexible.
Règles de segmentation
Une règle de partitionnement encode la manière dont nous propageons une opération. Étant donné que Shardy dépend désormais de StableHLO, il définit des règles de partitionnement pour chaque opération stablehlo. De plus, Shardy fournit ShardingRuleOpInterface, qui peut être utilisé par les propriétaires de dialectes dans leurs opérations pour définir des règles de partitionnement pour leurs propres opérations. Tant qu'une opération implémente cette interface, Shardy peut s'y propager.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Opérations de flux de données
Certaines opérations, par exemple les opérations basées sur les régions, nécessitent une approche différente, où les règles de fractionnement, qui ne décrivent que la correspondance entre les dimensions pour tous les opérandes et les résultats, ne suffisent pas. Dans ce cas, Shardy définit un ShardableDataFlowOpInterface afin que les propriétaires de dialectes puissent décrire la propagation du partitionnement via leurs opérations. Cette interface fournit des méthodes permettant d'obtenir les sources et les cibles de chaque arête de flux de données via leur propriétaire, ainsi que d'obtenir et de définir les fractionnements des propriétaires d'arêtes.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Consultez également Opérations de flux de données pour obtenir un aperçu général de la façon dont nous gérons les opérations de flux de données.
Interfaces pas encore implémentées
À l'avenir, d'autres interfaces et traits seront ajoutés pour rendre Shardy plus flexible et indépendant des dialectes. Nous les listons ci-dessous.
Division constante
La plupart des programmes de tenseur dans MLIR comportent une instance d'une constante réutilisée par toute opération qui a besoin de cette valeur. Cela est logique lorsque la constante requise est la même. Toutefois, pour un fractionnement optimal d'un programme, nous souhaitons permettre à chaque utilisation d'une constante d'avoir son propre fractionnement et de ne pas être affectée par la façon dont d'autres opérations utilisent cette constante.
Par exemple, dans l'image ci-dessous, si add est partitionné, cela ne devrait pas affecter la façon dont divide et subtract (dans différentes parties du calcul) sont partitionnés.
Nous appelons cela une fausse dépendance: comme les constantes sont peu coûteuses, il n'existe pas de véritable dépendance entre les opérations qui utilisent la même constante. Par conséquent, les utilisateurs peuvent décider du fractionnement de leurs opérations constantes (et semblables à des constantes). Chaque utilisation de cette constante peut ensuite avoir un fractionnement différent qui peut se propager de manière isolée à sa propre copie du sous-calcul constant.
Pour ce faire, les utilisateurs de Shardy doivent définir: - Un transfert your_dialect.constant -> sdy.constant ; - Un trait sdy::ConstantLike, tel que iota ; - Un trait mlir::Elementwise pour les opérations par élément, comme add et multiply ; - Un sdy::ConstantFoldable pour les opérations telles que slice/broadcast.
Ces opérations peuvent techniquement être calculées au moment de la compilation, si tous leurs opérandes/résultats sont des constantes.
Priorités d'opération
Dans GSPMD, les opérations élémentaires sont propagées en premier, suivies d'opérations telles que matmul.
Dans Shardy, nous souhaitons permettre aux utilisateurs de définir leurs propres priorités d'opérations, car nous ne connaissons pas leurs dialectes a priori. Nous leur demanderons donc de transmettre une liste d'opérations dans l'ordre dans lequel Shardy doit les propager.
La figure ci-dessous montre comment les priorités sont utilisées dans GSPMD pour propager les opérations dans le bon ordre.
Consultez l'article GSPMD pour en savoir plus sur l'importance des priorités d'opération.
Indépendance par rapport aux dialectes
Tant que vous implémentez les interfaces, les traits et le pass précédents, Shardy pourra fonctionner pour votre dialecte. Nous mettons tout en œuvre pour rendre Shardy plus flexible et indépendant des dialectes. Tenez-vous informé des nouveautés à venir.