ملخص تقني: AdaptiveLoad

1. بيان المشكلة

تواجه عملية تدريب نماذج محولات انتشار الفيديو (Video Diffusion Transformers) واسعة النطاق (مثل DiT و MMDiT) لنمذجة العالم (World Models) عقبتين حرجتين عند معالجة مجموعات البيانات مختلطة الأنماط التي تحتوي على أطوال تسلسل متباينة للغاية (تتراوح من الصور الثابتة إلى الفيديوهات طويلة المدة):

• عدم التوازن الحسابي بسبب التعقيد التربيعي: تعتمد خطوط أنابيب التدريب الموزعة الحالية على استراتيجيات "التجميع المتساوي للرموز" (Equal-token bucketing) (حيث يكون $B \times S = \text{constant}$). تفشل هذه الطريقة في مراعاة التعقيد الحسابي بمقدار $O(S^2)$ لآليات الانتباه الذاتي (Self-attention). ونتيجة لذلك، تصبح الحاويات (Buckets) التي تحتوي على تسلسلات طويلة "متخلفة حسابياً" (Computational stragglers)، مما يسبب عدم توازن شديد في الحمل. ويتم تحديد وقت المزامنة العالمي بواسطة أبطأ وحدة معالجة رسومية (GPU)، مما يؤدي إلى فترات خمول كبيرة (فقاعات مزامنة) للعقد الأخرى. وتشير البيانات التجريبية إلى وجود ارتباط ضعيف ($R \approx 0.35$) بين زمن تأخير الخطوة وإجمالي عدد الرموز، ولكن يوجد ارتباط قوي ($R \approx 0.92$) مع مربع طول التسلسل.
• جدار الذاكرة وعبء العمليات (Operator Overhead): في بنيات MMDiT، يتم استدعاء تطبيع الطبقة التكيفي (Adaptive Layer Normalization - AdaLN) مئات المرات في كل تكرار. تستخدم التنفيذات القياسية نوى CUDA منفصلة، مما يسبب رحلات ذهاب وإياب متكررة وغير ضرورية للموترات الوسيطة بين المعالجات متعددة الخيوط (SMs) والذاكرة ذات النطاق الترددي العالي (HBM). علاوة على ذلك، أثناء الانتشار العكسي (Backpropagation)، يعاني تقليل التدرج لمعلمات التعديل (Modulation parameters) من أنماط وصول غير فعالة للذاكرة ذات الخطوات (Strided access) (وصول غير متصل على طول بُعد التسلسل $N$)، مما يمنع الوصول المتماسك (Coalesced access) على مستوى الأجهزة ويؤدي إلى عدم استغلال عرض نطاق الذاكرة بشكل كافٍ.

2. المنهجية: إطار عمل AdaptiveLoad

يقترح البحث AdaptiveLoad، وهو إطار عمل تحسين كامل المكونات يدمج الجدولة عالية المستوى مع دمج النوى (Kernel Fusion) منخفض المستوى.

أ. موازنة الحمل التكيفية مزدوجة القيود

لمعالجة عدم التوازن الحسابي، استبدل المؤلفون افتراض "التساوي الخطي للرموز" بمنطق جدولة مزدوج القيود يفرض في آن واحد:

1. قيد الذاكرة الخطي: $B \times S \leq M_{mem}$ (لضمان ملاءمة حجم الدفعة لذاكرة وحدة معالجة الرسوميات).
2. القيد الحسابي متعدد الحدود: $B \times S^p \leq M_{comp}$ (لمراعاة التعقيد $O(S^2)$).

يتم تحديد حجم الدفعة النهائي للحاوية من خلال تقاطع هذين القيدين:
$$B_{shape} = \max(1, \min(\lfloor M_{mem}/S \rfloor, \lfloor M_{comp}/S^p \rfloor))$$

• نظام ملاءمة التكلفة (Cost Fitting System): لتحديد الأس التجريبي $p$ والحد الحسابي $M_{comp}$، يستخدم النظام سير عمل قياس أداء الشكل (Shape Benchmarking) القائم على البيانات. يقوم النظام بتنفيذ عمليات مسح بكسل اصطناعية في بيئة موزعة حية (مثل مسارات FSDP) لقياس زمن تأخير الخطوة مقابل حجم الدفعة وطول التسلسل. وباستخدام الانحدار غير الخطي، يقوم النظام بملاءمة نموذج تكلفة ($\text{step time} \approx a + b \times B \times S^p$) لمعايرة تكوينات الحاويات ديناميكياً، مما يلغي الحاجة إلى الضبط التجريبي اليدوي.

ب. نوى CUDA المدمجة لـ (LayerNorm-Modulate)

لمعالجة "جدار الذاكرة"، صمم المؤلفون نواة CUDA مدمجة لـ Fused AdaLN تجمع بين عمليات التطبيع والتعديل في إطلاق نواة واحدة.

