cp: fix(recipes): derive pretraining vocab from tokenizer (4996) into r0.6.0

docs/nemo2-migration-guide.md

@@ -693,19 +693,25 @@ tokenizer_config = TokenizerConfig(
 
 #### Vocab Size Priority
 
-In Megatron Bridge, vocabulary size can be specified in either the model provider or derived from the tokenizer. The priority order is:
+In Megatron Bridge, vocabulary size can be specified in the model provider or derived from the runtime tokenizer. The priority order is:
 
-1. **Model provider `vocab_size` is set**: Uses the model's vocab size
-   - Must be `>= tokenizer.vocab_size` (raises error if smaller)
-   - Sets `should_pad_vocab=False` (no automatic padding)
-   - Useful when you need a specific vocab size (e.g., for checkpoint compatibility)
+1. **`TokenizerConfig.use_tokenizer_vocab_size=True`**: Uses the tokenizer's vocab size
+   - Overrides a preset model-provider `vocab_size`.
+   - Sets `should_pad_vocab=True` (enables padding for efficient parallelism).
+   - Intended for from-scratch pretraining, where the dataset tokenizer defines the vocabulary.
+   - This policy remains active during checkpoint loading; disable it when checkpoint compatibility requires the explicit model vocabulary.
 
-2. **Model provider `vocab_size` is None**: Uses tokenizer's vocab size
+2. **Model provider `vocab_size` is set**: Uses the model's vocab size
+   - Must be `>= tokenizer.vocab_size` (raises an error if smaller).
+   - Sets `should_pad_vocab=False` (no automatic padding).
+   - Useful when a specific vocabulary is required for model or checkpoint compatibility.
+
+3. **Model provider `vocab_size` is `None`**: Uses the tokenizer's vocab size
    - Automatically derived from `tokenizer.vocab_size` after building the tokenizer.
-   - Sets `should_pad_vocab=True` (enables padding for efficient parallelism)
+   - Sets `should_pad_vocab=True`.
 
 ```python
-# Option 1: Let tokenizer determine vocab size
+# Option 1: Let tokenizer determine vocab size when the model has no preset
 config = ConfigContainer(
     model=GPTModelProvider(
         # vocab_size not set - will use tokenizer's vocab size
@@ -717,7 +723,19 @@ config = ConfigContainer(
     ),
 )
 
-# Option 2: Explicitly set vocab size in model
+# Option 2: Override a preset model vocab for from-scratch pretraining
+config = ConfigContainer(
+    model=GPTModelProvider(
+        vocab_size=128256,  # Ignored whenever the flag is enabled
+    ),
+    tokenizer=TokenizerConfig(
+        tokenizer_type="HuggingFaceTokenizer",
+        tokenizer_model="my-org/my-pretraining-tokenizer",
+        use_tokenizer_vocab_size=True,
+    ),
+)
+
+# Option 3: Explicitly set vocab size in model
 config = ConfigContainer(
     model=GPTModelProvider(
         vocab_size=128256,  # Explicitly set (must be >= tokenizer vocab size)
@@ -726,6 +744,17 @@ config = ConfigContainer(
 )
 ```
 
+Pretraining recipes enable `use_tokenizer_vocab_size` by default. For a new run, use an empty checkpoint directory so the runtime tokenizer defines the model vocabulary. A checkpoint created by that policy can be resumed with the same tokenizer and recipe configuration.
+
+Checkpoints created before a recipe enabled `use_tokenizer_vocab_size` may have used the model provider's larger explicit vocabulary. To preserve their embedding and output tensor shapes, disable the new policy and retain the vocabulary used to create the checkpoint:
+
+```python
+config.tokenizer.use_tokenizer_vocab_size = False
+config.model.vocab_size = 128256  # The vocabulary used to create the checkpoint
+```
+
+Do not change this setting partway through a run. Switching vocabulary policies changes model tensor shapes and is not a checkpoint migration mechanism.
+
 
 ### Parallelism Configuration Migration
 In NeMo 2.0, parallelism settings were configured on `MegatronStrategy`. In Megatron Bridge, these are set directly on the model provider:

scripts/performance/run_recipe.py

@@ -129,18 +129,29 @@ def _apply_training_argparse_overrides(config, args):
         # Tokenizer configuration
         from megatron.bridge.training.config import TokenizerConfig
 
+        use_tokenizer_vocab_size = config.tokenizer.use_tokenizer_vocab_size
         if args.tokenizer_type == "NullTokenizer":
-            config.tokenizer = TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=args.vocab_size)
+            config.tokenizer = TokenizerConfig(
+                tokenizer_type="NullTokenizer",
+                vocab_size=args.vocab_size,
+                use_tokenizer_vocab_size=use_tokenizer_vocab_size,
+            )
         elif args.tokenizer_type == "HuggingFaceTokenizer":
             if not args.tokenizer_model:
                 raise ValueError("--tokenizer-model is required when using HuggingFaceTokenizer")
             tokenizer_model = args.tokenizer_model
