Du kan skrive GPU-kode i Python med NVIDIAs cuTile

NVIDIA og Jane Street har skrevet sekvensmodellen Mamba-2 med cuTile, et Python-værktøj til at lave små programmer på grafikkortet. Det betyder noget, fordi Python-koden overlader styringen af grafikkortets regneenheder til compileren (oversætteren) og kan gøre specialkode lettere at lave. Nu bruges tre specialprogrammer til Mamba-2’s fremadrettede beregning, men hvor meget hurtigere den kører, er ikke oplyst.

NVIDIA og handelsvirksomheden Jane Street har skrevet sekvensmodellen Mamba-2 ved hjælp af cuTile. En sekvensmodel er en model, der læser tekst i rækkefølge. cuTile er et værktøj, hvor man skriver beregninger til grafikkortet i programmeringssproget Python. Tre små specialiserede beregningsdele, som man kalder kerner, står for modellens fremadrettede beregning. Det er den del, hvor input bliver regnet om til output.

Pointen med cuTile er, at selve koden skrives i Python, mens et oversætterprogram, en compiler, selv styrer de mange tråde, altså små arbejdsopgaver der kører samtidigt, og rammer grafikkortets Tensor Cores. Tensor Cores er regneenheder, der er bygget til matrixmultiplikation, altså at gange tal i et fast regnearkslignende mønster. Ifølge et blogindlæg fra NVIDIA er Jane Street blandt de første til at tage værktøjet i brug.

Mamba-2 er en sekvensmodel med en fast intern tilstand. Det er en slags hukommelse, der hele tiden har samme størrelse, og som bliver opdateret stykke for stykke. Den første kerne, chunk_cumsum, lægger henfaldstallene sammen for hvert ord og hvert hoved, altså en mindre del af modellen, i logaritmisk skala. Henfaldstal er tal, der beskriver, hvor hurtigt noget aftager, og logaritmisk skala er en måde at regne med tal på, hvor store spring bliver mere overskuelige. Så kan de næste kerner genbruge resultatet. Den anden kerne, fwd_pass_state, flytter den interne tilstand frem gennem tekststykkerne ved hjælp af Tensor Cores. Den tredje kerne, compute_out, samler den medbragte tilstand, den interaktionsmatrix som er henfaldet og maskeret, og springforbindelsen i én kerne. En interaktionsmatrix er en tabel over, hvordan ordene må påvirke hinanden, hvor nogle felter er lukket af, og en springforbindelse er en direkte vej, hvor information kan springe over et trin. På den måde bliver mellemregningerne liggende i grafikkortets registre, som er meget hurtige små lagerpladser.

Ingeniørerne fremhæver også flere optimeringer. Det er tip om ventetid, der hjælper med at beskrive trafikken til hukommelsen. Det er ombyttede operander, altså de tal der indgår i en beregning. Det er transponerede hukommelseslayouts, hvor rækker og kolonner bytter plads i hukommelsen. Og det er en autotuner, et program der selv prøver sig frem, for at finde den bedste blokstørrelse, altså hvor store stykker arbejdet deles op i, for hver inputform. Blogindlægget oplyser ikke, hvor meget hurtigere Mamba-2 kører med kernerne.