feat(scripts): build a draft vocabulary mapping without regenerating features - #765
Open
fg11991 wants to merge 2 commits into
Open
feat(scripts): build a draft vocabulary mapping without regenerating features#765fg11991 wants to merge 2 commits into
fg11991 wants to merge 2 commits into
Conversation
…features Choosing a draft vocabulary currently requires a full hidden-state capture first: process_token_dict_to_mappings runs inside prepare_hidden_states, so trying a second draft_vocab_size, or building a mapping for a corpus whose features were captured before the mapping was needed, means paying for the capture again. On a 600k-sample corpus that is hours of GPU time to answer a question about token frequencies. scripts/build_vocab_mapping.py derives the mapping from either source, with no GPU involved: - --data-path counts loss-bearing tokens from the source JSONL by applying the same chat template and truncation the trainer would; - --hidden-states-path counts them from already prepared features. Both stream: the tally is accumulated per shard instead of materializing a token column for the whole corpus, and the per-corpus counts are cached under a key that includes the corpus identity, so sweeping several draft_vocab_size values re-reads nothing. Two supporting changes: - specforge/training/vocab_mapping.py moves to specforge/data/. Scripts may not import specforge.training (tests/test_runtime/ test_package_architecture.py enforces it), and counting tokens in prepared features is a data concern that the trainer merely happens to be the first caller of. The module is unchanged. - process_token_dict_to_mappings builds t2d by scattering into a zeroed mask instead of testing membership in a Python list per target id. The old form is quadratic in the vocabulary: tens of seconds at V=248320 and K=64000, inside a phase that prints nothing, for a result that takes milliseconds. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
fg11991
requested review from
FlamingoPg,
FrankLeeeee,
shuaills,
sleepcoo and
zyksir
as code owners
August 12, 2026 08:47
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Choosing a draft vocabulary currently requires a full hidden-state capture first:
process_token_dict_to_mappingsruns insideprepare_hidden_states.py. Trying a seconddraft_vocab_size, or building a mapping for a corpus whose features were captured before a mapping was needed, means paying for the capture again — hours of GPU time on a 600k-sample corpus to answer a question about token frequencies.What's in it
scripts/build_vocab_mapping.pyderives the mapping from either source, with no GPU involved:--data-pathcounts loss-bearing tokens from the source JSONL, applying the same chat template and truncation the trainer would;--hidden-states-pathcounts them from already prepared features.Both stream: the tally is accumulated per shard instead of materializing a token column for the whole corpus, and per-corpus counts are cached under a key that includes the corpus identity, so sweeping several
draft_vocab_sizevalues re-reads nothing.Two supporting changes:
specforge/training/vocab_mapping.py→specforge/data/. Scripts may not importspecforge.training(tests/test_runtime/test_package_architecture.pyenforces this), and counting tokens in prepared features is a data concern that the trainer merely happens to be the first caller of. The module itself is unchanged.process_token_dict_to_mappingsbuildst2dby scattering into a zeroed mask instead of testing membership in a Python list per target id. The old form is quadratic in the vocabulary: tens of seconds atV=248320, K=64000, inside a phase that prints nothing, for a result that takes milliseconds.Testing
tests/test_scripts/test_build_vocab_mapping.pycovers both input paths, the cache-identity key, gzipped features, and — the part that actually matters — that the written file loads into a real EAGLE3 draft model's buffers and keeps the frequent tokens.tests/test_data/test_vocab_mapping_construction.pycovers the scatter rewrite against the previous list-based result.Full suite run on CPU (no CUDA available on the dev box): no test fails that does not already fail on
mainat the same commit.