‏مهاجرت vLLM از V0 به V1: قبل از تنظیم RL، درستی رو ثابت کن

تیم ServiceNow در پروژه PipelineRL از vLLM به‌عنوان موتور استنتاج برای تولید rollout استفاده می‌کنه. این موتور توکن‌ها رو sample می‌کنه و token logprobها رو برمی‌گردونه؛ trainer هم از این مقادیر برای محاسبه policy ratio، KL، clip rate، entropy و reward استفاده می‌کنه. وقتی تیم از نسخه 0.8.5 (V0) به 0.18.1 (V1) مهاجرت کرد، اختلاف‌های مشخصی در همین معیارها ظاهر شد که نشون‌دهنده یه train-inference mismatch بود.

هدف مهاجرت خیلی محدود و دقیق تعریف شده بود: اول ثابت کنن که V1 دقیقاً همون logprobهایی رو برمی‌گردونه که trainer انتظار داره، بعد نتایج رو با مرجع V0 مقایسه کنن، و فقط بعد از اثبات تطابق backend، تغییرات تابع هدف RL رو بررسی کنن. تیم سه لایه خرابی احتمالی رو شناسایی کرد: semantic mismatch (logprobها معنای متفاوتی دارن)، inference-path mismatch (تنظیمات پیش‌فرض متفاوت)، و objective mismatch (تابع هدف RL نیاز به اصلاح داره). اشتباه اولیه این بود که خیلی زود به دسته سوم مشکوک شدن.

اولین مشکل در سمانتیک logprob بود. vLLM V1 به‌صورت پیش‌فرض logprobها رو از خروجی خام مدل برمی‌گردونه، قبل از اعمال temperature scaling، penalties و top-k/p filtering. ولی PipelineRL انتظار داشت logprobها از توزیع پردازش‌شده‌ای که sampler استفاده می‌کنه بیان. راه‌حل ساده بود: تنظیم logprobs-mode=processed_logprobs.

دومین مشکل تنظیمات پیش‌فرض V1 بود. prefix caching و async scheduling در V1 به‌صورت پیش‌فرض فعال بودن، در حالی که V0 این‌طور نبود. در یه سیستم RL آنلاین که وزن‌های مدل مرتباً آپدیت می‌شن، prefix cache می‌تونه حالت‌هایی رو که قبل از آپدیت وزن محاسبه شدن دوباره استفاده کنه — که یه درجه آزادی اضافه و ناخواسته ایجاد می‌کنه. تنظیم صریح این مقادیر در config این مشکل رو حل کرد:

vllm_config:
  use_v1: true
  vllm_kwargs:
    logprobs-mode: processed_logprobs
    enable-prefix-caching: false
    async-scheduling: false

سومین مشکل مربوط به به‌روزرسانی وزن‌ها در حین اجرا بود. رفتار V0 چیزی شبیه به این بود: اجرا رو در یه مرز مشخص متوقف کن، وزن‌های جدید رو بارگذاری کن، بدون invalidate کردن cache از سر بگیر. معادل V1 این رفتار با این کد پیاده‌سازی شد:

await engine.pause_generation(mode="keep", clear_cache=False)
await engine_client.collective_rpc_async(
    "receive_weight_update",
    args=(request.model_dump_json(),),
)
await engine.resume_generation()

نکته کلیدی اینجاست که mode="keep" و clear_cache=False دقیقاً همون رفتار V0 رو شبیه‌سازی می‌کنن.

چهارمین مشکل دقت عددی لایه آخر مدل بود. trainer از یه lm_head با دقت fp32 برای projection نهایی استفاده می‌کرد، ولی backend rollout این رو match نمی‌کرد. تفاوت‌های کوچیک در logitها می‌تونن در policy ratio، KL و clipping به‌وضوح دیده بشن. جالبه که گزارش فنی MiniMax-M1 هم دقیقاً همین مشکل رو گزارش داده بود. بعد از اعمال هر چهار fix، منحنی‌های آموزش V1 با مرجع V0 کاملاً تطابق پیدا کردن.

نکات کلیدی:

• vLLM V1 به‌صورت پیش‌فرض logprob خام (قبل از پردازش) برمی‌گردونه؛ برای RL باید processed_logprobs فعال بشه
• در سیستم‌های RL آنلاین، prefix caching می‌تونه حالت‌های قبل از آپدیت وزن رو دوباره استفاده کنه و نتایج رو خراب کنه
• رفتار به‌روزرسانی وزن در حین اجرا باید صریحاً با mode="keep" و clear_cache=False تنظیم بشه تا با V0 مطابقت داشته باشه
• دقت fp32 در لایه lm_head بین trainer و inference backend باید یکسان باشه
• قبل از تغییر تابع هدف RL، باید مطمئن بشی backend به‌درستی کار می‌کنه