The Critical Bridge: Vocabulary Mapping in the EAGLE-3 Training Pipeline

A Pivot Point Between Data Preparation and Model Training

In the sprawling, multi-session effort to deploy speculative decoding for the Kimi-K2.5 model on 8x Blackwell GPUs, message [msg 2803] occupies a deceptively modest position. On its surface, it is a brief status update: the assistant runs a vocabulary mapping script, reports the results, and prepares to move to the next step. But this message is far more consequential than its length suggests. It represents the completion of the data preparation phase and the critical bridge between raw data and actual model training—a moment where the theoretical design of the EAGLE-3 draft model meets the concrete realities of the training data.

To understand why this message matters, one must appreciate the architecture it serves. The EAGLE-3 draft model is a lightweight auxiliary network that predicts the next several tokens the main model would generate, enabling speculative decoding to accelerate inference. Critically, the draft model operates with a much smaller vocabulary—32,000 tokens—compared to the full Kimi-K2.5 model's 163,840 tokens. This compression is what makes the draft model fast: fewer output logits to compute means faster autoregressive generation. But it introduces a fundamental problem: how do you map between the full model's token space and the draft model's token space? Every token the main model might predict must have a corresponding embedding in the draft model, and every draft prediction must be translatable back to the full vocabulary. This is the job of the vocabulary mapping, and message [msg 2803] documents its execution.

The Message in Full

The assistant writes:

1000 samples processed, 503K total tokens, 360K assistant tokens. Now let me run step 3 (vocab mapping) and then step 2 (hidden state extraction):

>

``` [bash] ssh root@10.1.230.174 '/root/ml-env/bin/python3 /root/eagle3-train/03_build_vocab_mapping.py \ --token-freq /root/eagle3-train/data_1k/prepared/token_freq.pt \ --output-dir /root/eagle3-train/data_1k/vocab_mapping \ --target-vocab-size 163840 \ --draft-vocab-size 32000 2>&1 | tail -20' Loading token frequency from /root/eagle3-train/data_1k/prepared/token_freq.pt... Shape: torch.Size([163840]) Non-zero tokens: 21270 Total frequency: 360421

>

Building vocab 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

>

Saved to /root/eagle3-train/data_1k/vocab_mapping/ t2d.pt, d2t.pt (torch), t2d.npy, d2t.npy (numpy... ```

