diff --git a/tests/test_qwen2.5_0.5B_opd_vllm.py b/tests/test_qwen2.5_0.5B_opd_vllm.py index 947670e4d..b43b268fe 100644 --- a/tests/test_qwen2.5_0.5B_opd_vllm.py +++ b/tests/test_qwen2.5_0.5B_opd_vllm.py @@ -129,7 +129,7 @@ def launch_teacher(): rm_args = ( "--custom-rm-path vime.rollout.on_policy_distillation.reward_func " "--custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards " - f"--rm-url http://{TEACHER_HOST}:{TEACHER_PORT}/v1/completions " + f"--rm-url http://{TEACHER_HOST}:{TEACHER_PORT}/inference/v1/generate " ) grpo_args = ( diff --git a/vime/rollout/on_policy_distillation.py b/vime/rollout/on_policy_distillation.py index 1ec269093..e25338ac6 100644 --- a/vime/rollout/on_policy_distillation.py +++ b/vime/rollout/on_policy_distillation.py @@ -1,87 +1,80 @@ -"""On-Policy Distillation (OPD): teacher logprobs via vLLM ``/inference/v1/generate``. - -The reward function sends the full (prompt + student response) token sequence to -an external vLLM teacher server's ``/inference/v1/generate`` endpoint with -``prompt_logprobs`` enabled; the post-process step reads the per-token teacher -logprobs out of the response and stores them on each sample for the OPD KL -penalty. - -Endpoint contract (vime vLLM disaggregated ``/inference/v1/generate``): - -- Request body:: - - { - "token_ids": [...], # full prompt+response token ids - "model": , # OPTIONAL on this endpoint; only sent when - # --opd-teacher-model is set (single-model - # teacher servers use their loaded model) - "sampling_params": { - "max_tokens": 1, # endpoint requires >=1; the generated - # token is ignored - "temperature": 0, - "prompt_logprobs": 1, # score every prompt token - "skip_special_tokens": False, - }, - } - -- Response body (``GenerateResponse``):: - - { - "choices": [{"token_ids": [...], "logprobs": {...}, ...}], - "prompt_logprobs": [ # TOP-LEVEL, aligned with the input token_ids - None, # position 0: no prior context - {: {"logprob": -3.21, "rank": 1, "decoded_token": "..."}, ...}, - ... - ], - } - -``prompt_logprobs[i]`` is a dict ``{token_id -> Logprob}`` for the token at -position ``i``. JSON serializes integer dict keys as strings, so we look up by -both ``int`` and ``str``. See ``vllm/entrypoints/serve/disagg/protocol.py`` -(``GenerateRequest`` / ``GenerateResponse``). -""" - from __future__ import annotations from typing import Any +from urllib.parse import urlsplit, urlunsplit -import aiohttp import torch +from vime.utils.http_utils import post from vime.utils.types import Sample +def _teacher_base_url(rm_url: str) -> str: + """Strip the path off ``--rm-url`` to get the teacher server base (scheme://host:port).""" + parts = urlsplit(rm_url) + return urlunsplit((parts.scheme, parts.netloc, "", "", "")) + + async def reward_func(args, sample, **kwargs): - if sample.multimodal_inputs and sample.multimodal_inputs.get("images"): - # ``/inference/v1/generate`` is token-only; a multimodal teacher requires the - # render -> generate flow (``/v1/chat/completions/render`` then attach - # ``features``), as in ``vime.rollout.vllm_rollout.generate``. Not yet wired - # for OPD — fail loudly rather than silently scoring a text-only sequence. + teacher_model = getattr(args, "opd_teacher_model", None) + sampling_params = { + "max_tokens": 1, + "temperature": 0, + "prompt_logprobs": 1, + "skip_special_tokens": False, + } + base = _teacher_base_url(args.rm_url) + mm = sample.multimodal_inputs or {} + images = mm.get("images") + # Only image multimodal is wired here (mirrors vllm_rollout.generate, which renders + # images only). Fail loud on other modalities rather than silently scoring text-only + # without their context. + unsupported = [k for k, v in mm.items() if k != "images" and v] + if unsupported and not images: raise NotImplementedError( - "OPD multimodal teacher scoring over /inference/v1/generate is not implemented; " - "wire the /v1/chat/completions/render -> /inference/v1/generate flow first." + "OPD teacher scoring over /inference/v1/generate supports only image multimodal; " + f"got unsupported modalities: {unsupported}" ) - payload: dict[str, Any] = { - "token_ids": sample.tokens, - "sampling_params": { - "max_tokens": 1, - "temperature": 0, - "prompt_logprobs": 1, - "skip_special_tokens": False, - }, - } - # ``model`` is optional on /inference/v1/generate. Only send it when the teacher - # model is explicitly named; never fall back to the *student* checkpoint - # (args.hf_checkpoint), which would mis-name a teacher!=student server. - teacher_model = getattr(args, "opd_teacher_model", None) + if images: + # Multimodal: render (preprocess images) to get token_ids + features, then + # score the student's full prompt+response token sequence with those features + # attached — mirrors vime.rollout.vllm_rollout.generate's disaggregated MM flow. + # Lazy import: these helpers live in vllm_rollout, which pulls in the rollout/ + # router stack the text-only reward path never needs — import only when an + # image sample is actually scored. + from vime.rollout.vllm_rollout import ( + _align_mm_feature_placeholders_to_tokens, + _coerce_flat_int_token_ids, + _mm_render_response_to_generate_body, + ) + from vime.utils.processing_utils import encode_image_for_rollout_engine + + content: list[dict[str, Any]] = [{"type": "text", "text": sample.prompt}] + for image in images: + content.append({"type": "image_url", "image_url": {"url": encode_image_for_rollout_engine(image)}}) + render_payload: dict[str, Any] = {"messages": [{"role": "user", "content": content}]} + if teacher_model: + render_payload["model"] = teacher_model + render_data = await post(f"{base}/v1/chat/completions/render", render_payload) + + body = _mm_render_response_to_generate_body(render_data, teacher_model or "") + canonical = _coerce_flat_int_token_ids(sample.tokens) + # Align the rendered feature placeholders to VIME's canonical ids BEFORE + # overriding token_ids (the aligner reads the render's token_ids first). + _align_mm_feature_placeholders_to_tokens(body, canonical) + body["token_ids"] = canonical + body["sampling_params"] = sampling_params + if teacher_model: + body["model"] = teacher_model + else: + body.pop("model", None) + return await post(args.rm_url, body) + + payload: dict[str, Any] = {"token_ids": sample.tokens, "sampling_params": sampling_params} if teacher_model: payload["model"] = teacher_model - - async with aiohttp.ClientSession() as session: - async with session.post(args.rm_url, json=payload) as resp: - resp.raise_for_status() - return await resp.json() + return await post(args.rm_url, payload) def _logprob_for_token(pos_entry: dict | None, token_id: int) -> float: