نقل الأوزان لمرحلة ما بعد تدريب التعلم المعزز في أقل من ثانيتين

لقد حققنا مؤخراً تحديثات للمعلمات عبر الأجهزة في غضون 1.3 ثانية لنموذج Kimi-K2 (بمعلمات تبلغ 1 تريليون)، حيث تم نقل الأوزان من 256 وحدة معالجة رسومية للتدريب (BF16) إلى 128 وحدة معالجة رسومية للاستدلال (FP8).

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

من خلال الاستفادة من الاتصال من نقطة إلى نقطة عبر تقنية RDMA، نجعل عملية نقل الأوزان سريعة للغاية دون تغيير محرك الاستدلال، مما يجعل الكود برمجيًا أسهل في الكتابة والصيانة.

يعتمد حلنا على خاصية كتابة RDMA، وهي وسيلة أساسية أحادية الجانب يقوم المصدر من خلالها بالتسجيل مباشرة في ذاكرة وحدة المعالجة الرسومية للوجهة.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

لن يتم إشعار جهة الوجهة بعملية النقل مطلقًا. يمنحنا هذا عمليات نقل منخفضة الكمون، وعالية الإنتاجية، وخالية من النسخ، مدفوعة بواسطة عقد التدريب دون أي منطق تحكم على عقد الاستدلال.

جمع البيانات الوصفية – يقوم وحدة التحكم بجمع البيانات الوصفية للمعلمات من جميع وحدات معالجة الرسومات الخاصة بالتدريب والاستدلال.

حساب الجدول الزمني – يقوم وحدة التحكم بحساب جدول زمني ثابت لنقل الأوزان، يحدد وحدة معالجة رسومات التدريب التي ترسل المعلمة إلى أي وحدة معالجة رسومات للاستدلال، وبالترتيب الذي يتم تحديده.

توزيع الجدول الزمني – يقوم وحدة التحكم بإرسال الجدول الزمني إلى كافة وحدات معالجة رسومات التدريب.

التنفيذ – بعد كل خطوة تدريب، يشير وحدة التحكم إلى وحدات معالجة رسومات التدريب لبدء عمليات النقل.

بعد تحديد سير العمل العام، يكمن التحدي الرئيسي في كيفية تنفيذ عمليات نقل الأوزان بكفاءة على نطاق تريليون معلمة. نوضح هنا تفاصيل مسار التنفيذ.

يتم توزيع المعلمات في التدريب وفقًا لتوزيعات FSDP. باستخدام full_tensor()، يمكن لجميع وحدات معالجة الرسومات في DeviceMesh إعادة بناء المعلمة الكاملة، وبالتالي يمكن أن تكون جميعها كمصدر لنقل الأوزان.

تُشكِّل شبكات DeviceMeshes المتعددة وغير المتداخلة مجموعة شبكة. ونظرًا لأن شبكات DeviceMeshes في نفس المجموعة غير متداخلة، فإن عمليات نقلها لا تتداخل ويمكن تشغيلها بالكامل بالتوازي. وبين مجموعات الشبكة، نقوم بإدراج حاجز عالمي فرض الترتيب.

تتراكب سلسلة نقل الأوزان مع نسخ البيانات من المضيف إلى الجهاز، وتجميع التنسورات، والدمج والكمّ، وعمليات النقل RDMA عبر MeshGroups.

نتعامل مع نقل كل موتر معلمة كمهمة. تستخدم عملية نقل الأوزان أنواعًا متعددة من مصادر الأجهزة، وبالتالي نقوم بتقسيم مهمة نقل الأوزان إلى مراحل مختلفة لخط الأنابيب تتداخل في الوقت:

نسخ الذاكرة من المضيف إلى الجهاز (Host to device memcpy) — إذا قام FSDP بتفريغ الأوزان إلى وحدة المعالجة المركزية (CPU)

إعداد المعلمات — إعادة بناء الوزن الكامل باستخدام full_tensor()، وتطبيق دمج الإسقاط، والقيام بالترميز الكمي إذا لزم الأمر.

نقل RDMA — كتابة خالية من النسخ إلى ذاكرة وحدة المعالجة الرسومية البعيدة للاستدلال

الحاجز العالمي — بعد اكتمال جميع استدعاءات full_tensor()، تتم المزامنة عبر مجموعات الشبكة باستخدام GLOO عبر شبكة إيثرنت.

في التنفيذ، نحافظ على طابور من نوع FIFO للمهام لكل مرحلة من مراحل خط الأنابيب. ومتى أكملت مهمة رأس الطابور المرحلة، يتم نقلها إلى ذيل طابور المرحلة التالية.

