C'è stata molta discussione dopo la presentazione del kernel MLA Jalapeño a HotChips e i successivi commenti di SemiAnalysis. Come team hardware di OpenAI, abbiamo appena sfiorato questa piccola pepita d'oro: il fatto che l'AI sta scrivendo i nostri kernel e che, quando lo fa, non abbiamo davvero bisogno di capire cosa fa il kernel riga per riga. Abbiamo clamorosamente tralasciato: come è possibile una cosa del genere? Qual è il modo giusto di pensare a questo, rispetto a un metodo più tradizionale di generazione del codice? Il kernel ottimizzato è solido quanto quello non ottimizzato?
Per quanto riguarda il mio background, ho lavorato su compilatori per acceleratori per ben oltre un decennio. Ho iniziato XLA, che è un'infrastruttura di compilazione eccellente con un team e uno sforzi stellari a livello interaziendale. Negli ultimi 2+ anni in OpenAI ho cercato di riconcettualizzare come dovrebbero funzionare i compilatori nell'era dell'AI. Le nuove formulazioni dei compilatori attingeranno ai punti di forza esistenti, ma è impossibile negare che ci sia un nuovo potente strumento da sfruttare nella cassetta degli attrezzi.
Sarà un po' un viaggio, ma spero di illuminare come l'AI viene utilizzata per l'automazione del miglioramento dei programmi informatici; cioè la compilazione ottimizzante. Credo davvero che, tramite l'AI, potremmo sperimentare qualcosa che consideriamo "compilatori 2.0". L'AI è fondamentalmente meno vincolata in ciò che può proporre, e ciò che propone è il risultato dell'addestramento e del contesto del modello; questo mi porta a classificarla come un "ottimizzatore stocastico" – questo può porre delle sfide ma, come vedremo, è anche una fonte di grandi punti di forza...
Molte ricerche accademiche e applicazioni industriali sono già dirette in questa direzione, e stanno rapidamente scoprendo il potenziale del coinvolgimento dell'AI nel regno del compilatore ottimizzante, ma siamo a un punto in cui merita una spiegazione a grandi linee.
Contesto
I compilatori prendono i programmi e restituiscono versioni tradotte o migliorate di quei programmi.
I programmi, sia in input che in output, hanno semantica che ci dice cosa significano i programmi, cosa potrebbero fare e come ragionare su quelle cose che potrebbero fare.
Quelli di noi che lavorano sui compilatori li considerano molto simili a funzioni pure – prendono una struttura dati e restituiscono una struttura dati che dovrebbe avere una semantica corrispondente.
A volte i nostri compilatori si concentrano sull'"abbassamento" o la "traduzione". Ad esempio, possono prendere il C e restituire assembly x86-64, che spesso consideriamo di "livello inferiore". Ma spesso fanno molto più della semplice traduzione come sottoparte di quel processo...
I nostri compilatori, in pratica, si concentrano sull'"ottimizzazione". Possono prendere una struttura dati che rappresenta il programma – nel nostro gergo una "Rappresentazione Intermedia" (IR) – e cercano di produrre una versione migliore di quel programma. A volte "migliore" significa che impiega meno cicli per essere eseguito, a volte significa che avrà meno codice non necessario, a volte significa specializzarsi per cose che possiamo dimostrare "devono essere vere" sul programma (valutazione parziale).
Ora, considera brevemente che gli LLM sono stati originariamente creati per tradurre testo umano da una lingua all'altra. Chiaramente la traduzione è nel loro campo di competenza. E possiamo vedere attraverso il nostro uso degli LLM nelle attività quotidiane che possono anche scrivere nuove soluzioni e migliorare quelle esistenti. Molti di noi programmatori hanno anche esperienza nel chiedere a un LLM "ottimizza questo frammento di codice" e loro notevolmente ci riescono. (Tuttavia, dobbiamo sapere che hanno ottimizzato il codice correttamente, cosa su cui torneremo!) Questo serve semplicemente a evidenziare che gli LLM hanno le capacità che cerchiamo in un compilatore ottimizzante.
Ottimizzazione e Ottimalità
I compilatori ottimizzanti, prevedibilmente, cercano di aumentare l'ottimalità del programma su cui stanno lavorando, secondo qualche obiettivo (di solito il tempo di esecuzione). Questo è così difficile da fare nel caso generale, per un programma arbitrario, che esiste un teorema chiamato il teorema della piena occupazione per gli ingegneri di compilatori. (L'ho scoperto solo dopo aver scelto di diventare un ingegnere di compilatori, ma mi ha comunque dato conforto!)
I "superottimizzatori" sono un meraviglioso sottocampo dei compilatori ottimizzanti. Immagina che esista un dato programma, e possiamo dire cosa fa tramite la semantica. Qual è il programma più ottimale che ha la stessa semantica? Questo è ciò che i superottimizzatori cercano di affrontare, ed è effettivamente un problema di ricerca...
Immagina di cercare di trovare il programma più breve che avesse la stessa semantica, e avessi un modo per chiedere se un programma candidato avesse la stessa semantica. Potrei, ipoteticamente, enumerare ogni programma in ordine obiettivo e scegliere il più piccolo che avesse la stessa semantica.
Tuttavia, enumerare ogni programma in ordine obiettivo sembra piuttosto intrattabile. Uno dei miei articoli accademici preferiti, realizzato nel 2013 intitolato "STOKE" (Superottimizzazione Stocastica), chiedeva: "beh, e se modificassimo casualmente i programmi più e più volte, alla fine osserveremmo il programma migliore?" Proponevano che attraverso una passeggiata casuale (e con il nostro amico OG del machine learning Markov Chain Monte Carlo / Metropolis-Hastings), alla fine avresti visto quel programma ottimale.
La modifica Monte Carlo è tipicamente stupida (scegli una modifica a caso), ma anche veloce. Gli LLM sono molto intelligenti (molti token di ragionamento), ma comparativamente lenti.
E se, invece della modifica Monte Carlo stupida/veloce, avessimo gli LLM a capire le direzioni in cui portare i programmi? Avremmo un ottimizzatore stocastico che fosse molto intelligente, che facesse camminare il nostro programma attraverso lo spazio dei programmi ottimizzati.
Intuizioni per l'Ottimizzazione
Facciamo un passo indietro. Considera la persona che conosci che meglio incarna "ottimizza al massimo i frammenti di codice". Per brevità chiamiamoli "Ottimizzino Ollie". Ollie probabilmente ha un'intuizione viscerale su quali tipi di modifiche al codice potrebbero dare frutti. Ollie probabilmente prova alcune cose per vedere se funzionano, e se non funzionano, le annulla e prova qualcos'altro. Ma hanno una certa intuizione su quali tipi di cose sono possibili e su come potrebbero essere in grado di battere il compilatore.
Queste intuizioni che Ollie ha sono spesso al di là di ciò che fanno i compilatori. Sebbene i moderni compilatori ottimizzanti siano piuttosto impressionanti nei loro risultati, si basano su regole ed euristiche abbastanza semplici. In gergo tecnico, si basano sull'idea di una trasformazione locale del flusso di dati eseguita fino a un punto fisso. Inoltre, ordiniamo in fasi le considerazioni; cioè costruiamo pipeline di compilazione per considerare A e poi B, ma non il problema composito AB. Gli scheduler e gli allocatori di registri sono un esempio notorio di questo, molti dottorati di ricerca sono stati tentati sull'allocatore di registri-scheduler composito (per ottenere i vantaggi del collasso dell'ordinamento delle fasi), ma sono stati difficili da far funzionare nella pratica.
Ecco perché l'esperienza di Ollie è preziosa. Spesso Ollie sa come bilanciare diversi problemi NP-completi con euristiche su misura per la situazione. Quindi c'è una maggiore consapevolezza e sensibilità contestuale su misura. Ollie è anche in grado di impiegare tecniche che i compilatori ottimizzanti potrebbero non applicare in modo redditizio, specialmente in combinazione, cose come l'outlining o la creazione di ABI personalizzate o trasformazioni per abilitare la vettorizzazione, o l'altra serie di cose che ci fanno borbottare "vorrei che il compilatore avesse un modo per fare proprio questo..."
Ora considera che l'AI, attraverso le sue capacità di ragionamento, potrebbe essere in grado di agire come un mini Ollie. Potrebbe non avere l'intuizione abbinata in termini di ciò che darà risultati, ma ha un'idea di ciò che può essere redditizio, e può fare molti, molti tentativi.
Con questo approccio, a differenza dell'articolo STOKE, non possiamo garantire che col tempo tendente all'infinito possiamo vedere il programma ottimale, ma poiché l'AI ha capacità di ragionamento "più simili a quelle umane", può effettivamente ottenere una trazione significativa simile a quella umana per unità di tempo.
Ricollegandoci: Kernel MLA
Lasciatemi iniziare dicendo: io non so quale codice di basso livello l'AI abbia prodotto per il nostro kernel MLA Jalapeño, ma so come scrivere il numpy per MLA.
Nel compilatore XLA su cui ho lavorato in precedenza, fondevamo quelle operazioni numpy insieme in grumi, e poi usavamo un metaprogramma chiamato "emettitore" per abbassarlo a loop, istruzioni e primitive di livello inferiore.
Quando il compilatore / programma emettitore XLA faceva questo, non mi importava quale assembly uscisse. Per il nostro ottimizzatore stocastico, l'AI concettualmente prende il posto del meta-programma emettitore – sia abbassa che ottimizza, e possiamo chiederle di ottimizzare ulteriormente e sempre più verso il roofline.
Spero che questo chiarisca dove si inserisce l'AI e come sia analoga a un componente in un sistema di compilazione ottimizzante esistente. È anche utile pensare: quale strato consideriamo "codice assembly" ora si sta spostando verso l'alto. Quando scrivi in normale C++ e lo compili a -O3 (il livello di ottimizzazione tipico più alto), non ti aspetti di capire l'assembly che esce, anche se capivi il C++ che avevi scritto. Stiamo facendo la cosa analoga qui, ma con una specifica di input di livello superiore e più matematica.
Ora, una domanda chiave è come controlliamo che il programma che otteniamo dall'AI sia effettivamente equivalente alla descrizione di livello superiore / numpy. Quel meccanismo di controllo stabilisce la solidità del processo di ottimizzazione stocastica dell'AI. Mi aspetto che un futuro post del blog possa entrare più in dettaglio su questo, ma per ora, basti dire che testare l'equivalenza semantica è possibile e lo facciamo. I programmi per acceleratori sono particolarmente adatti a contratti forti e completi che possiamo verificare "sono esattamente ciò che fa il programma ottimizzato dall'AI", poiché sono piuttosto matematici e orientati al flusso di dati nel loro contesto ampio.
Nota che molte tecniche rilevanti in quest'area sono state pionierizzate dagli sforzi nel sottocampo della sintesi di programmi. Mentre i compilatori ottimizzanti dicono, "ecco un programma con semantica, miglioralo ma con semantica equivalente!", la sintesi di programmi dice, "credo che esista un programma con questa semantica, per favore cerca di trovare il migliore che puoi". La sintesi di programmi è un problema più difficile della compilazione ottimizzante, ma è anche fondamentalmente meno vincolata. È effettivamente ciò che fanno gli umani come Ollie quando fanno meglio del compilatore ottimizzante, ed è qualcosa che l'AI può ora aiutarci ad automatizzare. L'AI può trarre "ispirazione" dal programma originale, ma non deve necessariamente eseguire solo piccole trasformazioni locali su di esso. I compilatori ottimizzanti classici non vedranno "oh, hai scritto un bubble sort" e, comprendendo il contratto, lo trasformeranno in un quick-sort, ma sia Ollie che l'AI sono in grado di farlo. Questo è ciò che ci colloca più nel regime di sintesi di programmi con ottimizzazione stocastica che nel regime classico del compilatore ottimizzante.
Tutto questo si riassume nel fatto che puoi iniziare con qualcosa che "non è molto lontano dal numpy", aspettare 48 ore, e avere un kernel ottimizzato con la stessa semantica, come abbiamo mostrato nel nostro talk a HotChips:
Come nota anche la slide, sulla nostra macchina siamo spesso in grado di osservare l'AI superare le prestazioni dei nostri esperti umani anche sui kernel che ritenevamo abbastanza ben ottimizzati. Spesso c'è ancora una percentuale raggiungibile decente semplicemente a causa delle molte varietà di combinazioni/permutazioni che potrebbero dover essere esplorate. Queste sono spesso intrattabilmente noiose per un ingegnere delle prestazioni umano.
Riepilogo e Conclusione
Un compilatore, in fin dei conti, è solo una funzione. Diamo il nostro programma a quella funzione, e otteniamo indietro una versione migliore del nostro programma. Il programma che otteniamo e il programma che inseriamo ci aspettiamo abbiano la stessa semantica.
I compilatori ottimizzanti tradizionali migliorano i programmi tramite regole del flusso di dati ed euristiche. Queste sono completamente comprensibili nella loro provenienza, ma possono anche essere più limitate nelle mosse che possono fare.
Al contrario, l'AI, come ottimizzatore stocastico, deve solo "pensare intensamente" e produrre qualcosa. Le sue mosse non sono fondamentalmente così limitate, rendendola più analoga al nostro ottimizzatore umano esperto. Abbiamo bisogno di modi per controllare che i programmi che produce siano solidi e implementino la stessa semantica che abbiamo inserito, e li abbiamo in atto. E questo tipo di ottimizzazione AI è particolarmente adatto per operazioni matematiche che hanno contratti molto forti. I contratti evitano la necessità di capire cosa fa il kernel riga per riga.
Ecco come abbiamo ottenuto il kernel MLA generato dall'AI!
