اگر قصد دارید مدل‌های زبانی بزرگ را روی سخت‌افزارهای محدود اجرا کنید، بزرگ‌ترین دشمن شما رشد درجه‌دوم (Quadratic) حافظه در مکانیزم توجه است. xFormers — ابزاری تخصصی برای بهینه‌سازی اجزای ترنسفورمر — این مشکل را با محاسبه توجه بدون ذخیره‌سازی کامل ماتریس امتیازات حل می‌کند.
در پیاده‌سازی‌های استاندارد، یک ماتریس عظیم با ابعاد [B, H, M, M] ایجاد می‌شود، جایی که با هر دو برابر شدن طول توالی (Sequence Length)، مصرف حافظه چهار برابر می‌شود. این گلوگاه معمولاً باعث توقف ناگهانی فرآیندهای آموزش یا محدود شدن پنجره متنی در مدل‌های عملیاتی می‌گردد. به نقل از راهنمای فنی Marktechpost، کتابخانه xFormers این رویکرد ساده‌لوحانه را با کرنل‌های بهینه جایگزین کرده تا رشد حافظه را به حالت خطی (Linear) تبدیل کند.
راه‌اندازی و اعتبارسنجی
برای شروع کار با این ابزار، محیط برنامه‌نویسی به نصب xformers و torch نیاز دارد. پیاده‌سازی با تأیید در دسترس بودن GPU و بررسی کرنل‌های توجه پشتیبانی‌شده توسط محیط از طریق دستور xformers.info آغاز می‌شود. برای مقایسه، یک تابع vanilla_attention به عنوان مرجع استفاده می‌شود تا ماتریس امتیازات [B, H, M, M] را با عملیات (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1]) ایجاد کند.
در یک مورد تست با اندازه دسته (B) برابر ۲، طول توالی (M) برابر ۵۱۲، تعداد سرها (H) برابر ۸ و بُعد سر (K) برابر ۶۴، توجه xFormers با این مرجع مقایسه شد. نتایج نشان داد که xFormers خروجی دقیقی (در محدوده خطاهای گرد کردن fp16) تولید می‌کند. این امر ثابت می‌کند که مسیر بهینه از نظر حافظه، صحت ریاضی را فدای سرعت نمی‌کند. این اعتبارسنجی حیاتی است تا اطمینان حاصل شود که عملگر xops.memory_efficient_attention دقیقاً مشابه فرمول استاندارد softmax-attention رفتار می‌کند.
بنچمارک کارایی حافظه
تست در برابر توجه استاندارد (Vanilla)، تضاد شدیدی را در مصرف منابع آشکار می‌کند. بنچمارک‌ها هم زمان اجرای CUDA (با استفاده از torch.cuda.Event برای دقت بالا، شامل ۵ تکرار گرم‌کردن و ۲۰ تکرار اندازه‌گیری) و هم اوج مصرف حافظه را از طریق torch.cuda.max_memory_allocated() می‌سنجند.
هنگام افزایش طول توالی از ۵۱۲ به ۴۰۹۶ توکن، نتایج کاملاً روشن است:
- توجه ساده‌لوحانه (Naive Attention): مصرف حافظه به صورت درجه‌دوم رشد می‌کند (تقریباً ۴ برابر به ازای هر دو برابر شدن طول توالی) زیرا باید کل ماتریس امتیازات BxHxMxM را ذخیره کند. این وضعیت اغلب منجر به بروز
RuntimeErrorدر طول پاس بازگشتی (Backward Pass) می‌شود. - xFormers: رشد حافظه تقریباً خطی باقی می‌ماند و زمان اجرا در تمام طول‌های تست شده، به طور قابل توجهی پایین‌تر است.
این کارایی هم در پاس رفت (Forward) و هم در پاس بازگشتی جاری است و تضمین می‌کند که گرادین‌ها در طول آموزش، حافظه VRAM گرافیک را اشباع نکنند. حتی در ۴۰۹۶ توکن، xFormers از رویدادهای Out-of-Memory (OOM) که در پیاده‌سازی‌های ساده رایج است، جلوگیری می‌کند و اجازه می‌دهد اندازه‌های دسته (Batch Size) بزرگ‌تر یا متون طولانی‌تری روی همان سخت‌افزار قرار گیرند.
پیاده‌سازی مکانیزم‌های پیشرفته توجه
فراتر از صرفه‌جویی در حافظه، xFormers چندین بهینه‌سازی معماری را که در مدل‌های پیشرفته (SOTA) استفاده می‌شود، فعال می‌کند:
ماسک‌گذاری علی (Causal Masking)
- ماسک‌گذاری ضمنی: این ابزار از یک
LowerTriangularMaskضمنی از طریقab.LowerTriangularMask()استفاده می‌کند. این بدان معناست که هیچ تانسور بولی MxM برای جلوگیری از نگاه مدل به توکن‌های آینده تخصیص نمی‌یابد و ردپای حافظه به حداقل می‌رسد. این پیاده‌سازی با مقایسه خروجی در برابر یک تابع توجه علی مرجع که ازtorch.triuبرای پر کردن مثلث بالایی با-infاستفاده می‌کند، تأیید شده است.
توالی‌های بسته‌بندی شده (Packed Sequences)
- اتلاف Padding صفر: با استفاده از
BlockDiagonalMask.from_seqlens()، توسعه‌دهندگان می‌توانند توالی‌هایی با طول متغیر را در یک تانسور واحد دسته‌بندی کنند. برای مثال، توالی‌هایی با طول‌های [۳۷، ۱۲۰، ۸، ۲۰۰] می‌توانند در یک توالی کلی ۳۶۵ توکنی ادغام شوند. - کارایی سبک vLLM: این روش اتلاف فضای Padding را حذف می‌کند؛ تکنیکی که هسته اصلی موتورهای با توان عملیاتی بالا مانند vLLM است. این ابزار همچنین از
BlockDiagonalCausalMaskبرای پاس‌های علی بسته‌بندی شده پشتیبانی می‌کند و بخش‌های اصلی را می‌توان با متدbias.split()بازیابی کرد.
توجه پرس‌وجوی گروهی (Grouped-Query Attention یا GQA)
- کاهش KV-Cache: xFormers از یک چیدمان ۵ بعدی [B, M, G, Hq, K] پشتیبانی می‌کند. در این ساختار، چندین سر پرس‌وجو (مثلاً ۸ سر) یک مجموعه کوچک‌تر از سرهای کلید-مقدار (مثلاً ۲ سر) را به اشتراک می‌گذارند.
- هم‌ترازی با مدل‌ها: این کار حجم KV-cache را به شدت کاهش می‌دهد و دقیقاً مشابه معماری مدل‌های Llama و Mistral در هنگام استنتاج است. شکل خروجی برای این عملیات [B, M, G, Hq, K] است.
بایاس ALiBi
- بایاس افزودنی سفارشی: این ابزار اجازه می‌دهد بایاس‌های موقعیتی افزودنی سفارشی اعمال شوند. پیاده‌سازی ALiBi شامل ایجاد شیب‌ها بر اساس تعداد سرها (مثلاً
2.0 ** (-8.0 / H)) و اعمال یک جریمه خطی برای توکن‌های دورتر با استفاده از یک ماتریس موقعیت نسبی است. - برون‌یابی متون طولانی: این بایاس با یک ماسک علی ترکیب شده و مستقیماً به عملگر توجه ارسال می‌شود تا برون‌یابی متون طولانی بهبود یابد. هرگونه بایاس افزودنی برای هر (سر، پرس‌وجو، کلید) را می‌توان به این روش ادغام کرد.
ساخت یک بلوک GPT قابل آموزش
برای اثبات هماهنگی این اجزا، یک مدل TinyGPT پیاده‌سازی شد. این مدل با واژگانی به اندازه ۶۴، طول توالی ۶۴ و ۳ لایه بلوک ترنسفورمر پیکربندی شده است. مدل شامل یک لایه جاسازی (Embedding) برای توکن‌ها و یک لایه مجزا برای کدگذاری موقعیتی است.
هر Block شامل موارد زیر است:
- نرمال‌سازی لایه (Layer Normalization): دو بار در هر بلوک (
n1وn2) قبل از مراحل توجه و پیش‌رو (Feed-forward) اعمال می‌شود. - پروجکشن QKV: یک لایه خطی (
nn.Linear(d, 3 * d)) که ورودی را به تانسورهای پرس‌وجو، کلید و مقدار تبدیل می‌کند و سپس به شکل [B, M, 3, H, K] تغییر اندازه می‌دهد. - توجه xFormers: با بهره‌گیری از
LowerTriangularMaskبرای رمزگشایی علی و یک لایه پروجکشن نهایی برای بازگرداندن تانسور به بُعد مدل. - لایه‌های پیش‌رو SwiGLU: مدل از
xops.SwiGLU(با ویژگی‌های ورودی، پنهان و خروجی) یا یک جایگزین دستی باF.silu(a) * bاستفاده می‌کند. این عملیات ادغام‌شده (Fused) از نظر محاسباتی بهینه‌تر از لایه‌های استاندارد ReLU است.
خط لوله آموزش از دقت ترکیبی خودکار (AMP) از طریق torch.amp.GradScaler و torch.autocast("cuda", dtype=torch.float16) بهره می‌برد. با استفاده از بهینه‌ساز AdamW و نرخ یادگیری 3e-3، مدل روی یک وظیفه مصنوعی پیش‌بینی توکن بعدی (شمارش صعودی به پیمانه اندازه واژگان) آموزش دید. داده‌های آموزشی توسط تابع make_batch ایجاد می‌شوند که توالی‌هایی را از یک توکن تصادفی شروع می‌کند.
طی ۴۰۰ گام، مدل همگرایی موفقیت‌آمیزی را نشان داد. حلقه آموزش، مقدار Loss و دقت پیش‌بینی توکن بعدی (میانگین پیش‌بینی‌های درست) را نظارت کرد و تأیید کرد که کرنل‌های بهینه حافظه، توانایی یادگیری مدل را از ابتدا تا انتها کاهش نمی‌دهند. این امر تأیید می‌کند که xFormers می‌تواند به عنوان جایگزینی مستقیم برای توجه استاندارد در یک حلقه کامل آموزش استفاده شود.
این تغییر در پیاده‌سازی به این معناست که توسعه‌دهندگان دیگر مجبور نیستند بین «پنجره‌های متنی عظیم» و «سخت‌افزارهای مقرون‌به‌صرفه» یکی را انتخاب کنند. با انتقال محاسبات به کرنل‌های تخصصی، صنعت از تخصیص حافظه به روش «زور خالص» در ترنسفورمرهای اولیه، به سمت یک رویکرد جراحی‌گونه در سطح کرنل حرکت می‌کند.
برای متخصصان، این یعنی مانع آموزش مدل‌های GPT سفارشی و کوچک‌مقیاس کمتر شده است. شما اکنون می‌توانید لایه‌های بیشتر یا توالی‌های طولانی‌تری را در همان فضای ۲۴ یا ۸۰ گیگابایتی GPU جای دهید، تنها با جایگزین کردن عملگر توجه.
برای مشاهده این بهینه‌سازی‌ها در عمل، می‌توانید پیاده‌سازی کامل را از طریق نوت‌بوک بررسی کرده و چیدمان GQA را در محیط CUDA خود تست کنید.
نتیجه‌گیری
در نهایت، ما درک عملی از نحوه بهبود کارایی ترنسفورمر توسط xFormers بدون تغییر در محاسبات بنیادی توجه به دست آوردیم. مشاهده کردیم که چگونه کرنل‌های بهینه حافظه هزینه توالی‌های طولانی را کاهش می‌دهند، در حالی که ماسک‌های علی، توالی‌های بسته‌بندی شده، توجه پرس‌وجوی گروهی و بایاس‌های افزودنی از جریان‌های کاری واقعی آموزش و استنتاج پشتیبانی می‌کنند. در نهایت با ادغام این قابلیت‌ها در یک مدل GPT فشرده و آموزش آن، پایه‌ای قوی برای به‌کارگیری xFormers در مدل‌های زبانی بزرگ‌تر و مجموعه‌داده‌های دشوارتر ایجاد کردیم.
گفتگو