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

يتطلب تدريب نموذج ذي 4 مليارات من المعلمات بدقة مختلطة مع حجم دفعة يبلغ 16 نحو 85 غيغابايت من ذاكرة وحدة معالجة الرسوميات. يساعدك فهم ما يشغل ذاكرة وحدة معالجة الرسوميات والعمليات الحسابية التي تحدث أثناء التدريب على تحديد الأماكن التي يمكنك فيها تقليل استهلاك الذاكرة.

┌─────────────────────────── TENSORS ──────────────────────────────────┐
│                                                                      │
│  MODEL WEIGHTS       ████████████████████████  6 bytes/param         │
│  (mixed precision)   ├── fp16 copy (2B) — forward/backward pass      │
│                      └── fp32 copy (4B) — stable weight updates      │
│                                                                      │
│  OPTIMIZER STATES    ████████████████████████████████  8 bytes/param │
│  (Adam)              ├── fp32 momentum  (4B)                         │
│                      └── fp32 variance  (4B)                         │
│                      ↳ quantized Adam (bitsandbytes) → 2 bytes/param │
│                                                                      │
│  GRADIENTS           ████████████████  4 bytes/param  (fp32)         │
│                      └── computed in backward pass, per parameter    │
│                                                                      │
│  ACTIVATIONS         ████ varies — batch × seq_len × depth × hidden  │
│  (forward cache)     └── cached for backward; can OOM even if model  │
│                          fits (long seqs / large batches)            │
│                          bf16/fp16 if using mixed precision          │
│                                                                      │
│  TEMPORARY TENSORS   ▓ short-lived (softmax, matmul scratch)         │
│                      └── peak spikes can cause OOM                   │
│                                                                      │
│  OTHER OVERHEAD      ▒ beam search caches, large embedding tables    │
└──────────────────────────────────────────────────────────────────────┘

تحتفظ ذاكرة وحدة معالجة الرسوميات (GPU) بفئتين من العناصر: الموترات المخزنة، والعمليات التي تعالجها.

المصفوفات

يتطلب التدريب ويولّد العديد من أنواع المصفوفات (tensors) التي تحتاج إلى تخزين.

• تُخزَّن أوزان النموذج على وحدة معالجة الرسوميات (GPU). في التدريب بدقة مختلطة، تكون هناك حاجة إلى نسختين من الأوزان - نسخة واحدة بدقة fp16 للتمرير الأمامي/الخلفي، ونسخة أخرى بدقة fp32 كـ"النسخة الرئيسية" لتحديثات الأوزان المستقرة. وهذا يعادل 6 بايت لكل معامل.

• تخزّن حالات المُحسِّنات مثل Adam تينسورين إضافيين لكل معامل، وهما الزخم والتباين، وكلاهما بصيغة fp32. وهذا يعني 8 بايتات إضافية لكل معامل.

  يمكنك استخدام مُحسِّن مختلف، مثل نسخة مُكمَّمة من Adam من bitsandbytes، لضغطه إلى بايتَين لكل معامل.

• تُحسب متجهات التدرج لكل معامل في المسار العكسي. وتُحفظ هذه المتجهات بدقة fp32، أي ما يعادل 4 بايت لكل معامل.

• تُحسب التفعيلات الأمامية (Forward activations) في المرور الأمامي وتُخزَّن مؤقتًا (cached) ليستخدمها المرور الخلفي في حساب التدرجات. وتختلف أحجام هذه التفعيلات تبعًا لحجم الدفعة (batch size) وطول التسلسل وعمق النموذج وحجم الطبقة المخفية. ولهذا السبب، قد يستنفد حجم الدفعة أو طول التسلسل ذاكرة وحدة معالجة الرسوميات (GPU) حتى لو كان النموذج نفسه يتسع في الذاكرة.

• تُنشأ المتجهات المؤقتة بواسطة عمليات مثل softmax وعمليات ضرب المصفوفات، وتُحرَّر بعد كل عملية. إذا كان الذروة في عملية واحدة مكثفة جداً، فقد تُسبب ارتفاعاً مؤقتاً في استهلاك الذاكرة يؤدي إلى نفادها.

• تضيف بعض الميزات تكاليف إضافية خاصة بها. فعلى سبيل المثال، يحافظ البحث بالحزمة (beam search) على مخرجات متعددة، كما يمكن أن تكون جداول التضمين (embedding tables) للمفردات الكبيرة ضخمة جداً.

عمليات التشغيل

توجد ثلاثة أنواع رئيسية من عمليات التدريب.

• تُعد عمليات ضرب المصفوفات (matmuls) النوع الرئيسي من العمليات، حيث تغطي الطبقات الخطية، وإسقاطات QKV، وإسقاطات مخرجات الانتباه، وطبقات FFN. أما المستهلكات الرئيسية للذاكرة فهي مصفوفات درجات الانتباه التي تنمو مع مربع طول التسلسل. ولهذا السبب تكون التسلسلات الطويلة مكلفة للغاية، وتوجد واجهات خلفية مختلفة للانتباه لتجنب إنشاء المصفوفة الكاملة في الذاكرة.

• تقرأ عمليات الاختزال، مثل softmax وlayer norm، موصلًا كاملًا، ثم تحسب إحصائية عبر بُعد، ثم تقرأ الموصل مرة أخرى لتطبيقها. وهذا يتطلب الوصول إلى الذاكرة عدة مرات لكل عملية.

• تُطبَّق عمليات العنصر-ب-العنصر، مثل دوال التنشيط وdropout، دالة على كل عنصر بشكل مستقل، ويكون استهلاكها للذاكرة متناسبًا مع حجم المصفوفة.

الخطوات التالية

• اطّلع على دليل تراكم التدرّجات لمعرفة كيفية محاكاة التدريب على حجم دفعة فعّال أكبر دون نفاد ذاكرة وحدة معالجة الرسوميات.
• اطّلع على دليل Gradient checkpointing لمعرفة كيفية تقليل استهلاك الذاكرة من خلال تخزين بعض التفعيلات الوسيطة فقط.
• اطّلع على دليل التدريب بدقة مختلطة لمعرفة كيفية استخدام أنواع بيانات بدقة أقل لتقليل الذاكرة وتسريع التدريب.
• اطّلع على دليل Kernels لمعرفة كيفية تسريع التدريب باستخدام نواة مدمجة مخصصة.