‏Cerebras چطوری یه مدل ۳ تریلیون پارامتری رو اجرا می‌کنه

تو این پست، تیم Cerebras به یه سؤال پرتکرار جواب داده: وقتی هر ویفر WSE-3T فقط ۴۴ گیگابایت حافظه‌ی روی چیپ (SRAM) داره و مدل‌های تریلیونی فقط برای وزن‌هاشون ترابایت‌ها حافظه می‌خوان، چطوری اجراشون می‌کنن؟ مثالشون مدل متن‌باز Kimi K3 هست با حدود ۳ تریلیون پارامتر، ۶۹ لایه‌ی linear-attention و ۲۴ لایه‌ی dense. این مدل تو Terminal-Bench 2.1 امتیاز ۸۸.۳٪ و تو BrowseComp امتیاز ۹۱.۲٪ گرفته.

فرض‌های محاسبه اینان: میانگین طول درخواست ۲۵۶ هزار توکن، کوانتیزه کردن وزن‌ها به ۴ بیت و KV cache هشت‌بیتی. با ۳ تریلیون پارامتر ۴ بیتی و تقسیم روی ۹۳ لایه، هر لایه حدود ۱۶ گیگابایت می‌شه.

حالا نوع لایه مهمه. لایه‌های dense به KV cache نیاز دارن، یعنی باید key و value همه‌ی توکن‌های هر درخواست فعال رو نگه دارن. برای همین هر لایه‌ی dense یه ویفر جدا می‌گیره و بقیه‌ی حافظه (حدود ۲۸ گیگ) برای کش می‌مونه. لایه‌های linear-attention فقط یه state کوچیک و ثابت برای هر درخواست دارن، پس دو تاشون (۳۲ گیگ) رو یه ویفر جا می‌شن. نتیجه: ۳۵ ویفر برای لایه‌های linear و ۲۴ تا برای dense، یعنی ۵۹ ویفر.

اجرا به شکل pipeline هست. چون هر لایه کامل روی یه چیپ جا می‌شه، وزن‌ها هیچ‌وقت جابه‌جا نمی‌شن و فقط activationها (یعنی خروجی میانی هر لایه) از طریق RDMA روی Ethernet به ویفر بعدی می‌رن، با حدود ۲ میکروثانیه برای هر پرش. ویفر ۱ لایه‌های ۱ و ۲ رو اجرا می‌کنه، نتیجه رو به ویفر ۲ می‌ده و همین‌طور تا آخر که توکن بعدی تولید بشه.

نویسنده یه تصور اشتباه رایج رو هم رد می‌کنه: این‌که ویفرها منتظر هم بیکار می‌مونن. در واقع هر مرحله هم‌زمان روی یه درخواست متفاوت کار می‌کنه، مثل خط مونتاژ. هر مرحله می‌تونه یه micro-batch (یه گروه کوچیک درخواست) هم پردازش کنه، پس تعداد درخواست‌های در جریان تقریباً تعداد مرحله‌ها ضربدر اندازه‌ی micro-batch هست. micro-batch کوچیک‌تر یعنی سرعت بیشتر برای هر درخواست، و بزرگ‌ترش یعنی throughput کل بیشتر.

سقف درخواست‌های هم‌زمان رو حافظه‌ی آزاد برای KV cache روی ویفرهای dense تعیین می‌کنه. به گفته‌ی نویسنده، معیار درست برای مقایسه‌ی سیستم‌های inference تعداد توکن در ثانیه به ازای هر مگاوات برقه، اون هم با سرعت تولید یکسان. همین روش برای مدل‌های بزرگ‌تر از ۳ تریلیون هم کار می‌کنه، فقط ویفر بیشتری لازمه.

نکات کلیدی:

• هر ویفر WSE-3T حدود ۴۴ گیگابایت SRAM روی چیپ داره
• Kimi K3 با وزن ۴ بیتی حدود ۱۶ گیگ برای هر لایه از ۹۳ لایه‌ش لازم داره
• هر لایه‌ی dense یه ویفر می‌گیره، ولی دو لایه‌ی linear رو یه ویفر جا می‌شن
• کل مدل روی ۵۹ ویفر پخش می‌شه و سرعت ۱۰۰۰ تا ۲۰۰۰ توکن در ثانیه می‌گیره
• انتقال activation بین ویفرها با RDMA حدود ۲ میکروثانیه طول می‌کشه
• micro-batch کوچیک سرعت هر درخواست رو بالا می‌بره، بزرگش throughput کل رو