Ring Attention

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

آلية عمل الانتباه الحلقي#

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

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

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

تطبع Matches full attention: True: لم يحتفظ كل جهاز في أي وقت إلا بكتلة درجات بحجم 8 × 8، وقد أنتجت الأجهزة مجتمعةً انتباهًا على المصفوفة الكاملة ذات الأبعاد 32 × 32. كما تطرح نوى الإنتاج قيمة عظمى تراكمية قبل رفع الدرجات أُسًّا للحفاظ على استقرار الأرقام. وعند استخدام قناع سببي، تتخطى الأجهزة التي تحتفظ بكتل الاستعلامات المبكرة معظم كتل المفاتيح وتبقى خاملة، وتعالج الانتباه المخطّط ذلك بمنح كل جهاز رموزًا موزعة بانتظام على التسلسل، ما يرفع إنتاجية التدريب من البداية إلى النهاية حتى 1.45 مرة عند 256K رمزًا.

التطبيقات#

يكتسب الانتباه الحلقي أهميته كلما كان أحد المدخلات أطول مما يستطيع جهاز واحد التعامل معه، والفيديو أوضح مثال على ذلك. واستخدمت المجموعة نفسها من بيركلي هذا الأسلوب لتدريب نموذج العالم الكبير (LWM)، وهو عائلة من نماذج الرؤية واللغة (VLMs) ذات 7 مليارات معامل، ازداد سياقها تدريجيًا من 4K إلى مليون رمز من النصوص والفيديو المتداخلين. يشفّر LWM كل إطار على هيئة 256 رمزًا، لذا يبلغ طول ساعة من الفيديو مأخوذة بمعدل إطار واحد في الثانية 921,600 رمز، وقد أجاب النموذج عن أسئلة حول مقاطع مجمّعة من YouTube تمتد ساعة كاملة، وهي مهمة تتطلب فهم الفيديو على مدى زمني طويل. ويواجه تدريب نماذج اللغة الكبيرة (LLMs) التي تقرأ قواعد شيفرة كاملة أو مجموعات مستندات في مطالبة واحدة الحد نفسه. وتستبدل واجهة context_parallel() في PyTorch انتباه جداء النقاط ذي القياس بالانتباه الحلقي ضمن سياقها، واستخدم فريق TorchTitan هذا الأسلوب لتدريب Llama 3 8B على تسلسلات بطول مليون رمز عبر 32 وحدة H100 GPU.

الانتباه الحلقي مقارنةً بأساليب التوازي الأخرى#

• توازي السياق: الاستراتيجية العامة لتقسيم تسلسل واحد على عدة أجهزة. الانتباه الحلقي هو الخوارزمية الكامنة وراء تنفيذ PyTorch، أما DeepSpeed Ulysses فهو بديل يستخدم تبادلات شاملة بين الجميع لمنح كل جهاز جميع الرموز لمجموعة فرعية من رؤوس الانتباه.
• توازي الموترات وتوازي خطوط الأنابيب: تقسيم أوزان النموذج داخل الطبقات أو بينها، ما يفيد عندما تتجاوز المعاملات سعة ذاكرة جهاز واحد.
• الانتباه السريع: نواة تعمل على جهاز واحد. ويشغّل الانتباه الحلقي نواة كتلية من هذا النوع على كل جهاز لمعالجة كتلته المحلية.

يقسّم توازي البيانات الدفعة بدلًا من التسلسل، وهو النمط الذي يستخدمه Ultralytics YOLO26 في التدريب باستخدام وحدات GPU متعددة: يؤدي تمرير device=[0, 1] إلى تشغيل DistributedDataParallel في PyTorch، الذي يحتفظ بنسخة كاملة من النموذج على كل GPU، ويمنح كل GPU حصة متساوية من كل دفعة من الصور، ويحسب متوسط التدرجات بعد كل خطوة.