Google ના Gemma-4 31B મોડેલને AWS Inferentia2 inf2.24xlarge પર પોર્ટ કરવાથી CPU રેફરન્સ સાથે પરફેક્ટ ટોકન-ફોર-ટોકન મેચ મળ્યો—છતાં પણ દરેક જનરેટ થયેલું વાક્ય અર્થહીન (gibberish) હતું. "મેચિંગ" અને "કામ કરવા" વચ્ચેનું આ અંતર હવે એવા કોઈપણ વ્યક્તિ માટે ચેતવણી સમાન છે જે વિશાળ LLMs ને Amazon ના કસ્ટમ ઇન્ફરન્સ ચિપ્સ પર ચલાવવાનો પ્રયાસ કરી રહ્યા છે.

ટોકન-બાય-ટોકન મેચ શા માટે પૂરતો નથી

ડેવલપરે Inferentia ડિવાઇસના દરેક આઉટપુટ ટોકનની સરખામણી મોડેલના CPU રન દ્વારા ઉત્પાદિત ટોકન સાથે કરી. સ્ટ્રીમ્સ સમાન હતા, તેથી એવું લાગતું હતું કે હાર્ડવેરે રેફરન્સ ઇમ્પ્લીમેન્ટેશનને બરાબર રીતે ફરીથી બનાવ્યું છે. વાસ્તવમાં, બંને સ્ટ્રીમ્સમાં મોડેલના ચેટ ટેમ્પલેટ વગરના અને ખોટા ટર્ન માર્કર્સ સાથેના ખામીયુક્ત પ્રોમ્પ્ટ મોકલવામાં આવ્યા હતા. ટેમ્પલેટના અભાવે મોડેલ અનંત લૂપમાં ફસાઈ ગયું અને અર્થહીન લખાણ આપવા લાગ્યું. હાર્ડવેરે તેનું કામ કર્યું—તેણે રેફરન્સ કોડમાં રહેલી ભૂલને જ ફરીથી રજૂ કરી.

બોધપાઠ સરળ છે: SEQ_MATCH (ક્રમિક ટોકન સમાનતા) નો અર્થ સાચું પરિણામ (correctness) નથી. જો રેફરન્સ ઇમ્પ્લીમેન્ટેશન જ બગડેલું હોય, તો હાર્ડવેરનું સચોટ રેપ્લિકા પણ તે જ નિષ્ફળતા વારસામાં મેળવે છે. વેરિફિકેશન ટોકન-લેવલની સમાનતાથી આગળ વધવું જોઈએ; તેમાં યોગ્ય રીતે ફોર્મેટ કરેલા ઇનપુટ્સ સાથે એન્ડ-ટુ-એન્ડ ફંક્શનલ ચેક્સ હોવા જરૂરી છે.

પેરામીટર્સના વેશમાં છુપાયેલા બફર્સ (Buffers)

લોડ ફેઝ દરમિયાન, મોડેલ લોડરે layer_scalar નામના ઘટકને સ્કીપ કરી દીધો. કોડમાં PyTorch મોડેલ ડેફિનેશનમાં આ ઓબ્જેક્ટને parameter ને બદલે buffer તરીકે રજિસ્ટર કરવામાં આવ્યો હતો. બફર્સ એ સ્ટેટિક ટેન્સર છે જેને ટ્રેનિંગ દરમિયાન અપડેટ કરવામાં આવતા નથી, અને ઘણા લોડર્સ Neuron-સુસંગત ફોર્મેટમાં રૂપાંતર કરતી વખતે તેમને અવગણે છે. તેને સ્કીપ કરવાથી ઘણા લેયર્સના સ્કેલિંગ ફેક્ટર્સ તેમના ડિફોલ્ટ પર રહી ગયા, જેના કારણે સમગ્ર નેટવર્કનું ગણિત ખોરવાઈ ગયું. કોઈ એરર (error) આવી નહીં; મોડેલ કમ્પાઈલ થયું અને ઇન્ફરન્સ પાઇપલાઇન ચાલી, પરંતુ ગાણિતિક પરિણામો ખોટા હતા.

Inferentia પર મોટા મોડેલ્સ લઈ જતા કોઈપણ વ્યક્તિએ દરેક નોન-પેરામીટર ટેન્સરનું ઓડિટ કરવું જોઈએ. ભલે ટેન્સર લર્ન કરવા માટે ન હોય, તેમ છતાં તે સાચા ફોરવર્ડ-પાસ કમ્પ્યુટેશન માટે આવશ્યક હોઈ શકે છે. બફર સમાવેશનું મેન્યુઅલી વેરિફિકેશન કરવાથી સાયલન્ટ સ્કેલ એરર્સ (silent scale errors) ને રોકી શકાય છે, જે અન્યથા શોધવી મુશ્કેલ હોય છે.

સ્પોટ-ઇન્સ્ટન્સની અસ્થિરતા અને 39-મિનિટનું કમ્પાઈલ

