فلش‌اتنشن چگونه سد حافظه پردازنده‌های گرافیکی را می‌شکند تا هوش مصنوعی با حافظه طولانی خلق شود

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

با طولانی‌تر شدن این توالی، حافظه مورد نیاز برای ذخیره این مقایسه‌ها به صورت توانی (مربعی) افزایش می‌یابد. یک توالی ۲۰۰۰ توکنی قابل مدیریت است، اما دو برابر کردن آن به ۴۰۰۰ توکن، فضای حافظه اشغالی را چهار برابر می‌کند. زمانی که مدل تلاش می‌کند یک کتاب ۶۴ هزار توکنی را بخواند، ماتریس توجه میانی آن‌قدر بزرگ می‌شود که دیگر در حافظه پنهان فوق‌سریع روی تراشه پردازنده گرافیکی، موسوم به حافظه دسترسی تصادفی ایستا (SRAM)، جا نمی‌گیرد. در عوض، سخت‌افزار مجبور می‌شود داده‌ها را مدام در حافظه بزرگ‌تر اما کندترِ با پهنای باند بالا (HBM) بخواند و بنویسد.

این جابه‌جایی مداوم داده‌ها یک ترافیک سنگین ایجاد می‌کند. به گفته پژوهشگران دانشگاه استنفورد، پیاده‌سازی‌های استاندارد مکانیزم توجه روی پردازنده گرافیکی انویدیا A100، بخش اعظم زمان خود را صرفاً در انتظار انتقال حافظه می‌گذرانند و تنها به ۲۵ تا ۴۰ درصد از حداکثر ظرفیت محاسباتی نظری سخت‌افزار دست می‌یابند. در حالی که گذرگاه حافظه برای همگام شدن دست و پا می‌زند، هسته‌های محاسباتی بی‌کار می‌مانند.[2]

در ماه مه ۲۰۲۲، پژوهشگرانی به نام‌های تری دائو (Tri Dao) و دن فو (Dan Fu) راه حلی را منتشر کردند که نحوه آموزش مدل‌های زبانی در این صنعت را از اساس تغییر داد. الگوریتم آن‌ها که فلش‌اتنشن (FlashAttention) نام گرفت، نه با تغییر ریاضیات، بلکه با تغییر مکان انجام محاسبات ریاضی، از گلوگاه حافظه عبور کرد. دائو در این باره توضیح داد: «فلش‌اتنشن الگوریتمی است که محاسبات توجه را بازآرایی می‌کند و با بهره‌گیری از تکنیک‌های کلاسیک (مانند کاشی‌کاری و محاسبه مجدد)، سرعت آن را به میزان قابل‌توجهی افزایش داده و مصرف حافظه را از حالت مربعی به خطی نسبت به طول توالی کاهش می‌دهد.»[1]

نوآوری اصلی فلش‌اتنشن تکنیکی به نام کاشی‌کاری (Tiling) است. این الگوریتم به جای بارگذاری کل توالی در حافظه HBM، محاسبه ماتریس عظیم توجه و بازنویسی آن، ورودی را به بلوک‌های کوچک‌تر یا همان کاشی‌ها تقسیم می‌کند. سپس یک کاشی منفرد را از حافظه ۸۰ گیگابایتی HBM به حافظه کوچک اما فوق‌العاده سریع ۲۰ مگابایتی SRAM که پهنای باندی برابر با ۱۹ ترابایت بر ثانیه دارد، منتقل می‌کند.[1][3]

به محض اینکه کاشی وارد SRAM می‌شود، پردازنده گرافیکی محاسبات توجه را به صورت محلی انجام داده و خروجی نهایی را مستقیماً به‌روزرسانی می‌کند. ماتریس‌های میانی عظیم اصلاً در HBM نوشته نمی‌شوند. با حذف این خواندن و نوشتن‌های غیرضروری در حافظه، نسخه اولیه الگوریتم فلش‌اتنشن توانست به ۱۵ درصد افزایش سرعت کلی در مدل‌های BERT-large و ۳ برابر افزایش سرعت در GPT-2 دست یابد. مهم‌تر از آن، این روش نیاز به حافظه را از حالت مربعی به خطی کاهش داد و ناگهان آموزش مدل‌ها روی توالی‌های ۱۶ هزار توکنی یا بیشتر را ممکن ساخت.[1]

