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

بلوک ترنسفورمر (Transformer Block) کوچک‌ترین واحد ساختاری مدل‌های زبانی امروزی نظیر GPT و LLaMA است. در نگاه اول، سرهم‌کردن لایه‌های اتنشن، شبکه‌ی پیش‌خور و نرمال‌سازی ساده به نظر می‌رسد؛ اما کوچک‌ترین جابه‌جایی در ترتیب لایه‌ها، مثل جایگاه نرمال‌سازی لایه‌ای (LayerNorm) یا نحوه‌ی اعمال ماسک علّی (Causal Mask)، می‌تواند پایداری مدل را در لایه‌های عمیق به‌کلی از بین ببرد و فرایند آموزش را به واگرایی بکشاند.

در این درس یاد می‌گیرید که چطور یک بلوک ترنسفورمر استاندارد را بر پایه‌ی PyTorch و بدون تکیه بر کتابخانه‌های آماده از صفر بسازید. هم‌چنین دو معماری پرکاربرد Pre-LN و Post-LN را در یک پشته‌ی 12 لایه‌ای روبه‌روی هم قرار می‌دهیم تا به‌صورت تجربی ببینید چرا مدل‌های مدرن همگی به سمت Pre-LN رفته‌اند و چگونه جریان گرادیان در این معماری بدون نیاز به تنظیمات پیچیده‌ی پیش‌گرمایش (Warmup) پایدار می‌ماند.

کدها و محاسبات این پیاده‌سازی کاملاً متکی بر پکیج‌های پایه‌ای PyTorch هستند و به‌راحتی روی پردازنده‌های معمولی (CPU) یا شتاب‌دهنده‌های سبک شخصی اجرا می‌شوند و نیازی به سرورها یا GPUهای پرهزینه‌ی ابری نخواهید داشت.

هر بلوک، واحد سازنده‌ی تمام مدل‌های زبانی بزرگ مبتنی بر رمزگشایی (Decoder LLM) امروزی است: نرمال‌سازی لایه‌ای (LayerNorm)، اتنشن چندسر (Multi-Head Attention)، اتصال باقیمانده (Residual Connection)، شبکه‌ی MLP و اتصال باقیمانده‌ی مجدد. نسخه‌ی Pre-LN بدون نیاز به وارم‌آپ (پیش‌گرمایش) با پایداری آموزش می‌بیند. نسخه‌ی Post-LN همان ساختاری است که در مقاله‌ی اصلی ترنسفورمر ارائه شد. در این درس هر دو نسخه را شانه‌به‌شانه می‌سازیم و بررسی می‌کنیم کدام‌یک در یک پشته‌ی 12 لایه‌ای با نرخ‌های یادگیری معمول دوام می‌آورد.

اهداف یادگیری

• ساخت یک بلوک ترنسفورمر در PyTorch با استفاده از چهار جزء متحرک: لایه‌ی LayerNorm، اتنشن چندسر علّی (Causal Multi-Head Attention)، اتصالات باقیمانده و MLP مستقل از موقعیت (Position-wise MLP).
• چیدمان لایه‌های LayerNorm در دو پیکربندی مختلف (Pre-LN و Post-LN) و توضیح این‌که چرا یکی از آن‌ها بدون وارم‌آپ با پایداری آموزش می‌بیند.
• پیاده‌سازی ماسک علّی (Causal Mask) درون اتنشن چندسر، به‌طوری که توکن i نتواند توکن‌های j > i را ببیند.
• ردگیری جریان گرادیان در هر دو نسخه روی یک پشته‌ی 12 لایه‌ای و خواندن نتایج بدون ابهام یا ساده‌انگاری.
• استفاده‌ی مجدد از این بلوک به‌عنوان یک واحد آماده (Drop-in) برای درس بعد، جایی که یک مدل GPT با 124 میلیون پارامتر سرهم می‌شود.

صورت مسئله

