گوگل کا Tunix سسٹم اس بڑی رکاوٹ کو ختم کرتا ہے جس نے بڑے پیمانے پر ایجنٹک ری انفورسمنٹ لرننگ (RL) کو TPU کا مؤثر طریقے سے استعمال کرنے سے روکا ہوا تھا۔ ٹریجیکٹری جنریشن کے کام کو پالیسی اپ ڈیٹ کرنے کے کام سے الگ کر کے، Tunix TPU کے استعمال (utilization) کو سنگل ڈیجٹ فیصد سے بڑھا کر تقریباً مکمل صلاحیت تک لے جاتا ہے، جس سے کمپیوٹ کے ضیاع میں نمایاں کمی آتی ہے۔

ایجنٹک RL میں رکاوٹ

ایجنٹک RL، روایتی "next-token" لینگویج ماڈل ٹریننگ سے مختلف ہے۔ ایک ایجنٹ کو API کالز بھیجنی پڑتی ہیں، کوڈ چلانا پڑتا ہے، یا کسی سمولیٹڈ ماحول (simulated environment) میں قدم اٹھانا پڑتا ہے، اور پھر نتیجے پر ردعمل دینا ہوتا ہے۔ اس لیے ٹریننگ لوپ synchronous ہوتا ہے: ماڈل ایک ایکشن پیدا کرتا ہے، ماحول چلتا ہے، نتیجہ واپس آتا ہے، اور اس کے بعد ہی ماڈل کو گریڈینٹ اپ ڈیٹ موصول ہوتا ہے۔ جب ایک واحد انوائرمنٹ سٹیپ میں کئی سیکنڈ لگتے ہیں، تو مہنگا TPU ہارڈ ویئر بیکار بیٹھا رہتا ہے، اور رپورٹ شدہ یوٹیلائزیشن 10% سے بھی نیچے گر سکتی ہے۔ یہ عدم کارکردگی براہ راست کلاؤڈ بلز میں اضافے اور تحقیق کے عمل میں تاخیر کا باعث بنتی ہے۔

Tunix کا ڈی کپلڈ آرکیٹیکچر

Tunix اس مسئلے کو دو مراحل—ٹریجیکٹری جنریشن اور پالیسی آپٹیمائزیشن—کو الگ الگ ہارڈ ویئر پولز پر منتقل کر کے حل کرتا ہے۔

  • Asynchronous actors سستے CPUs یا GPUs پر چلتے ہیں۔ ہر ایکٹر مسلسل اپنے تفویض کردہ ماحول کے ساتھ تعامل کرتا ہے، ایکشنز اور مشاہدات ریکارڈ کرتا ہے، اور حاصل ہونے والی ٹریجیکٹریز کو ایک مشترکہ اسٹور میں اسٹریم کرتا ہے۔
  • Continuous learners مخصوص TPU Pods پر کام کرتے ہیں۔ لرنر سینٹرل بفر سے بیچز (batches) لیتا ہے اور کسی بھی ایک ایکٹر کے رول آؤٹ مکمل ہونے کا انتظار کیے بغیر گریڈینٹ اپ ڈیٹس کرتا ہے۔
  • High-throughput buffer درمیان میں واقع ہوتا ہے، جو ٹریجیکٹریز کے لیے اسٹیجنگ ایریا کے طور پر کام کرتا ہے۔ چونکہ لرنر اتنی ہی تیزی سے ڈیٹا پڑھ سکتا ہے جتنی تیزی سے بفر ڈیٹا فراہم کر سکتا ہے، اس لیے TPU کبھی رکتا نہیں ہے۔

اس کا مجموعی اثر ایک ایسا ٹریننگ پائپ لائن ہے جہاں TPUs تقریباً ہر وقت مصروف رہتے ہیں، اور یوٹیلائزیشن کو 100% کے قریب لے جاتے ہیں۔

تکنیکی رکاوٹیں اور Tunix ان پر کیسے قابو پاتا ہے

مختلف لمبائی کے ایپی سوڈز اور XLA ری کمپائلیشن

JAX کا XLA کمپائلر فکسڈ ٹینسر شیپس (fixed tensor shapes) کے لیے آپٹیمائز کرتا ہے۔ تاہم، ایجنٹک ٹاسک مختلف لمبائی کے تسلسل (sequences) پیدا کرتے ہیں، جو عام طور پر مہنگی ری کمپائلیشنز کا باعث بنتے ہیں۔ Tunix مختصر تسلسل کو ایک ساتھ پیک کرتا ہے اور ملتے جلتے لمبائی کے ایپی سوڈز کو بکیٹس (buckets) میں گروپ کرتا ہے، جس سے شیپس اتنی دیر تک مستحکم رہتی ہیں کہ XLA کمپائل شدہ کرنلز کو دوبارہ استعمال کر سکے۔ اس کا نتیجہ بغیر کسی کمپائلر اوور ہیڈ کے مستقل تھرو پٹ کی صورت میں نکلتا ہے جو ورنہ کارکردگی کو مفلوج کر سکتا تھا۔

بہت سے TPU چپس پر بڑے ماڈلز کو اسکیل کرنا

70 بلین سے زیادہ پیرامیٹرز والے ایجنٹس کی ٹریننگ کے لیے ویٹس (weights) اور ڈیٹا کو متعدد TPU نوڈز پر پھیلانا ضروری ہوتا ہے۔ Tunix ماڈل پیرامیٹرز اور ایکٹیویشنز دونوں کو شارڈ (shard) کرنے کے لیے JAX کے ShardMap پرائمٹیو کا استعمال کرتا ہے، جس سے لرنر کو پورے ماڈل کو میموری میں رکھنے اور ساتھ ہی اسے تیز رفتاری سے ڈیٹا فراہم کرنے کی اجازت ملتی ہے۔ یہ شارڈنگ حکمت عملی ان ماڈلز کی ٹریننگ کو ممکن بناتی ہے جو پہلے ایک سنگل TPU پوڈ کے لیے پہنچ سے باہر تھے۔

ڈی کپلڈ پائپ لائنز سے پرانے (stale) گریڈینٹس

جب ایکٹرز لرنر سے آگے نکل جاتے ہیں، تو ان کا فراہم کردہ ڈیٹا موجودہ پالیسی کے لحاظ سے "stale" (پرانا یا غیر متعلقہ) ہو سکتا ہے۔ Tunix اس فرق کو دو طریقوں سے کم کرتا ہے: امپورٹنس سیمپلنگ (importance-sampling) پرانے نمونوں کو ان کی اہمیت کے مطابق دوبارہ وزن (reweight) دیتی ہے، اور ایک کنفیگر ایبل اسٹیلنس تھریش ہول (staleness threshold) ان ٹریجیکٹریز کو مسترد کر دیتا ہے جو مقررہ عمر سے زیادہ ہو چکی ہوں۔ یہ دونوں مل کر سیکھنے کے عمل کو مستحکم رکھتے ہیں چاہے پائپ لائن غیر ہم آہنگ (asynchronously) طریقے سے چل رہی ہو۔

صارفین کو کن باتوں کا خیال رکھنا چاہیے

  • لیٹنسی آڈٹ (Latency audit) – ڈی کپلڈنگ کا فائدہ انوائرمنٹ کے رسپانس ٹائم پر منحصر ہے۔ ٹیموں کو اینڈ ٹو اینڈ لیٹنسی کی پیمائش کرنی چاہیے اور اس بات کو یقینی بنانا چاہیے کہ ایکٹر پولز کا سائز ایسا ہو کہ بفر ہمیشہ بھرا رہے۔
  • ورکر پول ڈیزائن – سستے CPUs یا GPUs بہت سے ایکٹرز کو ہوسٹ کر سکتے ہیں، لیکن ان پر ضرورت سے زیادہ بوجھ ڈالنے سے نیٹ ورک یا اسٹوریج پر تناؤ پیدا ہو سکتا ہے۔ ایک متوازن پول جو بفر کی انگیسٹ ریٹ (ingest rate) سے میل کھاتا ہو، ضروری ہے۔
  • بفر کی مضبوطی (Buffer robustness) – سینٹرل اسٹور کو نیا بوٹل نیک (bottleneck) بنے بغیر ہائی رائٹ اور ریڈ ریٹس کو سنبھالنا چاہیے۔ کم ٹیل لیٹنسی (tail latency) اور کافی بینڈوتھ والے اسٹوریج سسٹم کا انتخاب کرنا اس آرکیٹیکچر کا ایک لازمی حصہ ہے۔

ممکنہ نقصانات

اس تقسیم شدہ آرکیٹیکچر میں زیادہ متحرک اجزاء شامل ہوتے ہیں: الگ الگ ہارڈ ویئر فلیٹس، ایک مستقل بفر، اور اسٹیلنس لمٹ کو نافذ کرنے کے لیے کوآرڈینیشن لاجک۔

خلاصہ

Tunix یہ ظاہر کرتا ہے کہ ایجنٹک RL میں اصل خرچہ ماڈل کا نہیں بلکہ سنکرونس انٹرایکشن لوپس کی وجہ سے ہونے والے فارغ وقت (idle time) کا ہے۔ رول آؤٹ کے کام کو سستے ہارڈ ویئر پر منتقل کر کے اور ہائی تھرو پٹ بفر کے ذریعے مسلسل سیکھنے والے TPU پوڈ کو ڈیٹا فراہم کر کے، گوگل نے 10% سے کم یوٹیلائزیشن کے مسئلے کو تقریباً مکمل صلاحیت والے ورک فلو میں بدل دیا ہے۔