يعمل نظام Tunix من Google على إزالة نقطة الاختناق التي منعت التعلم التعزيزي الوكيل (agentic RL) واسع النطاق من استخدام وحدات TPU بكفاءة. من خلال فصل عمل توليد بيانات التفاعل عن عمل تحديث السياسة، يرفع Tunix نسبة استخدام TPU من خانة الآحاد إلى ما يقرب من السعة الكاملة، مما يقلل من هدر الحوسبة بشكل كبير.

نقطة الاختناق في التعلم التعزيزي الوكيل (agentic RL)

يختلف التعلم التعزيزي الوكيل عن تدريب النماذج اللغوية المألوف القائم على "الرمز التالي" (next-token). يجب على الوكيل إرسال استدعاءات API، أو تنفيذ كود، أو المرور عبر بيئة محاكاة، ثم التفاعل مع النتيجة. وبالتالي، تكون حلقة التدريب متزامنة: ينتج النموذج إجراءً، وتعمل البيئة، وتعود النتيجة، وفقط بعد ذلك يتلقى النموذج تحديث التدرج (gradient update). عندما تستغرق خطوة واحدة في البيئة عدة ثوانٍ، تظل أجهزة TPU المكلفة خاملة، ويمكن أن ينخفض الاستخدام المسجل إلى أقل من 10%. ويترجم هذا عدم الكفاءة مباشرة إلى فواتير سحابة أعلى ودورات بحث أبطأ.

بنية Tunix المنفصلة

يعالج Tunix المشكلة عن طريق فصل المرحلتين — توليد المسارات (trajectory generation) وتحسين السياسة (policy optimization) — على مجموعات أجهزة منفصلة.

  • الممثلون غير المتزامنين (Asynchronous actors) يعملون على وحدات CPU أو GPU غير مكلفة. يتفاعل كل ممثل باستمرار مع البيئة المخصصة له، ويسجل الإجراءات والملاحظات، ويبث المسارات الناتجة إلى مخزن مشترك.
  • المتعلمون المستمرون (Continuous learners) يشغلون وحدات TPU Pods مخصصة. يسحب المتعلم الدفعات من المخزن المؤقت المركزي ويقوم بتحديثات التدرج دون انتظار انتهاء أي ممثل من عملية التدريب (rollout).
  • مخزن مؤقت عالي الإنتاجية (High-throughput buffer) يتوسط العملية، حيث يعمل كمنطقة تجهيز للمسارات. ولأن المتعلم يمكنه القراءة بنفس سرعة تزويد المخزن المؤقت للبيانات، فإن TPU لا يتوقف أبداً.

النتيجة الصافية هي خط أنابيب تدريب تظل فيه وحدات TPU مشغولة طوال الوقت تقريباً، مما يدفع الاستخدام نحو 100%.

العقبات التقنية وكيف يتغلب عليها Tunix

الحلقات (episodes) متغيرة الطول وإعادة تجميع XLA

يقوم مترجم XLA الخاص بـ JAX بالتحسين لأشكال الموترات (tensor shapes) الثابتة. ومع ذلك، تنتج المهام الوكيلية تسلسلات بأطوال مختلفة، مما قد يؤدي عادةً إلى عمليات إعادة تجميع مكلفة. يقوم Tunix بتجميع التسلسلات الأقصر معاً ووضع الحلقات ذات الأطوال المتشابهة في مجموعات (buckets)، مما يحافظ على استقرار الأشكال لفترة كافية لتمكين XLA من إعادة استخدام النوى المجمعة (compiled kernels). والنتيجة هي إنتاجية ثابتة دون العبء الإضافي للمترجم الذي قد يعيق الأداء.

توسيع النماذج الضخمة عبر العديد من رقائق TPU

يتطلب تدريب الوكلاء الذين يمتلكون أكثر من 70 مليار معلمة (parameter) توزيع الأوزان والبيانات عبر عقد TPU متعددة. يستخدم Tunix أداة ShardMap الأساسية في JAX لتقسيم كل من معاملات النموذج والتنشيطات (activations)، مما يسمح للمتعلم بالاحتفاظ بالنموذج بأكمله في الذاكرة مع الاستمرار في تغذيته بالبيانات بسرعة عالية. تجعل استراتيجية التقسيم هذه من الممكن تدريب نماذج كانت في السابق بعيدة المنال عن وحدة TPU pod واحدة.

التدرجات القديمة (Stale gradients) من خطوط الأنابيب المنفصلة

عندما يسبق الممثلون المتعلم، يمكن أن تصبح البيانات التي يزودونها "قديمة" (stale) بالنسبة للسياسة الحالية. يخفف Tunix من هذا الانحراف باستخدام آليتين: إعادة وزن عينات الأهمية (importance-sampling) للعينات القديمة لتعكس مدى صلتها، وعتبة تقادم قابلة للتكوين تتخلص من المسارات التي تتجاوز عمراً محدداً مسبقاً. معاً، يحافظان على استقرار التعلم حتى أثناء تشغيل خط الأنابيب بشكل غير متزامن.

ما يجب على المتبنين مراقبته

  • تدقيق زمن الاستجابة (Latency audit) – تعتمد فائدة الفصل على وقت استجابة البيئة. يجب على الفرق قياس زمن الاستجابة من البداية إلى النهاية والتأكد من أن مجمعات الممثلين بحجم مناسب لإبقاء المخزن المؤقت ممتلئاً جيداً.
  • تصميم مجمع العمال (Worker pool design) – يمكن لوحدات CPU أو GPU الرخيصة استضافة العديد من الممثلين، ولكن الإفراط في تخصيصها قد يسبب تنافساً على الشبكة أو التخزين. من الضروري وجود مجمع متوازن يطابق معدل إدخال المخزن المؤقت.
  • متانة المخزن المؤقت (Buffer robustness) – يجب أن يتعامل المخزن المركزي مع معدلات كتابة وقراءة عالية دون أن يصبح نقطة اختناق جديدة. يعد اختيار نظام تخزين يتميز بزمن استجابة ذيل (tail latency) منخفض ونطاق ترددي كافٍ جزءاً لا يقبل التفاوض من البنية التحتية.

السلبيات المحتملة

تقدم البنية المنفصلة المزيد من الأجزاء المتحركة: أساطيل أجهزة منفصلة، ومخزن مؤقت مستمر، ومنطق تنسيق لفرض حدود التقادم.

الخلاصة

يوضح Tunix أن التكلفة المهيمنة في التعلم التعزيزي الوكيل ليست النموذج نفسه، بل وقت الخمول الناتج عن حلقات التفاعل المتزامنة. من خلال نقل عمل التدريب (rollout) إلى أجهزة رخيصة وتغذية وحدة TPU pod تتعلم باستمرار من مخزن مؤقت عالي الإنتاجية، حولت Google مشكلة استخدام أقل من 10% إلى سير عمل بكامل السعة تقريباً.