یک ترنسفورمر چیزی نیست جز تکرار یک بلوک یکسان. اگر این بلوک را حتی یک بار اشتباه پیاده کنید و دوازده بار تکرارش کنید، مدلی تحویل می‌دهید که در همان دور (Epoch) اول واگرا می‌شود یا مجبور می‌شوید در ادامه‌ی مسیر دست به دامن ترفندهای وارم‌آپ شوید. دو حالت شکستی (Failure modes) که در این درس می‌بینید اصلاً عجیب و غریب نیستند؛ این‌ها اولین خطاهایی هستند که هر تازه‌کاری هنگام چیدن ناآگاهانه‌ی بلوک‌ها روی هم با آن‌ها روبه‌رو می‌شود. یکی از آن‌ها توجه لایه‌ی اتنشن به آینده است، و دیگری قرار گرفتن لایه‌ی LayerNorm در نقطه‌ای که نمی‌تواند سیگنال پسماند را در لایه‌های عمیق مهار کند.

وقتی اصل قضیه را ببینید، چاره‌اش کاملاً مکانیکی و سرراست است. بلوک دقیقاً دو مسیر باقیمانده و دقیقاً دو جایگاه برای نرمال‌سازی دارد. اگر این جایگاه‌ها را درست انتخاب کنید، بقیه‌ی کار پشته فقط حساب‌وکتاب داده‌ها خواهد بود.

مفاهیم پایه

هر بلوک ترنسفورمر صرفاً رمزگشا (Decoder-only)، تابعی است که تنسوری با ابعاد (batch, sequence, embedding) می‌گیرد و تنسوری با همان ابعاد تحویل می‌دهد. درون این بلوک، دو زیرلایه کار اصلی را پیش می‌برند.

این نسخه‌ی Pre-LN است. لایه‌ی LayerNorm داخل شاخه‌ی (branch) باقیمانده و قبل از زیرلایه می‌نشیند. اتصال باقیمانده، سیگنال نرمال‌نشده را مستقیم به جلو منتقل می‌کند.

در نسخه‌ی Post-LN، لایه‌ی LayerNorm به بعد از جمع باقیمانده منتقل می‌شود.

ابعاد خروجی کاملاً یکسان است، اما رفتار هنگام آموزش خیر. در Post-LN، گرادیانی که در مسیر باقیمانده به عقب برمی‌گردد باید از لایه‌ی LayerNorm عبور کند. در عمق دوازده لایه و با نرخ یادگیری 3e-4، این گرادیان آن‌قدر سریع تضعیف می‌شود که چاره‌ای جز استفاده از برنامه‌ی زمان‌بندی وارم‌آپ باقی نمی‌ماند. اما Pre-LN مسیر باقیمانده را بدون نرمال‌سازی رها می‌کند؛ در نتیجه گرادیان‌ها تر‌وتمیز به لایه‌ی امبدینگ می‌رسند. دقیقاً به همین دلیل است که از GPT-2 به بعد، همه‌ی مدل‌ها با ساختار Pre-LN عرضه می‌شوند.

اتنشن چندسر علّی

زیرلایه‌ی اتنشن، ورودی را به سه شکل روی تنسورهای پرسش (Query)، کلید (Key) و مقدار (Value) تصویر می‌کند (Projection). هر کدام از این تنسورها از شکل (B, T, D) به شکل (B, H, T, D/H) تغییر بعد می‌دهند که در آن H تعداد سرها است. اتنشن ضرب داخلی مقیاس‌شده (Scaled Dot-Product Attention) مقدار softmax(Q K^T / sqrt(d_k)) را برای هر سر محاسبه می‌کند، مثلث بالایی را با بی‌نهایت منفی ماسک می‌زند، ماسک را از طریق سافت‌مکس اعمال می‌کند و سپس حاصل را در V ضرب می‌کند. سرها دوباره در قالب یک تنسور واحد (B, T, D) به هم متصل می‌شوند و یک بار دیگر تصویر می‌شوند. این ماسک تنها قطعه‌ای است که مدل را علّی نگه می‌دارد. اگر ماسک را فراموش کنید، مدلی آموزش می‌دهید که تقلب می‌کند.

