O objetivo a longo prazo é tornar o Shardy um componente totalmente independente, capaz de funcionar com qualquer dialeto MLIR. No momento, o Shardy depende diretamente do StableHLO, mas estamos fazendo progressos para melhorar isso com várias abstrações e interfaces para tornar o Shardy mais flexível.
Regras de fragmentação
Uma regra de fragmentação codifica como propagamos uma operação. Como o Shardy
agora depende do StableHLO, ele define regras de fragmentação para cada operação de stablehlo. Além
disso, o Shardy fornece o
ShardingRuleOpInterface,
que pode ser usado pelos proprietários de dialeto nas operações para definir regras de fragmentação
para as próprias operações. Enquanto uma operação implementar essa interface,
o Shardy poderá se propagar por ela.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Operações de fluxo de dados
Algumas operações, como as baseadas em região, exigem uma abordagem diferente, em que as regras de fragmentação, que descrevem apenas a correspondência entre dimensões em todos os operandos e resultados, não são suficientes. Nesses casos, o Shardy define um
ShardableDataFlowOpInterface
para que os proprietários de dialeto possam descrever a propagação do sharding nas
operações. Essa interface fornece métodos para receber as origens e os destinos de cada borda
de fluxo de dados pelo proprietário e também para receber e definir os shardings dos proprietários
de borda.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Consulte também Operações de fluxo de dados para ter uma visão geral de alto nível de como processamos as operações de fluxo de dados.
Interfaces ainda não implementadas
No futuro, mais interfaces e características serão adicionadas para tornar o Shardy mais flexível e independente do dialeto. Confira a lista abaixo.
Divisão constante
A maioria dos programas de tensor no MLIR tem uma instância de uma constante que é reutilizada por qualquer operação que precise desse valor. Isso faz sentido quando a constante necessária é a mesma. No entanto, para o particionamento ideal de um programa, gostaríamos de permitir que cada uso de uma constante tenha o próprio particionamento e não seja afetado pela forma como outras operações usam essa constante.
Por exemplo, na figura abaixo, se o add for dividido em fragmentos, isso não afetará
como o divide e o subtract (em diferentes partes da computação) são
divididos em fragmentos.
Chamamos isso de dependência falsa: como as constantes são baratas, não há uma dependência real entre as operações que usam a mesma constante. Assim, os usuários podem decidir sobre o sharding das operações constantes (e semelhantes a constantes). Cada uso dessa constante pode ter um sharding diferente que pode se propagar de forma isolada para a própria cópia da subcomputação constante.
Para isso, os usuários do Shardy precisam definir: - Um passe your_dialect.constant ->
sdy.constant. - Uma característica sdy::ConstantLike, como
iota. - Uma característica mlir::Elementwise
para operações element-wise, como add e
multiply. - Um
sdy::ConstantFoldable para operações como
slice/broadcast.
Tecnicamente, essas operações podem ser calculadas no momento da compilação, se todos os
operandos/resultados forem constantes.
Prioridades de operação
No GSPMD, as operações element-wise são propagadas primeiro, seguidas por operações como matmul.
No Shardy, queremos permitir que os usuários definam as próprias prioridades de operação, já que não
sabemos a priori sobre os dialetos deles. Por isso, vamos pedir que eles transmitam uma lista
de operações na ordem em que querem que o Shardy as propague.
A figura abaixo mostra como as prioridades são usadas no GSPMD para propagar operações na ordem correta.
Consulte o artigo da GSPMD para uma discussão sobre por que as prioridades de operação são importantes.
Ser independente de dialeto
Desde que você implemente as interfaces, os traços e a transmissão anteriores, o Shardy poderá funcionar para seu dialeto. Estamos trabalhando para tornar o Shardy mais flexível e independente do dialeto. Fique de olho nas próximas atualizações.