Google-ന്റെ Tunix സിസ്റ്റം, ഏജന്റിക് റൈൻഫോഴ്‌സ്‌മെന്റ് ലേണിംഗ് (RL) വലിയ തോതിൽ നടത്തുമ്പോൾ TPU-കൾ കാര്യക്ഷമമായി ഉപയോഗിക്കുന്നതിന് തടസ്സമായി നിൽക്കുന്ന ഘടകത്തെ നീക്കം ചെയ്യുന്നു. ഇന്ററാക്ഷൻ ഡാറ്റാ ജനറേറ്റ് ചെയ്യുന്ന ജോലിയെയും പോളിസി അപ്‌ഡേറ്റ് ചെയ്യുന്ന ജോലിയെയും വേർതിരിക്കുന്നതിലൂടെ, Tunix TPU ഉപയോഗക്ഷമത ഒറ്റയക്ക ശതമാനത്തിൽ നിന്ന് പൂർണ്ണ ശേഷിക്കടുത്തേക്ക് ഉയർത്തുകയും കമ്പ്യൂട്ട് പാഴാകുന്നത് ഗണ്യമായി കുറയ്ക്കുകയും ചെയ്യുന്നു.

ഏജന്റിക് RL-ലെ തടസ്സങ്ങൾ

Agentic RL എന്നത് നമുക്ക് പരിചിതമായ "next-token" ലാംഗ്വേജ് മോഡൽ ട്രെയിനിംഗിൽ നിന്ന് വ്യത്യസ്തമാണ്. ഒരു ഏജന്റ് API കോളുകൾ അയക്കുകയോ, കോഡ് പ്രവർത്തിപ്പിക്കുകയോ, അല്ലെങ്കിൽ ഒരു സിമുലേറ്റഡ് എൻവയോൺമെന്റിലൂടെ കടന്നുപോവുകയോ ചെയ്ത ശേഷം അതിന്റെ ഫലത്തോട് പ്രതികരിക്കേണ്ടതുണ്ട്. അതിനാൽ ട്രെയിനിംഗ് ലൂപ്പ് സിൻക്രണസ് (synchronous) ആണ്: മോഡൽ ഒരു ആക്ഷൻ എടുക്കുന്നു, എൻവയോൺമെന്റ് പ്രവർത്തിക്കുന്നു, ഫലം തിരികെ ലഭിക്കുന്നു, അതിനുശേഷം മാത്രമേ മോഡലിന് ഒരു ഗ്രേഡിയന്റ് അപ്‌ഡേറ്റ് ലഭിക്കുകയുള്ളൂ. ഒരു എൻവയോൺമെന്റ് സ്റ്റെപ്പിന് ഏതാനും സെക്കൻഡുകൾ എടുക്കുമ്പോൾ, വിലകൂടിയ TPU ഹാർഡ്‌വെയർ ഉപയോഗശൂന്യമായി ഇരിക്കുകയും ഉപയോഗക്ഷമത 10% ന് താഴെയാവുകയും ചെയ്യുന്നു. ഈ കാര്യക്ഷമതയില്ലായ്മ നേരിട്ട് ഉയർന്ന ക്ലൗഡ് ബില്ലുകളിലേക്കും ഗവേഷണ വേഗത കുറയുന്നതിലേക്കും നയിക്കുന്നു.

Tunix-ന്റെ ഡി കപ്‌ൾഡ് ആർക്കിടെക്ചർ

ട്രജക്റ്ററി ജനറേഷൻ (trajectory generation), പോളിസി ഒപ്റ്റിമൈസേഷൻ (policy optimization) എന്നീ രണ്ട് ഘട്ടങ്ങളെ വ്യത്യസ്ത ഹാർഡ്‌വെയർ പൂളുകളിലേക്ക് മാറ്റിക്കൊണ്ട് Tunix ഈ പ്രശ്നത്തെ നേരിടുന്നു.

  • Asynchronous actors: ഇവ കുറഞ്ഞ ചിലവുള്ള CPU-കളിലോ GPU-കളിലോ പ്രവർത്തിക്കുന്നു. ഓരോ ആക്ടറും അതിന്റെ നിശ്ചിത എൻവയോൺമെന്റുമായി നിരന്തരം സംവദിക്കുകയും, ആക്ഷനുകളും ഒബ്സർവേഷനുകളും രേഖപ്പെടുത്തുകയും, resulting trajectories ഒരു ഷെയർഡ് സ്റ്റോറിലേക്ക് സ്ട്രീം ചെയ്യുകയും ചെയ്യുന്നു.
  • Continuous learners: ഇവ പ്രത്യേക TPU Pod-കളിൽ പ്രവർത്തിക്കുന്നു. ലേണർ സെൻട്രൽ ബഫറിൽ നിന്ന് ബാച്ചുകൾ എടുക്കുകയും ഒരു ആക്ടറും അതിന്റെ റോളൗട്ട് പൂർത്തിയാക്കാൻ കാത്തുനിൽക്കാതെ തന്നെ ഗ്രേഡിയന്റ് അപ്‌ഡേറ്റുകൾ നടത്തുകയും ചെയ്യുന്നു.
  • High-throughput buffer: ഇത് നടുവിൽ ഒരു സ്റ്റേജിംഗ് ഏരിയയായി പ്രവർത്തിക്കുന്നു. ബഫറിന് ഡാറ്റ നൽകാൻ കഴിയുന്ന വേഗതയിൽ ലേണറിന് അത് വായിക്കാൻ കഴിയുന്നതിനാൽ, TPU ഒരിക്കലും തടസ്സപ്പെടില്ല.

