مدل‌های متراکم در برابر مدل‌های پراکنده در معماری تراشه‌های تخصصی

تصور کنید منطق ریاضی و توزیع داده‌های یک مدل در تست‌های CPU به‌طور کامل درست عمل کند، اما به محض اجرا روی سخت‌افزار، خروجی مدل کاملاً خالی باشد. این شکست فنی دقیقاً همان نقطه‌ای بود که در مسیر انتقال مدل ترکیب خبره‌ها (Mixture-of-Experts یا MoE) Gemma-4 26B-A4B به سخت‌افزار AWS Inferentia2 رخ داد. طبق مستندات منتشرشده توسط توسعه‌دهنده xbill9، ریشه مشکل در نحوه مدیریت توزیع تانسورها در زمان اجرا بود که نیاز به یک مکانیزم خاص برای شناسایی رتبه‌بندی (SPMD-rank) داشت تا تداخل در فعال‌سازی‌های زمان اجرا برطرف شود. این پروژه قصد داشت ثابت کند که مدلی با این پیچیدگی می‌تواند به‌طور کامل روی پشته (Stack) Neuron ردیابی (Trace) و اجرا شود.

این تلاش بخشی از یک پروژه گسترده‌تر برای انتقال تمام خانواده Gemma-4 — شامل نسخه‌های متراکم E2B، E4B، 12B و 31B — به زیرساخت Inferentia2 است. همان‌طور که در تحلیل‌های پیشین ما درباره مدل‌های متراکم اشاره کردیم، نسخه 26B-A4B چالش متفاوتی دارد؛ زیرا این مدل «پراکنده» است و از ۱۲۸ خبره با مسیریابی Top-8 استفاده می‌کند. در دنیای شتاب‌دهنده‌های سخت‌افزاری، عملیات پراکنده به‌شدت دشوار هستند زیرا تبدیل آن‌ها به گراف‌های استاتیک برای کامپایلر پیچیدگی‌های زیادی دارد و اغلب نیازمند ایجاد تعادل بین کارایی محاسباتی و سازگاری با کامپایلر است. به نقل از گزارش‌های فنی، برای مدل google/gemma-4-26B-A4B-it، این چالش دوچندان بود چون تا پیش از این هیچ مدل MoE از خانواده Gemma-4 روی پشته AWS ردیابی نشده بود.

تلهٔ حافظه و معماری

نام‌گذاری مدل 26B-A4B می‌تواند در تخمین حافظه گمراه کند. عبارت «A4B» به معنای حدود ۴ میلیارد پارامتر فعال برای هر توکن (Token) است. اگرچه این موضوع محاسبات در مرحله پیش‌رو (Forward Pass) را کم می‌کند، اما اثر مستقیمی روی حجم حافظه ندارد و اثر آن را کاهش نمی‌دهد. برای پردازش یک توکن، تمام ۱۲۸ خبره — که مجموعاً حدود ۴۹ گیگابایت وزن‌ها هستند — باید در حافظه پهنای‌باند بالا (HBM) مستقر باشند. مسیریابی Top-8 فقط تعداد عملیات اعشاری (FLOPs) مورد نیاز برای محاسبه را کاهش می‌دهد، نه فضای ذخیره‌سازی مورد نیاز برای پارامترها را.

به همین دلیل، این مدل به ۱۹۲ گیگابایت HBM موجود در نمونه‌های inf2.24xlarge (که دارای ۱۲ هسته NeuronCores است) نیاز دارد. این موضوع، نیازهای حافظه‌ای این مدل را در همان دسته مدل متراکم 31B قرار می‌دهد؛ مدلی که از نظر تعداد کل پارامترها دو برابر بزرگ‌تر است اما سقف HBM یکسانی دارد. این یک محدودیت حیاتی برای کسانی است که می‌خواهند مدل را روی سخت‌افزارهای کوچک‌تر اجرا کنند؛ چراکه تعداد «پارامتر فعال» یک عدد مربوط به توان محاسباتی است، نه عددی برای تعیین حجم حافظه.

