🤖 machine learning

rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment

تقدم هذه الورقة البحثية rl-triton، وهي مكتبة مفتوحة المصدر تستفيد من إطار عمل مسح ترابطي موحد مُنفذ في Triton لتسريع سبعة خوارزميات متميزة لتعيين الائتمان في التعلم التعزيزي على وحدات معالجة الرسومات، محققةً تسريعات تتراوح بين 1.6 و5.70 ضعفاً مقارنة بالنماذج المرجعية المتجهة عن طريق تقليل عبء الذاكرة وتمكين الحوسبة المتوازية بمقياس $O(\log T)$ .

المؤلفون الأصليون: Lars Simon Zehnder

نُشر 2026-08-20

📖 5 دقيقة قراءة🧠 قراءة متعمّقة

CC BY 4.0

المؤلفون الأصليون: Lars Simon Zehnder

البحث الأصلي مرخَّص بموجب CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ هذا شرح مولَّده بالذكاء الاصطناعي للبحث أدناه. لم يكتبه المؤلفون ولم يصادقوا عليه. وللتحقق من الدقة التقنية، يرجى الرجوع إلى البحث الأصلي. اقرأ إخلاء المسؤولية الكامل

في عالم الذكاء الاصطناعي، هناك صراع مستمر لتعليم الحواسيب كيفية اتخاذ قرارات جيدة. تخيل روبوتًا يتعلم المشي أو برنامجًا يتعلم لعب لعبة ما. للتحسن، يجب على النظام معرفة أي إجراءات محددة أدت إلى النجاح وأيها أدى إلى الفشل. تسمى هذه العملية "تخصيص الائتمان" (credit assignment). إنها عملية النظر إلى الوراء في تسلسل من الأحداث وتحديد ما إذا كانت "هذه الخطوة كانت جيدة" أو "تلك الخطوة كانت سيئة"، حتى يتمكن النظام من تعديل سلوكه المستقبلي. وبينما قد يقضي الروبوت معظم وقته في استكشاف العالم أو إجراء حسابات معقدة لتقرير ما سيفعله بعد ذلك، فإنه عند اللحظة التي يحتاج فيها إلى التعلم من أخطائه، يجب عليه إجراء نوع محدد من الرياضيات. تتضمن هذه الرياضيات النظر إلى قائمة طويلة من الخطوات وربط النقاط بينها، حيث تعتمد قيمة خطوة واحدة على الخطوة التي تليها. ولفترة طويلة، كان القيام بهذه الرياضيات على شرائح حاسوبية قوية تسمى وحدات معالجة الرسومات (GPUs) بطيئًا لأن الكمبيوتر كان يضطر لمعالجة القائمة خطوة بخمة تلو الأخرى، مثل قراءة كتاب صفحة بصفحة، رغم أن الأجهزة كانت قادرة على قراءة صفحات عديدة في وقت واحد.

لقد طور باحث يدعى لارس سيمون زيندر أداة جديدة تسمى rl-triton تحل هذه العقبة. الأداة هي مجموعة من التعليمات الحاسوبية عالية الكفاءة المصممة خصيصًا لمهمة تخصيص الائتمان في التعلم التعزيزي. وبدلاً من إجبار الكمبيوتر على معالجة قائمة الخطوات كسلسلة متتالية بطيئة، تعيد الطريقة الجديدة تنظيم العمل بحيث يمكن حساب آلاف الخطوات في وقت واحد. الفكرة الجوهرية هي التعامل مع تسلسل الأحداث بأكمله كبنية رياضية موحدة يمكن تفكيكها وحلها بالتوازي. ومن خلال القيام بذلك، يمكن للكمبيوتر إنهاء الحساب في جزء ضئيل من الوقت الذي كان يستغرقه سابقًا، خاصة عند التعامل مع آلاف السيناريوهات المختلفة التي تحدث في نفس الوقت.

اختبر الباحثون هذا النهج الجديد مقابل الطرق القياسية المستخدمة حاليًا في المجال. ووجدوا أنه بالنسبة لأكثر السيناريوهات شيوعًا والأكثر تطلبًا — حيث يتم محاكاة آلاف البيئات في آن واحد — فإن الأداة الجديدة أسرع بشكل ملحوظ. وفي بعض الحالات، أكملت المهمة بسرعة تقارب ست مرات أسرع من أفضل طريقة سابقة. يأتي هذا التسريع من تغيير ذكي في كيفية انتقال البيانات عبر ذاكرة الكمبيوتر. في الطريقة القديمة، كان على الكمبيوتر أن يتوقف باستمرار ويجلب البيانات من بنك الذاكرة الرئيسي لكل خطوة في التسلسل، مما خلق ازدحامًا مروريًا. أما الطريقة الجديدة فتُبقي البيانات قريبة من محرك الحساب، مما يسم يسمح للكمبيوتر بالعمل خلال التسلسل بأكمله دون تلك التوقفات المستمرة. وهذا أمر مهم بشكل خاص لتدريب الذكاء الاصطناعي الحديث، حيث قد تقوم الأنظمة بتشغيل آلاف عمليات المحاكاة بالتوازي، كل منها يتكون من مئات الخطوات.

تفصل الورقة البحثية كيف يعمل هذا لسبعة أنواع مختلفة من خوارزميات التعلم، والتي تشترك جميعها في نفس النمط الرياضي الأساسي. تتعامل الأداة الجديدة معها جميعًا ضمن إطار عمل موحد. كما تولي اهتمامًا دقيقًا للواقع الفوضوي للبيانات في العالم الحقيقي، مثل عندما ينتهي "الإبيد" (episode) بشكل مفاجئ أو عندما يتم قطع المحاكاة. وقد أثبت الباحثون أن طريقتهم تعالج هذه الحدود بشكل صحيح، مما يضمن توقف إشارة التعلم في المكان الصحيح وعدم تسربها بالخطأ من سيناريو إلى آخر. لقد تحققوا من نتائجهم من خلال مقارنة الأداة الجديدة بكل من الطريقة القديمة البطيئة، والنسخة الأكثر حداثة وتحسينًا التي تستخدم أدوات برمجة قياسية. وقد تفوقت الأداة الجديدة باستمرار على كليهما، مما أظهر أن مكاسب السرعة كانت حقيقية وليست مجرد نتيجة لخدع برمجية أفضل.

من أكثر النتائج إثارة للاهتمام هي كيفية تغير ميزة السرعة اعتمادًا على حجم المشكلة. عندما تكون تسلسلات الخطوات قصيرة، تظل الأداة الجديدة أسرع، لكن الفرق يكون أصغر. ومع ذلك، كلما زادت طول التسلسلات، زادت الميزة. وذلك لأن الطرق القديمة تضطر لتكرار عملية جلب البيانات من الذاكرة مرات أكثر مع طول القائمة، بينما تتوسع الطريقة الجديدة بكفاءة أكبر بكثير. كما نظر الباحثون في كيفية تأثير ذلك على عملية التدريب الكاملة لوكيل الذكاء الاصطناعي. ووجدوا أنه بينما أصبح خطوة تخصيص الائتمان نفسها أسرع بكثير، فإن تسريع التدريب الإجمالي كان أحيانًا متواضعًا. وذلك لأن تخصيص الائتمان هو جزء واحد فقط من خط التدريب الكامل؛ فإذا كانت بقية العملية بطيئة، فإن تسريع جزء واحد منها لن يجعل العملية برمتها تعمل بشكل أسرع بشكل دراماتيكي. ومع ذلك، في إعدادات محددة حيث يستغوذ تخصيص الائتمان جزءًا أكبر من الوقت الإجمالي، تحسنت سرعة التدريب الإجمالية بشكل ملحوظ.

كما يسلط العمل الضوء على بعض القيود. بالنسبة للتسلسلات الطويلة جدًا، تواجه نوعًا معينًا من الخوارزميات يسمى "ريتريس" (Retrace) قيدًا متعلقًا بالأجهزة، حيث ينفد من شريحة الكمبيوتر نوع معين من مساحة التخزين السريعة، مما يسبب تباطؤًا. حدد الباحثون هذه المشكلة وأشاروا إلى أنها مقايضة معروفة في التصميم. كما ذكروا أن أداتهم الحالية تعمل بشكل أفضل مع تنسيقات البيانات القياسية وأن بعض المتغيرات المتخصصة قد تتطلب مزيدًا من التطوير. وبالرغم من هذه الحدود، تقدم الورقة حلًا واضحًا وعمليًا لمشكلة مستمرة في تدريب الذكاء الاصطناعي. ومن خلال تحويل الحساب المتسلسل خطوة بخطوة إلى حساب متوازٍ ومتزامن، أظهر الباحثون أنه من الممكن جعل التعلم التعزيزي أكثر كفاءة بشكل كبير. وهذه الكفاءة أمر بالغ الأهمية مع نمو أنظمة الذكاء الاصطناعي وزيادة تعقيدها، مما يتطلب منها التعلم من كميات هائلة من البيانات في أوقات أقصر. الأداة متاحة الآن للآخرين لاستخدامها، مما يوفر وسيلة لتسريع تدريب الأنظمة الذكية دون الحاجة إلى تغيير الطريقة الأساسية التي تتعلم بها.

ملخص تقني: rl-triton

بيان المشكلة

تتطلب خطوط أنابيب تدريب التعلم المعزز العميق (Deep Reinforcement Learning - RL) حساب كميات تخصيص الائتمان (credit assignment)—مثل المزايا (advantages)، وأهداف القيمة (value targets)، وتقديرات العائد (return estimates)—عند كل خطوة تحديث. تعمل هذه الحسابات على تنسورات (tensors) ذات شكل $[num\_envs, seq\_len]$ وتتضمن عادةً تكرارات خطية من الدرجة الأولى (مثل: تقدير الميزة المعمم GAE، وV-Trace، وRetrace( $\lambda$ )).

التحدي الجوهري يكمن في التبعية التسلسلية المتأصلة في هذه التكرارات. تتطلب التنفيذات القياسية $T$ من الخطوات المتسلسلة لحساب المتتالية، مما يخلق سلسلة زمنية صارمة تمنع التوازي الكامل عبر الخطوات الزمنية، حتى على وحدات معالجة الرسوميات (GPUs) التي تحتوي على آلاف النوى. وبينما يعمل torch.compile على إزالة عبء مفسر بايثون (Python interpreter)، إلا أنه لا يستطيع كسر تبعية البيانات التي تفرضها الخوارزمية نفسها، مما يترك العمق الحسابي عند $O(T)$ . علاوة على ذلك، فإن التنفيذات المتجهة (vectorized) الموجودة التي تعيد هيكلة المشكلة كمسح ترابطي متوازٍ (parallel associative scan) غالبًا ما تعاني من حركة مرور عالية للذاكرة: حيث تقوم بتجسيد النتائج الوسيطة في ذاكرة النطاق الترددي العالي (HBM) عند كل خطوة مضاعفة في عملية الاختزال $O(\log T)$ ، مما يتسبب في تكاليف كبيرة في زمن الاستجابة وعرض النطاق الترددي.

المنهجية

يقدم البحث rl-triton، وهي مكتبة مفتوحة المصدر لتنفيذ نوى (kernels) عالية الأداء لوحدات معالجة الرسوميات لعمليات تخصيص الائتمان في التعلم المعزز باستخدام لغة Triton. تستند المنهجية إلى ثلاث ركائز:

1. إطار جبري موحد

