سرهم‌بندی مدل GPT: پیاده‌سازی معماری 124M و تولید متن

در درس‌های گذشته اجزای درونی ترنسفورمر، از مکانیزم توجه گرفته تا لایه‌های MLP و نرمال‌سازی را به‌تفکیک بررسی کردیم. اما این قطعات تا زمانی که در قالب یک معماری یکپارچه سیم‌کشی نشوند، قادر به پردازش توالی‌های زبانی یا تولید کلمات نیستند. در این درس، تمام این بلوک‌ها را کنار امبدینگ توکن و موقعیت می‌گذاریم تا معماری نام‌آشنای GPT-2 با 124 میلیون پارامتر متولد شود.

درک نحوه‌ی اتصال این بخش‌ها، به‌ویژه تکنیک اشتراک وزن‌ها (Weight Tying) میان امبدینگ و لایه‌ی نهایی، نه‌تنها از مصرف بیهوده‌ی ده‌ها میلیون پارامتر حافظه جلوگیری می‌کند، بلکه استانداردی پایه‌ای برای پیاده‌سازی شبکه‌های مدرن‌تر مانند خانواده‌ی LLaMA به شمار می‌رود. علاوه بر ساخت مدل، پایپ‌لاین تولید خودبازگشتی متن با ابزارهای کنترلی دما و هرس top-k را نیز از صفر می‌سازیم.

برای اجرای این درس به پردازنده‌ی ابری یا سخت‌افزارهای گران‌قیمت خارجی نیازی ندارید؛ تمام مراحل تعریف معماری، بررسی تعداد پارامترها و آزمایش نمونه‌برداری روی پیکربندی‌های کوچک‌تر به‌راحتی روی CPU سیستم شخصی شما قابل اجراست.

روی هم چیدن 12 بلوک، یک امبدینگ توکن، یک امبدینگ موقعیت (position embedding) یادگیری‌پذیر، یک LayerNorm نهایی و یک سر مدل زبانی متصل (tied). کل مدل 124 میلیون پارامتری GPT همین است. در این درس این قطعات را کنار هم می‌گذاریم تا به یک کلاس کامل و آماده‌به‌کار برسیم، پارامترها را می‌شماریم تا مطمئن شویم ابعاد مدل با نسخه‌ی مرجع 124M مو نمی‌زند و در نهایت با نمونه‌برداری چندجمله‌ای (multinomial sampling)، دمای نمونه‌برداری (temperature) و top-k متن تولید می‌کنیم.

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

• سرهم کردن بلوک ترنسفورمر از درس 34 و ساخت مدل کامل GPT: امبدینگ توکن، امبدینگ موقعیت، N بلوک، LayerNorm نهایی و سر مدل زبانی (language model head).
• پیاده‌سازی دقیق پیکربندی 124 میلیون پارامتری: اندازه‌ی واژگان 50257، طول زمینه 1024، بعد امبدینگ 768، 12 سر اتنشن و 12 لایه.
• اشتراک ماتریس وزن‌های سر مدل زبانی با امبدینگ توکن (Weight Tying) و توضیح این‌که چرا این کار در این مقیاس حدود 38 میلیون پارامتر صرفه‌جویی می‌کند.
• تولید متن از روی پرامپت با نمونه‌برداری چندجمله‌ای، تنظیم دما و هرس top-k، همراه با حفظ طول زمینه به کمک پنجره‌ی لغزان (sliding window).
• سنجش تعداد پارامترها و هزینه‌ی محاسباتی گذر رو به جلو (forward pass) در مقایسه با هدف 124M.

صورت مسئله

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

ابعاد و ساختار مدل هم بسیار مهم است. مدل مرجع GPT-2 Small دقیقاً با همین پیکربندی بالا 124 میلیون پارامتر دارد. این اعداد تصادفی نیستند: حاصل‌ضرب واژگان 50257 در بعد امبدینگ 768، اندازه‌ی جدول توکن‌هاست. طول زمینه 1024 در 768، جدول موقعیت‌ها را می‌سازد. 12 بلوک که هر کدام حدود 7 میلیون پارامتر دارند، روی‌هم‌رفته 84 میلیون پارامتر می‌شوند. سر نهایی مدل هم با اشتراک وزن‌ها، دوباره از همان جدول توکن‌ها استفاده می‌کند. این تکه‌ها را جمع بزنید، دقیقاً به 124 میلیون پارامتر می‌رسید. اگر مدلی بسازید که تعداد پارامترهایش با این مرجع هم‌خوانی نداشته باشد، مطمئن باشید گوشه‌ای از کار را اشتباه سیم‌کشی کرده‌اید.

