Skip to content

feat(scripts): build a draft vocabulary mapping without regenerating features - #765

Open
fg11991 wants to merge 2 commits into
sgl-project:mainfrom
fg11991:feat/vocab-mapping-tooling
Open

feat(scripts): build a draft vocabulary mapping without regenerating features#765
fg11991 wants to merge 2 commits into
sgl-project:mainfrom
fg11991:feat/vocab-mapping-tooling

Conversation

@fg11991

@fg11991 fg11991 commented Aug 12, 2026

Copy link
Copy Markdown

Motivation

Choosing a draft vocabulary currently requires a full hidden-state capture first: process_token_dict_to_mappings runs inside prepare_hidden_states.py. Trying a second draft_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.py derives the mapping from either source, with no GPU involved:

  • --data-path counts loss-bearing tokens from the source JSONL, 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 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.pyspecforge/data/. Scripts may not import specforge.training (tests/test_runtime/test_package_architecture.py enforces 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_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, K=64000, inside a phase that prints nothing, for a result that takes milliseconds.

Testing

tests/test_scripts/test_build_vocab_mapping.py covers 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.py covers 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 main at the same commit.

…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>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant