رویکرد سخت‌افزاری انویدیا در برابر پیاده‌سازی‌های استاندارد PyTorch

اگر امروز برای آموزش مدل‌های زبانی هزینه می‌کنید، احتمالاً بخش بزرگی از بودجه شما صرف مدیریت گلوگاه‌های حافظه می‌شود. آموزش مدل‌های عمومی اغلب در رسیدن به عملکرد سطح تولید (Production-grade) ناکام می‌ماند، اما بهینه‌سازی‌های سخت‌افزار-آگاه می‌توانند این شکاف عظیم را پر کنند. موتور ترنسفورمر (Transformer Engine یا TE) انویدیا با ادغام هسته‌های GPU و بهره‌گیری از دقت FP8، بارهای کاری ترنسفورمر را روی سخت‌افزارهای سازگار شتاب می‌بخشد تا مدل‌ها سریع‌تر و با مصرف حافظه کمتر آموزش ببینند.

در حالی که پیاده‌سازی‌های سنتی ترنسفورمر بر عملیات‌های متوالی تکیه دارند که اغلب در رابط حافظه GPU متوقف می‌شوند، TE چندین عملیات را در قالب تک‌هسته‌های گرافیکی ادغام می‌کند. این تغییر، بازتابی از تلاش‌های گسترده‌تر برای حل گلوگاه‌های متوالی است؛ چالشی که پیش‌تر در پوشش خود درباره‌ی نحوه برخورد PyTorch Transformers با محدودیت‌های RNN بررسی کردیم. حذف توالی‌های کند، کلید دستیابی به مقیاس است.

برای توسعه‌دهندگان، این یعنی تفاوت بین مدلی که به‌سختی در حافظه VRAM جای می‌گیرد و مدلی که با سرعت خیره‌کننده آموزش می‌بیند و ردپای حافظه (Memory Footprint) آن کاهش یافته است. این سیستم به‌گونه‌ای طراحی شده که انعطاف‌پذیر باشد و برای GPUهای قدیمی‌تر، بازگشت‌های خودکار (Automatic Fallbacks) به دقت‌های BF16 یا FP32 را فراهم کند، در حالی که در معماری‌های Ampere و Blackwell قدرت کامل خود را آزاد می‌کند.

هسته‌های ادغام‌شده و بهره‌وری حافظه

طبق مستندات انویدیا، هسته اصلی این افزایش عملکرد، ماژول‌های ادغام‌شده (Fused Modules) هستند. به‌جای اجرای جداگانه هر هسته GPU برای هر عملیات، Transformer Engine اجزای یکپارچه‌ای را ارائه می‌دهد تا سربار اجرای هسته‌ها (Kernel Launch Overhead) کاهش یابد:

• te.Linear: یک لایه خطی ادغام‌شده که برای توان عملیاتی (Throughput) بالا بهینه شده و از انواع مختلف داده‌های پارامتری، از جمله torch.bfloat16 پشتیبانی می‌کند.
• te.LayerNorm: نرمال‌سازی با کارایی بالا که ترافیک حافظه را به‌طور چشمگیر کاهش می‌دهد.
• te.LayerNormLinear و te.LayerNormMLP: بلوک‌های پیشرفته‌ای که نرمال‌سازی را با عملیات خطی یا MLP در یک گام اجرایی ترکیب می‌کنند. این بلوک‌ها تغییر ابعادی مانند $H$ (768) به $3 imes H$ یا $4 imes H$ را تسهیل می‌کنند.
• te.TransformerLayer: یک بلوک کامل ترنسفورمر که توجه (Attention) و شبکه‌های پیش‌خور (Feed-forward Networks) را به‌عنوان یک واحد مدیریت می‌کند و از ماسک‌های توجه علی (Causal Attention Masks) و شماره‌گذاری خاص لایه‌ها پشتیبانی می‌کند.

این عملیات‌های ادغام‌شده، تعداد دفعاتی که داده‌ها باید بین ثبات‌های GPU (Registers) و حافظه سراسری (Global Memory) جابه‌جا شوند را به حداقل می‌رسانند. با استفاده از te.TransformerLayer شبکه توالی استاندارد شامل LayerNorm، MultiheadAttention و MLP را با یک فراخوانی واحد و بهینه شده برای سخت‌افزار جایگزین می‌کند. این رویکرد باعث می‌شود گلوگاه‌های انتقال داده که در پیاده‌سازی‌های معمولی PyTorch رایج است، تا حد زیادی از بین بروند.

مکانیسم اجرای FP8