The message opens with a summary of the preceding data preparation step: "1000 samples processed, 503K total tokens, 360K assistant tokens." This is the assistant confirming that Step 1—dataset preparation via 01_prepare_dataset.py—has completed successfully. The numbers are significant: 503,000 total tokens across 1,000 conversation samples from the mlabonne/open-perfectblend dataset, with 360,421 of those being assistant tokens (the ones the model actually needs to learn to predict). The 58 skipped samples (visible in [msg 2802]'s metadata) likely represent conversations that exceeded the 2048-token maximum sequence length.

The Reasoning Behind the Ordering

A subtle but important decision is visible in the assistant's phrasing: "Now let me run step 3 (vocab mapping) and then step 2 (hidden state extraction)." This ordering is a deliberate optimization. In [msg 2801], the assistant had laid out a four-step plan where Step 2 (hidden state extraction) was identified as the bottleneck, requiring a 22-minute model load on all 8 GPUs. Steps 1 and 3 were described as CPU-only and fast. By running Step 3 before Step 2, the assistant ensures that all data preparation is complete before committing to the expensive GPU-based extraction. If the vocab mapping had failed—say, because the token frequency file was malformed or the vocabulary sizes were mismatched—the assistant would have discovered the problem without wasting 22 minutes of GPU time loading the 547GB verifier model.

This is a classic optimization pattern in ML engineering: fail fast on cheap operations before committing to expensive ones. The assistant is implicitly treating the vocab mapping as a validation gate for the entire data pipeline. If the token frequencies look wrong, or if the mapping produces degenerate results, the assistant can diagnose the issue before sinking time into hidden state extraction.

The Vocabulary Mapping Mechanics

The script 03_build_vocab_mapping.py performs a conceptually simple but technically critical operation. It takes the token frequency histogram (a tensor of shape [163840]—one count per token in the full vocabulary) and constructs two mapping tensors:

  1. t2d (target-to-draft): A boolean mask of shape [163840] indicating which tokens in the full vocabulary are included in the draft vocabulary. Each position is True if the corresponding token ID is within the selected range.
  2. d2t (draft-to-target): An integer mapping of shape [32000] that maps each draft token ID to its corresponding full vocabulary token ID. This is the reverse lookup: given a draft model's output token (0–31999), which token in the full vocabulary does it represent? The results are revealing. The assistant reports "Non-zero tokens: 21270" out of 163840—meaning that in 503,000 tokens of training data, only 21,270 distinct token types actually appear. This is a remarkably sparse distribution, typical of large-vocabulary LLMs where the vast majority of the vocabulary consists of rare tokens that may never appear in any given dataset. The draft vocabulary of 32,000 slots is therefore more than sufficient: it has room for all 21,270 observed tokens plus headroom for unseen tokens.

The Critical Result: 100% Coverage

The most important line in the entire message is: "Draft vocab covers 100.0% of token frequency." This means that every token that appears in the training data falls within the selected token ID range [0, 163607], and therefore has a corresponding slot in the draft vocabulary. No token in the training data is lost or needs to be mapped to a special "unknown" token.

This 100% coverage is not guaranteed by default. The draft vocabulary of 32,000 tokens is much smaller than the full vocabulary of 163,840. The mapping strategy of selecting a contiguous range of token IDs [0, 163607] works because the Kimi-K2.5 tokenizer's vocabulary is organized with the most frequently used tokens at lower IDs. The special tokens—including token 163606 (<thinking>) and token 163607 (<response>), which are critical for the model's reasoning format—fall within this range. If the training data had required tokens above 163607, the mapping would have had to use a more sophisticated selection strategy, such as picking the top-K most frequent tokens rather than a contiguous range.

The selection of [0, 163607] as the range is itself a design decision worth examining. The draft vocabulary has 32,000 slots, so the range covers tokens 0 through 163607 inclusive—a span of 163,608 tokens. Wait: 163,608 tokens mapped into 32,000 slots? That seems impossible unless the mapping is not one-to-one. In fact, the t2d tensor is a boolean mask: it marks which of the 163,840 full-vocabulary tokens are "covered" by the draft vocabulary. The draft vocabulary only has 32,000 actual embeddings, but those 32,000 embeddings are assigned to the most important token IDs (0 through 163607). The remaining tokens (163608 through 163839) are not representable in the draft model—but since they don't appear in the training data, this is acceptable.

Wait—let me re-examine. The output says "Selected token ID range: [0, 163607]" and the d2t tensor has shape [32000]. This means the draft vocabulary of 32,000 tokens maps to a subset of the full vocabulary. The range [0, 163607] is the set of full-vocabulary token IDs that are covered—they have corresponding entries in the draft vocabulary. But 163,608 distinct token IDs cannot fit into 32,000 slots. So the mapping must be that the draft vocabulary covers the first 32,000 tokens (0–31999) plus some additional mechanism, or the "range" is describing something else.

Actually, looking more carefully: the t2d tensor is a boolean mask of shape [163840] where True means "this token ID is representable in the draft vocabulary." The d2t tensor of shape [32000] maps draft token IDs to full token IDs. The range [0, 163607] likely means that all tokens from 0 to 163607 are covered (t2d[i] = True for i in [0, 163607]), and the d2t mapping is the identity mapping for those first 32,000 tokens (d2t[i] = i for i in [0, 31999]), with the remaining tokens in the range being unmapped or handled differently.

This is a subtle point that the message does not fully clarify, but it doesn't need to—the 100% coverage result is what matters for the pipeline's correctness.

Assumptions Embedded in the Message

The assistant makes several assumptions in this message, most of which are reasonable but worth examining:

Assumption 1: The contiguous range strategy is optimal. By selecting [0, 163607] as the covered token range, the assistant assumes that the tokenizer's ID assignment correlates with frequency—that lower IDs are more important. This is generally true for most LLM tokenizers (BPE tokenizers assign IDs in order of discovery during training), but it's not guaranteed to be optimal for any specific dataset. The 100% coverage validates this assumption empirically for this particular dataset.

Assumption 2: The vocab mapping from 10 samples generalizes to 1000 samples. The assistant had previously built a vocab mapping for the 10-sample test run. The 1000-sample run produces a new mapping based on the larger dataset's token frequencies. The assistant implicitly assumes that the mapping stabilizes with more data and doesn't introduce regressions.

Assumption 3: The hidden state extraction will work with this vocab mapping. The vocab mapping is consumed by the extraction script (Step 2) and the training script (Step 4). The assistant assumes that the format—t2d.pt, d2t.pt, and their NumPy equivalents—is compatible with both downstream steps. This assumption is reasonable given that the same pipeline worked for 10 samples.

Assumption 4: The 100% coverage will hold for the test set. The vocab mapping is built from the training data's token frequencies. If the test set (or the model's actual generation targets during inference) contains tokens outside [0, 163607], the draft model won't be able to predict them. The assistant assumes that the training data is representative of the full token distribution.

Input Knowledge Required

