El objetivo a largo plazo es hacer de Shardy un componente completamente independiente, capaz de trabajar con cualquier dialecto de MLIR. Actualmente, Shardy depende directamente de StableHLO, pero estamos avanzando para eliminarlo a través de varias abstracciones e interfaces para que Shardy sea más flexible.
Reglas de fragmentación
Una regla de fragmentación codifica cómo propagamos a través de una operación. Dado que Shardy ahora depende de StableHLO, define reglas de fragmentación para cada operación de StableHLO. Además, Shardy proporciona ShardingRuleOpInterface, que los propietarios de dialectos pueden usar en sus operaciones para definir reglas de fragmentación para sus propias operaciones. Siempre que una operación implemente esta interfaz,
Shardy podrá propagarse a través de ella.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Operaciones de flujo de datos
Algunas operaciones, p.ej., las operaciones basadas en regiones, requieren un enfoque diferente en el que las reglas de fragmentación, que solo describen la correspondencia entre las dimensiones en todos los operandos y resultados, no son suficientes. En estos casos, Shardy define un ShardableDataFlowOpInterface para que los propietarios de dialectos puedan describir la propagación del particionamiento a través de sus operaciones. Esta interfaz proporciona métodos para obtener las fuentes y los destinos de cada borde del flujo de datos a través de su propietario, y también para obtener y configurar los particionamientos de los propietarios de los bordes.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Consulta también Operaciones de flujo de datos para obtener una descripción general de alto nivel sobre cómo controlamos las operaciones de flujo de datos.
Interfaces aún no implementadas
En el futuro, se agregarán más interfaces y atributos para que Shardy sea más flexible y no dependa del dialecto. A continuación, te las mencionamos.
División constante
La mayoría de los programas de tensores en MLIR tienen una instancia de una constante que reutiliza cualquier operación que necesite ese valor. Esto tiene sentido cuando la constante necesaria es la misma. Sin embargo, para lograr un fragmentación óptima de un programa, nos gustaría permitir que cada uso de una constante tenga su propio fragmentación y no se vea afectado por la forma en que otras operaciones usan esa constante.
Por ejemplo, en la siguiente imagen, si add está particionado, no debería afectar la forma en que se particionan divide y subtract (en diferentes partes del procesamiento).
Llamamos a esto una dependencia falsa: como las constantes son económicas, no hay una dependencia real entre las operaciones que usan la misma constante. Por lo tanto, los usuarios pueden decidir el fragmentación de sus operaciones constantes (y similares a constantes). Cada uso de esa constante puede tener un fragmento diferente que se puede propagar de forma aislada a su propia copia del subcálculo constante.
Para lograrlo, los usuarios de Shardy deben definir lo siguiente: - Un pase your_dialect.constant -> sdy.constant; - Un atributo sdy::ConstantLike, como iota; - Un atributo mlir::Elementwise para operaciones por elemento, como add y multiply; - Un sdy::ConstantFoldable para operaciones como slice/broadcast.
Estas operaciones se pueden calcular técnicamente en el tiempo de compilación, si todos sus operandos o resultados son constantes.
Prioridades de las operaciones
En GSPMD, primero se propagan las operaciones por elemento, seguidas de operaciones como matmul.
En Shardy, queremos permitir que los usuarios establezcan sus propias prioridades de operaciones, ya que no conocemos sus dialectos a priori. Por lo tanto, les pediremos que pasen una lista de operaciones en el orden en que quieren que Shardy las propague.
En la siguiente figura, se muestra cómo se usan las prioridades en GSPMD para propagar las operaciones en el orden correcto.
Consulta el documento de GSPMD para obtener una explicación sobre por qué son importantes las prioridades de las operaciones.
Ser independiente del dialecto
Siempre que implementes las interfaces, los atributos y el pase anteriores, Shardy poderá funcionar para tu dialecto. Estamos trabajando para que Shardy sea más flexible y no dependa del dialecto, así que no te pierdas las novedades.