Google-এর Tunix সিস্টেম সেই প্রতিবন্ধকতা দূর করে যা বৃহৎ আকারের এজেন্টিক রিইনফোর্সমেন্ট লার্নিং (RL)-কে TPU-কে দক্ষতার সাথে ব্যবহার করতে বাধা দিচ্ছিল। ইন্টারঅ্যাকশন ডেটা জেনারেশনের কাজ এবং পলিসি আপডেটের কাজকে আলাদা করার মাধ্যমে, Tunix TPU ইউটিলাইজেশন একক অংকের শতাংশ থেকে প্রায় পূর্ণ ক্ষমতায় নিয়ে আসে, যা কম্পিউট অপচয় নাটকীয়ভাবে কমিয়ে দেয়।
এজেন্টিক RL-এর প্রতিবন্ধকতা
এজেন্টিক RL প্রচলিত “নেক্সট-টোকেন” ল্যাঙ্গুয়েজ-মডেল ট্রেনিং থেকে আলাদা। একটি এজেন্টকে API কল পাঠাতে হয়, কোড এক্সিকিউট করতে হয়, অথবা একটি সিমুলেটেড এনভায়রনমেন্টের মধ্য দিয়ে যেতে হয় এবং তারপর ফলাফলের ওপর ভিত্তি করে প্রতিক্রিয়া জানাতে হয়। ফলে ট্রেনিং লুপটি সিনক্রোনাস হয়: মডেল একটি অ্যাকশন তৈরি করে, এনভায়রনমেন্ট চলে, ফলাফল ফিরে আসে এবং তারপরেই মডেল একটি গ্রেডিয়েন্ট আপডেট পায়। যখন একটি একক এনভায়রনমেন্ট স্টেপ নিতে কয়েক সেকেন্ড সময় লাগে, তখন মূল্যবান TPU হার্ডওয়্যার অলস বসে থাকে এবং রিপোর্ট করা ইউটিলাইজেশন ১০%-এর নিচে নেমে যেতে পারে। এই অদক্ষতা সরাসরি ক্লাউড বিল বৃদ্ধি এবং গবেষণার গতি কমিয়ে দেয়।
Tunix-এর ডিসাকাপলড আর্কিটেকচার
Tunix দুটি পর্যায়—ট্র্যাজেক্টরি জেনারেশন এবং পলিসি অপ্টিমাইজেশন—কে আলাদা হার্ডওয়্যার পুলে ভাগ করে এই সমস্যার সমাধান করে।
- অ্যাসিনক্রোনাস অ্যাক্টরস (Asynchronous actors) সস্তা CPU বা GPU-তে চলে। প্রতিটি অ্যাক্টর ক্রমাগত তার নির্ধারিত এনভায়রনমেন্টের সাথে ইন্টারঅ্যাক্ট করে, অ্যাকশন এবং অবজারভেশন রেকর্ড করে এবং প্রাপ্ত ট্র্যাজেক্টরিগুলো একটি শেয়ার্ড স্টোরে স্ট্রিম করে।
- কন্টিনিউয়াস লার্নারস (Continuous learners) ডেডিকেটেড TPU Pod দখল করে থাকে। লার্নার সেন্ট্রাল বাফার থেকে ব্যাচগুলো সংগ্রহ করে এবং কোনো একক অ্যাক্টরের রোলআউট শেষ হওয়ার জন্য অপেক্ষা না করেই গ্রেডিয়েন্ট আপডেট সম্পন্ন করে।
- হাই-থ্রুপুট বাফার (High-throughput buffer) মাঝখানে অবস্থান করে, যা ট্র্যাজেক্টরিগুলোর জন্য একটি স্টেজিং এরিয়া হিসেবে কাজ করে। যেহেতু লার্নার বাফার ডেটা সরবরাহ করার গতিতে ডেটা পড়তে পারে, তাই TPU কখনোই থমকে যায় না।
এর ফলে একটি ট্রেনিং পাইপলাইন তৈরি হয় যেখানে TPU প্রায় সব সময় ব্যস্ত থাকে এবং ইউটিলাইজেশন ১০০%-এর কাছাকাছি পৌঁছায়।
প্রযুক্তিগত বাধা এবং Tunix যেভাবে সেগুলো অতিক্রম করে
পরিবর্তনশীল দৈর্ঘ্যের এপিসোড এবং XLA রিকম্পাইলেশন
JAX-এর XLA কম্পাইলার ফিক্সড টেন্সর শেপের জন্য অপ্টিমাইজ করা থাকে। তবে, এজেন্টিক টাস্কগুলো ভিন্ন ভিন্ন দৈর্ঘ্যের সিকোয়েন্স তৈরি করে, যা সাধারণত ব্যয়বহুল রিকম্পাইলেশন ট্রিগার করে। Tunix ছোট সিকোয়েন্সগুলোকে একসাথে প্যাক করে এবং একই দৈর্ঘ্যের এপিসোডগুলোকে বাকেটে গ্রুপ করে রাখে, যা XLA-কে কম্পাইল করা কার্নেল পুনরায় ব্যবহার করার জন্য যথেষ্ট স্থিতিশীল শেপ প্রদান করে। এর ফলে কম্পাইলারের অতিরিক্ত বোঝা ছাড়াই স্থিতিশীল থ্রুপুট পাওয়া যায়, যা অন্যথায় পারফরম্যান্সকে মারাত্মকভাবে কমিয়ে দিত।
অনেক TPU চিপের মাধ্যমে বিশাল মডেল স্কেল করা
৭০ বিলিয়নেরও বেশি প্যারামিটার বিশিষ্ট এজেন্টদের প্রশিক্ষণ দিতে হলে একাধিক TPU নোডে ওয়েট (weights) এবং ডেটা ছড়িয়ে দেওয়া প্রয়োজন। Tunix মডেল প্যারামিটার এবং অ্যাক্টিভেশন উভয়কেই শার্ড (shard) করার জন্য JAX-এর ShardMap প্রিমিটিভ ব্যবহার করে, যা লার্নারকে উচ্চ গতিতে ডেটা সরবরাহ করার পাশাপাশি পুরো মডেলটি মেমরিতে রাখতে সাহায্য করে। এই শার্ডিং কৌশলটি এমন মডেল প্রশিক্ষণের সম্ভাবনা তৈরি করে যা আগে একটি একক TPU পডের পক্ষে অসম্ভব ছিল।
ডিসাকাপলড পাইপলাইন থেকে স্টেল গ্রেডিয়েন্ট (Stale gradients)
যখন অ্যাক্টররা লার্নারের চেয়ে এগিয়ে চলে, তখন তাদের সরবরাহ করা ডেটা বর্তমান পলিসির তুলনায় “স্টেল” বা পুরনো হয়ে যেতে পারে। Tunix দুটি মেকানিজমের মাধ্যমে এই বিচ্যুতি প্রশমিত করে: ইম্পর্টেন্স-স্যাম্পলিং (importance-sampling) পুরনো স্যাম্পলগুলোর প্রাসঙ্গিকতা প্রতিফলিত করতে সেগুলোকে রিওয়েট করে, এবং একটি কনফিগারযোগ্য স্টেলনেস থ্রেশহোল্ড (staleness threshold) নির্দিষ্ট বয়সের বেশি ট্র্যাজেক্টরিগুলোকে বাদ দিয়ে দেয়। এই দুটি পদ্ধতি পাইপলাইনটি অ্যাসিনক্রোনাসভাবে চললেও লার্নিং স্থিতিশীল রাখে।
ব্যবহারকারীদের যা খেয়াল রাখা প্রয়োজন
- ল্যাটেন্সি অডিট (Latency audit) – ডিসাকাপলিং-এর সুবিধা এনভায়রনমেন্টের রেসপন্স টাইমের ওপর নির্ভর করে। টিমগুলোর এন্ড-টু-এন্ড ল্যাটেন্সি পরিমাপ করা উচিত এবং নিশ্চিত করা উচিত যে বাফারটি পর্যাপ্তভাবে পূর্ণ রাখার জন্য অ্যাক্টর পুলের আকার সঠিক আছে।
- ওয়ার্কার পুল ডিজাইন (Worker pool design) – সস্তা CPU বা GPU অনেক অ্যাক্টর হোস্ট করতে পারে, কিন্তু অতিরিক্ত সাবস্ক্রিপশন নেটওয়ার্ক বা স্টোরেজেת সংঘর্ষ (contention) সৃষ্টি করতে পারে। বাফারের ইনজেস্ট রেটের সাথে সামঞ্জস্যপূর্ণ একটি ভারসাম্যপূর্ণ পুল থাকা অপরিহার্য।
- বাফার রোবাস্টনেস (Buffer robustness) – সেন্ট্রাল স্টোরটিকে নতুন কোনো প্রতিবন্ধকতা না হয়ে উচ্চ রাইট এবং রিড রেট সামলাতে হবে। লো টেইল ল্যাটেন্সি (low tail latency) এবং পর্যাপ্ত ব্যান্ডউইথ সম্পন্ন একটি স্টোরেজ সিস্টেম নির্বাচন করা এই আর্কিটেকচারের একটি অপরিহার্য অংশ।
সম্ভাব্য অসুবিধা
বিচ্ছিন্ন আর্কিটেকচারটি আরও অনেক জটিলতা বা মুভিং পার্টস (moving parts) নিয়ে আসে: আলাদা হার্ডওয়্যার ফ্লিট, একটি পারসিস্টেন্ট বাফার, এবং স্টেলনেস লিমিট কার্যকর করার জন্য কোঅর্ডিনেশন লজিক।
সারসংক্ষেপ
Tunix দেখায় যে এজেন্টিক RL-এর প্রধান খরচ মডেলের জন্য নয়, বরং সিনক্রোনাস ইন্টারঅ্যাকশন লুপের কারণে সৃষ্ট অলস সময়ের জন্য। রোলআউট কাজগুলোকে সস্তা হার্ডওয়্যারে স্থানান্তর করে এবং একটি হাই-থ্রুপুট বাফার থেকে একটি কন্টিনিউয়াসলি লার্নিং TPU পডকে ডেটা সরবরাহ করার মাধ্যমে, Google ১০%-এর নিচের ইউটিলাইজেশন সমস্যাকে প্রায় পূর্ণ ক্ষমতাসম্পন্ন একটি ওয়ার্কফ্লোতে রূপান্তরিত করেছে।