Google-ன் Tunix அமைப்பு, பெரிய அளவிலான ஏஜென்டிக் ரீஇன்ஃபோர்ஸ்மென்ட் லேர்னிங் (RL) முறைகள் TPUs-களை திறம்பட பயன்படுத்துவதைத் தடுத்து வந்த நெரிசல் புள்ளியை (choke point) நீக்குகிறது. தரவுகளை உருவாக்கும் பணியையும், பாலிசியைப் புதுப்பிக்கும் பணியையும் தனித்தனியாகப் பிரிப்பதன் மூலம், Tunix TPU பயன்பாட்டை ஒற்றை இலக்க சதவீதத்திலிருந்து முழுத் திறனுக்கு அருகாமையில் கொண்டு வருகிறது, இதன் மூலம் கணக்கீட்டு வீணாவதைக் (compute waste) கடுமையாகக் குறைக்கிறது.

ஏஜென்டிக் RL-ல் உள்ள நெரிசல்

ஏஜென்டிக் RL என்பது வழக்கமான “அடுத்த-டோக்கன்” (next-token) மொழி மாதிரிப் பயிற்சியிலிருந்து மாறுபட்டது. ஒரு ஏஜென்ட் API அழைப்புகளை அனுப்ப வேண்டும், குறியீட்டை (code) இயக்க வேண்டும் அல்லது ஒரு உருவகப்படுத்தப்பட்ட சூழலில் (simulated environment) செயல்பட வேண்டும், பின்னர் அதன் முடிவிற்கு எதிர்வினையாற்ற வேண்டும். எனவே, பயிற்சியின் சுழற்சி ஒத்திசைவானது (synchronous): மாதிரி ஒரு செயலைச் செய்கிறது, சூழல் இயங்குகிறது, முடிவு திரும்புகிறது, அதன் பின்னரே மாதிரிக்கு ஒரு கிரேடியன்ட் அப்டேட் (gradient update) கிடைக்கிறது. ஒரு சூழல் படி (environment step) பல வினாடிகள் எடுக்கும்போது, விலை உயர்ந்த TPU வன்பொருள் செயலற்ற நிலையில் (idle) உள்ளது, மேலும் அதன் பயன்பாடு 10%-க்கும் கீழே குறையக்கூடும். இந்தத் திறமையின்மை நேரடியாக அதிக கிளவுட் கட்டணங்களுக்கும், மெதுவான ஆராய்ச்சி சுழற்சிகளுக்கும் வழிவகுக்கிறது.

Tunix-ன் பிரிக்கப்பட்ட கட்டமைப்பு

டிராக்டரி உருவாக்கம் (trajectory generation) மற்றும் பாலிசி ஆப்டிமைசேஷன் (policy optimization) ஆகிய இரண்டு நிலைகளையும் தனித்தனி வன்பொருள் தொகுதிகளுக்கு (hardware pools) பிரிப்பதன் மூலம் Tunix இந்தப் பிரச்சனையைத் தீர்க்கிறது.

  • அசிங்க்ரோனஸ் ஆக்டர்கள் (Asynchronous actors) மலிவான CPU அல்லது GPU-களில் இயங்குகின்றன. ஒவ்வொரு ஆக்டரும் தனக்கு ஒதுக்கப்பட்ட சூழலுடன் தொடர்ந்து தொடர்பு கொண்டு, செயல்கள் மற்றும் அவதானிப்புகளைப் பதிவு செய்து, resulting trajectories-களை ஒரு பகிரப்பட்ட சேமிப்பகத்திற்கு (shared store) அனுப்புகிறது.
  • தொடர்ச்சியான லேர்னர்கள் (Continuous learners) பிரத்யேக TPU Pods-களை ஆக்கிரமிக்கின்றன. லேர்னர் மையப் பஃபரிலிருந்து (central buffer) தொகுதிகளைப் பெற்று, எந்தவொரு தனிப்பட்ட ஆக்டரும் தனது வேலையை முடிக்கும் வரை காத்திருக்காமல் கிரேடியன்ட் அப்டேட்களைச் செய்கிறது.
  • அதிகத் திறன் கொண்ட பஃபர் (High-throughput buffer) நடுவில் அமர்ந்து, டிராக்டரிகளுக்கான ஒரு இடைநிலைத் தளமாகச் செயல்படுகிறது. பஃபர் தரவை வழங்கும் வேகத்திலேயே லேர்னரால் படிக்க முடிவதால், TPU ஒருபோதும் தடையின்றி இயங்குகிறது.