يوضح المؤلف أن سبع خوارزميات تقدير مختلفة في التعلم المعزز—وهي GAE، وV-Trace، وRetrace( $\lambda$ )، وعوائد TD( $\lambda$ )، والعوائد المخصومة (discounted returns)، وآثار الأهلية (eligibility traces)، والجمع التراكمي للحلقات (episodic prefix sums)—هي جميعها حالات لـ تكرار خطي من الدرجة الأولى:
$A_t = \alpha_t + \beta_t \cdot A_{t+1}$
(أو مرآتها في الزمن الأمامي لآثار الأهلية).
تتم إعادة صياغة هذه الخوارزميات كـ مسوحات ترابطية متوازية (parallel associative scans) فوق أزواج $(\alpha_t, \beta_t)$ باستخدام عامل ترابطي محدد $\oplus$ :
$(\alpha_B, \beta_B) \oplus (\alpha_A, \beta_A) = (\alpha_B + \beta_B \alpha_A, \beta_A \beta_B)$
هذا التحويل يسمح بحل التكرار في $O(\log T)$ من الخطوات المتوازية بدلاً من $O(T)$ من الخطوات التسلسلية.

2. نوى Triton المدمجة (Fused Triton Kernels)

على عكس تنفيذات PyTorch القياسية التي تجسد تنسورات المعاملات الوسيطة ( $\alpha$ و $\beta$ ) في ذاكرة HBM قبل إجراء المسح، يستخدم rl-triton نوى مدمجة.

البناء على الشريحة (On-Chip Construction): تقوم النوى بتحميل مدخلات التدفق الخام (المكافآت، القيم، أعلام الانتهاء own، إلخ) وتبني معاملات التكرار ( $\alpha_t, \beta_t$ ) مباشرة في السجلات (registers).
المسح على الشريحة (On-Chip Scan): يتم إجراء المسح الترابطي بالكامل على الشريحة (باستخدام السجلات والذاكرة المشتركة) عبر tl.associative_scan. لا يتم كتابة حالات المسح الوسيطة أبدًا في HBM حتى الوصول إلى المخرج النهائي.
كفاءة الذاكرة: يقلل هذا النهج من رحلات الذهاب والإياب إلى HBM. بالنسبة لمتتالية بطول $T$ ، تقوم النواة المدمجة بعملية قراءة وكتابة واحدة فقط في HBM، بينما يتطلب المسح المجسد $O(\log T)$ من الرحلات.

3. دلالات الحدود الصريحة

تعالج المكتبة صراحةً شروط حدود الحلقة المعقدة التي غالبًا ما تسبب أخطاء في التنفيذات القياسية:

الانتهاء مقابل الاقتطاع (Termination vs. Truncation): تميز بين نهايات الحلقة الطبيعية (حيث تكون القيمة اللاحقة صفرًا) والاقتطاعات (حيث يجب استخدام قيمة استكمال/bootstrap مقدمة).
نوافذ التدفق (Rollout Windows): تدير بشكل صحيح الحد عند نهاية مخزن التدفق ( $t=T-1$ )، مما يضمن دمج قيم الاستكمال ( $\text{bootstrap values}$ ) في التحديث المحلي ( $\alpha_{T-1}$ ) أو في حمل المسح ( $A_T$ ) دون احتساب مزدوج، اعتمادًا على الخوارزمية المحددة.

المساهمات الرئيسية

إطار عمل موحد: إثبات أن سبع خوارزميات قياسية لتخصيص الائتمان في التعلم المعزز هي حالات لعامل مسح أفيني (affine scan operator) واحد.
تنفيذ مدمج: نوى Triton خاصة بكل خوارزمية تقوم ببناء المعاملات على الشريحة وتنفيذ المسح دون تجسيد وسيط في HBM.
معالجة حدود قوية: منطق صريح ومتحقق منه للخطوات المنتهية، والخطوات المقتطعة، وحدود نافذة التدفق.
اختبارات الأداء: تقييم شامل يظهر تسريعًا مقارنة بالحلقات التسلسلية القياسية ونماذج الأساس المتجهة torch.compile.
إصدار مفتوح المصدر: حزمة بايثون مع اختبارات الصحة، واختبارات تراجع الأداء، والتوثيق.

النتائج

