Google యొక్క Tunix సిస్టమ్, భారీ స్థాయి ఏజెంటిక్ రీఇన్‌ఫోర్స్‌మెంట్ లెర్నింగ్ (RL) ప్రక్రియలో TPUsని సమర్థవంతంగా ఉపయోగించకుండా అడ్డుకుంటున్న అడ్డంకిని (choke point) తొలగిస్తుంది. ఇంటరాక్షన్ డేటాను రూపొందించే పనిని మరియు పాలసీని అప్‌డేట్ చేసే పనిని వేరు చేయడం ద్వారా, Tunix TPU వినియోగాన్ని (utilization) ఒక అంకె శాతాల నుండి దాదాపు పూర్తి సామర్థ్యం వరకు పెంచుతుంది, తద్వారా కంప్యూట్ వృధాను గణనీయంగా తగ్గిస్తుంది.

ఏజెంటిక్ RLలో అడ్డంకి (bottleneck)

ఏజెంటిక్ RL అనేది మనకు బాగా తెలిసిన "next-token" లాంగ్వేజ్-మోడల్ ట్రైనింగ్ కంటే భిన్నమైనది. ఒక ఏజెంట్ API కాల్స్‌ను పంపాలి, కోడ్‌ను అమలు చేయాలి లేదా ఒక సిమ్యులేటెడ్ ఎన్విరాన్‌మెంట్‌లో అడుగులు వేయాలి, ఆపై ఫలితానికి అనుగుణంగా స్పందించాలి. అందువల్ల, ట్రైనింగ్ లూప్ సింక్రోనస్ (synchronous) గా ఉంటుంది: మోడల్ ఒక చర్యను (action) చేస్తుంది, ఎన్విరాన్‌మెంట్ రన్ అవుతుంది, ఫలితం తిరిగి వస్తుంది, ఆ తర్వాతే మోడల్‌కు గ్రేడియంట్ అప్‌డేట్ అందుతుంది. ఒకే ఎన్విరాన్‌మెంట్ స్టెప్ తీసుకోవడానికి కొన్ని సెకన్ల సమయం పట్టినప్పుడు, ఖరీదైన TPU హార్డ్‌వేర్ ఖాళీగా (idle) ఉంటుంది, దీనివల్ల వినియోగం (utilization) 10% కంటే తక్కువకు పడిపోవచ్చు. ఈ అసమర్థత నేరుగా అధిక క్లౌడ్ బిల్లులకు మరియు నెమ్మదైన పరిశోధన చక్రాలకు దారితీస్తుంది.

Tunix యొక్క డికప్లడ్ ఆర్కిటెక్చర్ (decoupled architecture)

Tunix సమస్యను పరిష్కరించడానికి రెండు దశలను—ట్రాజెక్టరీ జనరేషన్ (trajectory generation) మరియు పాలసీ ఆప్టిమైజేషన్ (policy optimization)—వేర్వేరు హార్డ్‌వేర్ పూల్స్‌లోకి మారుస్తుంది.

  • Asynchronous actors తక్కువ ఖరీదైన CPUs లేదా GPUsలపై నడుస్తాయి. ప్రతి యాక్టర్ తనకు కేటాయించిన ఎన్విరాన్‌మెంట్‌తో నిరంతరం ఇంటరాక్ట్ అవుతుంది, చర్యలు మరియు పరిశీలనలను (observations) రికార్డ్ చేస్తుంది, మరియు ఫలితంగా వచ్చే ట్రాజెక్టరీలను ఒక షేర్డ్ స్టోర్‌లోకి స్ట్రీమ్ చేస్తుంది.
  • Continuous learners ప్రత్యేక TPU Podsని ఉపయోగిస్తాయి. లెర్నర్ సెంట్రల్ బఫర్ నుండి బ్యాచ్‌లను తీసుకుంటుంది మరియు ఏ ఒక్క యాక్టర్ రోల్‌అవుట్ (rollout) పూర్తయ్యే వరకు వేచి ఉండకుండానే గ్రేడియంట్ అప్‌డేట్‌లను చేస్తుంది.
  • High-throughput buffer మధ్యలో ఉండి, ట్రాజెక్టరీల కోసం స్టేజింగ్ ఏరియాగా పనిచేస్తుంది. బఫర్ డేటాను ఎంత వేగంగా అందిస్తుందో, లెర్నర్ కూడా అంతే వేగంగా చదవగలదు కాబట్టి, TPU ఎప్పుడూ ఆగిపోదు (stalls).