تؤدي full_tensor() والعمليات الأخرى لوحدة معالجة الرسومات إلى استهلاك إضافي لذاكرتها. لتجنب خطأ نفاد الذاكرة، نبدأ في تنفيذ المهمة فقط إذا كانت المهام الجارية حالياً تشغل ذاكرة مؤقتة لوحدة معالجة الرسومات تقل عن الحد المرجعي القابل للتكوين.

تجعل العديد من خيارات التصميم نظامنا أسرع بكثير في التشغيل وأسهل في الصيانة مقارنة بالحلول الشائعة مفتوحة المصدر.

يجمع النقل القائم على Rank0 الأوزان قبل توزيعها على عقدة الاستدلال.

ينقل الاتصال من نقطة إلى نقطة الأوزان من عقدتي التدريب إلى واجهات الشبكة الخاصة بعقدة الاستدلال.

النمط الشائع هو توجيه جميع المعلمات عبر وحدات معالجة الرسومات ذات الترتيب 0 (rank-0): التجميع على ترتيب 0 للتدريب، والإرسال إلى ترتيب 0 للاستدلال، ثم التوزيع مرة أخرى. يصبح هذا بسرعة نقطة اختناق، محدودة بعرض النطاق الترددي لـ PCIe وبطاقة واجهة الشبكة (NIC) الخاصة بوحدة معالجة رسومات واحدة (مثل 400 جيجابت في الثانية ≈ 50 جيجابايت/ثانية).

في المقابل، يتيح إعدادنا من نقطة إلى نقطة لكل وحدة معالجة رسومات تدريبية الإرسال مباشرة إلى كل وحدة معالجة رسومات استدلالية، مما يشبع البنية الأساسية للشبكة بأكملها بدلاً من ارتباط واحد.

يستخدم النقل القائم على RPC مخزنًا مؤقتًا على جانب الاستدلال واستدعاءات العاملين، بينما يكتب النقل أحادي الجانب الأوزان مباشرةً من دون مشاركة عامل الاستدلال.

تعتمد بعض الأنظمة على استدعاء طريقة update_weight() الخاصة بمحرك الاستدلال لكل موتر. وهذا يعني إجراء تغييرات متداخلة على كود الاستدلال، بالإضافة إلى العبء الناتج عن استدعاءات الإجراءات البعيدة (RPC)، والتسلسل، وتنسيق طبقة التحكم.

مع استخدام خاصية كتابة RDMA، نقوم بتحديث الأوزان بصمت في ذاكرة وحدة المعالجة الرسومية للاستدلال، دون عمليات نسخ إضافية. لا توجد رسائل لطبقة التحكم ولا يوجد منطق تحكم لوحدة المعالجة المركزية. ولا يتطلب الأمر أي تعديل على محرك الاستدلال.

يمكن لعملية نقل الأوزان الاستفادة من أربعة أنواع من موارد الأجهزة: (1) نقل البيانات بين المضيف والجهاز (2) حساب وحدة معالجة الرسومات لدمج الإسقاط والترميز الكمي (3) شبكة RDMA لبيانات مستوى البيانات (4) إيثرنت لمستوى التحكم.

يقسم تصميمنا مهام نقل الأوزان إلى مراحل خط الأنابيب، مما يسمح بالتداخل السهل عبر موارد الأجهزة المختلفة.

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

يتم حساب جدولنا الزمني مرة واحدة عند التهيئة. يقوم كل تكرار تدريبي ببساطة بإعادة تشغيل الخطة: يُصدر وحدة التحكم إشارة "انطلاق"، وتتبع وحدات معالجة الرسومات مساراتها المحددة مسبقًا. التنفيذ يتسم بالموثوقية وخفة الوزن.

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

في هندستنا، نفصل هذه الخطوات كمكونات فردية. يمكن اختبار كل مكون على حدة، وتحليله، وتحسينه بشكل مستقل.

يعد نقل الأوزان السريع والموثوق لبنة أساسية للضبط الدقيق للتعلم المعزز واسع النطاق. من خلال الجمع بين خاصية كتابة RDMA، وجدول النقل الثابت، والتنفيذ عبر خطوط الأنابيب, قمنا بتقليص تحديثات التريليون معلمة إلى 1.3 ثانية فقط على Kimi-K2. النهج سهل الفهم، وسهل الصيانة، ويتجنب الاختناقات في التصاميم التقليدية.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

كتابة RDMA: عمليات نقل أحادية الجانب

سير العمل العام

تنفيذ نقل الأوزان

DeviceMesh ومجموعات الشبكة (Mesh Groups)

خط أنابيب المهام

التحكم في استخدام ذاكرة وحدة معالجة الرسومات (GPU)

لماذا هو سريع وبسيط

الالاتصال من نقطة إلى نقطة

نقل البيانات أحادي الجانب

خط الأنابيب (Pipelining)

جدول زمني ثابت

فصل نظيف

الخاتمة

اقتباس