‏سرعت‌بخشی به آموزش تقویتی مدل‌های زبانی با Speculative Decoding

آموزش تقویتی (Reinforcement Learning) مدل‌های زبانی مرزی (frontier LLMs) این روزا به یه بخش حیاتی از چرخه‌ی توسعه‌ی هوش مصنوعی تبدیل شده. اما یه چالش جدی وجود داره: مرحله‌ی تولید rollout که مدل باید توکن‌ها رو یکی‌یکی و به شیوه‌ی خودرگرسیو (autoregressive) تولید کنه، خیلی کند و پرهزینه‌ست. این گلوگاه باعث میشه کل فرآیند post-training کُند بشه و اسکیل‌پذیری رو محدود کنه.

روش‌های رایج برای حل این مشکل معمولاً رژیم rollout یا بهینه‌سازی رو تغییر می‌دن؛ مثلاً از off-policy execution یا replay استفاده می‌کنن. اما این رویکردها اغلب توزیع خروجی مدل اصلی رو تغییر می‌دن و ممکنه کیفیت یادگیری رو تحت‌تأثیر قرار بدن.

محققان در این مقاله پیشنهاد می‌دن از Speculative Decoding به عنوان یه primitive شتاب‌دهنده‌ی بدون افت (lossless) استفاده بشه. ایده‌ی اصلی Speculative Decoding اینه که یه مدل پیش‌نویس (draft model) کوچیک‌تر چند توکن رو با سرعت بیشتری پیشنهاد می‌ده، بعد مدل اصلی (target model) اون پیشنهادها رو تأیید یا رد می‌کنه. نتیجه‌ی نهایی دقیقاً همون چیزیه که مدل اصلی تولید می‌کرد، فقط خیلی سریع‌تر.

این روش در فریم‌ورک NeMo-RL با بک‌اند vLLM پیاده‌سازی شده. پشتیبانی از هر دو حالت synchronous و asynchronous وجود داره و می‌شه از مکانیزم‌های مختلف speculation استفاده کرد؛ از جمله MTP heads از پیش آموزش‌دیده، مدل‌های draft خارجی کوچیک، یا حتی تکنیک پیشرفته‌ای مثل Eagle3 که معمولاً بعد از مرحله‌ی RL اعمال می‌شه. این یعنی جدیدترین تکنیک‌های speculative decoding الان می‌تونن مستقیماً داخل حلقه‌ی آموزش RL استفاده بشن.

نتایج آزمایش‌ها روی یه workload استدلالی (reasoning post-training) با مدل ۸ میلیارد پارامتری نشون می‌ده که speculative decoding می‌تونه throughput تولید rollout رو تا ۱.۸ برابر بهبود بده. این عدد فقط برای مرحله‌ی rollout هست و بقیه‌ی pipeline رو تغییر نمی‌ده.

فراتر از آزمایش‌های واقعی، تیم تحقیقاتی با یه شبیه‌ساز با دقت بالا (high-fidelity performance simulator) پروجکشن‌هایی برای مقیاس‌های بزرگ‌تر زدن. تخمین‌ها نشون می‌ده که اگه speculative decoding با asynchronous RL ترکیب بشه، می‌شه در مدل ۲۳۵ میلیارد پارامتری تا ۲.۵ برابر سرعت کل فرآیند آموزش رو بالا برد؛ عددی که در این مقیاس بسیار قابل‌توجهه.

نکته‌ی مهم اینه که این روش کاملاً lossless هست؛ یعنی توزیع خروجی مدل target هیچ تغییری نمی‌کنه. برخلاف روش‌هایی که با off-policy learning یا precision reduction کار می‌کنن، اینجا کیفیت یادگیری حفظ میشه. این ویژگی این تکنیک رو به یه بلوک ساختاری ایده‌آل برای زیرساخت‌های RL post-training تبدیل می‌کنه.

نکات کلیدی:

• گلوگاه اصلی RL post-training، تولید خودرگرسیو rollout هست که کند و پرهزینه‌ست
• Speculative Decoding با استفاده از یه مدل draft کوچیک‌تر، سرعت تولید توکن رو بدون هیچ افت کیفیتی بالا می‌بره
• این روش در NeMo-RL با بک‌اند vLLM پیاده‌سازی شده و از Eagle3 هم پشتیبانی می‌کنه
• در مدل ۸B، throughput rollout تا ۱.۸ برابر بهبود پیدا کرده
• ترکیب با asynchronous RL در مقیاس ۲۳۵B، تا ۲.۵ برابر سرعت end-to-end آموزش رو بیشتر می‌کنه
• کاملاً lossless هست و توزیع خروجی مدل اصلی رو دست نمی‌زنه