انتقال از دقت BF16 به دقت ۸ بیتی (FP8) — شبیه به تبدیل یک عکس باکیفیت و سنگین به فرمتی فشرده‌تر که همچنان جزئیات اصلی را دارد — اجازه می‌دهد هسته‌های تنسور سریع‌تر اجرا شوند و مصرف حافظه کاهش یابد. با این حال، دقت ۸ بیتی مستعد ناپایداری عددی است. برای حل این مشکل، TE از دستورالعمل «مقیاس‌بندی تأخیری» از طریق recipe.DelayedScaling استفاده می‌کند.

به نقل از منابع فنی، این مکانیسم تاریخچه مقادیر حداکثر مطلق (amax) را در طول تکرارها ردیابی می‌کند. با حفظ این تاریخچه (که در آموزش‌ها با طول ۱۶ تکرار و الگوریتم محاسباتی "max" پیکربندی شده است)، موتور می‌تواند فاکتورهای مقیاس‌بندی برای تنسورها را به‌طور پویا تنظیم کند. این امر تضمین می‌کند که محدوده دینامیکی محدود FP8 منجر به محوشدن گرادیان (Gradient Vanish) یا انفجار گرادیان نشود.

این سیستم از یک فرمت ترکیبی (Hybrid) پشتیبانی می‌کند که بین E4M3 (۴ بیت برای نمایی و ۳ بیت برای مانتیس) و E5M2 جابه‌جا می‌شود. این کار تعادلی بین دقت در جایی که برای وزن‌ها لازم است و توان عملیاتی در جایی که برای فعال‌سازها (Activations) نیاز است، ایجاد می‌کند. تنظیم recipe.Format.HYBRID به موتور اجازه می‌دهد تا بهینه‌ترین فرمت را برای مراحل مختلف گذر ترنسفورمر انتخاب کند.

پیاده‌سازی و سازگاری

برای ادغام این سیستم، ابتدا باید قابلیت محاسباتی GPU از طریق torch.cuda.get_device_properties(0) شناسایی شود. سیستم نسخه‌های major و minor را بررسی می‌کند تا دسترسی را تعیین کند:

• پشتیبانی از TE: نیازمند قابلیت محاسباتی ۸.۰ یا بالاتر (معماری Ampere).
• پشتیبانی از FP8: نیازمند قابلیت محاسباتی ۸.۹ یا بالاتر (GPUهای L4، H100، Ada Lovelace یا Blackwell).

برای کاربران سخت‌افزارهای قدیمی‌تر مانند T4، سیستم به‌طور خودکار به مسیر پشتیبان خالص PyTorch بازمی‌گردد. این مدیریت توسط تابع te.fp8.check_fp8_support() انجام می‌شود تا اطمینان حاصل شود سخت‌افزار واقعاً قادر به اجرای هسته‌های با دقت کاهش‌یافته است. اگر TE در دسترس نباشد، یک کلاس Block_PT (Plain-PyTorch) برای بازسازی منطق te.TransformerLayer با استفاده از nn.MultiheadAttention و MLPهای nn.Sequential به کار گرفته می‌شود تا کد بدون تغییر در سخت‌افزارهای قدیمی اجرا شود.

در یک حلقه آموزش، این موتور از یک مدیر زمینه (Context Manager) به نام fp8_autocast استفاده می‌کند. این پوشش، تبدیل تنسورها به FP8 و مدیریت دستور مقیاس‌بندی را بدون نیاز به تبدیل دستی هر لایه توسط کاربر مدیریت می‌کند. این ابزار در کنار بهینه‌سازهایی مثل AdamW (که در تست‌های ما با نرخ یادگیری 3e-4 تنظیم شده بود) برای حفظ پایداری کار می‌کند و تضمین می‌کند که تبدیل‌ها در هر گام به‌درستی رخ دهند.

بنچمارک و دستاوردهای عملیاتی

در آزمایش‌های عملی روی یک مدل زبانی کوچک علی (Mini-causal Language Model) با حدود ۱۰ میلیون پارامتر روی توالی‌های سنتتیک، اثر FP8 فوری بود. هایپرپارامترهای معماری این مدل به شرح زیر تعریف شده است:

• واژگان (Vocabulary): ۹۶ توکن
• اندازه پنهان (D_MODEL): ۷۶۸
• سرهای توجه (Attention Heads): ۱۲
• اندازه FFN: ۳۰۷۲
• تعداد لایه‌ها: ۴
• طول توالی: ۲۵۶

