Shardy è un sistema di partizione di tensori basato su MLIR per tutti i dialetti. Realizzato grazie alla collaborazione dei team GSPMD e PartIR, incorpora il meglio di entrambi i sistemi e l'esperienza condivisa di team e utenti.
Vantaggi
- Maggiore controllo e prevedibilità per gli utenti grazie alla combinazione della propagazione di GSPMD con il partizionamento incrementale di PartIR.
- Nuove funzionalità basate sull'esperienza condivisa, ad esempio il nuovo supporto per le ristrutturazioni che notoriamente generano comunicazioni aggiuntive, a meno che gli utenti non sappiano come gestirle.
- Usabilità e possibilità di debug migliori per aumentare la velocità dell'utente finale, ad esempio utilizzando una rappresentazione dello sharding basata sugli assi.
- Una base di codice open source semplice che utilizza MLIR, con un insieme più ampio di collaboratori attivi (interni, esterni e in vari fusi orari) per supportare gli utenti.
Componenti
- Rappresentazione dello sharding: una rappresentazione dello sharding basata sugli assi associata a una mesh logica specifica (tra più mesh potenziali) e supporta gli assi e gli sharding delle dimensioni con limitazioni, la suddivisione degli assi per operazioni come la trasformazione, le priorità per la partizione incrementale e altro ancora.
- API del compilatore: un insieme di componenti del compilatore che possono essere utilizzati insieme alla rappresentazione dello sharding per influenzare la propagazione dello sharding.
- Sharding di input/output: associa uno sharding a un input o un output della funzione principale per indicare in che modo il tensore di input/output deve essere suddiviso quando viene assegnato alla funzione o restituito dalla funzione.
- Vincolo di suddivisione: associa una suddivisione a un tensore intermedio (ad es. il risultato di una moltiplicazione matriciale) per indicare in che modo deve essere suddiviso il tensore o un sottoinsieme dei relativi utilizzi.
- Shard As/Like: raggruppa più tensori in base a un ID per indicare che devono essere suddivisi nello stesso modo.
- Calcolo manuale: racchiude un sottocalcolo suddiviso manualmente utilizzando un sottoinsieme di assi della maglia, in cui gli shard lungo questi assi manuali sono specificati per tutti gli input e le uscite e all'interno del sottocalcolo i tipi di tensori sono locali rispetto a questi shard.
- Propagazione dello sharding: un algoritmo di propagazione che combina le priorità degli utenti e i vincoli di sharding con i modelli di costo e le strategie di compilatore:
- Priorità definite dall'utente, ad esempio parallelismo batch e poi ZeRO
- Priorità basate sulle operazioni, ad esempio operazioni elementari prima, poi moltiplicazioni matriciali e così via.
- Ipereuristiche più granulari, ad esempio preferiscono le dimensioni batch.
- Partizionatore SPMD: un componente che riduce le decisioni di propagazione dello sharding suddividendo il programma in un programma SPMD, aggiungendo le operazioni di movimento/formattazione dei dati e collettive necessarie durante il processo.
- A breve termine, l'implementazione iniziale utilizzerà il partitioning SPMD attuale di GSPMD.
- A lungo termine, prevediamo di creare un nuovo partizionatore SPMD basato su MLIR.
Repository di codice
Il progetto Shardy è in fase di sviluppo attivo e stiamo cercando feedback dalla community open source. Il codice di Shardy è disponibile all'indirizzo https://github.com/openxla/shardy