Google’s Gemma-4 31B మోడల్‌ను AWS Inferentia2 inf2.24xlargeకి పోర్ట్ చేసినప్పుడు, CPU రిఫరెన్స్‌తో టోకెన్-టు-టోకెన్ ఖచ్చితమైన మ్యాచ్ వచ్చింది—కానీ ప్రతి జనరేటెడ్ వాక్యం అర్థం లేనిదిగా (gibberish) ఉంది. "మ్యాచ్ అవ్వడం" మరియు "పని చేయడం" మధ్య ఉన్న ఈ తేడా, భారీ LLMలను Amazon యొక్క కస్టమ్ ఇన్ఫరెన్స్ చిప్‌లపైకి తీసుకురావాలనుకునే ఎవరికైనా ఒక హెచ్చరికగా నిలుస్తుంది.

టోకెన్-బై-టోకెన్ మ్యాచ్ ఎందుకు సరిపోదు

డెవలపర్ Inferentia పరికరం నుండి వచ్చిన ప్రతి అవుట్‌పుట్ టోకెన్‌ను, మోడల్ యొక్క CPU రన్‌ ద్వారా ఉత్పత్తి చేయబడిన టోకెన్‌తో పోల్చారు. స్ట్రీమ్స్ రెండూ ఒకేలా ఉన్నాయి, కాబట్టి హార్డ్‌వేర్ రిఫరెన్స్ ఇంప్లిమెంటేషన్‌ను ఖచ్చితంగా పునరుత్పత్తి చేసినట్లు కనిపించింది. వాస్తవానికి, రెండు స్ట్రీమ్స్ కూడా చాట్ టెంప్లేట్ లేని మరియు తప్పు టర్న్ మార్కర్లతో ఉన్న ఒక మల్ఫార్మ్డ్ ప్రాంప్ట్‌ను మోడల్‌కు అందించాయి. టెంప్లేట్ లేకపోవడం వల్ల మోడల్ ఒక ఇన్ఫినిట్ లూప్‌లోకి వెళ్లి, అర్థం లేని సమాచారాన్ని (nonsense) విడుదల చేసింది. హార్డ్‌వేర్ తన పనిని చేసింది—రిఫరెన్స్ కోడ్‌లో ఉన్న బగ్‌ను అది కూడా పునరుత్పత్తి చేసింది.

పాఠం చాలా సరళం: SEQ_MATCH (సీక్వెన్షియల్ టోకెన్ ఈక్వాలిటీ) అంటే కరెక్ట్‌నెస్ కాదు. రిఫరెన్స్ ఇంప్లిమెంటేషన్ లోపభూయిష్టంగా ఉంటే, హార్డ్‌వేర్ యొక్క నమ్మకమైన ప్రతిరూపం కూడా అదే వైఫల్యాన్ని పొందుతుంది. వాలిడేషన్ అనేది కేవలం టోకెన్-లెవల్ పారిటీకి మాత్రమే పరిమితం కాకూడదు; సరిగ్గా ఫార్మాట్ చేసిన ఇన్‌పుట్‌లతో ఎండ్-టు-ఎండ్ ఫంక్షనల్ చెక్‌లు అవసరం.

పారామీటర్ల రూపంలో ఉన్న బఫర్‌లు

లోడ్ ఫేజ్ సమయంలో, మోడల్ లోడర్ layer_scalar అనే కాంపోనెంట్‌ను స్కిప్ చేసింది. PyTorch మోడల్ డెఫినిషన్‌లో ఈ ఆబ్జెక్ట్‌ను కోడ్ ఒక పారామీటర్ కంటే బఫర్ గా రిజిస్టర్ చేసింది. బఫర్‌లు అంటే స్టాటిక్ టెన్సర్‌లు, వీటిని ట్రైనింగ్ సమయంలో అప్‌డేట్ చేయరు, మరియు Neuron-కంప్యాటబుల్ ఫార్మాట్‌లకు కన్వర్ట్ చేసేటప్పుడు చాలా లోడర్‌లు వీటిని విస్మరిస్తాయి. దీనిని స్కిప్ చేయడం వల్ల కొన్ని లేయర్‌ల స్కేలింగ్ ఫ్యాక్టర్లు వాటి డిఫాల్ట్ విలువల్లోనే ఉండిపోయాయి, దీనివల్ల మొత్తం నెట్‌వర్క్ అంతటా గణిత లెక్కలు (math) తప్పుగా మారాయి. ఎటువంటి ఎర్రర్ రాదు; మోడల్ కంపైల్ అయ్యింది మరియు ఇన్ఫరెన్స్ పైప్‌లైన్ నడిచింది, కానీ గణిత ఫలితాలు (numerical results) తప్పుగా వచ్చాయి.

Inferentiaకి పెద్ద మోడళ్లను తరలిస్తున్న ఎవరైనా, ప్రతి నాన్-పారామీటర్ టెన్సర్‌ను ఆడిట్ చేయండి. ఒక టెన్సర్ నేర్చుకోవడానికి (learned) ఉద్దేశించినది కాకపోయినా, సరైన ఫార్వార్డ్-పాస్ కంప్యూటేషన్ కోసం అది అవసరం కావచ్చు. బఫర్ ఇన్‌క్లూజన్‌ను మాన్యువల్‌గా వెరిఫై చేయడం ద్వారా, గుర్తించడం కష్టమైన సైలెంట్ స్కేల్ ఎర్రర్లను నివారించవచ్చు.

స్పాట్-ఇన్‌స్టాన్స్ అస్థిరత మరియు 39-నిమిషాల కంపైల్

