L'obiettivo a lungo termine è rendere Shardy un componente completamente autonomo, in grado di lavorare con qualsiasi dialetto MLIR. Attualmente, Shardy dipende direttamente da StableHLO, ma stiamo facendo progressi verso la rimozione di questa dipendenza tramite varie astrattizioni e interfacce per rendere Shardy più flessibile.
Regole di suddivisione in parti
Una regola di sharding codifica il modo in cui propaghiamo un'operazione. Poiché ora Shardy dipende da StableHLO, definisce le regole di sharding per ogni operazione stablehlo. Inoltre, Shardy fornisce il comando ShardingRuleOpInterface che può essere utilizzato dai proprietari di dialetti nelle loro operazioni per definire le regole di sharding per le proprie operazioni. Finché un'operazione implementa questa interfaccia,
Shardy potrà propagarsi al suo interno.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Operazioni di flusso di dati
Alcune operazioni, ad esempio quelle basate su regioni, richiedono un approccio diverso in cui le regole di sharding, che descrivono solo la corrispondenza tra le dimensioni in tutti gli operandi e i risultati, non sono sufficienti. In questi casi, Shardy definisce un
ShardableDataFlowOpInterface
in modo che i proprietari dei dialetti possano descrivere la propagazione dello sharding tramite le loro operazioni. Questa interfaccia fornisce metodi per recuperare le origini e le destinazioni di ogni angolo del flusso di dati tramite il relativo proprietario, nonché per recuperare e impostare gli shard dei proprietari di angolo.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Consulta anche Operazioni di flusso di dati per una panoramica generale di come gestiamo le operazioni di flusso di dati.
Interfacce non ancora implementate
In futuro verranno aggiunte altre interfacce e caratteristiche per rendere Shardy più flessibile e indipendente dal dialetto. Li elenchiamo di seguito.
Suddivisione costante
La maggior parte dei programmi tensori in MLIR ha un'istanza di una costante che viene riutilizzata da qualsiasi operazione che ne ha bisogno. Questo ha senso quando la costante necessaria è la stessa. Tuttavia, per uno sharding ottimale di un programma, vorremmo consentire a ogni utilizzo di una costante di avere il proprio sharding e di non essere influenzato dal modo in cui altre operazioni utilizzano la costante.
Ad esempio, nella figura seguente, se add è suddiviso in parti, questo non dovrebbe influire sul modo in cui divide e subtract (in parti diverse del calcolo) sono suddivisi in parti.
Si tratta di una dipendenza falsa: poiché le costanti sono economiche, non esiste una dipendenza reale tra le operazioni che utilizzano la stessa costante. Di conseguenza, gli utenti possono decidere lo sharding delle operazioni costanti (e simili). Ogni utilizzo di questa costante può avere un'organizzazione in parti diverse che può propagarsi in modo isolato alla propria copia del sottocalcolo costante.
Per farlo, gli utenti di Shardy devono definire: - Un passaggio your_dialect.constant ->
sdy.constant; - Un tratto sdy::ConstantLike, ad esempio
iota; - Un tratto mlir::Elementwise
per operazioni elementari come add e
multiply; - Un
sdy::ConstantFoldable per operazioni come
slice/broadcast.
Tecnicamente, queste operazioni possono essere calcolate in fase di compilazione, se tutti i loro operandi/risultati sono costanti.
Priorità operative
In GSPMD, le operazioni elementari vengono propagate per prime, seguite da operazioni come matmul.
In Shardy, vogliamo consentire agli utenti di impostare le proprie priorità di operazione, poiché non siamo a priori a conoscenza dei loro dialetti. Pertanto, gli chiederemo di passare un elenco di operazioni nell'ordine in cui vuole che Shardy le propaghi.
La figura seguente mostra come le priorità vengono utilizzate in GSPMD per propagare le operazioni nell'ordine corretto.
Consulta il documento GSMPD per una discussione sul perché le priorità op sono importanti.
Essere indipendente dal dialetto
Purché tu implementi le interfacce, i tratti e i pass precedenti, Shardy potrà lavorare per il tuo dialetto. Stiamo lavorando per rendere Shardy più flessibile e indipendente dal dialetto, quindi continua a seguirci per altri aggiornamenti.