Skip to content

[RFC]: Use vLLM Render Endpoint in preprocessing for Tokenization Alignment #652

Description

@shanjiaz

Motivation.

Speculators preprocesses training data by applying chat templates and tokenizing with HuggingFace's AutoProcessor/AutoTokenizer, then sends the resulting token IDs to vLLM's Completions API for hidden-state extraction. This assumes the HF tokenizer and vLLM's internal preprocessing produce identical token sequences for the same conversation. They sometimes don't, and the divergence occurs at two distinct layers. Tokenizer-level divergence. vLLM maintains a tokenizer registry (vllm/tokenizers/registry.py) with custom tokenizer backends that replace AutoTokenizer entirely when tokenizer_mode="auto". When vLLM auto-detects one of these models, it loads the custom tokenizer instead of HF's. Preprocessing-level divergence. Even when the underlying tokenizer is the same (e.g., gpt-oss uses HF's tokenizer for raw text encoding), vLLM applies additional conversation-level preprocessing that apply_chat_template() does not replicate. For gpt-oss specifically, vLLM uses Harmony's Rust encoder (openai_harmony.load_harmony_encoding) to render conversations into token sequences via render_for_completion(), which injects system preambles, formats tool descriptions, and structures messages in a way that is different from the Jinja2 chat template. vLLM also applies tool parsers and reasoning parsers that further transform the token sequence.

Two code paths for text-only vs. multimodal. Speculators currently maintains two separate data flows depending on content type. Text-only data is tokenized locally with HF and sent as pre-tokenized input_ids to the Completions API. Multimodal data sends raw messages to the Chat Completions API, letting vLLM handle tokenization, because the Completions API cannot carry image/video/audio references. This split creates its own problems: the text-only path is vulnerable to the tokenizer divergence described above, while the multimodal path relies on vLLM and HF inserting placeholder tokens at the same positions to keep the locally-built loss mask aligned. A single path that delegates all tokenization to vLLM would eliminate this branching and make our preprocessing step more maintainable.

Proposed Change.

Skip the local preprocessing step entirely and use a single vLLM instance for both tokenization and hidden-state extraction. The vLLM server already tokenizes inputs through its full preprocessing pipeline (chat template rendering, Harmony encoding, tool/reasoning parsing) before running the model. Speculators should send raw conversations to the Chat Completions API and let vLLM handle all tokenization, rather than tokenizing locally with HF and sending pre-tokenized IDs to the Completions API. This eliminates the tokenizer divergence problem and simplify the training process. Requests must forward chat_template_kwargs (e.g., enable_thinking for Gemma 4) so the rendered token sequence matches what vLLM produces at inference time with the same settings. Currently, preprocessing truncates input_ids to seq_length before sending to vLLM. With raw messages, truncation must happen at the conversation level (dropping early turns) or via max_model_len on the vLLM server side, since the Completions API pre-truncated path is no longer used.

Open Questions

  • How to efficiently return loss masks from the hidden-state extraction endpoint without a second pass
  • Currently speculators truncates input_ids to seq_length locally before sending them to vLLM. If we switch to sending raw messages and letting vLLM tokenize, we lose that truncation point. We need a strategy for enforcing length limits when speculators no longer owns the tokenization step.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions