SkewAdam Mixture-of-Experts کی ٹریننگ میموری کو 60% سے زیادہ کم کر دیتا ہے اور درستگی (accuracy) میں نمایاں اضافہ فراہم کرتا ہے، جس سے ڈویلپرز 6.78 بلین پیرامیٹر والے MoE ماڈل کو ایک ہی 40 GB GPU پر چلا سکتے ہیں۔

MoE ٹریننگ میموری کی حد (memory wall) کیوں پہنچ جاتی ہے

Mixture-of-Experts آرکیٹیکچرز ایک ڈینس بیک بون (dense backbone) کو برقرار رکھتے ہوئے ہر ان پٹ کو "ایکسپرٹ" (expert) سب نیٹ ورکس کے ایک چھوٹے سے حصے کے ذریعے گزارتے ہیں۔ اس کا فائدہ یہ ہے کہ ماڈل کمپیوٹیشن میں متناسب اضافے کے بغیر اربوں پیرامیٹرز تک پھیل سکتا ہے۔ عملی طور پر، آپٹیمائزر—خاص طور پر AdamW ورژن جسے زیادہ تر ٹیمیں استعمال کرتی ہیں—GPU RAM کا بڑا حصہ استعمال کر لیتا ہے۔ 6.78 B-پیرامیٹر ماڈل کے لیے، صرف آپٹیمائزر کے مومینٹم (momentum) اور ویریئنس (variance) ٹینسرز کو ہی تقریباً 50 GB کی ضرورت ہوتی ہے۔ اگر اس میں ایکٹیویشنز (activations) اور ماڈل ویٹس (weights) شامل کر دیے جائیں، تو عارضی میموری (peak memory) 81.4 GB تک پہنچ جاتی ہے، جس کی وجہ سے ملٹی-GPU سیٹ اپ یا سست ٹریننگ شیڈول پر مجبور ہونا پڑتا ہے۔

SkewAdam کیا مختلف کرتا ہے

SkewAdam کوئی نیا لرننگ رول ایجاد نہیں کرتا۔ یہ صرف اس جگہ کو تبدیل کرتا ہے جہاں Adam کے مومینٹس (moments) محفوظ ہوتے ہیں:

  • ڈینس بیک بون مکمل پریسیژن مومینٹم کو برقرار رکھتی ہے، تاکہ ڈینس لیئرز کے لیے ضروری ہموار اپ ڈیٹس (smooth updates) محفوظ رہیں۔
  • ایکسپرٹ بینک ویریئنس کا ایک فیکٹرڈ اپروکسیمیشن (factored approximation) اسٹور کرتا ہے، جس سے ہر ایکسپرٹ کے لیے محفوظ کیے جانے والے ڈیٹا میں کمی آتی ہے۔
  • روٹر (router)، جو فیصلہ کرتا ہے کہ کن ایکسپرٹس کو ایکٹیویٹ کرنا ہے، دوسرے مومنٹ (second-moment) کا درست تخمینہ برقرار رکھتا ہے۔

ان تینوں اجزاء کو الگ الگ "ٹیرز" (tiers) کے طور پر استعمال کرتے ہوئے، آپٹیمائزر وہاں غیر ضروری پریسیژن کو ختم کر دیتا ہے جہاں اس کی اہمیت کم ہوتی ہے اور اسے وہاں برقرار رکھتا ہے جہاں اس کی سب سے زیادہ ضرورت ہوتی ہے۔

پیمائش کے قابل اثرات

اسی 6.78 B-پیرامیٹر MoE ماڈل کو SkewAdam کے ساتھ چلانے سے درج ذیل نتائج حاصل ہوتے ہیں:

  • عارضی GPU میموری: 31.3 GB (81.4 GB سے کم)
  • آپٹیمائزر اسٹیٹ کا استعمال: 1.29 GB (50 GB سے کم)

کم شدہ اسٹیٹ (reduced state) ایک ہی 40 GB ایکسلریٹر کے اندر آسانی سے سما جاتی ہے۔

صرف بچت نہیں، بلکہ درستگی میں اضافہ

میموری میں کٹوتی اکثر ماڈل کے معیار کو نقصان پہنچاتی ہے، لیکن SkewAdam پرپلیکسیٹی (perplexity)—جو کہ لینگویج ماڈل کا ایک معیاری میٹرک ہے—کو 126.8 (AdamW) سے بہتر بنا کر 108.4 تک لے آتا ہے۔ یہ اسی ٹاسک پر Muon (120.2) اور Lion (393.7) سے بھی بہتر کارکردگی دکھاتا ہے۔ مصنفین کا کہنا ہے کہ یہ بہتری تینوں ٹیرز میں مومینٹم کو برقرار رکھنے کی وجہ سے آتی ہے؛ وہ طریقے جو مومینٹم کو مکمل طور پر ختم کر دیتے ہیں، جیسے کہ Adafactor، پیچھے رہ جاتے ہیں۔

کھلے سوالات

SkewAdam کے نتائج 6.78 B-پیرامیٹر ماڈل پر دکھائے گئے ہیں۔ یہ اب بھی غیر واضح ہے کہ کیا یہی میموری-اسٹیٹ تناسب اس سے کئی گنا بڑے ماڈلز یا لینگویج ماڈلنگ سے ہٹ کر دیگر کاموں کے لیے بھی برقرار رہے گا۔

کن چیزوں پر نظر رکھنی چاہیے

  • بڑے MoE کنفیگریشنز (دہائیوں اربوں پیرامیٹرز) پر بینچ مارکس۔
  • اوپن سورس فریم ورکس کی جانب سے اسے اپنانا اور مقبول ٹریننگ اسکرپٹس میں اس کا شامل ہونا۔
  • پوشیدہ توازن (trade-offs) پر کمیونٹی کا فیڈ بیک، جیسے کہ مختلف لرننگ ریٹ شیڈول کے تحت کنورجنس کی رفتار یا استحکام۔

ماخذ: ڈویلپر کی پوسٹ جس میں آپٹیمائزر کے ڈیزائن اور تجرباتی نتائج کی تفصیل دی گئی ہے۔