ഇതിന്റെ ഫലമായി TPUs മിക്കവാറും എല്ലാ സമയത്തും സജീവമായിരിക്കുന്ന ഒരു ട്രെയിനിംഗ് പൈപ്പ്‌ലൈൻ ലഭിക്കുന്നു, ഇത് ഉപയോഗക്ഷമത 100% ലേക്ക് എത്തിക്കുന്നു.

സാങ്കേതിക വെല്ലുവിളികളും Tunix അവയെ എങ്ങനെ മറികടക്കുന്നു എന്നതും

വ്യത്യസ്ത ദൈർഘ്യമുള്ള എപ്പിസോഡുകളും XLA റീകംപൈലേഷനും

JAX-ന്റെ XLA കംപൈലർ നിശ്ചിത ടെൻസർ ഷേപ്പുകൾക്കായി (fixed tensor shapes) ഒപ്റ്റിമൈസ് ചെയ്യുന്നു. എന്നാൽ ഏജന്റിക് ടാസ്ക്കുകൾ വ്യത്യസ്ത ദൈർഘ്യമുള്ള സീക്വൻസുകൾ ഉൽപ്പാദിപ്പിക്കുന്നു, ഇത് സാധാരണഗതിയിൽ ചിലവേറിയ റീകംപൈലേഷനുകൾക്ക് കാരണമാകും. Tunix ചെറിയ സീക്വൻസുകളെ ഒന്നിച്ച് പാക്ക് ചെയ്യുകയും സമാനമായ ദൈർഘ്യമുള്ള എപ്പിസോഡുകളെ ബക്കറ്റുകളായി ഗ്രൂപ്പ് ചെയ്യുകയും ചെയ്യുന്നു. ഇത് XLA-യ്ക്ക് കംപൈൽ ചെയ്ത കേർണലുകൾ (compiled kernels) വീണ്ടും ഉപയോഗിക്കാൻ ആവശ്യമായ സ്ഥിരത നൽകുന്നു. ഇതിലൂടെ കംപൈലർ മൂലമുണ്ടാകുന്ന കാലതാമസം ഇല്ലാതെ മികച്ച പ്രകടനം കാഴ്ചവെക്കാൻ സാധിക്കുന്നു.

നിരവധി TPU ചിപ്പുകളിലായി വൻകിട മോഡലുകളെ സ്കെയിൽ ചെയ്യുക

70 ബില്യൺ പാരമീറ്ററുകളിൽ കൂടുതൽ ഉള്ള ഏജന്റുകളെ പരിശീലിപ്പിക്കാൻ മോഡൽ പാരമീറ്ററുകളും ഡാറ്റയും ഒന്നിലധികം TPU നോഡുകളിലായി വിന്യസിക്കേണ്ടതുണ്ട്. Tunix, JAX-ന്റെ ShardMap പ്രിമിറ്റീവ് ഉപയോഗിച്ച് മോഡൽ പാരമീറ്ററുകളെയും ആക്ടിവേഷനുകളെയും ഷാർഡ് (shard) ചെയ്യുന്നു. ഇത് ഉയർന്ന വേഗതയിൽ ഡാറ്റ നൽകുന്നതിനോടൊപ്പം തന്നെ മുഴുവൻ മോഡലും മെമ്മറിയിൽ നിലനിർത്താൻ ലേണറെ അനുവദിക്കുന്നു. ഈ ഷാർഡിംഗ് തന്ത്രം മുമ്പ് ഒരു സിംഗിൾ TPU പോഡിൽ എത്തിപ്പിടിക്കാൻ കഴിയാതിരുന്ന മോഡലുകളെ പരിശീലിപ്പിക്കാൻ സഹായിക്കുന്നു.

ഡി കപ്‌ൾഡ് പൈപ്പ്‌ലൈനുകളിൽ നിന്നുള്ള സ്റ്റേൽ ഗ്രേഡിയന്റുകൾ