لایه‌ی MLP

لایه‌ی MLP مستقل از موقعیت، یک شبکه‌ی دولایه‌ای یکسان را به‌طور مستقل روی تک‌تک توکن‌ها اعمال می‌کند. عرض لایه‌ی پنهان چهار برابر عرض لایه‌ی امبدینگ است، تابع فعال‌سازی GELU بوده و بعد از لایه‌ی خطی دوم یک دراپ‌اوت قرار می‌گیرد. هیچ توکنی درون MLP با توکن دیگر تبادل اطلاعات نمی‌کند؛ همه‌ی آمیختگی توکن‌ها در اتنشن اتفاق می‌افتد.

نقش دوگانه‌ی اتصالات باقیمانده

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

پیاده‌سازی

فایل code/main.py موارد زیر را پیاده‌سازی می‌کند:

• کلاس LayerNorm همراه با مقیاس و جابه‌جایی یادگرفتنی، پارامتر eps اریب، که روی بردار هر توکن اعمال می‌شود.
• کلاس MultiHeadAttention با پارامترهای num_heads، بعد سرها به‌صورت head_dim = d_model // num_heads، تصویرسازی ادغام‌شده‌ی QKV، بافر ثبت‌شده‌ی ماسک علّی، و دراپ‌اوت برای اتنشن و اتصال باقیمانده.
• کلاس FeedForward با دو لایه‌ی خطی، تابع فعال‌سازی GELU و دراپ‌اوت.
• کلاس TransformerBlock همراه با پرچم pre_ln برای جابه‌جایی بین دو نسخه.
• یک دمو که دو پشته‌ی 6 لایه‌ای، یکی Pre-LN و دیگری Post-LN، با ورودی‌های کاملاً یکسان می‌سازد و موارد زیر را چاپ می‌کند: (الف) ابعاد خروجی، (ب) نرم گرادیان در لایه‌ی امبدینگ پس از یک گام رو به عقب.

اجرای کد:

python3 code/main.py

خروجی: بررسی ابعاد در هر دو پشته و مقایسه‌ی شانه‌به‌شانه‌ی نرم گرادیان‌ها. نرم گرادیان امبدینگ در پشته‌ی Pre-LN یک مرتبه بزرگی (Order of magnitude) بیشتر از پشته‌ی Post-LN در همان نرخ یادگیری است؛ این همان نشانه‌ی تجربی است که ثابت می‌کند Pre-LN بدون نیاز به وارم‌آپ آموزش می‌بیند.

ابزارها

• پکیج torch برای محاسبات تنسور، مشتق‌گیری خودکار (Autograd) و اتصالات لایه‌های nn.Module.
• بدون transformers و بدون وزن‌های از پیش‌آموزش‌دیده؛ این بلوک کاملاً از اجزای سازنده (Primitives) پیاده‌سازی شده است.

الگوهای بهینه‌سازی در محیط عملیاتی

سه الگو وجود دارد که بلوک متنی کتاب‌ها را به چیزی تبدیل می‌کند که واقعاً بشود در محیط عملیاتی (پروداکشن) استفاده کرد.

تصویرسازی ادغام‌شده‌ی QKV (Fused QKV projection). سه لایه‌ی خطی جداگانه به معنای سه بار اجرای کرنل (Kernel launch) و سه ضرب ماتریسی است. یک لایه‌ی خطی با عرض 3 * d_model همین کار را در یک‌بار اجرای کرنل انجام می‌دهد و بعد خروجی را در امتداد بعد آخر می‌شکند. مسیر ادغام‌شده روی تمام شتاب‌دهنده‌ها سریع‌تر است و با آن‌چه پیاده‌سازی‌های مرجع GPT-2، LLaMA و Mistral استفاده می‌کنند مطابقت دارد.