સ્પોટ ઇન્સ્ટન્સ પર 31-બિલિયન-પેરામીટર ધરાવતું મોડેલ ચલાવવું સસ્તું લાગે છે, પરંતુ આ બચત અનિશ્ચિત રીક્લેમ ઇવેન્ટ્સ (reclaim events) સાથે આવે છે. ડેવલપરનો કમ્પાઈલ સમય—મોડેલને Neuron-સુસંગત કોડમાં રૂપાંતરિત કરવા માટે લગભગ 39 મિનિટ—ત્યારે વ્યર્થ ગયો જ્યારે AWS એ ઇન્સ્ટન્સ પાછું ખેંચી લીધું (reclaimed). વિક્ષેપોથી બચવા માટે તેમણે ત્રણ-સ્તરીય સુરક્ષા કવચ બનાવ્યું:

  • ModelBuilder એ મેમરી વપરાશને 384 GB હોસ્ટ લિમિટની અંદર રાખ્યો, જેથી ક્રેસ (crash) ન થાય અને રીસ્ટાર્ટ કરવાની જરૂર ન પડે.
  • રો (raw) વેઇટ ફાઇલો અને કમ્પાઈલ કરેલી “neffs” (Neuron executable files) બંનેનું તાત્કાલિક S3 મિરરિંગ કરવાથી નવું ઇન્સ્ટન્સ બરાબર ત્યાંથી શરૂ કરી શક્યું જ્યાં અગાઉનું ઇન્સ્ટન્સ અટકી ગયું હતું.
  • એક મલ્ટી-રીજન પોલર એ ઉપલબ્ધ સ્પોટ કેપેસિટી માટે AWS રીજન્સને સ્કેન કર્યું અને જેવું નવું ઇન્સ્ટન્સ મળ્યું કે તરત જ તેને લોન્ચ કર્યું.

આ પગલાઓએ એક નાજુક, સિંગલ-પોઇન્ટ કમ્પાઈલને એક લવચીક (resilient) પાઇપલાઇનમાં બદલી નાખ્યું જે સ્પોટ માર્કેટના ઉતાર-ચઢાવમાં ટકી શકે છે.

મિશ્રડ એટેન્શન લેઆઉટ સાથે શાર્ડિંગની મુશ્કેલીઓ

Gemma-4 31B બે એટેન્શન કોન્ફિગરેશનનો ઉપયોગ કરે છે. કેટલાક લેયર્સ ચાર કી-વેલ્યુ (KV) હેડ્સનો ઉપયોગ કરે છે, જ્યારે અન્યમાં અલગ સંખ્યા હોય છે. જ્યારે કોઈ લેયરના KV હેડ કાઉન્ટને આઠ સમાંતર રેન્ક્સમાં સમાન રીતે વહેંચી ન શકાય, ત્યારે મોડેલને આઠ રેન્ક્સમાં સમાન રીતે વિભાજિત (split) કરવું નિષ્ફળ જાય છે. 4-હેડ ધરાવતા લેયરને આઠ રેન્ક્સમાં શાર્ડ કરવાનો પ્રયાસ કરવાથી દરેક રેન્કને અડધા હેડને હેન્ડલ કરવું પડે—જે ગાણિતિક રીતે અશક્ય છે અને તેનાથી શેપ મિસમેચ (shape mismatches) અને રનટાઇમ એરર્સ આવે છે.

તેનો ઉકેલ એ હતો કે ગ્લોબલી-શાર્ડેડ લેયર્સને (જેમના હેડ કાઉન્ટ સુસંગત હોય) તમામ રેન્ક્સમાં રિપ્લીકેટ કરવા અને માત્ર એવા “સ્લાઇડિંગ” લેયર્સને શાર્ડ કરવા જેના હેડ કાઉન્ટ સમાન વિભાજનની મંજૂરી આપે છે. આ હાઇબ્રિડ વ્યૂહરચનાએ ટેન્સર-પેરેલલ કાર્યક્ષમતા જાળવી રાખી અને KV હેડ્સના ગેરકાયદેસર વિભાજનને ટાળ્યું, જેનાથી અગાઉના પ્રયાસોમાં આવતી ટેન્સર-પેરેલલાઇઝેશન ભૂલો દૂર થઈ.

નિષ્કર્ષ (Takeaway)

એક વિશાળ LLM ને Inferentia પર પોર્ટ કરવું એ માત્ર કમ્પાઈલ-અને-રન કરવાની પ્રક્રિયા નથી. તે ટોકન સમાનતાથી આગળ વધીને કડક ફંક્શનલ ટેસ્ટિંગ, દરેક ટેન્સર—પેરામીટર અથવા બફર—યોગ્ય રીતે હેન્ડલ કરવામાં આવે છે તેની ઝીણવટભરી ચકાસણી, અને સ્પોટ-ઇન્સ્ટન્સ રીક્લેમેશનની અપેક્ષા રાખતી ડિપ્લોયમેન્ટ વ્યૂહરચનાની માંગ કરે છે. અંતે, શાર્ડિંગ મોડેલના આંતરિક એટેન્શન ભૂમિતિ (attention geometry) નો આદર કરતું હોવું જોઈએ; અન્યથા, જે પેરેલલિઝમ ઝડપ આપવાનું વચન આપે છે તે જ સાયલન્ટ ફેલ્યોરનું કારણ બની શકે છે.