O Shardy é um sistema de particionamento de tensor baseado em MLIR para todos os dialetos. Criada com a colaboração das equipes do GSPMD e do PartIR, ela incorpora o melhor dos dois sistemas e a experiência compartilhada de equipes e usuários.
Vantagens
- Mais controle e previsibilidade para os usuários ao combinar a propagação do GSPMD com o particionamento incremental do PartIR.
- Novos recursos motivados por experiências compartilhadas, por exemplo, suporte inovador para reformulações que geram comunicação extra, a menos que os usuários saibam como trabalhar com elas.
- Melhor usabilidade e capacidade de depuração para aumentar a velocidade do usuário final, por exemplo, usando uma representação de fragmentação baseada em eixo.
- Uma base de código simples e de código aberto que usa o MLIR, com um conjunto mais amplo de colaboradores ativos (internos, externos e em várias zonas horárias) para ajudar os usuários.
Componentes
- Representação de divisão: uma representação de divisão baseada em eixos que é vinculada a uma malha lógica específica (de várias malhas) e oferece suporte a divisões e eixos de restrição de dimensões, dividindo eixos para operações como remodelar, prioridades para particionamento incremental e muito mais.
- APIs do compilador: um conjunto de componentes de compilador que podem ser usados com a
representação de fragmentação para influenciar a propagação de fragmentação.
- Fragmentações de entrada/saída: associe uma fragmentação a uma entrada ou saída da função principal para indicar que é assim que o tensor de entrada/saída precisa ser fragmentado quando fornecido ou retornado pela função.
- Restrição de fragmentação: anexe uma fragmentação a um tensor intermediário (por exemplo, o resultado de um matmul) para indicar como esse tensor ou um subconjunto de usos dele deve ser fragmentado.
- Fragmentar como/como - agrupar vários tensores por um ID para indicar que eles precisam ser fragmentados da mesma maneira.
- Cálculo manual: inclui uma subcomputação que é particionada manualmente usando um subconjunto de eixos de malha, em que os particionamentos ao longo desses eixos manuais são especificados para todas as entradas e saídas. Dentro da subcomputação, os tipos de tensor são locais em relação a esses particionamentos.
- Propagação de fragmentação: um algoritmo de propagação que combina prioridades do usuário
e restrições de fragmentação com modelos de custo e heurísticas do compilador:
- Prioridades definidas pelo usuário, por exemplo, fazer paralelismo em lote e depois ZeRO
- Prioridades baseadas em operações, por exemplo, operações elementais primeiro, depois matmuls etc.
- Heurísticas mais detalhadas, por exemplo, preferência por dimensões de lote.
- Particionador SPMD: um componente que reduz as decisões de propagação de fragmentação
particionando o programa em um programa SPMD, adicionando o movimento/formatação de dados
necessário e as operações coletivas no processo.
- No curto prazo, a implementação inicial vai usar o particionador GSPMD SPMD atual.
- No longo prazo, planejamos criar um novo particionador SPMD baseado em MLIR.
Repositório do código
O projeto Shardy está em desenvolvimento ativo, e estamos buscando feedback da comunidade de código aberto. O código Shardy está disponível em https://github.com/openxla/shardy