Mindre minne, mer AI – metodutveckling för mindre språklabb

För ett mindre forskningslabb som AI Swedens NLU-grupp blir denna minnesbegränsning en kritisk flaskhals. Amarus forskning går därför bortom traditionell eskalering av resurser och fokuserar istället på att effektivisera processerna och den underliggande matematiken i algoritmerna som utgör grunden för språkmodeller. Hans forskning drivs av frågor som: Hur hanterar man effektivt stora mängder data? Vad kan man göra vid begränsat minne? Och vilka delar av processen kan förändras för att möjliggöra storskalighet?

GPU-minnets roll i LLM-träning

Högkvalitativa GPU:er är fundamentala för LLM-utveckling, och detsamma gäller strategierna för att utnyttja deras fulla kapacitet. De komplexa matematiska operationerna som modellerna bygger på är avhängiga GPU:ernas minne. Flaskhalsar uppstår i träningen om prestandan är otillräcklig, vilket skapar  fördröjningar – en typ av ovälkommen begränsning för en grupp som NLU, som jobbar i tidsbegränsade projekt med finansiering från EU och Vinnova, som EUs flaggskeppsprojekt OpenEuroLLM.

GPU:ns minne, eller VRAM, skiljer sig från traditionellt RAM-minne genom att det sitter direkt på GPU-kortet, vilket är nödvändigt för den massiva parallella databearbetningen. Under en träning måste alla modellens parametrar, tillsammans med data som behandlas och alla mellanresultat (som gradienter i backpropagation), lagras i detta minne. När en språkmodell växer till miljarder parametrar ökar kraven på hårdvaran enormt. Detsamma gäller när modellens aktiveringar växer snabbare än datans storlek, vilket till exempel sker vid naiv attention. I dessa lägen blir det begränsade VRAM-minnet snabbt en flaskhals som avgör om beräkningen överhuvudtaget går att genomföra.

Jakten på enkelhet

Kärnan i Amarus forskning handlar om att lösa problemet med ineffektiv  minnesanvändning som finns i neurala nätverkslager som attention. Tiofaldigandet av en sekvens hundrafaldigar minnesanvändningen i en naiv implementation. Hans arbete handlar här om att omformulera beräkningsproblemet. Det handlar om att flytta metoden från att hålla hela datamängden i minnet till att hantera det inkrementellt genom onlineberäkningar, likt hur man beräknar ett medelvärde av en miljon siffror utan att ladda in alla samtidigt. Hans teori heter FlashReduce och påminner om MapReduce, eftersom den går ut på att lösa delproblem och sedan kombinera lösningarna till ett slutresultat.

Det teoretiska ramverk Amaru har utvecklat formaliserar denna process. Genom att bryta ner stora, beräkningskrävande problem i små delproblem, en metod som kallas tiling (från ordet tiles, “brickor”). Utifrån detta har han utvecklat ett proof of concept som effektiviserar minnesanvändningen. Grundtanken tog sitt ursprung i FlashAttention, som kan ses som ett exempel på denna, mer generella, metod.

För forskaren ligger skönheten i att finna en lösning där man endast behöver information om slutresultatet och ett enskilt indatavärde för att beräkna gradienten, istället för att beakta alla enskilda delar. Detta är möjligt så länge ihopslagningsfunktionen av dellösningarna uppfyller vissa kriterier, vilka Amaru beskriver som "inte särskilt hårda", vilket ger utvecklarna stor frihet i designen.

Minskad implementationsbörda 

En av målsättningarna i FlashReduce är att minska den komplexa bördan för andra utvecklare och forskare. Precis som MapReduce hanterar de komplexa detaljerna för distribuerad compute, skulle hans ramverk kunna avlasta utvecklaren. Denne behöver då endast definiera hur man löser delproblemen och hur lösningarna kombineras, och låta ramverket identifiera hur dessa beståndsdelar på bästa sätt utförs på GPU:n. Detta gör det enklare att utveckla och experimentera med minneseffektiva kernels och är en möjlig metod för att skapa större och effektivare open source-modeller.

Detta arbetssätt visar på ett förhållningssätt där "enkelhet och elegans i matematisk teori kan översättas till praktisk effektivitet", menar han. Även om ett proof of concept visar på framsteg, är den återstående tröskeln hur man skapar en tillräckligt användarvänlig och effektiv implementering på GPU:er. Till skillnad från det teoretiska ramverket, som är enkelt och avgränsar forskarens sökrymd genom att visa vilka lager som kan göras minneseffektiva, ligger den verkliga utmaningen i de ständigt skiftande och komplexa hårdvarubegränsningarna. Målet är att ramverket ska hantera dessa detaljer helt automatiskt. På så sätt slipper programmeraren hantera de ”frihetsgrader” som annars krävs, vilket säkerställer optimal effektivitet. “Det här är ett möjligt sätt att skapa större och effektivare modeller och också att minska implementationsbördan,” menar Amaru.

En möjliggörare för open source i Europa

Amarus forskning, om att minska implementationsbördan och möjliggöra större, mer effektiva modeller inom open source-kategorin får en särskild tyngd när man ser till det europeiska AI-landskapet. OpenEuroLLM, ett stort initiativ som finansieras av Europeiska Unionen, som har som mål att skapa transparent AI för Europa genom att släppa storskaliga språkmodeller under öppen licens är beroende av tekniska genombrott som Amarus. I Europa, där tillgången till beräkningsresurser historiskt sett har varit begränsad, blir strävan efter effektivitet en strategisk nödvändighet, snarare än bara en akademisk utmaning.

Amarus teoretiska ramverk – som gör det enklare att utveckla minneseffektiva kernels och därmed minskar implementationsbördan för utvecklare – är den typ av forskning som är avgörande för att europeiska open source-modeller kan mäta sig i storlek och prestanda med andra internationella alternativ. Forskningen adresserar alltså den centrala knäckfrågan: hur kan europeiska aktörer skapa och upprätthålla spetskompetens inom storskalig AI, utan att enbart förlita sig på uppskalning av resurser? Genom att forska fram nya ramverk och processer vill man på AI Sweden bidra till att skapa lösningar på problem som den hårda begränsningen i GPU-minne för att därmed stärka den europeiska förmågan att bygga storskalig och transparent AI under öppen licens.