چطور fal مدل تصویرسازی Ideogram V4 رو ۶ برابر سریعتر کرد
هوش مصنوعیخلاصهٔ کاملتر
تیم fal تو این پست فنی توضیح میده که چطور سرویسدهی مدل تصویرسازی Ideogram V4 رو خیلی سریعتر کرده: یه تولید تصویر تو رزولوشن ۱K که قبلاً ۲.۷۵ ثانیه طول میکشید، حالا تو ۰.۴۴ ثانیه انجام میشه، یعنی ۶ برابر سریعتر و بدون افت کیفیت دیدنی. چارچوب فکریشون سادهست: هزینهٔ یه تصویر برابره با تعداد forward passهایی که ترنسفورمر diffusion اجرا میکنه، ضربدر هزینهٔ هر pass. برای همین رو هر دو طرف این معادله کار کردن.
قدم اول اجرای ترنسفورمر تو فرمت NVFP4 هست؛ یه فرمت اعشاری ۴ بیتی که رو سختافزار Blackwell بهصورت بومی پشتیبانی میشه و matmulها رو خیلی ارزونتر میکنه. ولی به گفتهٔ نویسنده، ارزونشدنِ matmul وقتی به درد میخوره که واقعاً وقت اونجا صرف بشه؛ تو یه ترنسفورمر diffusion بعد از هر GEMM بزرگ یه عملیات کوچیک (مثل RMSNorm یا gated-SiLU) میاد که اگه fuse نشه، خروجی رو مجبور میکنه یه رفتوبرگشت کامل تا حافظهٔ HBM بزنه و همون سرعتی که بهدست اومده رو میخوره.
برای همین اونها این عملیاتهای کوچیک رو با خودِ GEMM ادغام کردن؛ به این کار epilogue fusion میگن. یعنی عملیات درست قبل از ذخیرهشدنِ نتیجه رو HBM اعمال میشه تا خروجی فقط یهبار به حافظه بره. برای RMSNorm با head_dim=256 این کار سخت بود، چون یه head تو یه fragment جا نمیشه؛ راهحلشون یه زمانبندی «بازدید دوباره» بود که نیمهٔ اول head رو تو رجیسترها نگه میداره تا بعد از رسیدن نیمهٔ دوم نرمالسازی رو تموم کنه.
الگوی رایج دیگه gated-SiLU (یا SwiGLU) هست که یه GEMM خروجی میده و بعد اون رو نصف میکنه و در هم ضرب میکنه:
C = A @ B
output = SiLU(C[:, :N//2]) * C[:, N//2:]برای اینکه دو نیمه تو fragmentهای جدا نیفتن، ستونهای ماتریس B رو موقع بستهبندی وزنها جابهجا میکنن تا زوجها کنار هم بشینن و یه visitor سفارشی همونجا SiLU و ضرب رو انجام بده.
اما FP4 یه مشکل کیفیت داشت: رنگها، مخصوصاً قرمزها، بیرمق و کماشباع میشدن. چون مدل diffusion تصویر رو تو چند قدم پشتسرهم denoise میکنه، یه خطای کوچیک تو یه قدم، قدم بعدی رو بدتر میکنه و خطا مثل گلولهبرفی بزرگ میشه. اصلاح رنگ بعد از تولید هم جواب نداد، چون خطا تو خودِ latentها پخته شده بود. راهحل تکنیکی به اسم quantization-aware distillation (QAD) بود: یه teacher با دقت کامل bf16 به یه student کوانتیزهشده در FP4 آموزش میده، و گرادیانها با کمک straight-through estimator از میون رُندکردنِ FP4 عبور میکنن تا وزنها یاد بگیرن با نویز کوانتیزاسیون کنار بیان.
نویسنده میگه یه قدم فراتر از مقالهٔ اصلی QAD هم رفتن: مکانیزم CFG (که معمولاً ترنسفورمر رو تو هر قدم دو بار اجرا میکنه) رو توی همون distillation حل کردن. حالا teacher هر دو شاخه رو اجرا میکنه ولی student فقط شاخهٔ شرطی رو یاد میگیره، پس موقع سرویسدهی بهجای دو pass فقط یه pass لازمه.
بزرگترین اهرم اما timestep distillation بود: کمکردن تعداد قدمهای denoising. برای این کار بهجای رگرسیون ساده از خانوادهٔ روشهای Distribution Matching Distillation (DMD) استفاده کردن که بهجای تطبیق نمونهبهنمونه، توزیعِ تصاویر student رو با توزیع teacher تطبیق میده. اول بدون GAN یه student چندقدمیِ خوب ساختن و بعد یه discriminator اضافه کردن تا آخرین ذره وضوح و جزئیات رو برگردونه. نتیجهٔ نهایی یه مدل چندقدمی، تکشاخه و ۴ بیتیه که همون خروجی مدل کامل رو با کسری از محاسبات میده.
نکات کلیدی:
- تولید تصویر Ideogram V4 تو ۱K از ۲.۷۵ به ۰.۴۴ ثانیه رسیده، ۶ برابر سریعتر
- اجرای ترنسفورمر تو FP4 روی سختافزار Blackwell بههمراه epilogue fusion، هر pass رو ارزون میکنه
- QAD کیفیت و رنگی که FP4 از دست میداد رو با آموزش teacher به student برمیگردونه
- فولد کردن CFG تعداد شاخهها رو نصف میکنه
- timestep distillation با روشهای خانوادهٔ DMD تعداد قدمهای denoising رو به چندتا کاهش میده