مفهوم

شناسه‌های توکن به بردارهای توکن تبدیل می‌شوند. شناسه‌های موقعیت مکانی هم به بردارهای موقعیت بدل می‌گردند. این دو بردار با هم جمع شده و به درون دسته‌ی لایه‌ها فرستاده می‌شوند. بخش LayerNorm نهایی، تنها جزء بیرون از بلوک‌هاست که در تمام معماری‌های مدرن امروزی هم دست‌نخورده باقی مانده است. سر مدل زبانی (LM head) از همان ماتریس امبدینگ توکن دوباره استفاده می‌کند؛ معنای اشتراک ماتریس وزن‌ها دقیقاً همین است.

اشتراک ماتریس وزن‌ها (Weight tying)

امبدینگ توکن ابعادی معادل (vocab, d_model) دارد. سر مدل زبانی باید داده‌ها را از d_model دوباره روی vocab تصویر کند. این دو ماتریس ترانهاده‌ی (transpose) یکدیگرند. اشتراک این دو یعنی دقیقاً استفاده از یک تنسور پارامتر یکسان در دو جای مختلف. با واژگان 50257 و بعد مدلی 768، این ماتریس به‌تنهایی 38 میلیون پارامتر دارد. اگر آن‌ها را جدا نگه دارید، هزینه‌ی حافظه‌ی آن را دو بار می‌پردازید. وقتی آن‌ها را هم‌رسان (tied) می‌کنید، یک بار هزینه می‌دهید و از طرفی سیگنال گرادیان تمیزتری هم می‌گیرید، چون امبدینگ و سر خروجی هم‌زمان و هماهنگ با هم به‌روزرسانی می‌شوند.

امبدینگ موقعیت یادگیری‌پذیر است، نه سینوسی

مدل GPT-2 با امبدینگ موقعیت یادگیری‌پذیر (learned position embedding) کار می‌کند. جدول موقعیت، تنسوری با ابعاد (1024, 768) از پارامترهای یادگیری‌پذیر است. مدل در هر گذر رو به جلو، موقعیت‌های 0 تا T-1 را جست‌وجو (lookup) کرده و حاصل را به امبدینگ توکن اضافه می‌کند. این شیوه ساده‌ترین راهکار مدیریت موقعیت است (روش‌های دیگری مثل RoPE ،ALiBi و بایاس نسبی T5 گزینه‌های جایگزین هستند) و مدل مرجع 124M دقیقاً از همین روش استفاده می‌کند.

تولید متن: دما، top-k و نمونه‌برداری چندجمله‌ای

تولید متن به‌صورت خودبازگشتی (autoregressive) انجام می‌شود. مدل در هر گام، برای تک‌تک موقعیت‌ها لاجیت‌هایی روی کل واژگان برمی‌گرداند. شما فقط لاجیت‌های آخرین موقعیت را برمی‌دارید، آن را بر دما تقسیم می‌کنید، در صورت نیاز همه‌ی مقادیر را به‌جز k لاجیت برتر به منفی بی‌نهایت تبدیل می‌کنید (ماسک کردن)، با سافت‌مکس احتمال‌ها را به دست می‌آورید و در نهایت یک توکن از توزیع حاصل نمونه‌برداری می‌کنید.

سه پیچ تنظیم، سه رفتار کاملاً متفاوت می‌سازند: دمای نزدیک به صفر، رفتار مدل را به حالت حریصانه (greedy) تقلیل می‌دهد. دمای 1 توزیع طبیعی خود مدل را حفظ می‌کند. مقدار top-k برابر با 1 خروجی کاملاً حریصانه می‌دهد، در حالی که top-k برابر با 40 دنباله‌ی کم‌احتمال توزیع را فیلتر می‌کند. ترکیب این‌ها تفاوت زیادی در خروجی رقم می‌زند؛ در درس بعدی که به آموزش مدل اختصاص دارد، از فرایند تولید متن به عنوان یک سیگنال ارزیابی کیفی استفاده خواهیم کرد.

