diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5b9f798bdc7..a97312756cb 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -60,6 +60,7 @@ Changelog - LiLiCorr now trains offline and in streaming mode, not only online: the target logits its distractor penalty reads are reconstructed from the captured final hidden state, the same path self-logit distillation uses. - Add ``dflash_lilicorr_w_cal`` (default ``0.0``, off): an optional LiLiCorr proposal-calibration term, a KL divergence from the target's distribution over each slot's candidates to the reranker's candidate distribution, usable in place of the distractor penalty. - ``examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py`` ``--api_mode openai`` works against any OpenAI-compatible server, including vLLM, and gains ``--media_inline`` to send media as base64 ``data:`` URIs. Without ``--media_url_base`` or ``--media_inline`` the media path is still passed through verbatim, which only servers that read their own filesystem resolve (SGLang native does; vLLM needs ``--allowed-local-media-path``); that case now warns. +- ``server_generate_vlm_sglang.py`` gains ``--enable_thinking``, which sends ``chat_template_kwargs`` so a reasoning target actually reasons; without it the template's guard leaves thinking off and nothing in the output records which mode produced it. Requests that exceed the model's context window now retry with a halved output budget down to ``--min_max_tokens`` instead of being dropped, since the requests that overflow are the longest-context ones. *Megatron Framework (M-LM / M-Bridge)* diff --git a/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py b/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py index d6d40ca11a4..fc3e4a5a852 100755 --- a/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py +++ b/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py @@ -357,6 +357,16 @@ def _prompt_with_vision_token(prompt: str, media_type: str, token_format: str) - return f"{QWEN_IMAGE_TOKEN}\n{prompt}" +def _is_context_overflow(body: str) -> bool: + """Whether a non-200 body reports the request exceeding the model's context window. + + Matched on the message rather than the status code: a 400 covers many unrelated + problems, and retrying those with a smaller budget would just burn requests. + """ + lowered = body.lower() + return "maximum context length" in lowered or "longer than the maximum" in lowered + + def _openai_generate_one( request: dict[str, Any], url: str, @@ -368,6 +378,8 @@ def _openai_generate_one( media_root: str | None, input_root: str | None, media_inline: bool = False, + enable_thinking: bool = False, + min_max_tokens: int = 256, ) -> tuple[str, str | None]: try: import requests @@ -412,20 +424,38 @@ def _openai_generate_one( "temperature": temperature, "max_tokens": max_tokens, } - + if enable_thinking: + # Reasoning models gate their prefill on a chat-template variable, and the + # usual guard is ``enable_thinking is defined and enable_thinking`` -- so leaving it + # unset is indistinguishable from setting it false, and the target answers without + # reasoning. Nothing in the response says which mode produced it. + payload["chat_template_kwargs"] = {"enable_thinking": True} + + budget = max_tokens try: - response = requests.post(url, json=payload, timeout=timeout) - if response.status_code != 200: - return "", f"HTTP {response.status_code}: {response.text[:1000]}" - data = response.json() - choices = data.get("choices") or [] - if not choices: - return "", f"No choices in OpenAI response: {data}" - message = choices[0].get("message") or {} - answer = _coerce_text(message.get("content")) - if not answer: - return "", f"Could not extract assistant content from OpenAI response: {data}" - return answer, None + while True: + payload["max_tokens"] = budget + response = requests.post(url, json=payload, timeout=timeout) + if response.status_code != 200: + body = response.text[:1000] + # The server REJECTS an over-budget request, it does not truncate: prompt + # tokens plus max_tokens must fit the context window or it is a 400. + # Returning the error here would drop the sample, and the samples that + # overflow are the longest-context ones -- long videos -- which is exactly + # the data a multimodal draft model needs to have seen. Shrink and retry. + if _is_context_overflow(body) and budget > min_max_tokens: + budget = max(min_max_tokens, budget // 2) + continue + return "", f"HTTP {response.status_code}: {body}" + data = response.json() + choices = data.get("choices") or [] + if not choices: + return "", f"No choices in OpenAI response: {data}" + message = choices[0].get("message") or {} + answer = _coerce_text(message.get("content")) + if not answer: + return "", f"Could not extract assistant content from OpenAI response: {data}" + return answer, None except Exception: return "", traceback.format_exc() @@ -507,6 +537,20 @@ def main() -> None: help="Send media as base64 data: URIs instead of paths/URLs. Self-contained, but " "inflates every request body -- prefer --media_url_base for video.", ) + parser.add_argument( + "--enable_thinking", + action="store_true", + help="Ask the chat template to prefill the reasoning block, for targets that gate " + "it on an enable_thinking template variable. Reasoning traces are far longer than " + "direct answers, so raise --max_tokens accordingly.", + ) + parser.add_argument( + "--min_max_tokens", + type=int, + default=256, + help="Floor when shrinking the output budget to fit the context window. Below this " + "a completion is too stunted to be worth keeping.", + ) parser.add_argument( "--api_mode", choices=("openai", "native"), @@ -602,6 +646,8 @@ def main() -> None: args.media_root, args.input_root, args.media_inline, + args.enable_thinking, + args.min_max_tokens, ): meta for request_args, meta in zip(batch_args, batch_meta) } diff --git a/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py b/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py index 42a63bdadf9..9d5955fde0b 100644 --- a/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py +++ b/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py @@ -16,6 +16,7 @@ """Tests for the multimodal SGLang generation client.""" import base64 +import copy import importlib.util from pathlib import Path @@ -173,3 +174,145 @@ def test_openai_media_value_prefers_remote_urls_over_inlining(tmp_path): assert ( server_generate_vlm_sglang._as_openai_media_value(remote, None, None, None, True) == remote ) + + +def test_context_overflow_is_recognised_only_from_the_message(): + """A 400 covers many problems; only the context message should trigger a retry.""" + + assert server_generate_vlm_sglang._is_context_overflow( + "This model's maximum context length is 32768 tokens. However, you requested 40000" + ) + assert server_generate_vlm_sglang._is_context_overflow( + "Input is longer than the maximum supported length" + ) + assert not server_generate_vlm_sglang._is_context_overflow("Invalid value for 'top_k'") + assert not server_generate_vlm_sglang._is_context_overflow("") + + +def _fake_requests(monkeypatch, responses): + """Install a requests stub that returns ``responses`` in order, recording payloads.""" + + sent = [] + + class _Resp: + def __init__(self, status, body): + self.status_code = status + self.text = body + self._body = body + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + queue = list(responses) + + class _Requests: + @staticmethod + def post(url, json=None, timeout=None): + # Snapshot: the caller mutates one payload dict across retries, so recording + # the reference would make every entry show the final budget. + sent.append(copy.deepcopy(json)) + return _Resp(*queue.pop(0)) + + monkeypatch.setitem(__import__("sys").modules, "requests", _Requests) + return sent + + +def test_enable_thinking_is_passed_as_chat_template_kwargs(monkeypatch): + """Without this the target answers in non-thinking mode and nothing says so.""" + + sent = _fake_requests(monkeypatch, [(200, "")]) + request = {"prompt": "hi", "image_path": "https://example.com/a.jpg"} + answer, err = server_generate_vlm_sglang._openai_generate_one( + request, + "http://x/v1/chat/completions", + "m", + 0.7, + 2048, + 60, + None, + None, + None, + False, + enable_thinking=True, + ) + + assert err is None and answer == "ok" + assert sent[0]["chat_template_kwargs"] == {"enable_thinking": True} + + +def test_thinking_is_absent_unless_requested(monkeypatch): + sent = _fake_requests(monkeypatch, [(200, "")]) + request = {"prompt": "hi", "image_path": "https://example.com/a.jpg"} + server_generate_vlm_sglang._openai_generate_one( + request, "http://x/v1/chat/completions", "m", 0.7, 2048, 60, None, None, None, False + ) + + assert "chat_template_kwargs" not in sent[0] + + +def test_budget_shrinks_until_the_request_fits(monkeypatch): + """An over-budget request is retried smaller, not dropped. + + The samples that overflow are the longest-context ones -- long videos -- so dropping + them biases the dataset against exactly the data a multimodal drafter needs. + """ + + overflow = (400, "This model's maximum context length is 32768 tokens.") + sent = _fake_requests(monkeypatch, [overflow, overflow, (200, "")]) + request = {"prompt": "hi", "video_path": "https://example.com/a.mp4"} + + answer, err = server_generate_vlm_sglang._openai_generate_one( + request, + "http://x/v1/chat/completions", + "m", + 0.7, + 8192, + 60, + None, + None, + None, + False, + enable_thinking=False, + min_max_tokens=256, + ) + + assert err is None and answer == "ok" + assert [p["max_tokens"] for p in sent] == [8192, 4096, 2048] + + +def test_budget_shrinking_stops_at_the_floor(monkeypatch): + """Below the floor a completion is too stunted to keep, so report instead of looping.""" + + overflow = (400, "This model's maximum context length is 32768 tokens.") + sent = _fake_requests(monkeypatch, [overflow, overflow]) + request = {"prompt": "hi", "video_path": "https://example.com/a.mp4"} + + answer, err = server_generate_vlm_sglang._openai_generate_one( + request, + "http://x/v1/chat/completions", + "m", + 0.7, + 512, + 60, + None, + None, + None, + False, + enable_thinking=False, + min_max_tokens=256, + ) + + assert answer == "" and "maximum context length" in err + assert [p["max_tokens"] for p in sent] == [512, 256] + + +def test_non_context_errors_are_not_retried(monkeypatch): + sent = _fake_requests(monkeypatch, [(400, "Invalid value for 'top_k'")]) + request = {"prompt": "hi", "image_path": "https://example.com/a.jpg"} + + answer, err = server_generate_vlm_sglang._openai_generate_one( + request, "http://x/v1/chat/completions", "m", 0.7, 8192, 60, None, None, None, False + ) + + assert answer == "" and "top_k" in err + assert len(sent) == 1