پیچیدگی مسیر دوگانه FFN

معماری این مدل فراتر از یک جایگزینی ساده MLP به MoE است. با بررسی ماژول transformers-5.13 روی یک دستگاه Meta پیش از نوشتن کدهای انتقال، مشخص شد که شبکه پیش‌خور (Feed-Forward Network) در واقع یک سیستم مسیر-دوگانه (Dual-path) است. هر یک از ۳۰ لایه، یک MLP متراکم مشترک را به‌طور موازی با MoE ۱۲۸ خبره اجرا می‌کند.

این دو مسیر ترکیب شده و از چهار لایه نرمال‌سازی (Layernorm) با ترتیب زیر عبور می‌کنند:

• باقی‌مانده (Residual) پس از مرحله توجه (Attention)
• مسیر متراکم: pre_feedforward_layernorm $
  ightarrow$ mlp (2112) $
  ightarrow$ post_feedforward_layernorm_1
• مسیر MoE: pre_feedforward_layernorm_2 $
  ightarrow$ router $
  ightarrow$ 128 experts (704) $
  ightarrow$ post_feedforward_layernorm_2
• مجموع پنهان: hidden = dense + moe
• گذر نهایی: post_feedforward_layernorm $
  ightarrow$ + residual $
  ightarrow$ × layer_scalar

تحلیل عمیق مسیریابی و خبره‌ها

بخش Gemma4TextRouter دارای RMSNorm، یک مقدار مقیاس (Scale) و یک مقیاس اختصاصی برای هر خبره (per_expert_scale) است. مکانیزم مسیریابی با استفاده از یک تابع softmax(128, fp32)، هشت خبره برتر (Top-8) را انتخاب کرده و سپس آن‌ها را دوباره نرمال‌سازی کرده و در per_expert_scale ضرب می‌کند تا وزن نهایی مشارکت هر خبره تعیین شود.

بخش Gemma4TextExperts از یک لایه ادغام‌شده gate_up_proj با ابعاد [128, 1408, 2816] و یک down_proj با ابعاد [128, 2816, 704] استفاده می‌کند. در پیاده‌سازی استاندارد Hugging Face، مرحله پیش‌رو یک حلقه جمع‌آوری/پراکندگی پراکنده (Sparse gather/scatter) است که از torch.where و index_add_ برای محاسبه تنها خبره‌های انتخاب‌شده استفاده می‌کند.

حل مسئلهٔ ردیابی (Traceability)

عملیات مسیریابی MoE استاندارد بر اساس اینکه کدام خبره برای هر توکن انتخاب شود، یک حلقه جمع‌آوری/پراکندگی پراکنده است. با این حال، این عملیات وابسته به داده (Data-dependent) نمی‌تواند به یک گراف HLO استاتیک در Neuron ردیابی شود. برای دور زدن این مشکل، در این پورت از تکنیک «وزن‌دهی متراکم» (Dense-weighting) استفاده شد: سیستم تمام ۱۲۸ خبره را برای هر توکن محاسبه می‌کند و سپس آن‌ها را بر اساس خروجی مسیریاب وزن‌دهی می‌کند.

از آنجا که خبره‌های انتخاب‌نشده وزن صفر می‌گیرند ($0 \times expert(x) = 0$)، نتیجه ریاضی این کار با مسیریابی Top-8 کاملاً یکسان است. اگرچه این روش از نظر FLOPs اسراف است — زیرا ۱۲۸ خبره را فقط برای استفاده از ۸ مورد محاسبه می‌کند — اما یک توالی با شکل ثابت (Fixed-shape) از ضرب ماتریس‌ها ایجاد می‌کند که کامپایلر آن را به‌شدت می‌پسندد. این اجازه می‌دهد خبره‌ها را به‌صورت لایه‌های خطی موازی استاندارد از طریق ModelBuilder توزیع کرد:

