The Vocab Mapping Script: A Critical Bridge in the EAGLE-3 Training Pipeline

Introduction

In the sprawling, multi-session effort to deploy and optimize a 1-trillion-parameter Kimi-K2.5 model on 8× RTX PRO 6000 Blackwell GPUs, a pivotal moment arrived when profiling revealed that AllReduce communication consumed 51.5% of decode time. The assistant and user pivoted from further low-level optimization to investigate speculative decoding — a technique where a smaller "draft" model generates candidate tokens that the large "target" model verifies in parallel, trading computation for reduced latency. After researching options and ruling out n-gram speculation (which proved 9–26% slower than baseline for this MoE reasoning model), the team settled on training a custom EAGLE-3 draft head, following the approach pioneered by Baseten. This led to the message under analysis: a brief but consequential step in building the training pipeline.

The subject message, <msg id=2526>, is deceptively simple:

[assistant] Now the vocab mapping script: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/03_build_vocab_mapping.py Wrote file successfully.

Beneath this terse surface lies a critical architectural decision. The vocabulary mapping script is the third of five pipeline components being constructed in rapid succession, and it solves a fundamental problem unique to the EAGLE-3 architecture: the draft model does not predict tokens from the base model's full vocabulary, but rather from a compressed subset. This article unpacks why this message was written, the reasoning behind it, the assumptions embedded in the approach, and how it fits into the broader speculative decoding strategy.

The Pipeline Under Construction

To understand <msg id=2526>, one must first understand the pipeline it belongs to. The assistant was building a complete EAGLE-3 training pipeline from scratch, consisting of five components, each written as a standalone Python script in the eagle3-train/ directory:

  1. draft_config.json — The configuration file defining the draft model architecture, matching the AQ-MedAI/Kimi-K2-Instruct-eagle3 structure with a 32K draft vocabulary, 4-layer MLP head, and specific layer indices for feature extraction.
  2. 01_prepare_dataset.py — A dataset preparation script that downloads samples from HuggingFace, tokenizes them using the Kimi-K2.5 tokenizer, computes token frequencies, and saves the prepared data in the format expected by the speculators training infrastructure.
  3. 02_extract_hidden_states.py — The most technically challenging component: a script that stops the running vLLM server, uses speculators' VllmHiddenStatesGenerator to run the base model in offline mode, attaches PyTorch hooks to capture hidden states from specific layers, and saves the extracted features alongside the token sequences.
  4. 03_build_vocab_mapping.py — The subject of this article: a vocabulary mapping script that creates a bidirectional mapping between the base model's full vocabulary (163,840 tokens for Kimi-K2.5) and the draft model's compressed vocabulary (32,000 tokens).
  5. 04_train.py — The training script that loads the prepared hidden states and vocab mapping, initializes the EAGLE-3 draft model from scratch, and runs the actual training loop using speculators' Eagle3Trainer.
  6. run_pipeline.sh — A shell orchestrator that ties all steps together with proper error handling and logging. The assistant had already written the draft config (<msg id=2523>), the dataset preparation script (<msg id=2524>), and the hidden state extraction script (<msg id=2525>). Message 2526 continues this assembly-line approach, adding the vocabulary mapping component before moving on to the training script (<msg id=2527>) and orchestrator (<msg id=2528>).

Why Vocabulary Mapping Matters

The need for vocabulary mapping arises from a fundamental asymmetry in EAGLE-3's architecture. The draft model in EAGLE-3 does not predict tokens from the base model's full vocabulary. Instead, it predicts from a compressed "draft vocabulary" — typically 32,000 tokens versus the base model's 163,840. This compression is possible because in practice, only a fraction of the full vocabulary is used frequently enough to warrant dedicated prediction slots.

The draft model's output is a probability distribution over these 32,000 draft tokens. During speculative decoding, each draft token must be mapped back to a real token ID in the base model's vocabulary before it can be fed to the target model for verification. Conversely, when preparing training data, the ground-truth next token (from the base model's vocabulary) must be mapped to the draft vocabulary so the draft model can learn to predict it.

This bidirectional mapping — token-to-draft (t2d) and draft-to-token (d2t) — is what 03_build_vocab_mapping.py computes. The mapping is frequency-based: the most commonly occurring tokens in the training data are assigned slots in the draft vocabulary, while rare or never-seen tokens are either excluded or mapped to a fallback.

Technical Implementation Details

From the successful test run visible in <msg id=2532>, we can reconstruct what the script does:

# Loading token frequency from prepared data
token_freq = torch.load("/root/eagle3-train/data_test/prepared/token_freq.pt")
# Shape: torch.Size([163840]) — one count per token ID
# Non-zero tokens: 899 (out of 163,840)
# Total frequency: 2875 (across 10 samples, 512 tokens each)

# Building mapping: 163840 -> 32000
# Draft vocab covers 100.0% of token frequency
# Selected token ID range: [0, 163607]
# t2d shape: torch.Size([163840]), dtype: torch.bool
# d2t shape: torch.Size([32000]), dtype: torch.int64

The script outputs four files: t2d.pt and d2t.pt (PyTorch tensors) plus t2d.npy and d2t.npy (NumPy arrays for interoperability). The t2d tensor is a boolean mask of shape [163840] — for each token in the base vocabulary, a True value indicates it has a slot in the draft vocabulary. The d2t tensor is an integer mapping of shape [32000] — for each draft token ID, it stores the corresponding base vocabulary token ID.

A critical detail: the test showed that the draft vocabulary covered 100% of token frequency with only 899 non-zero tokens out of 163,840. This means that with just 10 samples of 512 tokens each, only 899 unique token types appeared, and all fit comfortably within the 32,000-slot draft vocabulary. In a full-scale training run with millions of tokens, the coverage would be lower, and the mapping would need to prioritize the most frequent tokens — but the architecture supports this gracefully.

Assumptions and Decisions

Several important assumptions are baked into this message and the script it writes:

Assumption 1: The draft vocabulary size of 32,000 is appropriate. This matches the AQ-MedAI/Kimi-K2-Instruct-eagle3 configuration and is the standard for EAGLE-3. The assumption is that 32,000 tokens provide sufficient coverage of the base model's vocabulary for effective speculative decoding. If the acceptance rate were too low, one might need a larger draft vocabulary — but 32K is the established baseline.

Assumption 2: Frequency-based mapping is sufficient. The script selects draft vocabulary slots based on token frequency in the training data. This assumes that the training data's token distribution is representative of the deployment distribution. For a general-purpose reasoning model, this is a reasonable but not guaranteed assumption.

Assumption 3: The speculators training infrastructure expects this exact mapping format. The assistant had studied the speculators codebase thoroughly (<msg id=2507>, <msg id=2517>, <msg id=2518>, <msg id=2519>, <msg id=2520>) and confirmed that the training data loader (speculators/train/data.py) expects t2d and d2t tensors in the format produced by this script. This was a deliberate design decision to maximize compatibility.

Assumption 4: The pipeline will be ported to B200/B300 hardware. The user explicitly stated this in <msg id=2505>: "on our existing machine with lowered numbers, then should be easy to port to much more expensive b300 machine." The scripts are designed to be parameterized (via command-line arguments) so that the hero run can scale up seamlessly.

The Thinking Process Visible in Context

While the subject message itself is terse, the surrounding messages reveal extensive reasoning. In <msg id=2516>, the assistant explicitly weighed multiple approaches to hidden state extraction before settling on the speculators-based approach. In <msg id=2519>, we see a stream of consciousness as the assistant iterates through four different strategies:

  1. "Write a standalone hidden state extraction script that loads the model directly using HuggingFace transformers + register_forward_hook"
  2. "Use our running vLLM server for inference + write a separate lightweight model"
  3. "Use vLLM offline LLM with the worker extension"
  4. "Write a script that uses vLLM's offline LLM class with the speculators worker extension" This iterative refinement — each option considered and rejected for specific technical reasons — demonstrates the careful architectural thinking that preceded the actual implementation. The final choice to use speculators' VllmHiddenStatesGenerator was pragmatic: it leverages existing, tested code rather than writing a custom extraction mechanism that might have subtle bugs. Similarly, the decision to write 03_build_vocab_mapping.py as a separate script rather than folding it into the training script reflects a modular design philosophy. Each step can be tested independently, re-run if needed, and the outputs inspected before proceeding. This is visible in the test sequence: the assistant ran Step 1 (<msg id=2531>), then Step 3 (<msg id=2532>), before tackling the more complex Step 2 (<msg id=2533>).

Mistakes and Incorrect Assumptions

The most significant incorrect assumption — one that would surface in subsequent messages — was that the speculators library's VllmHiddenStatesGenerator would work seamlessly with vLLM 0.16.0rc2. The assistant had noted the version mismatch in &lt;msg id=2511&gt; ("speculators requires vllm<=0.15.0 for data generation which is a problem — our vLLM is 0.16.0") but proceeded anyway, hoping that the imports would work and runtime issues could be patched. This assumption proved partially correct: imports succeeded, but runtime failures emerged due to API changes in SchedulerConfig, KV cache utilities, and the Kimi-K2.5 multimodal model wrapper architecture (model.language_model.model.layers vs. model.model.layers).

The vocabulary mapping script itself, however, had no such issues — it ran successfully on the first attempt because it depends only on PyTorch tensor operations and the token frequency data produced by Step 1, not on any vLLM internals. This independence was a deliberate design choice that paid off.

Conclusion

Message &lt;msg id=2526&gt; is a small but essential piece of a much larger puzzle. In writing the vocabulary mapping script, the assistant completed the data preparation layer of the EAGLE-3 training pipeline, enabling the critical transformation between the base model's full vocabulary and the draft model's compressed vocabulary. The script embodies several key design decisions: modular pipeline architecture, compatibility with the speculators training framework, frequency-based token selection, and parameterization for scaling to production hardware. While the message itself is only a few lines, the reasoning behind it spans dozens of preceding messages — research into two competing frameworks (SpecForge and Speculators), analysis of the existing K2 EAGLE-3 model, empirical testing of n-gram speculation, and careful consideration of the vLLM version compatibility constraints. It stands as a testament to the methodical, research-driven approach that characterized this entire optimization effort.