تأثیر این دستاورد بر صنعت هوش مصنوعی فوری بود. پلتفرم هاگینگ فیس (Hugging Face) این الگوریتم را در کتابخانه هسته ترانسفورمرز خود ادغام کرد و خاطرنشان ساخت که «با موازی‌سازی مضاعف محاسبات توجه در طول توالی، می‌تواند سرعت استنتاج را به طور چشمگیری افزایش دهد.» مدل‌های متن‌باز بزرگی از جمله لاما (Llama) و فالکون (Falcon) نیز آن را به عنوان استانداردی برای آموزش و استنتاج پذیرفتند.[4]

با این حال، نسخه اول هنوز جای پیشرفت داشت. اگرچه این نسخه ترافیک حافظه را کاهش داد، اما نتوانست کار محاسباتی را به شکلی بی‌نقص بین رشته‌های (Threads) داخلی پردازنده گرافیکی تقسیم کند. در ژوئیه ۲۰۲۳، دائو فلش‌اتنشن-۲ را منتشر کرد که در آن کرنل برای بهینه‌سازی تقسیم کار بازنویسی شده بود. این به‌روزرسانی، بار کاری را با کارایی بیشتری در میان وارپ‌های (Warps) پردازنده گرافیکی - گروه‌های ۳۲ تایی از رشته‌ها که دستورالعمل‌ها را همزمان اجرا می‌کنند - توزیع کرد.[2]

این تقسیم‌بندی، نیاز رشته‌ها به همگام‌سازی و اشتراک‌گذاری حافظه با یکدیگر را کاهش داد. نتایج حاصل، بهره‌وری سخت‌افزار را به سطوحی بی‌سابقه رساند. فلش‌اتنشن-۲ به سرعت آموزشی معادل ۲۲۵ ترافلاپس بر ثانیه روی پردازنده گرافیکی A100 رسید و توانست ۷۲ درصد از حداکثر بهره‌وری نظری فلاپ (FLOP) مدل را محقق کند. این نسخه تقریباً دو برابر سریع‌تر از نسخه اصلی بود.[2]

مقیاس‌پذیری خطی حافظه، قابلیت‌های کاملاً جدیدی را برای مدل‌های زبانی به ارمغان آورد. پیش از سال ۲۰۲۲، پنجره‌های زمینه عمدتاً به ۲۰۴۸ یا ۴۰۹۶ توکن محدود می‌شدند. پس از پذیرش گسترده الگوریتم‌های توجهِ آگاه از ورودی/خروجی (IO-aware)، توسعه‌دهندگان شروع به انتشار مدل‌هایی کردند که قادر به پردازش همزمان ۳۲ هزار، ۱۰۰ هزار و در نهایت یک میلیون توکن بودند. این امر به کاربران اجازه داد تا کل پایگاه‌های کد، گزارش‌های مالی یا رمان‌ها را در یک پرامپت واحد بارگذاری کنند.

رقابت برای بهینه‌سازی در کنار استقرار سخت‌افزارهای جدید همچنان ادامه دارد. در ژوئیه ۲۰۲۴، پژوهشگران فلش‌اتنشن-۳ را معرفی کردند که به طور خاص برای معماری هاپر (Hopper) انویدیا، از جمله پردازنده گرافیکی H100 طراحی شده بود. این نسخه از قابلیت‌های ناهمگام هسته‌های تنسور (Tensor Cores) جدید بهره می‌برد و به تراشه اجازه می‌دهد تا محاسبات و جابه‌جایی داده‌ها را به صورت همزمان و هم‌پوشانی‌شده انجام دهد.

فلش‌اتنشن-۳ با درهم‌آمیختن ضرب ماتریس و عملیات سافت‌مکس (Softmax)، پردازنده H100 را به ۷۵ درصد از حداکثر ظرفیت نظری خود می‌رساند و در دقت ۱۶ بیتی به سرعت ۷۴۰ ترافلاپس دست می‌یابد. هنگام استفاده از فرمت‌های ممیز شناور ۸ بیتی با دقت پایین‌تر، این الگوریتم به مرز ۱.۲ پتافلاپس نزدیک می‌شود.

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