దీని ఫలితంగా, TPUs దాదాపు ఎల్లప్పుడూ బిజీగా ఉండే ట్రైనింగ్ పైప్‌లైన్ ఏర్పడుతుంది, ఇది వినియోగాన్ని (utilization) 100% వైపు నడిపిస్తుంది.

సాంకేతిక అడ్డంకులు మరియు Tunix వాటిని ఎలా అధిగమిస్తుంది

వేరియబుల్-లెంగ్త్ ఎపిసోడ్లు మరియు XLA రీకంపిలేషన్ (recompilation)

JAX యొక్క XLA కంపైలర్ ఫిక్స్‌డ్ టెన్సర్ షేప్స్ (fixed tensor shapes) కోసం ఆప్టిమైజ్ చేస్తుంది. అయితే, ఏజెంటిక్ టాస్క్‌లు వేర్వేరు పొడవుల సీక్వెన్స్‌లను ఉత్పత్తి చేస్తాయి, ఇవి సాధారణంగా ఖరీదైన రీకంపిలేషన్‌లకు దారితీస్తాయి. Tunix చిన్న సీక్వెన్స్‌లను కలిపి ప్యాక్ చేస్తుంది మరియు ఒకే రకమైన పొడవు ఉన్న ఎపిసోడ్‌లను బకెట్లలో గ్రూప్ చేస్తుంది, తద్వారా XLA కంపైల్డ్ కెర్నల్స్‌ను (compiled kernels) తిరిగి ఉపయోగించుకోవడానికి వీలుగా షేప్స్‌ను స్థిరంగా ఉంచుతుంది. దీని ఫలితంగా, పనితీరును దెబ్బతీసే కంపైలర్ ఓవర్‌హెడ్ లేకుండా స్థిరమైన త్రూపుట్ (throughput) లభిస్తుంది.

అనేక TPU చిప్‌ల ద్వారా భారీ మోడళ్లను స్కేల్ చేయడం

70 బిలియన్ల కంటే ఎక్కువ పారామితులు ఉన్న ఏజెంట్‌లను ట్రైన్ చేయడానికి వెయిట్స్ (weights) మరియు డేటాను బహుళ TPU నోడ్ల ద్వారా విస్తరించాల్సి ఉంటుంది. Tunix, JAX యొక్క ShardMap ప్రిమిటివ్‌ను ఉపయోగించి మోడల్ పారామితులు మరియు యాక్టివేషన్‌లు రెండింటినీ షార్డ్ (shard) చేస్తుంది, దీనివల్ల లెర్నర్ హై-స్పీడ్‌లో డేటాను అందిస్తూనే మొత్తం మోడల్‌ను మెమరీలో ఉంచుకోగలదు. ఈ షార్డింగ్ వ్యూహం వల్ల, గతంలో ఒకే TPU పోడ్ పరిధికి మించిన భారీ మోడళ్లను కూడా ట్రైన్ చేయడం సాధ్యమవుతుంది.

డికప్లడ్ పైప్‌లైన్‌ల నుండి వచ్చే స్టేల్ గ్రేడియంట్లు (Stale gradients)