تم تقييم المكتبة على وحدات NVIDIA H100 (80GB HBM3) وRTX 2000 Ada عبر مختلف أعداد البيئات وأطوال المتتاليات.

التسريع مقارنة بالحلقات القياسية: في أنظمة المحاكاة ذات التوازي الهائل (مثل: 4,096 بيئة، 128 خطوة)، يحقق rl-triton تسريعًا يتراوح بين 16–317 ضعفًا مقارنة بالحلقات التسلسلية غير المجمعة. ويرجع ذلك أساسًا إلى تقليل عمق التبعية من $O(T)$ إلى $O(\log T)$ .
التسريع مقارنة بنماذج الأساس المتجهة: مقارنة بنموذج أساس متجه torch.compile (الذي يستخدم أيضًا مسح $O(\log T)$ $O (lo g T)$ ولكنه يجسد الوسائط)، يحقق rl-triton تسريعًا في كامل الاستدعاء يتراوح بين 1.6–5.70×.
- يكون التسريع أكثر وضوحًا للخوارزميات ذات التعقيد الحسابي الأقل لكل خطوة (مثل: العوائد المخصومة، وعوائد $\lambda$ ) وأطوال المتتاليات الطويلة، حيث يصبح تكرار رحلات الذهاب والإياب إلى HBM في النموذج الأساسي هو عنق الزجاجة.
- بالنسبة لـ Retrace( $\lambda$ )، تكون معدلات التسريع أقل (1.6–1.9×) بسبب الحسابات الأثقل لكل خطوة (تحميل احتمالات الإجراء) ومشكلات ضغط السجلات عند أطوال المتتاليات الطويلة جدًا ( $>2048$ ).
التأثير النهائي (End-to-End): في تحديث PPO اصطناعي، ترجم تسريع النواة إلى تسريع في التدريب الكامل يتراوح بين 1.11–1.16× عندما كان GAE يمثل 10–15% من إجمالي وقت التحديث. وعندما كانت حصة GAE أصغر (~~2.2%)، كان الربح النهائي ضئيلاً (~~1.02×)، وهو ما يتوافق مع قانون أمدال (Amdahl's Law).
المقارنة مع PufferLib: يتفوق rl-triton على نواة CUDA المكتوبة يدويًا في PufferLib في نظام المحاكاة ذي التوازي الهائل (تدفقات قصيرة، عدد بيئات مرتفع) في 31 من أصل 40 شكلًا تم اختبارها، لا سيما مع زيادة طول المتتالية.

الأهمية والادعاءات

يدعي البحث أن rl-triton يوفر حلاً موحدًا عالي الأداء للمسار الحرج لتخصيص الائتمان في التعلم المعزز. من خلال الاستفادة من إطار المسح الترابطي وقدرات الدمج في Triton، فإنه يقضي على اختناقات عرض نطاق HBM التي تعاني منها كل من التنفيذات القياسية ونماذج PyTorch المتجهة التقليدية.

يؤكد المؤلف أن المكتبة مهمة بشكل خاص لـ أنظمة المحاكاة ذات التوازي الهائل (آلاف البيئات، تدفقات قصيرة) الشائعة في التدريب الموزع الحديث وخطوط أنابيب RLHF، حيث يؤدي تقليل حركة مرور HBM إلى توفير كبير في وقت التنفيذ الفعلي. وبينما لا تدعي المكتبة تسريع خط أنابيب التدريب بأكمله (حيث غالبًا ما تهيمن عمليات تمرير الشبكة الأمامي والخلفي)، إلا أنها تقدم تحسينًا جوهريًا لمكون تخصيص الائتمان، وهو تبعية صارمة في كل خطوة تحديث. كما يسلط العمل الضوء على إمكانية تطبيق هذه المسوحات المدمجة على مجالات أخرى تتضمن التكرارات الأفينية المجزأة، مثل تدريب النماذج اللغوية الكبيرة (LLM) ذات السياق الطويل والشبكات العصبية النبضية (spiking neural networks).