AI i wafer-skala löser flaskhalsar för dataflöden i LLM

Integrering av AI på wafer-nivå (Wafer-Scale AI) adresserar flaskhalsar för dataflöden i stora språkmodeller (LLM) genom att konsolidera hundratusentals processorkärnor och massivt distribuerat minne på en enda kontinuerlig kiselyta. Denna arkitektoniska förändring eliminerar den ”bandbreddsvägg” som finns i traditionella multi-GPU-kluster, där kommunikationshastigheten sjunker avsevärt när data korsar chipgränser. Genom att använda wafer-on-wafer (WoW) hybrid-bondning har forskare möjliggjort interconnects med ultrahög bandbredd som tillåter ett sömlöst dataflöde över hela systemet, vilket underlättar den massiva parallella bearbetning som krävs för moderna stora språkmodeller (LLM).

Flaskhalsen för dataflöden i modern AI

Den primära flaskhalsen för skalning av stora språkmodeller har skiftat från rå beräkningskapacitet till de fysiska begränsningarna för dataflöden mellan enskilda chip. I konventionella högpresterande datorkluster, såsom de som använder NVIDIA DGX-system, begränsas enskilda GPU:er av de fysiska stiften och kablaget som krävs för att kommunicera med närliggande processorer. Detta skapar en ”kommunikationsvägg” där energikostnaden och latensen för att flytta data mellan chip blir de dominerande faktorerna för träningstid och inferenseffektivitet.

Latens i interconnects skapar betydande overhead under synkroniseringsfaserna vid träning av Transformer-baserade modeller. När modeller växer till biljoner parametrar anstränger behovet av frekventa ”all-reduce”-operationer – där gradienter delas mellan alla beräkningsnoder – traditionell nätverkshårdvara. Detta har lett till en forskningsera ”efter GPU:n”, där datavetare söker vägar att behandla en hel kisel-wafer som ett enda enhetligt superchip för att upprätthålla det höghastighetsflöde av data som är nödvändigt för nästa generation av generativ AI.

Vad är Cerebras wafer-scale-motor och hur fungerar den?

Cerebras wafer-scale engine (WSE) är ett massivt system-på-wafer som integrerar hundratusentals AI-optimerade kärnor och gigabyte av on-chip SRAM på en enda bit kisel. Genom att kringgå den traditionella processen att skära wafers i små enskilda chip bibehåller WSE det kablage med hög bandbredd som är inneboende i kisel-tillverkningsprocessen över hela ytan på 215 mm x 215 mm. Denna design möjliggör en mesh-arkitektur som stöder kommunikation med ultralåg latens, vilket är avgörande för att utföra minnesbundna operationer som inferens av stora språkmodeller i stor skala.

WSE fungerar genom att utnyttja specialiserade parallelliseringsstrategier, såsom MeshGEMM och MeshGEMV, för att optimera åtkomst till distribuerat minne. I en standard-GPU är högbandbreddsminne (HBM) placerat utanför processorn, vilket skapar en flaskhals; wafer-scale-metoden placerar dock minnet direkt intill beräkningskärnorna. Denna arkitektur gör det möjligt för systemet att hantera de massiva dataflöden som krävs för modell-prefilling och tokengenerering utan den strömkrävande och långsamma processen att hämta data från externa komponenter, vilket resulterar i storleksordningar snabbare prestanda för AI-arbetslaster.

Wafer-Scale-integrering och hybrid-bondning förklarat

Wafer-on-wafer (WoW) hybrid-bondning är en banbrytande tillverkningsteknik som staplar två kisel-wafers vertikalt för att skapa en högdensitets-3D-interconnect. Till skillnad från traditionella förpackningsmetoder som använder mikro-bumps eller lödtenn för att ansluta chip, smälter hybrid-bondning samman kopparytorna på två wafers direkt. Detta ger ett betydligt högre antal anslutningspunkter per kvadratmillimeter, vilket effektivt fördubblar eller tredubblar den tillgängliga kommunikationsytan mellan bearbetningselement (PE) och möjliggör mycket mer komplexa nätverkstopologier i Wafer-Scale AI-system.

• Ökad interconnect-densitet: Hybrid-bondning möjliggör miljontals vertikala anslutningar över wafer-ytan.
• Minskad parasitkapacitans: Det direkta koppar-till-koppar-gränssnittet sänker det elektriska motståndet och energin som krävs för att överföra bitar.
• Förbättrad termisk hantering: Stapling av wafers kan förbättra värmeavledningsvägarna när de kombineras med avancerad vätskekylning.
• Retikel-integrering: Processen möjliggör sömlös kommunikation mellan de rektangulära ”retiklarna” (basenheterna i chiptillverkning) som utgör wafern.

Hur påverkar fysisk retikelplacering nätverkstopologin i wafer-scale-chip?