To fully understand this message, one needs:

  1. Understanding of speculative decoding and EAGLE-3: The concept of a draft model with a smaller vocabulary that predicts multiple tokens per step, and the need to map between draft and full vocabularies.
  2. Knowledge of the Kimi-K2.5 model architecture: Specifically, its vocabulary size of 163,840 tokens, its use of special tokens like <thinking> (163606) and <response> (163607), and the structure of its tokenizer.
  3. Familiarity with the speculators library: The library's conventions for vocabulary mapping tensors (t2d, d2t), their shapes and dtypes, and how they're consumed by the training and inference pipelines.
  4. Context from the broader session: The assistant's multi-session effort to deploy Kimi-K2.5, the earlier struggles with flash-attn compilation, the pivot from GLM-5 to Kimi-K2.5, and the EAGLE-3 training pipeline development documented in [msg 2786] through [msg 2802].
  5. Understanding of token frequency distributions: Why 21,270 non-zero tokens out of 163,840 is expected behavior for a 503K-token dataset, and why a 32,000-token draft vocabulary is more than sufficient.

Output Knowledge Created

This message creates several pieces of actionable knowledge:

  1. Validation of the data pipeline: The vocab mapping completes successfully, confirming that the data preparation steps (tokenization, frequency counting) produced correct outputs.
  2. Token distribution statistics: The training data contains 21,270 unique tokens out of 163,840 possible, with 360,421 total assistant token occurrences. This sparsity confirms that the draft vocabulary of 32,000 tokens is appropriately sized.
  3. Mapping coverage confirmation: 100% of token frequency is covered by the selected token ID range [0, 163607], meaning no training data tokens will be lost in the draft vocabulary compression.
  4. Ready-to-use mapping tensors: The t2d.pt, d2t.pt, t2d.npy, and d2t.npy files are saved and ready for consumption by the hidden state extraction and training scripts.
  5. Pipeline ordering decision: The assistant establishes that Step 3 (vocab mapping) runs before Step 2 (hidden state extraction), serving as a lightweight validation gate before the expensive GPU-based operation.

The Thinking Process Visible in the Message

The assistant's reasoning is visible in the structure of the message itself. The opening line—"1000 samples processed, 503K total tokens, 360K assistant tokens"—is not just a status update. It's a self-check: the assistant is confirming that the preceding step produced reasonable numbers before proceeding. The 360K assistant tokens versus 503K total tokens (roughly 72% assistant, 28% user/system) is a typical ratio for instruction-following datasets, confirming that the data is well-formed.

The decision to pipe the output through tail -20 is also revealing. The assistant knows from experience (the spammy tokenizer output seen in [msg 2801]) that the script produces voluminous intermediate logging. By using tail -20, the assistant focuses on the final results—the coverage statistics and file outputs—while suppressing noise. This is a practical engineering judgment: the important information is the summary, not the per-token details.

The assistant's choice to report the tensor shapes and dtypes (torch.Size([163840]), torch.bool, torch.int64) reflects a debugging-oriented mindset. These details would be unnecessary if everything worked perfectly, but they serve as documentation in case something goes wrong downstream. If the training script fails with a shape mismatch, the assistant can immediately compare against these reported values.

Finally, the trailing ellipsis in "t2d.npy, d2t.npy (numpy..." suggests the output was truncated. The assistant didn't re-run the command to get the full output—a sign that the key information (the coverage percentage and the tensor shapes) was already captured, and the remaining details (likely just file size information) are not critical.

The Broader Significance

In the context of the full session ([msg 2786] through [msg 2802] and beyond), this message marks the transition from data preparation to the expensive GPU-based operations. The assistant has been building toward this moment for hours: first exploring the speculators library's API, rewriting the training script, validating on 10 samples, and verifying vLLM compatibility. Now, with the vocab mapping confirmed, the assistant can proceed to hidden state extraction—the 22-minute model load that will produce the actual training targets for the EAGLE-3 draft model.

The message also represents a quiet validation of the entire pipeline design. The 100% coverage result confirms that the contiguous-range mapping strategy works for this dataset and this tokenizer. If the coverage had been lower—say, 95%—the assistant would have needed to reconsider the mapping strategy, perhaps switching to a top-K frequency selection or a learned mapping. The clean result means the pipeline can proceed without architectural changes.

In the end, [msg 2803] is a message about readiness. The data is prepared, the mapping is built, and the assistant is about to commit to the expensive GPU operation that will produce the actual training data. It is the calm before the storm—the last cheap validation before the 22-minute model load, the 2.9-minute extraction at 2912 tok/s, and the 27.7-minute training run that will produce the first real EAGLE-3 draft model checkpoint. The vocabulary mapping, mundane as it seems, is the key that unlocks everything that follows.