تُقدّم المقالة بحثًا جديدًا بعنوان "CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs"، هدفه الأساسي هو تحسين كفاءة تدريب نماذج Transformer، خاصةً من خلال معالجة العمليات "المكثفة من حيث الذاكرة" التي تبدو منفصلة ولكنها تتراكم لتسبب تأخيرًا كبيرًا.

كاتب المقال، المصدر: Machine Learning China

في 22 مايو، قام تري داو بمشاركة تغريدة من هان غوو على وسائل التواصل الاجتماعي. كما كتب: "بعد إعادة صياغة رياضية، تبين أن كل شيء في Transformer هو سلسلة من GEMM + epilogue (ضرب المصفوفات زائد الخاتمة). وبفضل بعض الأوامر المحسّنة، يمكن للـ LLM (والمبتدئين) كتابة نوى بسرعة الضوء لجميع عمليات Transformer!"

تري داو هو أحد المؤلفين الأساسيين لسلسلة FlashAttention، وهذه التغريدة تشير إلى الورقة البحثية التي نشروها ذلك اليوم: CODA.

• عنوان البحث: CODA: إعادة كتابة كتل Transformer كبرامج GEMM-Epilogue
• رابط الورقة: https://arxiv.org/abs/2605.19269
• عنوان الكود: https://github.com/HanGuo97/coda-kernels

اسم هذا المشروع يُقرأ كـ"الخاتمة" ويُنطق كـ"CUDA". حاول باحثون من MIT وبرينستون وTogether AI وMeta استخدام مجموعة جديدة من التجريدات البرمجية لمعالجة منهجية جميع الحسابات المبعثرة التي تُهمل غالبًا ولكنها تستهلك وقتًا مستمرًا في تدريب Transformers.

ضريبة الكسل لتدريب النماذج الكبيرة

لفهم المشكلة التي تحلها CODA، يجب أولاً فهم أين تذهب وقت تدريب النماذج الكبيرة.

عند تدريب نموذج بحجم 1 مليار معلمة بأسلوب LLaMA-3 على وحدة NVIDIA H100 واحدة، سيظن معظم الناس بشكل غريزي أن الوقت يُنفق بالكامل على عمليات ضرب المصفوفات وحسابات الانتباه، لأنها تمثل "الحساب الحقيقي". هذه الغريزة صحيحة إلى حد كبير: عمليات ضرب المصفوفات (GEMM) وحسابات الانتباه تستهلك الجزء الأكبر من القدرة الحسابية.

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

هذا ما يُعرف بـ "ضيق عرض النطاق الترددي للذاكرة": مثل طاهٍ ماهر جدًا، لكنه يجب أن ينقل المكونات من مستودع بعيد كلما أراد إعداد طبق، ثم يعيدها بعد الانتهاء، بدلاً من وضعها على الطاولة بجانبه. حتى لو كانت يديه سريعتين جدًا، فإن الوقت الذي يُهدره في الانتظار للنقل هو وقت مهدور فعليًا.

الأكثر سوءًا، مع جعل تنسيقات الدقة المنخفضة مثل FP8 وFP4 من NVIDIA عمليات المصفوفة أسرع، فإن التكلفة النسبية لهذه العمليات "النقل" تزداد: فقد تسارعت ضرب المصفوفات، لكن تكلفة نقل التنسورات إلى الداخل والخارج لم تقل بنفس المعدل.

يحتوي البحث على مجموعة من البيانات البديهية: عند تدريب نموذج بـ 1 مليار معلمة على H100 باستخدام TorchTitan، تشغل العمليات غير ضرب المصفوفات نسبة كبيرة من وقت التشغيل الشامل، وستزداد هذه النسبة بشكل أكبر مع إدخال دقة FP8.

الإطارات البرمجية الحالية عاجزة تمامًا عن التعامل مع هذا. يقوم PyTorch بتمثيل حسابات Transformer كسلسلة من العوامل، مع حدود واضحة بين العوامل. هذه الحدود مواتية جدًا للتمايز التلقائي (autograd)، لكنها تمنع بالضبط تحسينات الدمج عبر العوامل: كل حدود عامل غالبًا ما تمثل كتابة غير ضرورية إلى الذاكرة.

CODA: الكنز مخفي في "الخاتمة"

انطلاقًا من ملاحظة بسيطة.

على GPU، يُقسَّم نواة ضرب المصفوفات عالية الأداء (GEMM) هيكلياً إلى جزأين: الحلقة الرئيسية (mainloop) التي تُنفِّذ عمليات الضرب والجمع الأساسية للمصفوفات المجزأة، ونهاية العملية (epilogue) التي تقوم ببعض المعالجات الختامية قبل كتابة النتائج إلى ذاكرة العرض، مثل إضافة التحيز، وتحويل النوع، والتقسيم البسيط.

إن معنى وجود الخاتمة هو أن مخرجات ضرب المصفوفات لا تزال "حية" في المسجلات المدمجة، ولم تُكتب بعد إلى ذاكرة الفيديو العالمية. هذا نافذة ذهبية قصيرة: إذا تم إجراء المزيد من الحسابات في هذه اللحظة، يمكن تجنب بالكامل عملية الكتابة ثم القراءة المتكررة لذاكرة الفيديو.

البصيرة الأساسية لـ CODA هي أن العديد من العمليات التي تستهلك الذاكرة في Transformer يمكن إعادة معلمتها جبريًا وتنفيذها داخل نافذة "الخاتمة" هذه.

هذا يتطلب بعض المهارات الرياضية. على سبيل المثال، النمط الأكثر شيوعًا GEMM-RMSNorm-GEMM: نتيجة ضرب مصفوفة، تليها إضافة متبقية، ثم تطبيع RMS، ثم ضرب مصفوفة آخر. في الممارسة التقليدية، يتم تنفيذ هذه العمليات الثلاث كعوامل منفصلة بالتسلسل، مع كتابة النتائج الوسيطة مرتين إلى ذاكرة GPU.

اكتشف فريق CODA أن عامل التحجيم الصفّي r في تطبيع RMS، نظرًا لأنه قياس مشترك لكل صف، فإنه يحقق خاصية التبادلية مع عملية ضرب المصفوفات اللاحقة: يمكن تأجيل تطبيق r من "قبل GEMM الثاني" إلى "نهاية GEMM الثاني". بعد التأجيل، لا يتطلب نهاية GEMM الأول سوى حساب "جذر متوسط المربعات الجزئي" (partial RMS)، والذي يُدمج بواسطة نواة تجميع خفيفة جدًا، بينما اختفت حسابات RMSNorm الكاملة.

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

خمسة "قطع بناء" و一套 "لغة ليجو"

CODA ليست نواة دمج محددة، بل مجموعة من التجريدات البرمجية.

إنه يُثبّت الحلقة الرئيسية المُحسّنة من قبل الخبراء لـ GEMM، ثم يُعرّض في نهاية العملية خمسة أنواع من الوحدات الأساسية القابلة للتركيب:

• التحولات عنصرًا بعنصر (الجمع المتبقي، دالة التنشيط، RoPE)
• تحميل وتخزين المتجهات (إذاعة أوزان RMSNorm)
• تحميل وتخزين المصفوفات المجزأة (حفظ التنشيطات الوسيطة للاستخدام في الانتشار العكسي)
• تقسيم التقليل (الجذر التربيعي المحلي المتوسط، تقسيم log-sum-exp)
• تغيرات ذات حالة (الإحصائيات المطلوبة للتوحيد عبر الإنترنت: الحد الأقصى وsum-exp)

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

الأمر الأكثر إثارة هو مرونة هذا الإجراء المجرد تجاه "من يكتب الكود". قام البحث بتقييم نمطين من التحقيق: أحدهما مكتوب بواسطة مبرمجين بشريين، والآخر مُنشأ بواسطة Claude Code — حيث يتم إعطاء وصفات CODA الأولية، وأمثلة متعددة، وسجلات التنفيذ، ويقوم الذكاء الاصطناعي بإنشاء معظم كود النواة، مع إشراف بشري خفيف.

أظهر كلا النمطين أداءً عاليًا. قال تري داو في تغريدة: "يمكن لـ LLM والمبتدئين كتابة نواة بسرعة الضوء"، وهو ما يعكس نتائج التجارب في الورقة البحثية على أرض الواقع.

نتائج التجربة

اختارت CODA منافسين صعبين في الاختبارات المرجعية: cuBLAS بالإضافة إلى torch.compile، وLiger Kernel وFlashInfer المُحسّنين خصيصًا لـ LLM.

قُيّمت كل نواة بتنفيذين: CODA (LLM) الذي تم إنشاؤه بواسطة Claude Code، حيث قدم الباحثون وصفًا لل primitives، وعدة أمثلة، وسجلًا مستمرًا للنصائح التنفيذية، وأكمل الذكاء الاصطناعي الكود الأساسي مع إشراف بشري خفيف؛ CODA (Human) الذي كتبه مبرمجون بشريون بشكل مستقل، باستخدام نفس فكرة إعادة المعلمة من المستوى العالي، لكن دون الاعتماد على مجموعة primitives الخاصة بـ CODA. وتم مقارنة نتائج المجموعتين مع مكتبات التحسين مثل cuBLAS + torch.compile وLiger Kernel وFlashInfer.

على مستوى المشغل الواحد، وعلى سبيل المثال النمط النموذجي GEMM-RMSNorm-GEMM، تفوق CODA على الأساسيات cuBLAS + PyTorch في أبعاد الخفاء الثلاثة للموديلات 1B و7B و70B. كما أظهرت مجموعات النهاية مثل SwiGLU وRoPE وCross-Entropy أداءً مشابهًا.

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

تُظهر عوائد التدرج العكسي تحسنًا ملحوظًا: يمكن أن يحقق نواة التدرج العكسي GEMM-Residual-PartialRMS-GEMM تسريعًا يتراوح بين 1.6 و1.8 مرة مقارنة بالخط الأساسي، كما يحقق SwiGLU العكسي تحسنًا يتراوح بين 1.4 و1.6 مرة. في هذا المجال، لا يزال الفرق بين نماذج LLM والتنفيذ اليدوي صغيرًا جدًا. وهذا لا يثير الدهشة: فالتدرج العكسي يتضمن بشكل طبيعي مزيدًا من الوصول إلى المصفوفات الوسيطة، مما يزيد من فائدة دمج النهايات؛ كما أن التصميم الأساسي لـ CODA واضح بما يكفي ليسمح للنماذج الذكية بإكمال التجميع بشكل صحيح.

في اختبارات النهاية إلى النهاية لطبقة Transformer الكاملة، يكون تسريع التقدم في CODA بين 5% و20% عبر المقاييس المختلفة، مع تحسن أكثر وضوحًا في أحجام النماذج الأكبر (المقابلة لأبعاد الإخفاء بحجم 70B).

من حيث دقة القيم، يُعيد تعيين CODA وقت تطبيق عامل التطبيع RMS، لكن التجارب تُظهر أن خطأها العددي مماثل لتنفيذ PyTorch المرجعي، وفي بعض التكوينات يكون الخطأ أقل — وذلك بفضل وجود مُجمّعات ذات دقة أعلى في الحلقة الأساسية GEMM.

ما يمكن لـ CODA فعله: قائمة مرجعية سريعة لتحديد حدود قدرات CODA قبل الانتقال إلى المنظور الأكبر.

• النطاق: تقريبًا جميع الحسابات باستثناء الانتباه وترميز الكلمات في التقدم والbackward لـ Transformer القياسي (مثل بنية LLaMA)، بما في ذلك RMSNorm، جمع المتبقي، تنشيط SwiGLU، ترميز الموضع الدوراني RoPE، خسارة الإنتروبيا المتقاطعة، وحسابات التدرج العكسي للعمليات المذكورة أعلاه.
• تأثير التسريع: في الأبعاد المخفية ضمن نطاق 1B إلى 70B، يُظهر المستوى الفردي للعامل تحسينات متفاوتة مقارنة بقاعدة cuBLAS + torch.compile، حيث يكون الربح في التدرج العكسي الأكثر وضوحًا (يمكن أن يصل بعض النوى إلى أكثر من 1.6 مرة)؛ ويصل التسريع الشامل للتقدم الأمامي لطبقة Transformer الكاملة إلى حوالي 5٪ إلى 20٪، مع تحسن أكثر وضوحًا في حجم النموذج الأكبر.
• يمكن لأي شخص استخدام: CODA، المُنفَّذ على أساس CuTeDSL (لغة برمجة مُخصصة لـ NVIDIA CUTLASS)، والذي يدعم طريقتين لكتابة النوى: بواسطة المبرمجين البشريين ونماذج الذكاء الاصطناعي، وكلتا الطريقتين تحققان أداءً عاليًا.
• القيود الحالية: يُدعم حاليًا فقط سيناريو GPU واحد، دون تضمين التدريب الموزع؛ يركز إعادة المعلمة على بنية Transformer القياسية، وقابلية تطبيقها على هياكل أخرى ما زالت قيد التحقق.

خاتمة

لا يُعد CODA عملًا منعزلًا. إنه تنفيذ ملموس لفئة من الأفكار: على وحدات معالجة الرسومات، فإن مساحة التحسين الحقيقية غالبًا لا تكمن في "ما يتم حسابه"، بل في "كيفية نقله".

يسمح FlashAttention بدمج حسابات الانتباه داخل الذاكرة على الشريحة، بينما يحاول CODA دمج وظائف التطبيع والتفعيل أيضًا داخلها. يقلل Triton من عتبة كتابة النوى المخصصة، وتوسع ThunderKittens وTileLang وغيرها في استكشاف هذا المجال على مستويات مختلفة. تشير هذه الجهود معًا إلى اتجاه واحد موحد: دمج سهولة التعبير في رسم خوارزميات PyTorch مع كفاءة التنفيذ القريبة من CUDA المكتوبة يدويًا ضمن إطار قابل للبرمجة واحد.

الجملة الأخيرة من تغريدة Tri Dao تستحق التأمل مرة أخرى: "يمكن للنماذج اللغوية الكبيرة والمبتدئين كتابة نوى بسرعة الضوء لجميع عمليات Transformer." وراء هذا منطق أعمق: عندما يتم تصميم التجريدات البرمجية بشكل جيد بما يكفي، يمكن للنماذج الذكية نفسها المشاركة في تحسين بنية التدريب الخاصة بها. هذه الدورة هي ما يجعل CODA الأكثر إثارة للتأمل.

من هذا المنظور، قد يكون اسم "CODA" ذا معنى أعمق. في الموسيقى الكلاسيكية، تُشير "Coda" إلى الجزء الختامي الذي يُنهي القطعة الموسيقية. هنا، إنها "الختام" لنموذج GEMM — وكتابة هذا الختام قد تكون بالضبط الفصل التالي المهم في تحسين كفاءة نظام تدريب Transformer.