Google’s Tunix system inaondoa kizuizi kikubwa kilichozuia agentic reinforcement learning (RL) ya kiwango kikubwa kutumia TPU kwa ufanisi. Kwa kutenganisha kazi ya kuzalisha data ya mwingiliano (interaction data) kutoka kwa kazi ya kusasisha sera (policy), Tunix inaongeza matumizi ya TPU kutoka asilimia za tarakimu moja hadi karibu uwezo kamili, ikipunguza upotevu wa rasilimali za kompyuta kwa kiasi kikubwa.

Kizuizi katika agentic RL

Agentic RL ni tofauti na mafunzo ya modeli ya lugha ya "next-token" ambayo ni ya kawaida zaidi. Wakala (agent) lazima atume API calls, atekeleze kodi, au apite katika mazingira ya simulation, kisha arekodi matokeo. Mzunguko wa mafunzo (training loop) kwa hivyo ni synchronous: modeli inazalisha kitendo, mazingira yanaendeshwa, matokeo yanarudi, na baada ya hapo ndipo modeli inapopokea gradient update. Wakati hatua moja ya mazingira inachukua sekunde kadhaa, vifaa vya TPU vya gharama kubwa vinakaa bila kufanya kazi, na matumizi yanayoripotiwa yanaweza kushuka chini ya 10%. Kutofanya kazi kwa ufanisi kunasababisha moja kwa moja gharama kubwa za cloud na mizunguko ya utafiti kuwa ya polepole.

Muundo wa Tunix uliotenganishwa (decoupled architecture)

Tunix inashughulikia tatizo hili kwa kuhamisha hatua mbili—uundaji wa trajectory na uboreshaji wa sera (policy optimization)—kwenye makundi tofauti ya vifaa (hardware pools).

  • Asynchronous actors hukimbia kwenye CPU au GPU zisizo na gharama kubwa. Kila actor hujihusisha mara kwa mara na mazingira aliyopangiwa, anarekodi vitendo na uchunguzi, na kutuma trajectory zinazozalishwa kwenye hifadhi ya pamoja.
  • Continuous learners huchukua TPU Pods maalum. Mwanafunzi (learner) huchukua batches kutoka kwenye buffer kuu na kufanya gradient updates bila kusubiri actor yeyote kumaliza rollout.
  • High-throughput buffer inakaa katikati, ikifanya kazi kama eneo la mpito kwa ajili ya trajectory. Kwa sababu mwanafunzi anaweza kusoma kwa kasi ile ile ambayo buffer inaweza kutoa data, TPU haisimami kamwe.

Athari yake ya jumla ni mtiririko wa mafunzo ambapo TPU zinabaki zikiwa na kazi karibu wakati wote, zikifikisha matumizi kuelekea 100%.

Vikwazo vya kiufundi na jinsi Tunix inavyovishinda

Vipindi vya urefu tofauti (Variable-length episodes) na XLA recompilation

Kichujio (compiler) cha XLA cha JAX kinatengeneza matokeo kulingana na tensor shapes zilizofungwa. Hata hivyo, kazi za agentic huzalisha mfuatano (sequences) zenye urefu tofauti, jambo ambalo kwa kawaida lingesababisha recompilations za gharama kubwa. Tunix inapakia mfuatano mifupi pamoja na kuweka vipindi vya urefu unaofanana katika makundi (buckets), ikifanya shapes ziwe thabiti kwa muda mrefu wa kutosha ili XLA itumie tena compiled kernels. Matokeo yake ni uwezo thabiti wa kupitisha data (throughput) bila mzigo wa compiler ambao ungeharibu utendaji.

Kukuza modeli kubwa kwenye chip nyingi za TPU

Kufundisha wakala (agents) wenye zaidi ya vigezo (parameters) bilioni 70 kunahitaji kusambaza uzito (weights) na data kwenye node nyingi za TPU. Tunix inatumia JAX ShardMap primitive kusambaza (shard) vigezo vya modeli na activations, ikimruhusu mwanafunzi kuweka modeli nzima kwenye kumbukumbu (memory) huku akiendelea kuilisha data kwa kasi kubwa. Mkakati huu wa usambazaji (sharding) unawezesha kufundisha modeli ambazo hapo awali hazikuweza kufikiwa na TPU pod moja.

Gradient zilizopitwa na wakati (Stale gradients) kutoka kwenye mifumo iliyotenganishwa

Wakati actors wanapokimbia mbele ya mwanafunzi, data wanayotoa inaweza kuwa "stale" (iliyopitwa na wakati) kulinganisha na sera ya sasa. Tunix inapunguza mkengeuko huu kwa njia mbili: importance-sampling inatunza sampuli za zamani ili kuonyesha umuhimu wake, na kizingiti cha staleness kinachoweza kurekebishwa kinatupa trajectory zinazozidi umri uliowekwa. Kwa pamoja, zinafanya mafunzo yawe thabiti hata wakati mtiririko unakimbia kwa njia ya asynchronous.

Mambo ambayo watumiaji wanapaswa kuzingatia

  • Ukaguzi wa latency – Faida ya kutenganisha inategemea muda wa majibu ya mazingira. Timu zinapaswa kupima latency ya mwisho hadi mwisho na kuhakikisha kuwa actor pools zimeundwa ili kuweka buffer ikiwa imejazwa vizuri.
  • Muundo wa worker pool – CPU au GPU rahisi zinaweza kuhifadhi actors wengi, lakini kuwapa kazi nyingi kupita kiasi kunaweza kusababisha msongamano kwenye mtandao au hifadhi. Worker pool iliyobalansi inayolingana na kasi ya buffer kuingiza data ni muhimu.
  • Uimara wa buffer – Hifadhi kuu lazima ihimili viwango vya juu vya kuandika na kusoma bila kuwa kizuizi kipya. Kuchagua mfumo wa hifadhi wenye latency ndogo na upana wa kutosha wa bandwidth ni sehemu isiyoweza kujadiliwa ya muundo huu.

Hasara zinazoweza kutokea

Muundo huu uliogawanyika unaleta sehemu nyingi zinazohusika: makundi ya vifaa tofauti, buffer ya kudumu, na mantiki ya uratibu ili kusimamia mipaka ya staleness.

Muhtasari

Tunix inaonyesha kuwa gharama kuu katika agentic RL si modeli yenyewe bali muda wa kukaa bila kazi unaosababishwa na mzunguko wa mwingiliano wa synchronous. Kwa kuhamishia kazi ya rollout kwenye vifaa rahisi na kuilisha TPU pod inayojifunza mfululizo kutoka kwenye buffer yenye uwezo mkubwa, Google imebadilisha tatizo la matumizi ya chini ya 10% kuwa mtiririko wa kazi wenye uwezo karibu kamili.