• ColumnParallelLinear: لایه gate_up_proj [128, 1408, 2816] به ابعاد [180224, 2816] تغییر شکل می‌یابد. در تنظیمات TP=8، رتبه r خبره‌های اندیس 16*r تا 16*r+15 را دریافت می‌کند.
• RowParallelLinear: لایه down_proj [128, 2816, 704] به ابعاد [2816, 90112] تغییر شکل یافته و از توزیع ورودی و عملیات all-reduce استفاده می‌کند.

با جایگزینی self.experts با یک پیاده‌سازی سفارشی DenseExperts و در حالی که مسیر پیش‌رو لایه دوگانه و مسیریاب اصلی بدون تغییر باقی ماندند، مدل ردیابی‌پذیر شد. مسیریاب در تمام رتبه‌ها تکثیر می‌شود و top_k_weights آن پیشاپیش شامل نرمال‌سازی و مقیاس مورد نیاز هر خبره است.

باگ «خروجی تهی»

قبل از استقرار روی سخت‌افزار، تمامی محاسبات و توزیع TP=8 روی CPU در برابر Hugging Face تست شدند که منجر به حداکثر اختلاف (MAXDIFF) حدود $2e-6$ و شباهت کسینوسی ۱.۰ شد. اما در اولین اجرای واقعی، اتفاق عجیبی افتاد: مدل کامپایل شد — و اولین ردیابی MoE از نوع خود با ۳۰ لایه روی Neuron ثبت شد (MB_TRACED) — اما خروجی مدل یک رشته خالی ('') بود. مدل بلافاصله توکن پایان-نوبت را صادر می‌کرد.

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

توسعه‌دهنده برای وزن‌دهی خبره‌ها در هر رتبه از scatter_to_tensor_model_parallel_region استفاده کرده بود. این تابع رتبه خود را با استفاده از get_tensor_model_parallel_rank() شناسایی می‌کند که یک عدد صحیح پایتونی است. چون ModelBuilder یک رتبه را کامپایل کرده و سپس آن گراف را برای تمام ۸ رتبه در زمان اجرا تکثیر می‌کند، رتبه ۰ در گراف «باکد» (Bake) شده بود (یعنی برش Wd[:, 0:16] در گراف ثبت شده بود). در زمان اجرا، تمام رتبه‌ها خبره‌های ۰ تا ۱۵ را وزن‌دهی می‌کردند، در حالی که لایه‌های gate_up آن‌ها در حال محاسبه خبره‌های متفاوتی بودند (مثلاً رتبه ۱ در حال محاسبه خبره‌های ۱۶ تا ۳۱ بود). این عدم همراستایی کامل، خروجی مدل را به زباله تبدیل می‌کرد و باعث می‌شد مدل با اطمینان کامل نوبت را به پایان برساند.

اصلاح و بهینه‌سازی

راه حل، پیاده‌سازی دقیق مکانیزمی بود که خود ماژول MoE در NxD استفاده می‌کند: enable_spmd_rank. این کار شامل ثبت یک ماژول SPMDRank است — یک پارامتر int32 با اندازه [1] که از طریق arange(TP) در چک‌پوینت بارگذاری می‌شود — تا هر رتبه شماره منحصربه‌فرد خود را دریافت کند. سپس کد به‌روز شد تا از نسخه آگاه از رتبه در زمان اجرا استفاده کند:

Wl = scatter_to_process_group_spmd(Wd, 1, self.spmd_rank.get_rank())

پس از این تغییر، خروجی سخت‌افزار دقیقاً با مرجع CPU تطبیق یافت (SEQ_MATCH True) و زمان پیش‌پر (Prefill) به ۷۷ میلی‌ثانیه رسید.

در ادامه، برای بهینه‌سازی خبره‌ها که حدود ۹۳٪ از وزن‌ها را تشکیل می‌دهند، یک نسخه int8 ساخته شد. این کار با استفاده از کوانتش متقارن در هر کانال (per-channel symmetric quantization) از طریق QuantizedColumnParallel و QuantizedRowParallel انجام شد. این مرحله نیازمند دو اصلاح در autograd مربوط به NxD بود: فراخوانی .clone() روی خروجی لایه down و تغییر تابع scale_dequantize به گونه‌ای که خارج از مکان (out-of-place) عمل کند، زیرا پیش‌تر عملیات x *= scale را به‌صورت in-place روی یک view ارتباطی نامتقارن انجام می‌داد.

نتایج به شرح زیر بود:

• دقت: نسخه int8 از نظر عددی کامل بود و توکن به توکن با نسخه fp32 یکسان بود (SEQ_MATCH True).
• اندازه: حجم فایل کامپایل شده (neff) از ۶۴.۶ گیگابایت به ۴۱.۸ گیگابایت کاهش یافت که نشان‌دهنده صرفه‌جویی دقیق ۲۲.۸ گیگابایتی در بخش خبره‌ها است.
• محدودیت سخت‌افزاری: مدل همچنان به نمونه 24xlarge نیاز دارد. روی باکس‌های ۲ هسته‌ای کوچک‌تر (inf2.8xlarge/inf2.xlarge با ۳۲ گیگابایت HBM و TP=2)، حجم خبره‌ها به ۱۱.۴ گیگابایت در هر رتبه می‌رسد. با این حال، لایه lm_head، گراف کامپایل‌شده و رزروهای زمان اجرای Neuron، نیاز حافظه را حدود ۳ تا ۴ گیگابایت بالاتر از محدودیت ۱۶ گیگابایتی هر هسته می‌برد. از آنجا که نسخه ۴ هسته‌ای وجود ندارد، جای‌گذاری مدل در باکس ۲ هسته‌ای مستلزم استفاده از خبره‌های fp4 است.

خلاصه موانع فنی

• عملیات وابسته به داده: عملیاتی که وابسته به داده باشد ردیابی نمی‌شود. راه حل این است که شکل آن را ثابت (Fixed-shape) کنیم. رویکرد «تمام خبره‌ها به‌صورت متراکم»، FLOPs را فدای گراف استاتیک و صحت مطلق می‌کند.
• رتبه‌های زمان اجرا: لایه‌های خطی موازی که به‌طور خودکار توزیع شده‌اند ایمن هستند زیرا وزن‌ها در زمان بارگذاری برش می‌خورند. اما توزیع‌های تانسوری مستقل روی فعال‌سازی‌های زمان اجرا باید از نسخه SPMD-rank استفاده کنند تا از «باکد شدن» داده‌های رتبه ۰ جلوگیری شود.
• بسته‌بندی: ایجاد یک حلقه آینه‌ای (Mirroring) به S3 در هنگام torch.jit.save می‌تواند باعث آپلود ناقص شود. در یک مورد، یک فایل neff با حجم ۶۴.۶ گیگابایت به ۲۱.۹ گیگابایت کاهش یافت و خطاهای PytorchStreamReader ... failed finding central directory ایجاد کرد. راه حل این است که پیش از آپلود نهایی، صحت فایل با zipfile.is_zipfile(...) تایید شود.

این انتقال ثابت می‌کند که اگرچه MoE محاسبات (FLOPs) را کاهش می‌دهد، اما حافظه (HBM) را نجات نمی‌دهد. مرز بعدی این پروژه پیاده‌سازی خبره‌های fp4 است — که نیازمند ادغام عمیق microscaling در NxD با uint16‌های بسته‌بندی شده و هوک‌های from_float است — تا در نهایت مدل روی نمونه‌های ۲ هسته‌ای Inferentia2 جای بگیرد.

آرتیفکت‌ها و منابع

• Docker Hub: نسخه‌های xbill9/gemma4-optb-26b (bf16) و xbill9/gemma4-optb-26b-int8
• Hugging Face: مدل xbill9/gemma-4-26B-A4B-it-inferentia2 (در دو حالت bf16 و int8)
• دستورالعمل: فایل‌های tp_mb_moe.py (برای DenseExperts و SPMDRank scatter) و tp_mb_moe_int8.py