بافر ثبت‌شده برای ماسک علّی. این ماسک فقط به حداکثر طول زمینه (Context length) وابسته است. کافی است آن را یک بار هنگام ساخت شیء با register_buffer تخصیص دهید، پنجره‌ی فعال را در هر گذر رو به جلو برش بزنید و از تخصیص حافظه‌ی تکراری در هر فراخوانی پرهیز کنید. فراموش کردن این نکته، ماسک را در طول زمینه‌های طولانی به یک گلوگاه تخصیص حافظه تبدیل می‌کند.

دراپ‌اوت در دو نقطه، نه سه نقطه. دراپ‌اوت باید بعد از سافت‌مکس اتنشن (دراپ‌اوت اتنشن) و بعد از لایه‌ی خطی دوم MLP (دراپ‌اوت باقیمانده) بنشیند. قرار دادن دراپ‌اوت روی خود اتصال باقیمانده، خاصیت همانیِ جمع‌پذیری را که ضامن عبور گرادیان در لایه‌های عمیق است از بین می‌برد. برخی پیاده‌سازی‌های اولیه این اشتباه را مرتکب شدند و تاوانش را با آموزش‌های شکننده و ناپایدار پس دادند.

کاربردها

• بلوک این درس مستقیماً و بدون تغییر در سرهم‌بندی GPT در درس 35 قرار می‌گیرد.
• نسخه‌ی Pre-LN همان ساختاری است که تمام مدل‌های متن‌باز امروزی به کار می‌برند. نسخه‌ی Post-LN ساختار مقاله‌ی اتنشن سال 2017 است. تسلط بر هر دو نسخه برای فهم هر معماری مبتنی بر رمزگشایی که با آن مواجه می‌شوید کافی است.
• اگر GELU را با SiLU عوض کنید، به تابع فعال‌سازی خانواده‌ی LLaMA می‌رسید. اگر LayerNorm را با RMSNorm جایگزین کنید، ساختار نرمال‌سازی خانواده‌ی LLaMA را خواهید داشت؛ اسکلت اصلی یکی است.

تمرین‌ها

1. پرچم bias=False را به تمام لایه‌های خطی بلوک اضافه کنید. مدل‌های زبانی متن‌باز امروزی لایه‌های خطی را بدون بایاس به کار می‌برند. حساب کنید در یک مدل 12 لایه‌ای با بعد 768، چه تعداد پارامتر صرفه‌جویی می‌شود.
2. لایه‌ی nn.LayerNorm را با یک لایه‌ی RMSNorm دست‌نویس جایگزین کنید و مطمئن شوید که ابعاد خروجی دست‌نخورده باقی می‌ماند.
3. پرچمی اضافه کنید که وزن‌های اتنشن را برای سر اول به‌صورت تنسور (B, T, T) برگرداند. مثلث بالایی آن را رسم کنید تا مطمئن شوید بعد از سافت‌مکس مقدارش صفر است.
4. یک آزمون اعتبارسنجی بسازید که تنسوری با ابعاد (2, 16, 384) و H=6 را به هر دو نسخه بفرستد و بررسی کند که وقتی وزن‌ها کاملاً یکسان مقداردهی اولیه شده‌اند و دراپ‌اوت صفر است، خروجی‌های گذر رو به جلو با هم تفاوت دارند (مثلاً با not torch.allclose).

اصطلاحات کلیدی

منابع برای مطالعه‌ی بیشتر

• بخش 7، درس 02 (خوداتنشنی از صفر) برای درک ریاضیات اتنشنِ زیربنای این بلوک.
• بخش 7، درس 05 (ترنسفورمر کامل) برای مشاهده‌ی نسخه‌ی رمزگذار-رمزگشا (Encoder-Decoder) از همین ساختار.
• بخش 10، درس 04 (پیش‌آموزش مینی‌جی‌پی‌تی) برای یادگیری فرایند آموزشی که این بلوک درون آن قرار می‌گیرد.
• بخش 19، درس 35 (همین مسیر آموزشی) که دوازده تا از این بلوک‌ها را در قالب یک مدل GPT سرهم می‌کند.