De porting van Google's Gemma-4 31B-model naar een AWS Inferentia2 inf2.24xlarge leverde een perfecte token-voor-token match op met de CPU-referentie—toch was elke gegenereerde zin onzin. Het gat tussen "overeenkomst" en "werkend zijn" dient nu als waarschuwing voor iedereen die probeert enorme LLM's op Amazons aangepaste inference-chips te persen.

Waarom een token-voor-token match niet genoeg is

De ontwikkelaar vergeleek elk output-token van het Inferentia-apparaat met het token dat werd geproduceerd door een CPU-run van het model. De streams waren identiek, dus het leek erop dat de hardware de referentie-implementatie exact had gereproduceerd. In werkelijkheid voedden beide streams een foutief geformatteerde prompt aan een model dat was ontdaan van zijn chattemplate en was voorzien van de verkeerde turn markers. De ontbrekende template stuurde het model in een oneindige lus, waardoor er onzin werd uitgespuwd. De hardware deed zijn werk—het reproduceerde een bug die al in de referentiecode aanwezig was.

De les is simpel: SEQ_MATCH (sequentiële token-gelijkheid) is niet gelijk aan correctheid. Als de referentie-implementatie defect is, neemt een getrouwe hardware-replica dezelfde fout over. Validatie moet verder gaan dan pariteit op tokenniveau; het vereist end-to-end functionele controles met correct geformatteerde inputs.

Buffers vermomd als parameters

Tijdens de laadfase sloeg de model loader een component genaamd layer_scalar over. De code registreerde dit object als een buffer in plaats van een parameter in de PyTorch-modeldefinitie. Buffers zijn statische tensoren die tijdens de training niet worden bijgewerkt, en veel loaders negeren ze bij het converteren naar Neuron-compatibele formaten. Door dit over te slaan, bleven de schaalfactoren voor verschillende lagen op hun standaardwaarden staan, wat de berekeningen in het hele netwerk verstoorde. Er werd geen foutmelding gegeven; het model compileerde en de inference-pipeline draaide, maar de numerieke resultaten waren onjuist.

Voor iedereen die grote modellen naar Inferentia verplaatst: controleer elke niet-parameter tensor. Zelfs als een tensor niet bedoeld is om te worden geleerd, kan deze nog steeds essentieel zijn voor de correcte forward-pass berekening. Het handmatig verifiëren van de inclusie van buffers kan stille schaalfouten voorkomen die anders moeilijk te diagnosticeren zijn.

Spot-instance volatiliteit en de 39-minuten compile

Het draaien van een model met 31 miljard parameters op een spot instance lijkt goedkoop, maar de besparingen gaan gepaard met onvoorspelbare reclaim-events. De compileertijd van de ontwikkelaar—ongeveer 39 minuten om het model te vertalen naar Neuron-compatibele code—vervloog toen AWS de instance terugvorderde. Om onderbrekingen te overleven, bouwden ze een drieledig vangnet:

  • ModelBuilder hield het geheugengebruik binnen de limiet van 384 GB van de host, waardoor crashes die een herstart zouden forceren werden voorkomen.
  • Onmiddellijke S3-mirroring van zowel de ruwe weight-bestanden als de gecompileerde “neffs” (Neuron executable files) zorgde ervoor dat een nieuwe instance precies kon voortgaan waar de vorige was gestopt.
  • Een multi-region poller scande AWS-regio's op beschikbare spot-capaciteit en lanceerde een nieuwe instance zodra er een beschikbaar was.

Deze stappen veranderden een kwetsbare, single-point compile in een veerkrachtige pipeline die de volatiliteit van spot-markten overleeft.

Sharding-valstrikken met gemengde attention-layouts

Gemma-4 31B gebruikt twee attention-configuraties. Sommige lagen maken gebruik van vier key-value (KV) heads, andere een ander aantal. Het gelijkmatig verdelen van het model over acht parallelle ranks mislukt wanneer het aantal KV-heads van een laag niet netjes deelbaar is. Het proberen te sharden van een laag met 4 heads over acht ranks zou elke rank dwingen om een halve head te verwerken—een wiskundige onmogelijkheid die shape mismatches en runtime-fouten veroorzaakt.

De oplossing was om de globally-sharded lagen (die met compatibele head counts) te repliceren over alle ranks en alleen de “sliding” lagen te sharden waarvan de head counts een gelijkmatige verdeling toelieten. Deze hybride strategie behield de tensor-parallel efficiëntie terwijl illegale verdeling van KV-heads werd vermeden, waardoor de tensor-parallelisatie-fouten die eerdere pogingen teisterden werden geëlimineerd.

Conclusie

Het porten van een gigantische LLM naar Inferentia is meer dan alleen een compile-and-run oefening. Het vereist rigoureuze functionele tests die verder gaan dan token-gelijkheid, nauwgezette verificatie dat elke tensor—parameter of buffer—correct wordt afgehandeld, en een deploymentstrategie die rekening houdt met de terugvordering van spot-instances. Ten slotte moet sharding de interne attention-geometrie van het model respecteren; anders wordt de parallelliteit die snelheid belooft een bron van stille fouten.