యాక్టర్లు లెర్నర్‌ కంటే ముందుగా నడుస్తున్నప్పుడు, వారు అందించే డేటా ప్రస్తుత పాలసీకి సంబంధించి "స్టేల్" (stale - పాతది) కావచ్చు. Tunix ఈ తేడాను రెండు పద్ధతుల ద్వారా తగ్గిస్తుంది: ఇంపార్టెన్స్-శాంప్లింగ్ (importance-sampling) పాత శాంపిల్స్ యొక్క ప్రాముఖ్యతను ప్రతిబింబించేలా వాటిని రీవెయిట్ చేస్తుంది, మరియు కాన్ఫిగర్ చేయదగిన స్టేల్‌నెస్ త్రెషోల్డ్ (staleness threshold) నిర్ణీత వయస్సును మించిన ట్రాజెక్టరీలను తొలగిస్తుంది. ఇవి రెండూ కలిసి, పైప్‌లైన్ అసమకాలికంగా (asynchronously) నడుస్తున్నప్పటికీ లెర్నింగ్‌ను స్థిరంగా ఉంచుతాయి.

వినియోగదారులు గమనించవలసిన అంశాలు

  • Latency audit – డికప్లింగ్ వల్ల కలిగే ప్రయోజనం ఎన్విరాన్‌మెంట్ యొక్క రెస్పాన్స్ టైమ్‌పై ఆధారపడి ఉంటుంది. బఫర్ నిండుగా ఉండేలా చూసుకోవడానికి, టీమ్‌లు ఎండ్-టు-ఎండ్ లేటెన్సీని (end-to-end latency) కొలవాలి మరియు యాక్టర్ పూల్స్ పరిమాణాన్ని సరిగ్గా నిర్ణయించాలి.
  • Worker pool design – తక్కువ ఖరీదైన CPUs లేదా GPUs అనేక యాక్టర్లను హోస్ట్ చేయగలవు, కానీ వాటిని అతిగా ఉపయోగించడం (oversubscribing) వల్ల నెట్‌వర్క్ లేదా స్టోరేజ్ సమస్యలు రావచ్చు. బఫర్ యొక్క ఇంజెస్ట్ రేట్‌కు (ingest rate) సరిపోయే సమతుల్యమైన పూల్ అవసరం.
  • Buffer robustness – సెంట్రల్ స్టోర్ కొత్త అడ్డంకిగా మారకుండా, అధిక రైట్ మరియు రీడ్ రేట్లను తట్టుకోగలగాలి. తక్కువ టెయిల్ లేటెన్సీ (tail latency) మరియు తగిన బ్యాండ్‌విడ్త్ ఉన్న స్టోరేజ్ సిస్టమ్‌ను ఎంచుకోవడం ఈ ఆర్కిటెక్చర్‌లో అత్యంత కీలకం.

సంభావ్య ప్రతికూలతలు (Potential downsides)

ఈ విడి ఆర్కిటెక్చర్ వల్ల మరిన్ని అంశాలను నిర్వహించాల్సి ఉంటుంది: వేర్వేరు హార్డ్‌వేర్ ఫ్లీట్లు, పర్సిస్టెంట్ బఫర్ మరియు స్టేల్‌నెస్ పరిమితులను అమలు చేయడానికి కోఆర్డినేషన్ లాజిక్.

ముగింపు (Takeaway)

ఏజెంటిక్ RLలో ప్రధాన ఖర్చు మోడల్ వల్ల కాకుండా, సింక్రోనస్ ఇంటరాక్షన్ లూప్‌ల వల్ల కలిగే ఖాళీ సమయం (idle time) వల్లేనని Tunix నిరూపిస్తుంది. రోల్‌అవుట్ పనిని తక్కువ ఖరీదైన హార్డ్‌వేర్‌కు బదిలీ చేయడం మరియు హై-త్రూపుట్ బఫర్ ద్వారా నిరంతరం నేర్చుకునే TPU పోడ్‌కు డేటాను అందించడం ద్వారా, Google 10% కంటే తక్కువగా ఉన్న వినియోగ సమస్యను దాదాపు పూర్తి సామర్థ్యం కలిగిన వర్క్‌ఫ్లోగా మార్చింది.