پیاده‌سازی از صفر

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

• کلاس دیتای class GPTConfig با مقادیر پیش‌فرض مدل 124M: شامل vocab_size=50257، context_length=1024، d_model=768، num_heads=12، num_layers=12، mlp_expansion=4، dropout=0.1، use_bias=True و weight_tying=True.
• کلاس class GPTModel شامل امبدینگ توکن، امبدینگ موقعیت، دراپ‌اوت امبدینگ، 12 بلوک TransformerBlock، لایه‌ی نهایی LayerNorm و یک lm_head که در صورت فعال بودن پرچم مربوطه، وزن‌هایش با امبدینگ توکن به اشتراک گذاشته می‌شود.
• تابع کمکی count_parameters که تعداد پارامترهای یکتا را می‌شمارد (تا اثر اشتراک وزن‌ها در شمارش نهایی درست محاسبه شود).
• تابع generate که دما، top-k، نمونه‌برداری چندجمله‌ای و پنجره‌ی لغزان برای کنترل زمینه را مدیریت می‌کند.
• یک نمونه‌ی اجرایی (demo) که مدل را می‌سازد، تعداد پارامترهایش را در کنار مدل مرجع 124M چاپ می‌کند و یک توالی کوتاه از روی پرامپتی ثابت تولید می‌کند تا کل پایپ‌لاین را از ابتدا تا انتها نشان دهد.

اجرای کد:

python3 code/main.py

خروجی: تعداد پارامترها در مقایسه با مرجع 124M، شناسه‌های توکن تولیدشده از پرامپت ثابت و تأییدیه‌ای مبنی بر این‌که در صورت فعال بودن اشتراک وزن‌ها، سر مدل زبانی و امبدینگ توکن از حافظه‌ی مشترک استفاده می‌کنند.

برای این‌که اجرای آزمایشی سریع بماند، اسکریپت یک پیکربندی بسیار کوچک (d_model=64, num_layers=2) را هم از ابتدا تا انتها اجرا کرده و توالی توکن‌های تولیدشده را در خروجی نمایش می‌دهد. پیکربندی 124M ساخته می‌شود، اما فقط تعداد پارامترهای آن محاسبه شده و یک گذر رو به جلو روی آن اجرا می‌گردد.

استک فنی

• کتابخانه‌ی torch برای محاسبات تانسوری، مشتق‌گیری خودکار (autograd) و ساختار ماژول‌ها.
• فایل code/main.py ساختار بلوک ترنسفورمر درس 34 را مجدداً به‌صورت محلی پیاده می‌کند.

الگوهای محیط عملیاتی

سه تفاوت فنی، مرز میان مدلی را که فقط اجرا می‌شود با مدلی که واقعاً آماده‌ی استقرار در محیط عملیاتی (پروداکشن) است مشخص می‌کنند:

مقداردهی اولیه‌ی کوچک برای ماتریس‌های تصویر پسماند. خروجی تصویر اتنشن و لایه‌ی خطی دوم در MLP، هر دو مستقیماً وارد یک جمع پسماند (residual add) می‌شوند. اگر این وزن‌ها را با همان انحراف معیاری مقداردهی اولیه کنید که بقیه‌ی لایه‌های خطی دارند، جریان پسماند (residual stream) با افزایش عمق شبکه مدام بزرگ‌تر می‌شود و لایه‌ی نهایی LayerNorm را وارد ناحیه‌ای ناپایدار و بیش‌ازحد داغ می‌کند. اگر انحراف معیار این دو تصویر را با ضریب 1 / sqrt(2 * num_layers) مقیاس کنید، جریان پسماند در تمام 12 لایه در محدوده‌ای منطقی و پایدار باقی می‌ماند.

کش کردن تنسور شناسه‌های موقعیت، نه محاسبه‌ی مجدد آن. اجرای torch.arange(T) در هر گذر رو به جلو، حافظه‌ی جدیدی تخصیص می‌دهد. این تنسور را یک بار در __init__ برای حداکثر طول زمینه بسازید، در هر فراخوانی T درایه‌ی اول آن را برش (slice) بزنید و از رفت‌وآمد بیهوده به تخصیص‌دهنده‌ی حافظه جلوگیری کنید.

اشتراک وزن‌ها در سطح پارامتر، نه صرفاً کپی کردن. نوشتن عبارت lm_head.weight = token_embedding.weight خود تنسور را به اشتراک می‌گذارد؛ در حالی که کپی کردن این کار را نمی‌کند. بهینه‌ساز (optimizer) باید فقط یک پارامتر را به‌روزرسانی کند و گراف مشتق‌گیری خودکار باید فقط یک بار تجمع گرادیان داشته باشد. اگر داده‌ها را صرفاً کپی کنید، با گذشت زمان سر خروجی از امبدینگ فاصله می‌گیرد و عملاً هیچ سودی از اشتراک وزن‌ها نمی‌برید.

کاربرد عملی

• کلاس مدلی که در این درس پیاده شد، دقیقاً همان ساختاری را دارد که در درس بعدی آموزش می‌دهیم.
• با جایگزینی امبدینگ موقعیت یادگیری‌پذیر با RoPE، بدون نیاز به دست زدن به ساختار بلوک‌ها یا سر خروجی، به معماری خانواده‌ی LLaMA می‌رسید.
• جایگزینی تابع فعال‌سازی GELU با SiLU و تعویض LayerNorm با RMSNorm، باقی تغییرات خانواده‌ی LLaMA را برایتان تکمیل می‌کند.
• تابع تولید متن این درس با هر منبعی از لاجیت‌ها کار می‌کند و مختص این مدل نیست. می‌توانید در درس 37 لاجیت‌ها را از فایل مدل ازپیش‌آموزش‌دیده‌ی GPT-2 بخوانید و از همین حلقه‌ی تولید متن استفاده کنید.

تمرین‌ها

1. حالت اشتراک وزن‌ها را بین سر مدل زبانی و امبدینگ توکن غیرفعال کنید و پارامترها را دوباره بشمارید. بررسی کنید که اختلاف دقیقاً حاصل 50257 در 768 یعنی 38 میلیون پارامتر باشد.
2. امبدینگ موقعیت یادگیری‌پذیر را با یک جدول سینوسی جایگزین کنید که در زمان ساخت مدل محاسبه می‌شود. تأیید کنید که گذر رو به جلو هنوز کار می‌کند و تعداد پارامترها 786432 عدد کاهش می‌یابد.
3. یک پرچم greedy=True به تابع تولید متن اضافه کنید که از مرحله‌ی نمونه‌برداری عبور کند و مستقیماً مقدار argmax را بردارد. مطمئن شوید خروجی در اجراهای مختلف کاملاً قطعی و تکرارپذیر است.
4. پارامتر تنظیمی جدیدی به نام repetition_penalty اضافه کنید که لاجیت هر توکن موجود در پرامپت یا متن تولیدشده تا این لحظه را پیش از سافت‌مکس بر مقداری ثابت تقسیم کند. روی یک پرامپت ثابت نشان دهید که مقادیر بزرگ‌تر از 1 چطور دفعات تکرار توکن‌ها را در خروجی کاهش می‌دهند.
5. نمونه‌برداری top_p (هسته‌ای یا nucleus) را در کنار top_k پیاده‌سازی کنید. یک بررسی دوخطی اضافه کنید که نشان دهد مجموع احتمال‌های توکن‌های باقی‌مانده از top_p بیشتر باشد.

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

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

• بخش 19 درس 34: بررسی بلوکی که در این مدل روی هم چیده شد.
• بخش 19 درس 36: حلقه‌ی آموزش مدل با تابع زیان (loss) انتروپی متقاطع.
• بخش 19 درس 37: بارگذاری وزن‌های ازپیش‌آموزش‌دیده‌ی GPT-2 روی همین معماری دقیق.
• بخش 7 درس 07 (مدل‌سازی علّی زبان GPT): مبانی ریاضیاتی پیش‌بینی توکن بعدی.
• بخش 10 درس 04 (پیش‌آموزش Mini GPT): فرایند اولیه‌ی آموزش روی همین معماری.