From b79f20e8cb89b7b86e125164a977c5287a7aa64c Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 14 Oct 2026 10:55:58 -0700 Subject: [PATCH] Let the multimodal generator enable thinking and fit its output budget Two gaps found generating drafter data against a reasoning VLM target, both of which fail quietly. Thinking could not be enabled at all. Reasoning models gate the 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. Every sample comes back as a direct answer and nothing in the output says which mode produced it. On the target that surfaced this, the same prompt answers in 201 characters without the flag and 5,315 with it, and 20,000 generated records contained no reasoning block at all before anyone noticed. --enable_thinking sends the chat_template_kwargs. Requests that exceed the context window were dropped. The server rejects an over-budget request rather than truncating it: prompt tokens plus max_tokens must fit or it is a 400. Returning that as an error loses 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. The budget now halves until the request fits, down to --min_max_tokens. Overflow is matched on the message rather than the status code, because a 400 covers many unrelated problems and retrying those with a smaller budget would just burn requests. Testing: 18 tests in tests/examples/speculative_decoding/ test_server_generate_vlm_sglang.py (12 existing, 6 new), covering the chat_template_kwargs payload with and without the flag, the shrink sequence 8192 -> 4096 -> 2048, stopping at the floor, and a non-context 400 not being retried. pre-commit clean on all three files. Signed-off-by: Ye Yu --- CHANGELOG.rst | 1 + .../server_generate_vlm_sglang.py | 72 +++++++-- .../test_server_generate_vlm_sglang.py | 143 ++++++++++++++++++ 3 files changed, 203 insertions(+), 13 deletions(-) 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