Skip to content

fix(recipes): derive pretraining vocab from tokenizer - #4996

Merged
malay-nagda merged 12 commits into
mainfrom
agent/pretraining-tokenizer-vocab
Jul 28, 2026
Merged

fix(recipes): derive pretraining vocab from tokenizer#4996
malay-nagda merged 12 commits into
mainfrom
agent/pretraining-tokenizer-vocab

Conversation

@cuichenx

@cuichenx cuichenx commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Summary

  • add an explicit TokenizerConfig.use_tokenizer_vocab_size policy for from-scratch pretraining
  • enable it in the shared language-model pretraining template
  • resolve model embedding/output vocabulary from the runtime tokenizer even when AutoBridge supplied a Hugging Face model vocabulary
  • use the matching Hugging Face tokenizer in canonical Qwen3-VL and Qwen3.5-VL pretraining recipes, so they follow the same policy without a recipe-specific override
  • keep NullTokenizer in the Qwen-VL performance layer, sized to the model vocabulary so it cannot change the benchmark model shape
  • document vocabulary precedence and the explicit compatibility override for checkpoints created under the previous policy

Root cause

AutoBridge maps the Hugging Face vocab_size into the Megatron provider even when recipes use load_weights=False. Training setup treated that value as authoritative when it was larger than the runtime tokenizer. A launcher replacing the recipe tokenizer could therefore keep the source model vocabulary instead of sizing the from-scratch model to the dataset tokenizer.

This caused Nemotron 3 Nano to use 131072 entries with the RP2 32K tokenizer instead of its previous 32000-entry behavior. Qwen3.5 text similarly retained 248320 entries with the RP2 32K tokenizer.

Qwen-VL initially needed an exception because its mock dataset gets token IDs from the Hugging Face processor while the recipe exposed a 32000-entry NullTokenizer. Enabling tokenizer-derived vocabulary with that placeholder would have incorrectly resized the language model to 32000 entries. Canonical Qwen-VL pretraining now uses the real tokenizer at the processor path; only performance recipes replace it with a model-sized null tokenizer.

Scope

All language-model recipes that stem from the shared pretraining template select the runtime tokenizer vocabulary. Standalone vanilla GPT needs no explicit flag override because its model vocabulary is unset, so setup already derives it from its synthetic tokenizer.

Canonical Qwen3-VL and Qwen3.5-VL mock-pretraining recipes now follow the same tokenizer-derived policy. Their performance wrappers preserve the source model shape with a model-sized NullTokenizer. Diffusion vocabularies remain unrelated to language-token embeddings.

SFT, PEFT, conversion, and checkpoint-compatible flows retain the existing explicit-model-vocabulary default.

Checkpoint compatibility

Pretraining recipes enable use_tokenizer_vocab_size for new from-scratch runs. The policy is unconditional, including during checkpoint loading, and must not be changed partway through a run.

To resume a checkpoint created before a recipe enabled this policy, set tokenizer.use_tokenizer_vocab_size=False and retain the explicit model.vocab_size used to create that checkpoint. This preserves its embedding/output shape without guessing across persistent, non-persistent, resharded, or custom-managed checkpoint formats.

Validation

  • 26.08 RC2 container: 840 targeted setup, Qwen-VL recipe, canonical recipe-factory, and performance recipe-factory tests passed
  • repository-wide pre-commit suite: passed
  • independent self-review after the compatibility amendments: no blockers

This is independent of #4854. Once combined, its Qwen3.5 text recipes inherit the tokenizer-derived vocabulary policy without recipe-specific overrides.

Signed-off-by: Chen Cui <chcui@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 21, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Signed-off-by: Chen Cui <chcui@nvidia.com>
Comment thread src/megatron/bridge/training/setup.py Outdated
@claude

claude Bot commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Review — feat: tokenizer-driven vocab size for pretraining

Solid, well-scoped change. The new use_tokenizer_vocab_size flag on TokenizerConfig is defaulted False (safe, backward-compatible), enabled centrally in _pretrain_common(), and explicitly disabled in the qwen_vl mock VLM recipes where a NullTokenizer is only a runtime placeholder. The logic change in _validate_and_set_vocab_size is minimal and correct, and both the from-scratch and checkpoint-compat paths are unit-tested.

Correctness checks (all passed):

  • _validate_and_set_vocab_size: use_tokenizer_vocab_size=True short-circuits to (tokenizer_vocab_size, True), ignoring a larger preset model vocab. False preserves prior behavior exactly. The default arg keeps the single call site and external callers backward-compatible.
  • Coverage sweep: the only GPT/text pretrain recipe that bypasses _pretrain_common() is vanilla_gpt.py, which sets the flag True directly — no gap. All standard text/MoE pretrain recipes route through _pretrain_common(). wan.py/flux.py keep the default False but are diffusion recipes (not GPTDatasetConfig), so that is correct.
  • All 3 pretrain factories in qwen3_vl.py and all 5 in qwen35_vl.py set use_tokenizer_vocab_size = False, matching the added test assertions 1:1.
  • Docs (nemo2-migration-guide.md) reordered the priority list accurately and added a clear checkpoint-migration warning. No other doc duplicates the old vocab-priority section, so no stale references.

Minor (non-blocking): the inline comment at setup.py:652 now only describes one of the two branch conditions; the branch also fires on use_tokenizer_vocab_size=True. Suggested tweak posted inline.

Nit: consider a unit test for the use_tokenizer_vocab_size=True + model_vocab_size=None combination to lock in that the flag composes with an unset model vocab (both existing cases use an explicit model vocab).

Suggested test cases
No perf tests impacted (no scripts/performance/configs/ recipe configs touched). The recipe/setup changes are covered by unit tests:

  • tests/unit_tests/training/test_setup.py::TestValidateAndSetVocabSize::test_pretraining_uses_tokenizer_vocab_over_larger_model_vocab
  • tests/unit_tests/training/test_setup.py::TestValidateAndSetVocabSize::test_checkpoint_compatibility_override_preserves_model_vocab
  • tests/unit_tests/recipes/test_all_recipe_factories.py::test_recipe_factory_builds_config
  • tests/unit_tests/recipes/test_all_perf_recipe_factories.py::test_perf_recipe_factory_builds_config

Signed-off-by: Chen Cui <chcui@nvidia.com>
@yaoyu-33 yaoyu-33 added area:recipe Training recipes and launch configs breaking-change Public behavior or API compatibility changes bug Something isn't working full-test-suite needs-review PR is ready for code review and waiting on a reviewer labels Jul 21, 2026
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
balasaajay
balasaajay previously approved these changes Jul 24, 2026
Signed-off-by: Chen Cui <chcui@nvidia.com>
yaoyu-33
yaoyu-33 previously approved these changes Jul 25, 2026
Signed-off-by: Chen Cui <chcui@nvidia.com>
@malay-nagda
malay-nagda merged commit af17edf into main Jul 28, 2026
156 checks passed
@malay-nagda
malay-nagda deleted the agent/pretraining-tokenizer-vocab branch July 28, 2026 07:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:recipe Training recipes and launch configs breaking-change Public behavior or API compatibility changes bug Something isn't working full-test-suite needs-review PR is ready for code review and waiting on a reviewer

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants