معماری FFN Fusion
معماری FFN Fusion

محققان هوش مصنوعی انویدیا تکنیک بهینه‌سازی نوین FFN Fusion را معرفی می‌کنند: نمایشی از موازی‌سازی مؤثر محاسبات متوالی در مدل‌های زبان بزرگ (LLM)

مدل‌های زبان بزرگ (LLMها) در حوزه‌های مختلفی حیاتی شده‌اند و امکان ایجاد برنامه‌های کاربردی با عملکرد بالا مانند تولید زبان طبیعی، تحقیقات علمی و دستیاران مکالمه‌ای را فراهم کرده‌اند. در زیربنای این پیشرفت‌ها، معماری ترانسفورمر قرار دارد که در آن لایه‌های متناوب مکانیزم‌های توجه و شبکه‌های پیشخور (FFN) به صورت متوالی ورودی توکن‌شده را پردازش می‌کنند. با این حال، با افزایش اندازه و پیچیدگی، بار محاسباتی مورد نیاز برای استنتاج به طور قابل توجهی افزایش می‌یابد و یک گلوگاه کارایی ایجاد می‌کند. استنتاج کارآمد اکنون یک نگرانی اساسی است و بسیاری از گروه‌های تحقیقاتی بر روی استراتژی‌هایی تمرکز کرده‌اند که می‌توانند تأخیر را کاهش دهند، توان عملیاتی را افزایش دهند و هزینه‌های محاسباتی را کاهش دهند، در حالی که عملکرد مدل را حفظ یا بهبود می‌بخشند.

در مرکز این مشکل کارایی، ساختار ذاتاً متوالی ترانسفورمرها قرار دارد. خروجی هر لایه به لایه بعدی وارد می‌شود و نیازمند ترتیب و همگام‌سازی دقیق است که به ویژه در مقیاس بزرگ مشکل‌ساز است. با گسترش اندازه‌های مدل، هزینه محاسبات متوالی و ارتباط بین پردازنده‌های گرافیکی (GPU) افزایش می‌یابد که منجر به کاهش کارایی و افزایش هزینه استقرار می‌شود. این چالش در سناریوهایی که نیاز به تولید سریع چند توکن دارند، مانند دستیاران هوش مصنوعی بلادرنگ، تشدید می‌شود. کاهش این بار متوالی ضمن حفظ قابلیت‌های مدل، یک مانع فنی کلیدی است. باز کردن قفل استراتژی‌های موازی‌سازی جدید که دقت را حفظ می‌کنند اما به طور قابل توجهی عمق محاسبات را کاهش می‌دهند، برای گسترش دسترسی و مقیاس‌پذیری LLMها ضروری است.

چندین تکنیک برای بهبود کارایی پدیدار شده‌اند. کوانتیزاسیون (Quantization) دقت نمایش‌های عددی را کاهش می‌دهد تا نیاز به حافظه و محاسبات را به حداقل برساند، اگرچه اغلب با خطر از دست دادن دقت همراه است، به ویژه در عرض بیت پایین. هرس کردن (Pruning) پارامترهای اضافی را حذف کرده و مدل‌ها را ساده می‌کند اما بدون دقت کافی، به طور بالقوه به دقت آسیب می‌زند. مدل‌های ترکیب کارشناسان (MoE یا Mixture-of-Experts) فقط زیرمجموعه‌ای از پارامترها را برای هر ورودی فعال می‌کنند و آن‌ها را برای حجم کاری خاص بسیار کارآمد می‌سازند. با این حال، ممکن است در اندازه‌های دسته متوسط به دلیل استفاده کم از سخت‌افزار، عملکرد ضعیفی داشته باشند. اگرچه این استراتژی‌ها ارزشمند هستند، اما دارای معاوضه‌هایی هستند که کاربرد جهانی آن‌ها را محدود می‌کند. در نتیجه، این حوزه به دنبال روش‌هایی است که بهبودهای گسترده کارایی را با سازش‌های کمتر ارائه دهند، به ویژه برای معماری‌های متراکم که آموزش، استقرار و نگهداری آن‌ها ساده‌تر است.

محققان در انویدیا یک تکنیک بهینه‌سازی معماری جدید به نام FFN Fusion (همجوشی FFN) را معرفی کردند که با شناسایی توالی‌های FFN که می‌توانند به صورت موازی اجرا شوند، گلوگاه متوالی در ترانسفورمرها را برطرف می‌کند. این رویکرد از این مشاهده ناشی شد که وقتی لایه‌های توجه با استفاده از ابزار Puzzle حذف می‌شوند، مدل‌ها اغلب توالی‌های طولانی از FFNهای متوالی را حفظ می‌کنند. این توالی‌ها وابستگی متقابل کمی از خود نشان می‌دهند و بنابراین می‌توانند به طور همزمان پردازش شوند. با تجزیه و تحلیل ساختار LLMهایی مانند Llama-3.1-405B-Instruct، محققان مدل جدیدی به نام Ultra-253B-Base را با هرس کردن و بازسازی مدل پایه از طریق FFN Fusion ایجاد کردند. این روش منجر به مدلی به طور قابل توجهی کارآمدتر می‌شود که عملکرد رقابتی خود را حفظ می‌کند.

FFN Fusion چندین لایه FFN متوالی را در یک FFN واحد و گسترده‌تر ادغام می‌کند. این فرآیند مبتنی بر هم‌ارزی ریاضی است: با الحاق وزن‌های چندین FFN، می‌توان یک ماژول واحد تولید کرد که مانند مجموع لایه‌های اصلی رفتار می‌کند اما می‌تواند به صورت موازی محاسبه شود. به عنوان مثال، اگر سه FFN به صورت متوالی روی هم چیده شوند که هر کدام به خروجی لایه قبلی وابسته است، همجوشی آن‌ها این وابستگی‌ها را با اطمینان از اینکه هر سه بر روی ورودی یکسان عمل می‌کنند و خروجی‌هایشان جمع می‌شوند، از بین می‌برد. مبنای نظری این روش نشان می‌دهد که FFN همجوشی شده همان ظرفیت نمایشی را حفظ می‌کند. محققان با استفاده از فاصله کسینوسی بین خروجی‌های FFN، تجزیه و تحلیل وابستگی را برای شناسایی مناطقی با وابستگی متقابل کم انجام دادند. این مناطق برای همجوشی بهینه تلقی شدند، زیرا تغییر حداقلی در جهت توکن بین لایه‌ها نشان‌دهنده امکان‌پذیری پردازش موازی بود.

اعمال FFN Fusion بر روی مدل Llama-405B منجر به Ultra-253B-Base شد که دستاوردهای قابل توجهی در سرعت و کارایی منابع به همراه داشت. به طور خاص، مدل جدید بهبودی ۱.۷۱ برابری در تأخیر استنتاج و کاهش ۳۵ برابری در هزینه محاسباتی به ازای هر توکن در اندازه دسته ۳۲ را به دست آورد. این کارایی به قیمت از دست دادن قابلیت‌ها تمام نشد. Ultra-253B-Base در آزمون MMLU امتیاز ۸۵.۱۷٪، در MMLU-Pro امتیاز ۷۲.۲۵٪، در Arena Hard امتیاز ۸۴.۹۲٪، در HumanEval امتیاز ۸۶.۵۸٪ و در MT-Bench امتیاز ۹.۱۹ را کسب کرد. این نتایج اغلب با مدل اصلی ۴۰۵ میلیارد پارامتری مطابقت داشت یا از آن فراتر رفت، حتی اگر Ultra-253B-Base فقط ۲۵۳ میلیارد پارامتر داشت. استفاده از حافظه نیز با کاهش ۲ برابری در نیازهای کش کلید-مقدار (kv-cache) بهبود یافت. فرآیند آموزش شامل تقطیر (distilling) ۵۴ میلیارد توکن در پنجره زمینه ۸ هزار، و سپس تنظیم دقیق مرحله‌ای (staged fine-tuning) در زمینه‌های ۱۶ هزار، ۳۲ هزار و ۱۲۸ هزار بود. این مراحل تضمین کردند که مدل همجوشی شده دقت بالایی را حفظ کرده و از اندازه کاهش یافته بهره‌مند می‌شود.

این تحقیق نشان می‌دهد که چگونه طراحی مجدد معماری متفکرانه می‌تواند دستاوردهای قابل توجهی در کارایی ایجاد کند. محققان نشان دادند که لایه‌های FFN در معماری‌های ترانسفورمر اغلب مستقل‌تر از آنچه قبلاً تصور می‌شد هستند. روش آن‌ها برای کمی‌سازی وابستگی بین لایه‌ها و تبدیل ساختارهای مدل، امکان کاربرد گسترده‌تر در مدل‌هایی با اندازه‌های مختلف را فراهم کرد. این تکنیک همچنین بر روی یک مدل ۷۰ میلیارد پارامتری تأیید شد و قابلیت تعمیم آن را اثبات کرد. آزمایش‌های بیشتر نشان داد که در حالی که لایه‌های FFN اغلب می‌توانند با حداقل تأثیر همجوشی شوند، موازی‌سازی کامل بلوک، از جمله توجه، به دلیل وابستگی‌های متقابل قوی‌تر، تخریب عملکرد بیشتری را به همراه دارد.

چند نکته کلیدی از تحقیق در مورد FFN Fusion:

  • تکنیک FFN Fusion با موازی‌سازی لایه‌های FFN با وابستگی کم، محاسبات متوالی در ترانسفورمرها را کاهش می‌دهد.
  • همجوشی با جایگزینی توالی‌های FFN با یک FFN واحد و گسترده‌تر با استفاده از وزن‌های الحاق شده به دست می‌آید.
  • مدل Ultra-253B-Base، برگرفته از Llama-3.1-405B، به استنتاج ۱.۷۱ برابر سریع‌تر و هزینه ۳۵ برابر کمتر به ازای هر توکن دست می‌یابد.
  • نتایج بنچمارک‌ها عبارتند از: ۸۵.۱۷٪ (MMLU)، ۷۲.۲۵٪ (MMLU-Pro)، ۸۶.۵۸٪ (HumanEval)، ۸۴.۹۲٪ (Arena Hard) و ۹.۱۹ (MT-Bench).
  • استفاده از حافظه به دلیل بهینه‌سازی کش کلید-مقدار (kv-cache) به نصف کاهش می‌یابد.
  • FFN Fusion در مقیاس‌های مدل بزرگتر مؤثرتر است و با تکنیک‌هایی مانند هرس کردن و کوانتیزاسیون به خوبی کار می‌کند.
  • موازی‌سازی کامل بلوک ترانسفورمر پتانسیل دارد اما به دلیل وابستگی‌های متقابل قوی‌تر، نیازمند تحقیقات بیشتر است.
  • یک روش سیستماتیک با استفاده از فاصله کسینوسی به شناسایی اینکه کدام توالی‌های FFN برای همجوشی ایمن هستند کمک می‌کند.
  • این تکنیک در اندازه‌های مختلف مدل، از جمله ۴۹ میلیارد، ۷۰ میلیارد و ۲۵۳ میلیارد پارامتر تأیید شده است.
  • این رویکرد پایه و اساس طراحی‌های LLM سازگارتر با موازی‌سازی و کارآمدتر از نظر سخت‌افزاری را بنا می‌نهد.

برای جزئیات بیشتر، به مقاله پژوهشی مراجعه کنید. تمام اعتبار این تحقیق متعلق به پژوهشگران این پروژه است.