Google’s Tunix system lifts the choke point that has kept large-scale agentic reinforcement learning (RL) from using TPUs efficiently. By separating the work of generating interaction data from the work of updating the policy, Tunix drives TPU utilization from single-digit percentages up to near-full capacity, cutting compute waste dramatically.

گلوگاه در agentic RL

agentic RL با آموزش مدل‌های زبانی متداول «توکن بعدی» متفاوت است. یک عامل (agent) باید فراخوانی‌های API ارسال کند، کد اجرا کند یا در یک محیط شبیه‌سازی‌شده گام بردارد و سپس به نتیجه واکنش نشان دهد. بنابراین، حلقه آموزش همگام (synchronous) است: مدل یک عمل انجام می‌دهد، محیط اجرا می‌شود، نتیجه بازمی‌گردد و تنها پس از آن است که مدل به‌روزرسانی گرادیان را دریافت می‌کند. وقتی هر گام در محیط چندین ثانیه طول می‌کشد، سخت‌افزار گران‌قیمت TPU بیکار می‌ماند و میزان بهره‌وری گزارش‌شده می‌تواند به زیر ۱۰٪ کاهش یابد. این ناکارآمدی مستقیماً به قبض‌های ابری بالاتر و چرخه‌های تحقیقاتی کندتر منجر می‌شود.

معماری مجزای Tunix

Tunix با جدا کردن دو مرحله — تولید مسیر (trajectory generation) و بهینه‌سازی سیاست (policy optimization) — و قرار دادن آن‌ها در مجموعه‌های سخت‌افزاری مجزا، با این مشکل مقابله می‌کند.

  • بازیگران ناهمگام (Asynchronous actors) روی CPU یا GPUهای ارزان‌قیمت اجرا می‌شوند. هر بازیگر به‌طور مداوم با محیط اختصاص یافته خود تعامل می‌کند، اقدامات و مشاهدات را ثبت می‌کند و مسیرهای (trajectories) حاصل را به یک ذخیره‌ساز مشترک ارسال می‌کند.
  • یادگیرندگان مداوم (Continuous learners) از TPU Podهای اختصاصی استفاده می‌کنند. یادگیرنده دسته‌های داده را از بافر مرکزی می‌گیرد و بدون منتظر ماندن برای اتمام اجرای هر بازیگر، به‌روزرسانی‌های گرادیان را انجام می‌دهد.
  • بافر با توان عملیاتی بالا (High-throughput buffer) در میان این دو قرار دارد و به عنوان یک منطقه آماده‌سازی برای مسیرها عمل می‌کند. از آنجایی که یادگیرنده می‌تواند با همان سرعتی که بافر داده‌ها را تأمین می‌کند، داده‌ها را بخواند، TPU هرگز متوقف نمی‌شود.

نتیجه نهایی، یک خط لوله (pipeline) آموزشی است که در آن TPUها تقریباً تمام مدت مشغول به کار هستند و بهره‌وری را به سمت ۱۰۰٪ سوق می‌دهند.

موانع فنی و نحوه غلبه Tunix بر آن‌ها

اپیزودهای با طول متغیر و بازکامپایل شدن XLA

کامپایلر XLA در JAX برای اشکال ثابت تنسور (fixed tensor shapes) بهینه‌سازی شده است. با این حال، وظایف عامل‌محور توالی‌هایی با طول‌های متفاوت تولید می‌کنند که در حالت عادی باعث بازکامپایل شدن‌های پرهزینه می‌شود. Tunix توالی‌های کوتاه‌تر را در کنار هم قرار می‌دهد و اپیزودهایی با طول مشابه را در دسته‌های (buckets) مشخص گروه‌بندی می‌کند تا اشکال (shapes) برای مدتی کافی پایدار بمانند تا XLA بتواند از کرنل‌های کامپایل‌شده مجدداً استفاده کند. نتیجه، توان عملیاتی پایدار بدون بار اضافی کامپایلر است که در غیر این صورت عملکرد را مختل می‌کرد.

مقیاس‌پذیری مدل‌های عظیم در چندین تراشه TPU

آموزش عامل‌هایی با بیش از ۷۰ میلیارد پارامتر مستلزم توزیع وزن‌ها و داده‌ها در چندین گره (node) TPU است. Tunix از پریمیتیو ShardMap در JAX برای تکه‌تکه کردن (sharding) پارامترهای مدل و فعال‌سازها (activations) استفاده می‌کند که به یادگیرنده اجازه می‌دهد کل مدل را در حافظه نگه دارد و در عین حال با سرعت بالا به آن داده تزریق کند. این استراتژی sharding، آموزش مدل‌هایی را که قبلاً برای یک TPU pod واحد دور از دسترس بودند، امکان‌پذیر می‌کند.

گرادیان‌های کهنه در خطوط لوله مجزا

وقتی بازیگران جلوتر از یادگیرنده اجرا می‌شوند، داده‌هایی که آن‌ها تأمین می‌کنند می‌تواند نسبت به سیاست فعلی «کهنه» (stale) شود. Tunix این انحراف را با دو مکانیسم کاهش می‌دهد: نمونه‌برداری اهمیت (importance-sampling) نمونه‌های قدیمی‌تر را برای بازتاب میزان مرتبط بودن آن‌ها بازوزن‌دهی می‌کند، و یک آستانه کهنگی (staleness threshold) قابل تنظیم، مسیرهایی را که از سن تعیین‌شده فراتر رفته‌اند، حذف می‌کند. این دو در کنار هم باعث می‌شوند یادگیری حتی زمانی که خط لوله به‌صورت ناهمگام اجرا می‌شود، پایدار بماند.

مواردی که پذیرندگان باید مراقب آن‌ها باشند

  • بررسی تأخیر (Latency audit) – مزیت جداسازی به زمان پاسخگویی محیط بستگی دارد. تیم‌ها باید تأخیر سرتاسری (end-to-end latency) را اندازه‌گیری کنند و اطمینان حاصل کنند که استخرهای بازیگر (actor pools) به گونه‌ای تنظیم شده‌اند که بافر را همیشه پر نگه دارند.
  • طراحی استخر کارگر (Worker pool design) – CPU یا GPUهای ارزان می‌توانند میزبان بازیگران زیادی باشند، اما تخصیص بیش از حد آن‌ها می‌تواند باعث ایجاد رقابت (contention) در شبکه یا ذخیره‌سازی شود. داشتن یک استخر متوازن که با نرخ ورود داده‌های بافر مطابقت داشته باشد، ضروری است.
  • استحکام بافر (Buffer robustness) – ذخیره‌ساز مرکزی باید نرخ‌های بالای نوشتن و خواندن را بدون تبدیل شدن به یک گلوگاه جدید مدیریت کند. انتخاب یک سیستم ذخیره‌سازی با تأخیر دم (tail latency) پایین و پهنای باند کافی، بخشی غیرقابل مذاکره از معماری است.

معایب احتمالی

معماری تقسیم‌شده اجزای متحرک بیشتری را معرفی می‌کند: ناوگان‌های سخت‌افزاری مجزا، یک بافر پایدار و منطق هماهنگی برای اعمال محدودیت‌های کهنگی.

نتیجه‌گیری

Tunix نشان می‌دهد که هزینه غالب در agentic RL خودِ مدل نیست، بلکه زمان بیکاری ناشی از حلقه‌های تعامل همگام است. گوگل با انتقال کار اجرای مسیر (rollout) به سخت‌افزارهای ارزان و تغذیه یک TPU pod که به‌طور مداوم در حال یادگیری است از طریق یک بافر با توان عملیاتی بالا، مشکل بهره‌وری زیر ۱۰٪ را به یک جریان کاری با ظرفیت تقریباً کامل تبدیل کرده است.