ആക്ടർമാർ ലേണറിനേക്കാൾ മുന്നോട്ട് പോകുമ്പോൾ, അവർ നൽകുന്ന ഡാറ്റ നിലവിലെ പോളിസിയെ അപേക്ഷിച്ച് "സ്റ്റേൽ" (stale) ആകാൻ സാധ്യതയുണ്ട്. Tunix ഇത് രണ്ട് രീതിയിൽ പരിഹരിക്കുന്നു: ഒന്ന്, ഇംപോർട്ടൻസ്-സാമ്പിളിംഗ് (importance-sampling) ഉപയോഗിച്ച് പഴയ സാമ്പിളുകളുടെ പ്രസക്തി കണക്കിലെടുത്ത് അവയെ റീവെയ്റ്റ് ചെയ്യുന്നു; രണ്ട്, നിശ്ചിത പ്രായപരിധി കഴിഞ്ഞ ട്രജക്റ്ററികളെ ഒഴിവാക്കാൻ ഒരു കോൺഫിഗറബിൾ സ്റ്റേൽനസ് ത്രെഷോൾഡ് (staleness threshold) ഉപയോഗിക്കുന്നു. ഇവ രണ്ടും ചേർന്ന് പൈപ്പ്‌ലൈൻ അസിൻക്രണസ് ആയി പ്രവർത്തിക്കുമ്പോഴും ലേണിംഗ് സ്ഥിരതയുള്ളതാക്കി നിലനിർത്തുന്നു.

സ്വീകരിക്കുന്നവർ ശ്രദ്ധിക്കേണ്ട കാര്യങ്ങൾ

  • Latency audit – ഡി കപ്‌ളിംഗിന്റെ ഗുണം എൻവയോൺമെന്റിന്റെ റെസ്പോൺസ് ടൈമിനെ ആശ്രയിച്ചിരിക്കുന്നു. ടീമുകൾ എൻഡ്-ടു-എൻഡ് ലേറ്റൻസി അളക്കേണ്ടതും ബഫർ എപ്പോഴും നിറഞ്ഞിരിക്കാൻ ആവശ്യമായ രീതിയിൽ ആക്ടർ പൂളുകൾ ക്രമീകരിക്കേണ്ടതും അത്യാവശ്യമാണ്.
  • Worker pool design – കുറഞ്ഞ ചിലവുള്ള CPU-കളോ GPU-കളോ നിരവധി ആക്ടർമാരെ ഉൾക്കൊള്ളနိုင်ുമെങ്കിലും, അവ അമിതമായി ഉപയോഗിക്കുന്നത് നെറ്റ്‌വർക്ക് അല്ലെങ്കിൽ സ്റ്റോറേജ് തടസ്സങ്ങൾക്ക് കാരണമായേക്കാം. ബഫറിന്റെ ഇൻജസ്റ്റ് റേറ്റുമായി (ingest rate) പൊരുത്തപ്പെടുന്ന ഒരു ബാലൻസ്ഡ് പൂൾ അത്യാവശ്യമാണ്.
  • Buffer robustness – സെൻട്രൽ സ്റ്റോറിന് പുതിയൊരു തടസ്സമായി മാറാതെ തന്നെ ഉയർന്ന റൈറ്റ്, റീഡ് നിരക്കുകൾ കൈകാര്യം ചെയ്യാൻ കഴിയണം. കുറഞ്ഞ ലേറ്റൻസിയും മതിയായ ബാൻഡ്‌വിഡ്ത്തും ഉള്ള ഒരു സ്റ്റോറേജ് സിസ്റ്റം തിരഞ്ഞെടുക്കുന്നത് ഈ ആർക്കിടെക്ചറിന്റെ ഒഴിച്ചുകൂടാനാവാത്ത ഭാഗമാണ്.

സാധ്യമായ പോരായ്മകൾ

ഈ സ്പ്ലിറ്റ് ആർക്കിടെക്ചർ കൂടുതൽ ഘടകങ്ങൾ introduces ചെയ്യുന്നു: പ്രത്യേക ഹാർഡ്‌വെയർ ഫ്ലീറ്റുകൾ, ഒരു പെർസിസ്റ്റന്റ് ബഫർ, കൂടാതെ സ്റ്റേൽനസ് പരിധികൾ നടപ്പിലാക്കുന്നതിനുള്ള കോർഡിനേഷൻ ലോജിക് എന്നിവ ഇതിൽ ഉൾപ്പെടുന്നു.

ചുരുക്കം

ഏജന്റിക് RL-ലെ പ്രധാന ചെലവ് മോഡൽ അല്ലെന്നും, മറിച്ച് സിൻക്രണസ് ഇന്ററാക്ഷൻ ലൂപ്പുകൾ മൂലമുണ്ടാകുന്ന ഉപയോഗശൂന്യമായ സമയമാണെന്നും Tunix കാണിച്ചുതരുന്നു. റോളൗട്ട് ജോലികൾ കുറഞ്ഞ ചിലവുള്ള ഹാർഡ്‌വെയറിലേക്ക് മാറ്റുന്നതിലൂടെയും, ഹൈ-ത്രൂപുട്ട് ബഫറിൽ നിന്ന് ഒരു കണ്ടിന്യൂസ് ലേണിംഗ് TPU പോഡിന് ഡാറ്റ നൽകുന്നതിലൂടെയും, Google 10%-ൽ താഴെയുള്ള ഉപയോഗക്ഷമതയെ പൂർണ്ണ ശേഷിക്കടുത്തുള്ള ഒരു വർക്ക്ഫ്ലോയായി മാറ്റിയിരിക്കുന്നു.