Nowy model językowy przewiduje wydajność kodu: przełom w optymalizacji i projektowaniu systemów

Sztuczna inteligencja coraz śmielej wkracza w obszary, które do niedawna wydawały się domeną wyłącznie ludzkiej ekspertyzy. Najnowszym przykładem jest model językowy RLM (Regression Language Model), stworzony przez badaczy z Cornell i Google. Potrafi on przewidywać szereg parametrów wydajnościowych kodu – od opóźnień w kernelach GPU, przez zużycie pamięci, aż po dokładność działania sieci neuronowych. Co istotne, robi to bez konieczności ręcznego definiowania cech czy stosowania skomplikowanych narzędzi.

Czym wyróżnia się RLM?

RLM to uniwersalny model regresyjny, który na podstawie surowego tekstu kodu (w językach takich jak Python, C/C++, a nawet w reprezentacji ONNX) potrafi przewidzieć:

• Szczytowe zużycie pamięci.
• Opóźnienia kernelów GPU w środowisku Triton.
• Dokładność i opóźnienia sieci neuronowych, uwzględniając specyfikę sprzętową.

Kluczowe jest to, że RLM nie wymaga żadnego wstępnego przetwarzania danych, inżynierii cech, ani użycia specjalizowanych koderów grafów. Działa bezpośrednio na tekście, co upraszcza cały proces.

Konkretne wyniki robią wrażenie

Model osiąga imponujące wyniki korelacji. Przykładowo, w zadaniu przewidywania zużycia pamięci dla kodu z platformy APPS LeetCode, RLM osiągnął współczynnik korelacji Spearmana ρ ≈ 0.93. Dla opóźnień kerneli Triton wynik to ρ ≈ 0.52, a dla różnych języków z CodeNet średnia korelacja przekracza 0.5. W testach na przestrzeniach NAS (Network Architecture Search) RLM osiągnął współczynnik Kendalla τ ≈ 0.46, co jest porównywalne, a czasem nawet lepsze, niż wyniki uzyskiwane przez bardziej złożone modele oparte na grafach.

Wielowymiarowe przewidywanie

Autoregresyjna natura dekodera RLM pozwala na uwzględnianie zależności między różnymi metrykami. Na przykład, model może najpierw przewidzieć dokładność sieci neuronowej, a następnie, na podstawie tej predykcji, oszacować opóźnienia na różnych urządzeniach. Umożliwia to uwzględnienie realistycznych kompromisów i optymalizację pod kątem wielu celów jednocześnie.

Dlaczego to jest ważne?

Obecne metody przewidywania wydajności w kompilatorach, systemach doboru kerneli GPU czy w NAS często opierają się na specyficznych cechach, drzewach składniowych lub enkoderach GNN. To sprawia, że są one wrażliwe na nowe operacje i języki. RLM, traktując regresję jako generowanie tekstu (ciągu cyfr), unifikuje cały proces. Kod źródłowy, Triton IR czy grafy ONNX są tokenizowane jako zwykły tekst, a następnie dekodowane na liczby. To obniża koszty utrzymania i ułatwia transfer wiedzy na nowe zadania poprzez fine-tuning.

Jak to działa?

RLM bazuje na architekturze encoder-decoder z inicjalizacją z modelu T5-Gemma (około 300 milionów parametrów). Wejście to surowy tekst (kod lub ONNX), a wyjście to ciąg cyfr reprezentujących liczbę. Specjalne ograniczenia podczas dekodowania zapewniają, że generowane są poprawne liczby i pozwalają na oszacowanie niepewności poprzez próbkowanie.

Badacze przeprowadzili szereg eksperymentów, które pokazały, że wstępne uczenie językowe przyspiesza konwergencję i poprawia przewidywanie opóźnień w Triton, a dekodowanie tekstu z ograniczeniami jest lepsze niż tradycyjne głowice regresyjne MSE. Użycie tokenizera wyspecjalizowanego w operatorach ONNX zwiększa efektywny kontekst, a dłuższy kontekst i skalowanie do większego enkodera Gemma dodatkowo poprawiają wyniki.

Implikacje dla przyszłości

Przedstawione wyniki sugerują, że traktowanie przewidywania wydajności jako problemu generowania tekstu otwiera nowe możliwości. Kompaktowy model RLM, inicjalizowany przez T5Gemma, potrafi analizować kod źródłowy (Python/C++), kernele Triton lub grafy ONNX i generować skalibrowane dane liczbowe poprzez dekodowanie z ograniczeniami. Uzyskane korelacje są na tyle wysokie, że mogą mieć realny wpływ na heurystyki kompilatorów, przycinanie kerneli oraz wielowymiarową optymalizację NAS. Dostępność otwartego zbioru danych i biblioteki ułatwia replikację wyników i obniża próg wejścia dla dostrajania modelu do nowego sprzętu lub języków.