-            config.tokenizer = TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model=tokenizer_model)
+            config.tokenizer = TokenizerConfig(
+                tokenizer_type="HuggingFaceTokenizer",
+                tokenizer_model=tokenizer_model,
+                use_tokenizer_vocab_size=use_tokenizer_vocab_size,
+            )
         elif args.tokenizer_type == "SentencePieceTokenizer":
             if not args.tokenizer_model:
                 raise ValueError("--tokenizer-model is required for SentencePieceTokenizer")
             config.tokenizer = TokenizerConfig(
-                tokenizer_type="SentencePieceTokenizer", tokenizer_model=args.tokenizer_model
+                tokenizer_type="SentencePieceTokenizer",
+                tokenizer_model=args.tokenizer_model,
+                use_tokenizer_vocab_size=use_tokenizer_vocab_size,
             )
     else:
         # Diffusion recipes (FLUX, WAN) keep their own dataset object (Wan/FluxDatasetConfig).

src/megatron/bridge/models/gemma/gemma3_provider.py

@@ -119,7 +119,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> "MCoreG
         if hasattr(model, "embedding"):
             model.embedding = Gemma3LanguageModelEmbedding(
                 config=self,
-                vocab_size=self.vocab_size,
+                vocab_size=model.vocab_size,
                 max_sequence_length=self.seq_length,
                 position_embedding_type=self.position_embedding_type,
                 scatter_to_sequence_parallel=self.scatter_embedding_sequence_parallel,

src/megatron/bridge/models/gemma/gemma4_provider.py

@@ -345,7 +345,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> "MCoreG
         if hasattr(model, "embedding"):
             model.embedding = Gemma3LanguageModelEmbedding(
                 config=self,
-                vocab_size=self.vocab_size,
+                vocab_size=model.vocab_size,
                 max_sequence_length=self.seq_length,
                 position_embedding_type=self.position_embedding_type,
                 scatter_to_sequence_parallel=self.scatter_embedding_sequence_parallel,

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py

@@ -55,6 +55,7 @@
     get_packed_seq_cp_partition_indices,
     get_packed_seq_q_cu_seqlens,
 )
+from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 
 
 def _is_mrope_position_ids(position_ids: torch.Tensor | None) -> bool:
@@ -324,10 +325,22 @@ def __init__(
                 pg_collection=pg_collection,
             )
         if self.add_decoder:
+            assert language_transformer_config.vocab_size is not None, (
+                "vocab_size must be configured before constructing the Qwen3-VL language model"
+            )
+            if language_transformer_config.should_pad_vocab:
+                language_model_vocab_size = calculate_padded_vocab_size(
+                    language_transformer_config.vocab_size,
+                    language_transformer_config.make_vocab_size_divisible_by,
+                    language_transformer_config.tensor_model_parallel_size,
+                )
+            else:
+                language_model_vocab_size = language_transformer_config.vocab_size
+
             self.language_model = Qwen3VLGPTModel(
                 config=language_transformer_config,
                 transformer_layer_spec=language_transformer_layer_spec,
-                vocab_size=language_transformer_config.vocab_size,
+                vocab_size=language_model_vocab_size,
                 max_sequence_length=language_transformer_config.language_max_sequence_length,
                 parallel_output=parallel_output,
                 position_embedding_type="mrope",

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py

@@ -28,6 +28,8 @@ class Qwen3VLTransformerConfig(TransformerConfig):
     """Configuration for Qwen3-VL transformer with vision and language components."""
 
     vocab_size: int = 64000
+    make_vocab_size_divisible_by: int = 128
+    should_pad_vocab: bool = False
     language_max_sequence_length: int = 4096
 
     patch_size: int = 16

src/megatron/bridge/perf_recipes/_common.py

@@ -33,8 +33,9 @@ def _benchmark_common(cfg: ConfigContainer, cross_entropy_impl: str = "te") -> N
     Intended for performance benchmark recipes only. Sets short training runs,
     disables checkpointing/eval, tunes scheduler, and enables perf-oriented kernels.
 
-    Must stay in sync with ``_set_common_perf_overrides`` in
-    ``scripts/performance/utils/overrides.py``.
+    This is the fixed-model-shape policy for flat performance recipes.
+    Canonical recipes launched with ``scripts/performance --use_recipes``
+    retain their own tokenizer vocabulary policy.
 
     Individual recipes may override any of these after calling this function
     (e.g. Kimi K2 sets ``grad_reduce_in_fp32 = True``).
@@ -44,6 +45,10 @@ def _benchmark_common(cfg: ConfigContainer, cross_entropy_impl: str = "te") -> N
     cfg.train.manual_gc = True
     cfg.train.manual_gc_interval = 100
 
+    # Performance recipes benchmark a fixed model shape. Synthetic or runtime
+    # tokenizers must not resize the embedding and output layers during setup.
+    cfg.tokenizer.use_tokenizer_vocab_size = False
+
     cfg.checkpoint.save = None
 
     cfg.logger.log_interval = 1

src/megatron/bridge/perf_recipes/qwen_vl/common.py

@@ -32,11 +32,24 @@
 from megatron.bridge.training.config import ConfigContainer
 
 
+def _use_model_vocab_null_tokenizer(cfg: ConfigContainer) -> None:
+    """Use a model-sized synthetic tokenizer for Qwen-VL performance runs."""
+    if cfg.model.vocab_size is None:
+        raise ValueError("Qwen-VL performance recipes require a model vocabulary size.")
+    cfg.tokenizer.tokenizer_type = "NullTokenizer"
+    cfg.tokenizer.tokenizer_model = None
+    cfg.tokenizer.vocab_size = cfg.model.vocab_size
+    # The synthetic tokenizer mirrors the fixed benchmark model shape; it does
+    # not define a new tokenizer-derived vocabulary for from-scratch training.
+    cfg.tokenizer.use_tokenizer_vocab_size = False
+
+
 def _qwen35_vl_common(cfg: ConfigContainer) -> None:
     """Apply VLM-specific performance benchmark settings for Qwen3.5-VL.
 
     Must be called before ``_benchmark_common`` and after setting precision.
     """
+    _use_model_vocab_null_tokenizer(cfg)
     cfg.model.bias_activation_fusion = True
     cfg.model.recompute_granularity = None
     cfg.model.recompute_method = None
@@ -82,6 +95,7 @@ def _qwen35_vl_post_clear_scope_with_overlap(cfg: ConfigContainer) -> None:
 
 def _finalize_qwen3_vl(cfg: ConfigContainer) -> None:
     """Apply Qwen3-VL perf defaults that must override generic benchmark defaults."""
+    _use_model_vocab_null_tokenizer(cfg)
     # _benchmark_common sets apply_rope_fusion=True; Qwen3-VL asserts it must be False
     # (per-token absolute positional frequencies are incompatible with TE's fused RoPE).
     cfg.model.apply_rope_fusion = False

src/megatron/bridge/recipes/common.py

@@ -111,6 +111,7 @@ def _pretrain_common() -> ConfigContainer:
         tokenizer=TokenizerConfig(
             tokenizer_type="HuggingFaceTokenizer",
             tokenizer_model=None,  # Must be set by each recipe
+            use_tokenizer_vocab_size=True,
         ),
         # Checkpoint config
         checkpoint=CheckpointConfig(

src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py

@@ -30,7 +30,6 @@
 from megatron.bridge.recipes.utils.dataset_utils import default_peft_config
 from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS
 from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing
-from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
 from megatron.bridge.training.config import ConfigContainer
 
 
@@ -72,8 +71,7 @@ def qwen35_vl_9b_pretrain_4gpu_h100_bf16_mock_config() -> ConfigContainer:
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False
@@ -121,8 +119,7 @@ def qwen35_vl_27b_pretrain_16gpu_h100_bf16_mock_config() -> ConfigContainer:
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False
@@ -171,8 +168,7 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_h100_bf16_mock_config() -> ConfigContainer:
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False
@@ -222,8 +218,7 @@ def qwen35_vl_122b_a10b_pretrain_128gpu_h100_bf16_mock_config() -> ConfigContain
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False
@@ -274,8 +269,7 @@ def qwen35_vl_397b_a17b_pretrain_512gpu_h100_bf16_mock_config() -> ConfigContain
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False

src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py

@@ -32,7 +32,6 @@
 from megatron.bridge.recipes.utils.dataset_utils import default_peft_config
 from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS
 from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing
-from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend
 
@@ -77,8 +76,7 @@ def qwen3_vl_8b_pretrain_4gpu_h100_bf16_mock_config() -> ConfigContainer:
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False
@@ -127,8 +125,7 @@ def qwen3_vl_30b_a3b_pretrain_8gpu_h100_bf16_mock_config() -> ConfigContainer:
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False
@@ -178,8 +175,7 @@ def qwen3_vl_235b_a22b_pretrain_256gpu_h100_bf16_mock_config() -> ConfigContaine
         persistent_workers=False,
         pad_to_max_length=True,
     )
-    cfg.tokenizer.tokenizer_type = "NullTokenizer"
-    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+    cfg.tokenizer.tokenizer_model = hf_path
     cfg.train.eval_interval = 500
     cfg.train.eval_iters = 32
     cfg.ddp.overlap_grad_reduce = False

src/megatron/bridge/training/setup.py

@@ -248,6 +248,7 @@ def setup(
     cfg.model.vocab_size, cfg.model.should_pad_vocab = _validate_and_set_vocab_size(
         model_vocab_size=cfg.model.vocab_size,
         tokenizer_vocab_size=tokenizer.vocab_size,
+        use_tokenizer_vocab_size=getattr(cfg.tokenizer, "use_tokenizer_vocab_size", False),
     )
 
     if hasattr(cfg.dataset, "tokenizer"):
@@ -630,12 +631,18 @@ def _apply_peft_transformation(peft, base_model: list[MegatronModule]) -> list[M
     return transformed_model
 
 
-def _validate_and_set_vocab_size(model_vocab_size: Optional[int], tokenizer_vocab_size: int) -> tuple[int, bool]:
+def _validate_and_set_vocab_size(
+    model_vocab_size: Optional[int],
+    tokenizer_vocab_size: int,
+    use_tokenizer_vocab_size: bool = False,
+) -> tuple[int, bool]:
     """Validate and determine the correct vocab size for the model.
 
     Args:
         model_vocab_size: Vocab size set in model config (can be None)
         tokenizer_vocab_size: Unpadded tokenizer vocab size
+        use_tokenizer_vocab_size: Ignore a preset model vocabulary and derive it
+            from the tokenizer. Intended for from-scratch pretraining recipes.
 
     Returns:
         tuple[int, bool]: The validated unpadded vocab size and padding flag
@@ -645,8 +652,9 @@ def _validate_and_set_vocab_size(model_vocab_size: Optional[int], tokenizer_voca
     Raises:
         ValueError: If model vocab size is invalid
     """
-    if model_vocab_size is None:
-        # If model vocab size is not set, use the tokenizer's vocab size
+    if use_tokenizer_vocab_size or model_vocab_size is None:
+        # Use the tokenizer's vocab size when the model vocab is unset, or when
+        # use_tokenizer_vocab_size forces it for from-scratch pretraining.
         # Enable padding since this came from tokenizer
         return tokenizer_vocab_size, True
     elif model_vocab_size < tokenizer_vocab_size:

src/megatron/bridge/training/tokenizers/config.py

@@ -33,6 +33,17 @@ class TokenizerConfig(MTrainTokenizerConfig):
     rank: int = 0
     """Distributed rank used by MCore tokenizer helper logging."""
 
+    use_tokenizer_vocab_size: bool = False
+    """Use the runtime tokenizer vocabulary size for the model.
+
+    Enable this for from-scratch pretraining, where the tokenizer selected for
+    the dataset defines the embedding and output vocabulary. Keep it disabled
+    when model or checkpoint compatibility requires an explicitly configured
+    model vocabulary size. This policy also applies during checkpoint loading;
+    disable it and configure the checkpoint's original model vocabulary when
+    resuming a run created with a different vocabulary policy.
+    """
+
     hf_tokenizer_kwargs: dict[str, Any] | None = field(default_factory=dict)
     """Additional keyword arguments to pass to HuggingFace AutoTokenizer.from_pretrained.