Federated learning i miljöer med blandad hårdvara – lärdomar från projektet NextGen

Hantera heterogen hårdvara

Även om traditionella ramverk för federerad inlärning (’federated learning’) har stöd för att blanda hårdvara från olika leverantörer kräver de ofta komplexa, stabila och ‘always-open’ nätverksanslutningar. NextGen placerar i stället enhetshanterare i containrar och använder ett asynkront kommunikationsmönster (’blackboard’) byggt på standard-API:er. Metoden gör det möjligt för ramverket att framgångsrikt koppla samman NVIDIA-GPU:er, Intel Gaudi-HPU:er och AMD-grafikkort, antingen placerade lokalt eller bakom en molntjänst, inom en och samma träningsloop.

Projektet leddes av AI Sweden och medfinansierades av Vinnova. Målet var att bygga en slags ’plug and play’-wrapper för infrastruktur som låter företag samarbeta säkert även om deras hårdvaruuppsättningar skiljer sig åt. Frågan har lyfts fram i ljuset av ett växande intresse för datasuveränitet och nya regelverk som EU:s AI-förordning (EU AI Act).

Det finns också ett starkt ekonomiskt incitament, eftersom nya studier visar att den totala ägandekostnaden (TCO) för lokal AI-beräkningskraft (on-premise) i allt högre grad trumfar tredjepartsmolntjänster.

Prövat mot farmaceutisk data

Bland projektdeltagarna fanns Intel, AstraZeneca, Zenseact, Aixia, NetApp och Proact. För att testa ramverket mot verkliga företagskrav lades fokus på ett komplext första use case: förmågor hos stora språkmodeller (LLM) för informationshämtning i ett farmaceutiskt RAG-system.

Projektet fokuserade på att finjustera inbäddningsmodellen i RAG-systemet och att använda den för att hämta relevanta delar från informationstäta, hundrasidiga farmaceutiska dokument. En separat konversationsinriktad språkmodell (LLM) – en chatbot – fungerade som generator och assistent för att svara på användarnas frågor baserat på den framtagna kontexten, precis som i en typisk RAG-uppsättning.

– Blackboard-mönstret vi använder är helt asynkront och flexibelt när klienter av någon anledning faller bort eller nya tillkommer. Modelluppdateringar 'pushas' aldrig till klienterna, de hämtas enbart. Detta ger klienterna själva kontrollen över informationsflödet, säger Mauricio Muñoz.  

Att blanda olika typer av hårdvara under träningen skapade dock en utmaning. Maskiner med mindre minne och beräkningskraft behövde kunna bearbeta mindre batchar av data, vilket spelar roll när man tränar modeller med kontrastiv inlärning ('contrastive learning'), den speciella typ av inlärningsalgoritm som användes.

– Bildligt talat kan noder med stort minne använda en 'högupplöst karta' över datan, medan noder med mindre minne måste navigera med en 'suddig kontur', förklarar Mauricio Muñoz. Att slå samman dessa inlärningssignaler till en enda modelluppdatering leder till intressant dynamik.

Eftersom kontrastiv inlärning i hög grad är beroende av batchstorlek för att lära sig effektivt, ledde utjämningen av dessa "högupplösta" och "lågupplösta" resultat till att mätvärdena för inhämtning av rådata sjönk något. Det slutgiltiga AI-systemet förblev dock mycket träffsäkert, vilket bevisar att ramverket för blandad hårdvara fungerar väl utan att försämra slutprodukten.

En grund av öppen källkod att bygga vidare på

Den huvudsakliga utkomsten från NextGen-projektet är en funktionell "plug and play"-motor med öppen källkod, som har publicerats i form av centrala kodarkiv (nextgen-framework, nextgen-train och nextgen-rag-eval) tillsammans med demos och omfattande teknisk dokumentation.

Genom att kombinera effektiva träningstekniker med smart dataaggregering ger ramverket en flexibel och säker grund för framtidens decentraliserade AI.