స్పాట్ ఇన్‌స్టన్స్‌పై 31-బిలియన్ పారామీటర్ల మోడల్‌ను రన్ చేయడం చౌకగా అనిపించవచ్చు, కానీ ఆ పొదుపుతో పాటు ఊహించని రీక్లెయిమ్ ఈవెంట్‌లు (reclaim events) కూడా వస్తాయి. మోడల్‌ను Neuron-కంప్యాటబుల్ కోడ్‌గా మార్చడానికి పట్టిన డెవలపర్ యొక్క కంపైల్ సమయం—సుమారు 39 నిమిషాలు—AWS ఇన్‌స్టన్స్‌ను తిరిగి తీసుకున్నప్పుడు (reclaimed) అంతా పోయింది. అంతరాయాలను తట్టుకోవడానికి వారు మూడు అంచెల భద్రతా వలయాన్ని (safety net) నిర్మించారు:

  • ModelBuilder మెమరీ వినియోగాన్ని 384 GB హోస్ట్ లిమిట్ లోపల ఉంచింది, తద్వారా రీస్టార్ట్ చేయాల్సిన పరిస్థితి వచ్చేలా చేసే క్రాష్‌లను నివారించింది.
  • రా రియల్ వెయిట్ ఫైల్స్ మరియు కంపైల్ చేసిన “neffs” (Neuron executable files) రెండింటినీ వెంటనే S3 మిర్రరింగ్ చేయడం వల్ల, కొత్త ఇన్‌స్టాన్స్ పాతది ఎక్కడ ఆగిపోయిందో అక్కడి నుండే తిరిగి ప్రారంభించగలిగింది.
  • ఒక మల్టీ-రీజియన్ పోలర్ అందుబాటులో ఉన్న స్పాట్ కెపాసిటీ కోసం AWS రీజియన్లను స్కాన్ చేసి, కొత్త ఇన్‌స్టాన్స్ కనిపించగానే దానిని లాంచ్ చేసింది.

ఈ దశలు బలహీనమైన, సింగిల్-పాయింట్ కంపైల్‌ను స్పాట్ మార్కెట్ల హెచ్చుతగ్గులను తట్టుకోగల దృఢమైన పైప్‌లైన్‌గా మార్చాయి.

మిక్స్‌డ్ అటెన్షన్ లేఅవుట్‌లతో షార్డింగ్ ఇబ్బందులు

Gemma-4 31B రెండు అటెన్షన్ కాన్ఫిగరేషన్‌లను ఉపయోగిస్తుంది. కొన్ని లేయర్‌లు నాలుగు కీ-వాల్యూ (KV) హెడ్‌లను ఉపయోగిస్తే, మరికొన్ని వేరే సంఖ్యను ఉపయోగిస్తాయి. ఒక లేయర్ యొక్క KV హెడ్ కౌంట్ సమానంగా విభజించబడనప్పుడు, మోడల్‌ను ఎనిమిది సమాంతర ర్యాంక్‌లలో (parallel ranks) సమానంగా విభజించడం (splitting) విఫలమవుతుంది. 4-హెడ్ లేయర్‌ను ఎనిమిది ర్యాంక్‌లలో షార్డ్ చేయడానికి ప్రయత్నిస్తే, ప్రతి ర్యాంక్ సగం హెడ్‌ను హ్యాండిల్ చేయాల్సి వస్తుంది—ఇది గణితపరంగా అసాధ్యం మరియు దీనివల్ల షేప్ మిస్‌మ్యాచ్‌లు (shape mismatches) మరియు రన్‌టైమ్ ఎర్రర్‌లు వస్తాయి.

దీనికి పరిష్కారం ఏమిటంటే, గ్లోబల్‌గా-షార్డ్ చేయబడిన లేయర్‌లను (సరిపోలే హెడ్ కౌంట్లు ఉన్నవి) అన్ని ర్యాంక్‌లలో రీప్లికేట్ చేయడం మరియు హెడ్ కౌంట్లు సమానంగా విభజించడానికి వీలైన “స్లైడింగ్” లేయర్‌లను మాత్రమే షార్డ్ చేయడం. ఈ హైబ్రిడ్ వ్యూహం టెన్సర్-ప్యారలల్ సామర్థ్యాన్ని (tensor-parallel efficiency) కాపాడుతూనే, KV హెడ్‌ల అక్రమ విభజనను నివారించింది, తద్వారా మునుపటి ప్రయత్నాలలో ఎదురైన టెన్సర్-ప్యారలలైజేషన్ ఎర్రర్‌లను తొలగించింది.

ముఖ్య అంశం (Takeaway)

ఒక భారీ LLMను Inferentiaకి పోర్ట్ చేయడం అనేది కేవలం కంపైల్-అండ్-రన్ చేసే ప్రక్రియ మాత్రమే కాదు. ఇది టోకెన్ ఈక్వాలిటీ కంటే మించిన కఠినమైన ఫంక్షనల్ టెస్టింగ్‌ను, ప్రతి టెన్సర్—పారామీటర్ లేదా బఫర్—సరైన పద్ధతిలో హ్యాండిల్ చేయబడిందని నిశితంగా వెరిఫై చేయడాన్ని మరియు స్పాట్-ఇన్‌స్టాన్స్ రీక్లెయిమెంట్‌ను ముందుగానే ఊహించే డిప్లాయ్‌మెంట్ వ్యూహాన్ని కోరుతుంది. చివరగా, షార్డింగ్ అనేది మోడల్ యొక్క అంతర్గత అటెన్షన్ జ్యామితిని (attention geometry) గౌరవించాలి; లేకపోతే, వేగాన్ని పెంచే ప్యారలలిజం అనేది సైలెంట్ ఫెయిల్యూర్లకు మూలమవుతుంది.