• المسار الأمامي (Forward Pass): تستخدم النواة السجلات (Registers) والذاكرة المشتركة (Shared Memory) لتخزين الإحصائيات الوسيطة (المتوسط والتباين)، وتنفذ عمليات تقليص على مرحلتين (Warp-level Shuffle و Block-level Reduction). تُطبق التحويلات الخطية مباشرة أثناء المسح الثالث لبُعد الميزة $D$، مما يلغي عمليات القراءة والكتابة المتكررة في الذاكرة (HBM).
• المسار الخلفي (D-tile Coalesced Reduction): لحل مشكلة الوصول عبر الخطوات في تقليل التدرج، قدم المؤلفون استراتيجية تقليل متماسك بنمط D-tile. بدلاً من التقليل على طول بُعد التسلسل $N$ (الذي يسبب وصولاً عبر الخطوات)، يتم تقسيم الشبكة (Grid) إلى تخطيط $(d_{tile}, n_{tile})$. يتم تثبيت الخيوط (Threads) على مؤشر ميزة $d$ وتقوم بالتراكم الرأسي على بلاطات التسلسل. يضمن هذا أن عمليات الوصول إلى الذاكرة داخل الـ Warp تكون متصلة، مما يسمح بمعاملات HBM متماسكة ويعظم استغلال عرض النطاق الترددي.
• تبسيط الرسم البياني (Graph Simplification): من خلال دمج سلسلة العمليات، يتم ضغط الرسم البيو الحسابي من 5-8 عقد منفصلة إلى عقدة ذرية واحدة. تظل الإحصائيات الوسيطة في السجلات، مما يتجاوز إدارة الذاكرة العالمية للإطار البرمجي ويقلل بشكل كبير من بصمة ذاكرة التنشيط (Activation memory).

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

1. تحديد ودرء عدم تطابق الحمل: يحدد البحث بدقة عدم تطابق الحمل في تدريب انتشار الفيديو ويقترح خوارزمية جدولة حاويات تكيفية مزدوجة القيود. يقلل هذا النهج من معدل عدم توازن الحمل بنسبة تصل إلى 50% من خلال الإدارة النشطة لتعقيد $O(S^2)$ للتسلسلات الطويلة.
2. نمذجة تكلفة قائمة على البيانات: صمم المؤلفون نظاماً متكاملاً ينقل تكوين التدريب الموزع من التوجيه بالحدس إلى التوجيه بالبيانات. ومن خلال الملاءمة الآلية لنموذج التكلفة عبر "قياس أداء الشكل"، يحسن النظام استغلال الموارد على مستوى الكتلة (Cluster) دون تدخل يدوي.
3. نوى مدمجة عالية الأداء: تم تطوير نواة Fused AdaLN CUDA متخصصة لبنيات MMDiT. من خلال تقديم تقليل D-tile المتماسك، تعالج النواة مشكلة الوصول عبر الخطوات في الانتشار العكسي، مما يزيد من إنتاجية التدريب بأكثر من 25% مع تقليل استخدام ذاكرة التنشيط في الوقت نفسه.

4. النتائج التجريبية

أُجريت التجارب على النموذج العالمي Wan 2.1 باستخدام مجموعة بيانات مختلطة تضم 10 ملايين عينة (WebDataset و Koala-36m) عبر مجموعات مكونة من 8 و 16 وحدة معالجة رسومية.

• تحسين الإنتاجية (Throughput):
  • مجموعة 8 وحدات معالجة رسومية: ارتفع متوسط إنتاجية التدريب من 14,383 إلى 18,069 رمز/ثانية (+25.6%).
  • مجموعة 16 وحدة معالجة رسومية: ارتفع متوسط الإنتاجية من 30,170 إلى 38,372 رمز/ثانية (+27.18%).
  • حافظ منحنى AdaptiveLoad على حد أدنى أعلى للإنتاجية مقارنة بالنموذج الأساسي، حتى في أسوأ حالات توزيع التسلسل.
• كفاءة موازنة الحمل:
  • انخفض معدل عدم التوازن الحسابي (المقاس بمعامل الاختلاف للحمل الحسابي) من 39% إلى 18.9%.
  • انخفضت تكلفة المزامنة (CV لزمن الخطوة) بشكل كبير، حيث شهد تكوين 16 وحدة معالجة رسومية انخفاضاً من متوسط أساسي قدره 18.7% إلى 10.4% باستخدام AdaptiveLoad.
• أداء النواة والذاكرة:
  • ذروة استخدام ذاكرة الفيديو (VRAM): تحسنت بنسبة 22.7%.
  • تسريع النواة: حقق المسار الأمامي المدمج تسريعاً يتراوح بين 3.21× إلى 3.39×. وحقق المسار الخلفي (D-tile reduction) تسريعاً يصل إلى 1.42× عند طول تسلسل 64 ألف.
  • توفير الذاكرة: قللت النواة المدمجة ذاكرة التنشيط بنسبة تقارب 61.9% (على سبيل المثال، توفير ~4.8 جيجابايت عند $N=64,000$).
  • سعة الحمل: دعم النظام المحسن أقصى طول تسلسل قدره 52,800 رمز (زيادة بنسبة 10% عن حد 48,000 رمز في النموذج الأساسي) دون زيادة في ذروة استخدام الذاكرة.
• استقرار التدريب: ظلت منحنيات الخسارة (Loss curves) لـ AdaptiveLoad متطابقة تماماً مع النموذج الأساسي، مما يثبت أن التحسينات حافظت على الدقة العددية وجودة التقارب.

5. الأهمية

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