இதன் இறுதி விளைவு என்னவென்றால், TPUs கிட்டத்தட்ட எப்போதும் வேலையில் இருக்கும் ஒரு பயிற்சிப் பாதையாகும் (training pipeline), இது பயன்பாட்டை 100%-ஐ நோக்கித் தள்ளுகிறது.

தொழில்நுட்பத் தடைகள் மற்றும் Tunix அவற்றை எவ்வாறு முறியடிக்கிறது

மாறுபட்ட நீளம் கொண்ட எபிசோட்கள் மற்றும் XLA ரீகம்பைலேஷன்

JAX-ன் XLA கம்பைலர் நிலையான டென்சர் வடிவங்களுக்காக (fixed tensor shapes) மேம்படுத்தப்பட்டுள்ளது. இருப்பினும், ஏஜென்டிக் பணிகள் மாறுபட்ட நீளம் கொண்ட தொடர்களை உருவாக்குகின்றன, இது வழக்கமாக அதிக செலவு பிடிக்கும் ரீகம்பைலேஷன்களைத் தூண்டும். Tunix குறுகிய தொடர்களை ஒன்றாக இணைத்து, ஒரே மாதிரியான நீளம் கொண்ட எபிசோட்களைத் தொகுப்புகளாக (buckets) குழுவாக்குகிறது, இதன் மூலம் XLA ஏற்கனவே கம்பைல் செய்யப்பட்ட கர்னல்களை (compiled kernels) மீண்டும் பயன்படுத்தும் அளவுக்கு வடிவங்களை நிலையாக வைத்திருக்கிறது. இதன் விளைவாக, செயல்திறனைப் பாதிக்கும் கம்பைலர் சுமை இன்றி சீரானத் திறன் கிடைக்கிறது.

பல TPU சிப்களில் பிரம்மாண்டமான மாதிரிகளை அளவிடுதல்

70 பில்லியன் பாராமீட்டர்களுக்கும் அதிகமான ஏஜென்ட்களைப் பயிற்றுவிக்க, பல TPU நோட்களில் (nodes) எடைகளையும் (weights) தரவுகளையும் பரப்ப வேண்டும். Tunix, JAX-ன் ShardMap primitive-ஐப் பயன்படுத்தி மாடல் பாராமீட்டர்கள் மற்றும் ஆக்டிவேஷன்கள் (activations) இரண்டையும் பிரிக்கிறது (shard), இது லேர்னர் முழு மாடலையும் நினைவகத்தில் (memory) வைத்திருக்கும் அதே வேளையில் அதிக வேகத்தில் தரவை வழங்க அனுமதிக்கிறது. இந்த ஷார்டிங் (sharding) உத்தி, இதற்கு முன்பு ஒரு தனி TPU போடால் செய்ய முடியாத மாடல்களைப் பயிற்றுவிப்பதைச் சாத்தியமாக்குகிறது.

பிரிக்கப்பட்ட குழாய்களில் இருந்து வரும் காலாவதியான கிரேடியன்ட்கள்