بر اساس گزارش Marktechpost، بنچمارک‌های مربوط به کل حلقه آموزش (شامل انتشار پیش‌رو، پس‌انتشار و به‌روزرسانی بهینه‌ساز) نشان داد که اجرای FP8 هم زمان هر گام (ms/step) و هم اوج حافظه GPU تخصیصی شده را در مقایسه با BF16 کاهش می‌دهد. برای اندازه دسته‌ای (Batch Size) ۳۲ و طول توالی ۲۵۶، بنچمارک‌ها تأخیر دقیق و سربار حافظه را با استفاده از توابع torch.cuda.max_memory_allocated() و torch.cuda.synchronize() اندازه‌گیری کردند تا اثرات دقیق همگام‌سازی GPU در محاسبات لحاظ شود.

این دستاوردها با مقیاس‌پذیر شدن مدل دراماتیک‌تر می‌شوند. برای مثال، افزایش ابعاد پنهان مدل به ۲۰۴۸ و تعداد لایه‌ها به ۱۲، شکاف عملکردی بین حالت FP8 و حالت‌های با دقت بالاتر را به‌طور قابل‌توجهی عمیق‌تر می‌کند و نشان می‌دهد که در مدل‌های بزرگ‌تر، صرفه‌جویی در حافظه و سرعت FP8 حیاتی‌تر می‌شود.

نظارت بر پایداری عددی

برای تأیید سلامت آموزش FP8، موتور متاداده‌ها را از طریق fp8_meta نمایش می‌دهد. توسعه‌دهندگان می‌توانند با بررسی ویژگی scaling_fwd در زیر-ماژول‌ها، موارد زیر را مشاهده کنند:

• فاکتورهای مقیاس‌بندی: ضرایبی که در حال حاضر برای گنجاندن تنسورها در محدوده FP8 اعمال می‌شوند.
• تاریخچه Amax: مقادیر حداکثری ثبت شده در چندین تکرار اخیر (مثلاً ۴ مورد اول تاریخچه) که برای محاسبه فاکتور مقیاس بعدی استفاده می‌شوند.

این شفافیت به متخصصان اجازه می‌دهد تشخیص دهند آیا مقیاس‌بندی تأخیری به‌درستی توزیع تنسورها را ثبت می‌کند یا خیر. در تست سنتتیک، مدل موفق شد یک الگوی حسابی (گام توکن ثابت) را در واژگان ۹۶ تایی یاد بگیرد، که ثابت می‌کند FP8 توانایی مدل برای همگرایی روی قوانین ساده را تخریب نمی‌کند. این موضوع از طریق تولید خودرگرسیو حریصانه (Greedy Autoregressive Generation) تأیید شد، جایی که توکن‌های تولید شده متوالی برای بررسی تفاوت گام‌های ثابت مقایسه شدند و صحت یادگیری مدل در دقت ۸ بیتی به اثبات رسید.

تحلیل: تعریف مجدد خط مبنای آموزش

این انتقال به FP8 و هسته‌های ادغام‌شده، نشان‌دهنده تغییری در تعریف «آموزش بهینه» است. ما در حال حرکت از بهینه‌سازی‌های صرفاً الگوریتمی به سمت «طراحی مشترک سخت‌افزار-نرم‌افزار» (Hardware-Software Co-design) هستیم. وقتی کامپایلر و هسته از معماری خاص هسته‌های تنسور آگاه باشند، پیک تئوری TFLOPS سخت‌افزار در واقع قابل دستیابی می‌شود.

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

با این حال، این موضوع یک وابستگی به سخت‌افزار (Hardware-lock) ایجاد می‌کند. قابل‌توجه‌ترین دستاوردهای عملکردی اکنون توسط جدیدترین معماری‌های NVIDIA کنترل می‌شوند و احتمالاً شکاف بین کسانی که به جدیدترین سیلیکون‌ها دسترسی دارند و کسانی که بر خوشه‌های قدیمی تکیه می‌کنند را بیشتر می‌کنند.

گام‌های عملی برای توسعه‌دهندگان:

• تنظیم recipe.Format.HYBRID برای یافتن تعادل بهینه بین سرعت و دقت.
• افزایش amax_history_len تا ۱۰۲۴ برای پایداری بیشتر در مدل‌های بسیار بزرگ.
• استفاده از تابع fp8_model_init() برای ذخیره خودِ وزن‌ها در قالب FP8 جهت استنتاج فوق‌سریع.

این رویکرد جامع، ترنسفورمر را از مجموعه‌ای از لایه‌های مجزا به یک خط لوله (Pipeline) بهینه شده و جریان‌یافته در سخت‌افزار تبدیل می‌کند. اما داستان سخت‌افزاری این تحول حتی شگفت‌انگیزتر است — به تحلیل ما درباره‌ی تراشه‌های Blackwell مراجعه کنید.