Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

إليك ملخص تقني مفصل لورقة البحث بعنوان "طي التوازي التنسوري والتسلسلي من أجل تدريب واستدلال نماذج المحولات بكفاءة عالية في الذاكرة" (Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference).

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

تواجه عمليات تدريب واستدلال النماذج اللغوية الكبيرة (LLMs) عقبتين أساسيتين في الذاكرة:

1. ذاكرة المعلمات (Parameter Memory): المساحة المطلوبة لأوزان النموذج، والتدرجات (gradients)، وحالات المُحسن (optimizer states).
2. ذاكرة التنشيط (Activation Memory): المساحة المطلوبة للتنشيطات الوسيطة، والتي تتناسب طردياً مع مربع طول التسلسل ($S^2$) في آليات الانتباه (attention mechanisms).

تعالج استراتيجيات التوازي الحالية هذه المشكلات بشكل منفصل أو مع وجود مقايضات:

• التوازي التنسوري (Tensor Parallelism - TP): يقوم بتجزئة أوزان النموذج عبر الأجهزة، مما يقلل ذاكرة المعلمات ولكنه يترك ذاكرة التنشيط كاملة ومكررة. ويصبح غير فعال عند أطوال التسلسل الطويلة بسبب ارتفاع ذاكرة التنشيط وحجم الاتصالات.
• التوازي التسلسلي (Sequence Parallelism - SP): يقوم بتجزئة مدخلات التسلسل عبر الأجهزة، مما يقلل ذاكرة التنشيط ولكنه يترك أوزان النموذج مكررة بالكامل. وهو غير فعال عندما تكون ذاكرة المعلمات هي العائق.
• الجمع التقليدي بين TP وSP: يجمع بينهما عبر تعيينهما لمحاور متعامدة في شبكة الأجهزة (على سبيل المثال، شبكة $T \times \Sigma$). ورغم فعاليته، إلا أنه يستهلك عدداً كبيراً من الأجهزة للنموذج المتوازي ($T \times \Sigma$)، مما يترك أجهزة أقل للتوازي البياناتي (Data Parallelism). علاوة على ذلك، فإنه غالباً ما يجبر مجموعات التوازي على الامتداد عبر روابط بين العقد (inter-node) أبطأ، مما يؤدي إلى عدم الاستفادة الكاملة من الاتصالات داخل العقدة الواحدة (intra-node) عالية النطاق الترددي (مثل NVLink أو AMD Infinity Fabric).

التحدي الجوهري: كيف يمكن تقليل كل من ذاكرة المعلمات وذاكرة التنشيط في آن واحد على محور جهاز واحد لتعظيم كفاءة الذاكرة وتوفير الأجهزة للتوازي البياناتي، دون تكبد تكاليف اتصال باهظة.

2. المنهجية: التوازي التنسوري والتسلسلي (TSP)

يقترح المؤلفون استراتيجية TSP، وهي استراتيجية تقوم بـ "طي" التوازي التنسوري والتسالي في محور جهاز واحد. بدلاً من تعيين TP وSP لمحاور متعامدة، تقوم TSP بتعيين جزء من أوزان النموذج وجزء من توكنات (tokens) التسلسل لكل رتبة (rank/GPU).

المكونات التصميمية الرئيسية:

• التجزئة الموحدة (Unified Sharding): يحمل كل جهاز $1/D$ من معلمات النموذج و$1/D$ من توكنات التسلسل. هذا يقلل من الذاكرة المرتبطة بالمعلمات والذاكرة المرتبطة بالتنشيط بمقدار $1/D$ على محور واحد.
• جداول التشغيل (Runtime Schedules): للتعامل مع حقيقة أن الأوزان والتنشيطات يتم تجزئتها بشكل مختلف، تستخدم TSP جداول اتصال محددة لطبقات مختلفة:
  • طبقات الانتباه (Attention Layers): تستخدم جدول البث (broadcast).
    • تمر الرتب (ranks) على أجزاء الأوزان (المجمعة $W_Q, W_K, W_V, W_O$).
    • في كل تكرار، تقوم رتبة واحدة ببث جزء الوزن الخاص بها إلى جميع الرتب الأخرى.
    • تقوم جميع الرتب بحساب الإسقاطات المحلية (local projections) على أجزاء التسلسل الخاصة بها.
    • تقوم عملية تجميع الكل (All-Gather) بإعادة بناء سياق المفتاح/القيمة (Key/Value) الكامل عبر بُعد التسلسل لتنفيذ خوارزمية FlashAttention.
    • تُستخدم طريقة تقسيم الزجزاج (Zigzag partitioning) لموازنة عبء عمل الانتباه السببي (ضمان عدم تحمل التوكنات المتأخرة لعمل إضافي غير متناسب).
  • طبقات Gated MLP: تستخدم جدول اتصال حلقي (ring communication).
    • تدور أجزاء الأوزان ($W_1, W_2, W_3$) حول الحلقة (point-to-point).
    • تقوم كل رتبة بتجميع المخرجات الجزئية محلياً.
    • هذا يلغي الحاجة إلى عملية All-Reduce (المستخدمة في TP القياسي)، حيث يظل التسلسل محلياً وتتحرك الأوزان فقط.
• الوعي بالعتاد (Hardware Awareness): تم تحسين التصميم ليتناسب مع بنيات مثل AMD MI300X (Infinity Fabric)، حيث يتوسع النطاق الترددي الجماعي مع عدد المعالجات المشاركة. تضمن TSP أن عمليات الاتصال (البث والحلقة) تشغل المجموعة الكاملة أو تتداخل مع العمليات الحسابية لتعظيم الكفاءة.

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

1. استراتيجية توازي مبتكرة: تقديم TSP، التي تدمج محوري التوازي التنسوري والتسلسلي في محور واحد، مما يسمح بالتقليل المتزامن لذاكرة المعلمات وذاكرة التنشيط.
2. التحليل النظري: نموذج رياضي صارم يقارن TSP مع DP وTP وSP وTP+SP فيما يتعلق بـ:
   • الذاكرة: TSP هي المخطط الوحيد الذي يقلل من الذاكرة المرتبطة بالوزن والذاكرة المرتبة بالتنشيط بمقدار $1/D$ دون الحاجة إلى $D^2$ من الأجهزة.
   • الاتصال: بينما تقدم TSP حجماً أكبر من اتصالات حركة الأوزان (التي يتجنبها TP+SP)، يثبت المؤلفون أنه بالنسبة للتسلسلات الطويلة أو أحجام الدفعات المتوسطة، فإن الاتصال المرتبط بالتنشيط هو المهيمن، مما يجعل TSP منافسة.
   • العمليات الحسابية (FLOPs): تحافظ TSP على نفس كفاءة الحوسبة لكل جهاز كما في TP وSP.
3. التنفيذ والتداخل (Implementation & Overlap): تنفيذ تفصيلي لكتل Attention وMLP في TSP مع استراتيجيات تداخل واضحة بين الاتصال والحوسبة (توصيل نقل الأوزان خلف عمليات GEMM).
4. تحسين الطوبولوجيا (Topology Optimization): إثبات أن TSP تسمح باستيعاب نسخة النموذج بالكامل داخل عقدة واحدة عالية النطاق (intra-node)، مما يتجنب عقوبات الأداء الناتجة عن الامتداد عبر الروابط البينية (inter-node) الأبطأ التي تتطلبها شبكات TP+SP ثنائية الأبعاد.

4. النتائج

قيم المؤلفون TSP على عنقود يصل إلى 1,024 من معالجات AMD MI300X (128 عقدة) لنموذج محول كثيف بـ 7 مليار معلمة.

• كفاءة الذاكرة:
  • حققت TSP باستمرار أدنى ذروة ذاكرة عبر جميع أطوال التسلسل المختبرة (من 16 ألف إلى 128 ألف توكن).
  • عند السياقات القصيرة، طابقت TSP نموذج TP (الذي يهيمن عليه وزن المعلمات).
  • عند السياقات الطويلة، تفوقت TSP بشكل كبير على TP وTP+SP (التي تهيمن عليها ذاكرة التنشيط)، مما منع حدوث أخطاء نفاذ الذاكرة (OOM) التي فشلت فيها الاستراتيجيات الأخرى.
• الإنتاجية (Throughput):
  • رغم زيادة حجم الاتصالات النظري (بسبب حركة الأوزان)، حافظت TSP على إنتاجية تنافسية مع TP+SP.
  • تم إخفاء الاتصالات الإضافية بنجاح خلف الحوسسة عبر تقنية الـ pipelining.
  • نمت مزايا الإنتاجية مع زيادة درجات التوازي ($D$) وأطوال التسلسل.
• القابلية للتوسع (Scalability):
  • سمحت TSP باستخدام دفعات دقيقة (micro-batches) أكبر بسبب تقليل بصمة الذاكرة، مما حسن الإنتاجية بشكل أكبر.
  • مكنت من استيعاب نسخ النموذج بالكامل داخل عقدة واحدة (8 معالجات)، مما أتاح المزيد من الرتب للتوازي البياناتي (DP) في العناقيد الأكبر.

5. الأهمية

• التوسع الواعي بالعتاد: توفر TSP حلاً عملياً لعناقيد الحوسبة عالية الأداء (HPC) الحديثة حيث يكون النطاق الترددي داخل العقدة عالياً بينما يكون النطاق بين العقد عائقاً. إنها تعظم الاستفادة من الاتصالات الداخلية المكلفة.
• قابلية السياق الطويل: مع توجه النماذج نحو نوافذ سياق أطول (100 ألف+ توكن)، تصبح ذاكرة التنشيط هي العائق الرئيسي. TSP مهيأة بشكل فريد للتعامل مع هذا دون الحاجة إلى التكاليف الهائلة للأجهزة المطلوبة في التوازي ثنائي الأبعاد.
• المرونة: يمكن دمج TSP مع مخططات توازي أخرى (مثل التوازي في خط الأنابيب Pipeline Parallelism، أو توازي الخبراء Expert Parallelism، أو التوازي البياناتي Data Parallelism)، مما يوفر "محوراً" مرناً لمصممي الأنظمة لتحسين الأداء بناءً على طوبولوجيا العتاد.
• كفاءة التكلفة: من خلال تقليل عدد الأجهزة المطلوبة لتوازي النموذج، تسمح TSP باستخدام المزيد من الأجهزة للتوازي البياناتي، مما قد يقلل من حجم العنقود الإجمالي المطلوب لمهمة تدريب معينة أو يزيد من حجم الدفعة الفعال.

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