ஆக்டர்கள் லேர்னரை விட வேகமாக இயங்கும்போது, அவை வழங்கும் தரவு தற்போதைய பாலிசியுடன் ஒப்பிடும்போது “காலாவதியாக” (stale) மாறக்கூடும். Tunix இந்த விலகலை (drift) இரண்டு வழிமுறைகள் மூலம் குறைக்கிறது: இம்பார்டன்ஸ்-சாம்பிளிங் (importance-sampling) பழைய மாதிரிகளின் பொருத்தத்தைப் பிரதிபலிக்கும் வகையில் அவற்றின் எடைகளை மாற்றியமைக்கிறது, மேலும் ஒரு குறிப்பிட்ட காலக்கெடுவைத் தாண்டும் டிராக்டரிகளை ஒரு குறிப்பிட்ட காலாவதி வரம்பின் (staleness threshold) மூலம் நீக்குகிறது. இவை இரண்டும் இணைந்து, பயிற்சிச் சுழற்சி அசிங்க்ரோனஸாக இயங்கினாலும் கற்றலை நிலையாக வைத்திருக்கின்றன.

பயன்படுத்துபவர்கள் கவனிக்க வேண்டியவை

  • லேட்டன்சி தணிக்கை (Latency audit) – பிரிக்கப்பட்ட கட்டமைப்பின் பயன் சூழலின் பதிலளிக்கும் நேரத்தைப் பொறுத்தது. குழுக்கள் எண்ட்-டு-எண்ட் லேட்டன்சியை அளவிட வேண்டும் மற்றும் பஃபர் எப்போதும் நிரம்பிக் கொண்டிருப்பதை உறுதி செய்ய ஆக்டர் பூல்களைச் சரியான அளவில் அமைக்க வேண்டும்.
  • வொர்க்கர் பூல் வடிவமைப்பு (Worker pool design) – மலிவான CPU அல்லது GPU-களில் பல ஆக்டர்களை இயக்க முடியும், ஆனால் அவற்றை அளவுக்கு அதிகமாகப் பயன்படுத்துவது நெட்வொர்க் அல்லது ஸ்டோரேஜில் சிக்கலை ஏற்படுத்தலாம். பஃபரின் தரவு உள்வாங்கும் வேகத்திற்கு (ingest rate) இணையான ஒரு சமநிலையான பூல் அவசியம்.
  • பஃபரின் உறுதித்தன்மை (Buffer robustness) – மையச் சேமிப்பகம் புதிய நெரிசல் புள்ளியாக மாறாமல், அதிக எழுத்து (write) மற்றும் வாசிப்பு (read) விகிதங்களைக் கையாள வேண்டும். குறைந்த லேட்டன்சி மற்றும் போதுமான பேண்ட்வித் (bandwidth) கொண்ட ஒரு ஸ்டோரேஜ் சிஸ்டத்தைத் தேர்ந்தெடுப்பது இந்த கட்டமைப்பின் தவிர்க்க முடியாத பகுதியாகும்.

சாத்தியமான குறைபாடுகள்

இந்தத் தனித்த கட்டமைப்பு அதிக இயங்கும் கூறுகளை அறிமுகப்படுத்துகிறது: தனித்தனி வன்பொருள் தொகுதிகள், ஒரு நிலையான பஃபர் மற்றும் காலாவதி வரம்புகளை நடைமுறைப்படுத்த ஒருங்கிணைப்பு தர்க்கம் (coordination logic).

முடிவுரை

ஏஜென்டிக் RL-ல் ஆதிக்கம் செலுத்தும் செலவு மாடல் மட்டுமல்ல, ஒத்திசைவான தொடர்புச் சுழற்சிகளால் ஏற்படும் செயலற்ற நேரமே என்பதை Tunix காட்டுகிறது. ரோல்அவுட் (rollout) பணிகளை மலிவான வன்பொருளுக்கு மாற்றுவதன் மூலமும், அதிகத் திறன் கொண்ட பஃபரிலிருந்து தொடர்ந்து கற்றுக் கொள்ளும் TPU போடிற்குத் தரவை வழங்குவதன் மூலமும், கூகுள் 10%-க்கும் குறைவான பயன்பாட்டுப் பிரச்சனையை முழுத் திறனுக்கு அருகாமையில் இயங்கும் ஒரு பணிப்பாய்வாக (workflow) மாற்றியுள்ளது.