Den fysiska retikelplaceringen dikterar den uppnåeliga nätverkstopologin genom att definiera de gränser där höghastighetskablage på chipet måste övergå till länkar mellan retiklar eller wafers. I en standarddesign för Wafer-Scale AI begränsar litografiprocessen storleken på en enskild funktionell enhet (retikeln), vilket kräver ett stegmönster för att täcka hela wafern. Forskarna, inklusive Luca Benini, Torsten Hoefler och Maciej Besta, fann att sättet dessa retiklar är inriktade – oavsett om de är perfekt staplade eller förskjutna – avgör vilken väg data måste ta, vilket direkt påverkar latens och genomströmning.

Sofistikerade placeringsstrategier kan avsevärt minska antalet ”hopp” (hop count) mellan avlägsna processorkärnor. Forskargruppen undersökte fyra primära konfigurationer: Aligned, Interleaved, Rotated och Contoured. Genom att strategiskt förskjuta eller rotera retiklarna under hybrid-bondningsprocessen kunde de skapa genvägar som skulle vara omöjliga i ett standard-2D-mesh. Till exempel tillåter en ”Contoured”-placering nätverket att följa waferns naturliga cirkulära form, vilket optimerar användningen av kisel i kantzonerna och skapar en mer effektiv global kommunikationsstruktur för djupinlärningsuppgifter.

Optimering av topologi: De fyra strategierna för retikelplacering

Baslinjen för nätverk på wafer-nivå är vanligtvis ett 2D-mesh, vilket begränsas av dess stela, rutnätsliknande anslutningsmöjligheter. Även om det är funktionellt lider 2D-meshet av hög latens när data måste färdas från ena sidan av wafern till den andra. För att hantera detta föreslår forskningen som presenteras av Hoefler och hans kollegor att den fysiska placeringen av wafers kan manipuleras för att skapa ”expressfiler” för data. Strategierna Aligned och Interleaved fokuserar på att förenkla de vertikala anslutningarna mellan två bondade wafers för att säkerställa att varje processorkärna har en direkt partner på det motsatta lagret.

Avancerade konfigurationer som Rotated och Contoured representerar ett genombrott inom högpresterande nätverksdesign. Rotated-strategin innebär att man roterar en wafer 90 grader före bondning, vilket omvandlar den fysiska layouten till ett high-radix-nätverk. Denna konfiguration möjliggör ”långväga” hopp över wafern utan att lägga till extra fysiskt kablage. Dessa optimerade topologier säkerställer att nätverket kan hantera de oregelbundna datamönster som ofta ses i glesa AI-modeller, vilket ger en flexibel infrastruktur som kan anpassas till olika arkitekturer för neurala nätverk.

Kvantifiering av vinsterna: Prestanda och effektivitet

Implementeringen av optimerade strategier för retikelplacering resulterade i en massiv ökning av nätverksgenomströmningen med 250 % jämfört med traditionella 2D-mesh-baslinjer. Detta prestandahopp beror främst på minskad nätverksträngsel; genom att tillhandahålla fler direkta vägar för data kan systemet flytta större volymer information samtidigt. För stora språkmodeller innebär detta att ”all-reduce”- och ”all-to-all”-kommunikationsmönster, som ofta bromsar träningen, kan slutföras på en bråkdel av tiden.

Utöver rå hastighet lyfter forskningen fram betydande förbättringar i latens och energimässig hållbarhet.

• Latensreduktion: Kommunikationsfördröjningar minskades med upp till 36 %, vilket möjliggör snabbare synkronisering av modellvikter.
• Energieffektivitet: Energin som krävs per överförd byte minskade med upp till 38 %, ett kritiskt mått för AI:s miljöpåverkan.
• Systemtillförlitlighet: Genom att minska antalet hopp och det elektriska avståndet som data måste färdas, förbättrade forskarna också den övergripande robustheten i kommunikationsstrukturen.

Branschkonsekvenser: Från Cerebras till framtidens superchip

Denna forskning ger en färdplan för framtiden inom halvledartillverkning och går bortom gränserna för enskilda GPU-chip. Medan företag som Cerebras Systems redan har bevisat livskraften hos wafer-scale-motorer, tyder arbetet av Torsten Hoefler och hans team vid ETH Zurich på att vi bara är i början av vad som är möjligt. Genom att använda wafer-on-wafer-bondning och avancerad topologidesign kan framtida superchip rymma biljoner transistorer i ett enda paket, vilket ger den hårdvara som krävs för att träna modeller som är betydligt större än GPT-4.

Skiftet mot integrerad Wafer-Scale AI kommer sannolikt att definiera om konkurrensen på marknaden för AI-hårdvara. Eftersom efterfrågan på LLM fortsätter att skjuta i höjden kommer förmågan att minimera energiförbrukningen samtidigt som genomströmningen maximeras att bli den främsta konkurrensfördelen. Övergången från att bygga ”kluster av chip” till ”wafer-scale-system” representerar en grundläggande förändring i hur vi ser på datorarkitektur och lovar en framtid där den fysiska storleken på vår hårdvara är den enda gränsen för våra modellers intelligens.