اجرای مدل‌های Gemma ۴ روی یک تراشه TPU v۵e با روش QAT

آموزش اجرای مدل‌های Gemma ۴ روی تراشه TPU v۵e با روش QAT و بهینه‌سازی وزن‌های int۴ برای افزایش سرعت پردازش و کاهش حافظه مصرفی در Google Cloud.

اجرای مدل‌های زبانی بزرگ روی سخت‌افزارهای محدود، چالشی است که با استفاده از وزن‌های QAT (Quantization-Aware Training) و بهینه‌سازی فرمت ذخیره‌سازی، به شکلی عملیاتی حل می‌شود. با بازآرایی این وزن‌ها در فرمت‌های int۴ و int۸ و استفاده از پچ‌های اختصاصی برای vLLM، می‌توان مدل‌هایی با ابعاد مختلف را روی یک تراشه Google Cloud TPU v۵e مستقر کرد. این رویکرد به کاربران اجازه می‌دهد بدون نیاز به خوشه‌های بزرگ، از توان پردازشی مدل‌های ۱۲B و حتی ۲۶B بهره‌مند شوند.

پیش‌نیازهای فنی برای راه‌اندازی در بستر ابری

برای شروع این فرآیند، ابتدا باید یک پروژه در Google Cloud با سهمیه TPU v۵e در منطقه us-west۴-a ایجاد کنید. همچنین داشتن یک مخزن ذخیره‌سازی ابری برای نگهداری چک‌پوینت‌ها و نتیجه‌های خروجی ضروری است. فرآیند اصلی شامل کپی‌برداری از مخزن gemma۴-dev و تنظیم توکن Hugging Face در Secret Manager است. پس از آماده‌سازی این زیرساخت، می‌توانید با استفاده از اسکریپت‌های موجود، وزن‌های QAT را به فرمت‌های بهینه تبدیل کنید.

گام‌های بازآرایی وزن‌ها

برای تبدیل وزن‌ها، باید از اسکریپت‌های اختصاصی جهت بازآرایی گروه‌های وزنی استفاده کنید تا جایگاه آن‌ها در شبکه ۱۶-سطحی حفظ شود. این کار باعث می‌شود مدل بدون نیاز به گرد کردن مجدد وزن‌ها، در قالب‌های فشرده‌ای قرار گیرد که vLLM به‌راحتی آن‌ها را پردازش می‌کند. برای نمونه، در مدل‌های ۱۲B، این روش باعث می‌شود وزن‌ها در فضای 11.31 گیگابایتی جای گیرند که با ظرفیت 15.75 گیگابایتی حافظه HBM در تراشه v۵e کاملاً سازگار است.

چه مدل‌هایی روی یک تراشه v۵e اجرا می‌شوند؟

نتایج نشان می‌دهد مدل‌های Gemma ۴ از نسخه E۲B تا ۱۲B با عملکردی هم‌تراز با مدل‌های bf۱۶ عمل می‌کنند. در مدل ۱۲B، سرعت خروجی به ۶۷۵ توکن در ثانیه می‌رسد که برای کاربردهای سازمانی بسیار کارآمد است. با این حال، مدل‌های بزرگ‌تر مانند ۲۶B محدودیت‌هایی در زمینه طول پنجره متن (Context) دارند و پس از ۲۱۷۶ توکن دچار محدودیت می‌شوند. هوش مصنوعی نهال با ارائه زیرساخت‌های مشابه، به سازمان‌ها کمک می‌کند تا این مدل‌ها را در محیط‌های داخلی و بدون پیچیدگی‌های فنی مستقر کنند.

نکات کلیدی

• مدل‌های Gemma ۴ تا اندازه ۱۲B با استفاده از فرمت int۸ و int۴ روی یک تراشه v۵e با کارایی کامل اجرا می‌شوند.
• استفاده از وزن‌های QAT، دقت مدل را تا 2.4 واحد نسبت به خروجی‌های استاندارد ۴-بیتی گوگل بهبود می‌دهد.
• محدودیت حافظه HBM در تراشه v۵e اجازه اجرای مدل‌های بزرگ‌تر از ۲۶B را نمی‌دهد.
• پچ‌های vLLM برای این سخت‌افزار شامل روش‌های W۸A۸ برای ضرب ماتریسی سریع‌تر و جداول جای‌گذاری (Embedding) فشرده است.

منبع: dev.to