Le più grandi aziende nel campo dello sviluppo di machine learning hanno presentato il progetto OpenXLA, volto a sviluppare collaborativamente strumenti per la compilazione e l'ottimizzazione dei modelli per i sistemi di machine learning. Il progetto ha integrato lo sviluppo di strumenti che consentono di unificare la compilazione dei modelli realizzati nei framework TensorFlow, PyTorch e JAX, per un'efficace formazione e esecuzione su diverse GPU e acceleratori specializzati. Hanno partecipato al progetto aziende come Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba e Amazon.
Si prevede che, grazie alla collaborazione delle principali team di ricerca e dei rappresentanti della comunità, si riuscirà a stimolare lo sviluppo dei sistemi di machine learning e risolvere il problema della frammentazione dell'infrastruttura per diversi framework e hardware. OpenXLA consente di implementare un supporto efficace per diverse attrezzature, indipendentemente dal framework su cui è stato creato il modello di machine learning. Si aspetta che, grazie a OpenXLA, si possa ridurre il tempo di addestramento dei modelli, aumentare la capacità, ridurre i ritardi, abbattere i costi delle risorse di calcolo e accelerare il time-to-market.
OpenXLA comprende tre componenti principali, il cui codice è distribuito sotto la licenza Apache 2.0:
- XLA (Accelerated Linear Algebra) — un compilatore che ottimizza i modelli di machine learning per un'esecuzione ad alte prestazioni su diverse piattaforme hardware, comprese GPU, CPU e acceleratori specializzati di vari produttori.
- StableHLO — la specifica e l'implementazione base di un insieme di operazioni ad alto livello (HLO, High-Level Operations) per l'uso nei modelli dei sistemi di machine learning. Funge da strato intermedio tra i framework di machine learning e i compilatori che trasformano il modello per l'esecuzione su hardware specifico. Gli strati per la generazione di modelli nel formato StableHLO sono stati preparati per i framework PyTorch, TensorFlow e JAX. Alla base di StableHLO è utilizzato un insieme di MHLO, che è stato ampliato con il supporto per la serializzazione e la versioning.
- IREE (Intermediate Representation Execution Environment) — compilatore e runtime che trasforma i modelli di machine learning in una rappresentazione intermedia universale, basata sul formato MLIR (Multi-Level Intermediate Representation) del progetto LLVM. Tra le caratteristiche si segnala la possibilità di compilazione anticipata (ahead-of-time), il supporto per la gestione del flusso, la possibilità di utilizzare elementi dinamici nei modelli e l'ottimizzazione per diverse CPU e GPU, con basse spese generali.
I principali vantaggi degli strumenti OpenXLA:
- Raggiungimento di performance ottimali senza la necessità di approfondire nella scrittura di codice specifico per determinati dispositivi. Fornitura di ottimizzazioni pronte, includendo la semplificazione delle espressioni algebraiche, l'assegnazione efficace in memoria, e la programmazione dell'esecuzione tenendo conto della riduzione dei picchi di consumo di memoria e dei sovraccarichi.
- Semplificazione della scalabilità e della parallelizzazione dei calcoli. Gli sviluppatori devono semplicemente aggiungere annotazioni per un sottoinsieme di tensori critici, su cui il compilatore può automaticamente generare codice per calcoli paralleli.
- Garanzia di portabilità attraverso il supporto per diverse piattaforme hardware, come GPU AMD e NVIDIA, CPU basate su architettura x86 e ARM, acceleratori ML TPU Google, IPU AWS Trainium Inferentia, Graphcore e Cerebras Wafer-Scale Engine.
- Supporto per l'integrazione di estensioni con implementazione di funzionalità aggiuntive, come il supporto per la scrittura di primitivi di machine learning profondo utilizzando CUDA, HIP, SYCL, Triton e altri linguaggi per il calcolo parallelo. Possibilità di tuning manuale dei colli di bottiglia nei modelli.
Fonte: opennet.ru
