From 7bc76ff49222732d467321fcda20bbb015122e8a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 00:30:20 -0700 Subject: [PATCH] fix(recipes): derive pretraining vocab from tokenizer (#4996) Signed-off-by: Chen Cui (cherry picked from commit af17edf52c514c58c79cd291b04a9b42bc5c57f3) --- docs/nemo2-migration-guide.md | 47 ++++++++++++++---- scripts/performance/run_recipe.py | 17 +++++-- .../bridge/models/gemma/gemma3_provider.py | 2 +- .../bridge/models/gemma/gemma4_provider.py | 2 +- .../qwen_vl/modelling_qwen3_vl/model.py | 15 +++++- .../modelling_qwen3_vl/transformer_config.py | 2 + src/megatron/bridge/perf_recipes/_common.py | 9 +++- .../bridge/perf_recipes/qwen_vl/common.py | 14 ++++++ src/megatron/bridge/recipes/common.py | 1 + .../bridge/recipes/qwen_vl/h100/qwen35_vl.py | 16 ++----- .../bridge/recipes/qwen_vl/h100/qwen3_vl.py | 10 ++-- src/megatron/bridge/training/setup.py | 14 ++++-- .../bridge/training/tokenizers/config.py | 11 +++++ .../ckpts/qwen3_4b/test_qwen3_4b_ckpt.py | 8 +++- .../models/gemma/test_gemma3_provider.py | 5 +- .../models/gemma/test_gemma4_provider.py | 23 +++++++++ .../qwen_vl/modelling_qwen3_vl/test_model.py | 48 +++++++++++++++++++ .../recipes/qwen_vl/test_qwen35_vl_recipes.py | 4 +- .../recipes/qwen_vl/test_qwen3_vl_recipes.py | 4 +- .../recipes/test_all_perf_recipe_factories.py | 8 +++- .../recipes/test_all_recipe_factories.py | 6 +++ .../unit_tests/recipes/test_gemma4_recipe.py | 16 ++++--- .../recipes/test_perf_recipe_environment.py | 4 +- .../performance/test_recipe_environment.py | 38 +++++++++++++++ .../training/test_modelopt_setup.py | 4 +- tests/unit_tests/training/test_setup.py | 22 +++++++++ 26 files changed, 297 insertions(+), 53 deletions(-) diff --git a/docs/nemo2-migration-guide.md b/docs/nemo2-migration-guide.md index 76464e3f57..dfd0f9ff49 100644 --- a/docs/nemo2-migration-guide.md +++ b/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: diff --git a/scripts/performance/run_recipe.py b/scripts/performance/run_recipe.py index 53344b071a..1954962376 100644 --- a/scripts/performance/run_recipe.py +++ b/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). diff --git a/src/megatron/bridge/models/gemma/gemma3_provider.py b/src/megatron/bridge/models/gemma/gemma3_provider.py index e52e9a4545..7718ff9d81 100644 --- a/src/megatron/bridge/models/gemma/gemma3_provider.py +++ b/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, diff --git a/src/megatron/bridge/models/gemma/gemma4_provider.py b/src/megatron/bridge/models/gemma/gemma4_provider.py index 85759f30dd..7b7b28fe0e 100644 --- a/src/megatron/bridge/models/gemma/gemma4_provider.py +++ b/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, diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py index 7d1e1b651d..7ac9ea1cb5 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +++ b/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", diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py index 41cbc08fa6..c8f2bec15d 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +++ b/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 diff --git a/src/megatron/bridge/perf_recipes/_common.py b/src/megatron/bridge/perf_recipes/_common.py index f21d54b18d..fb51246535 100644 --- a/src/megatron/bridge/perf_recipes/_common.py +++ b/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 diff --git a/src/megatron/bridge/perf_recipes/qwen_vl/common.py b/src/megatron/bridge/perf_recipes/qwen_vl/common.py index 4948781827..17f6d97576 100644 --- a/src/megatron/bridge/perf_recipes/qwen_vl/common.py +++ b/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 diff --git a/src/megatron/bridge/recipes/common.py b/src/megatron/bridge/recipes/common.py index 809a4e765f..6c620ad631 100644 --- a/src/megatron/bridge/recipes/common.py +++ b/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( diff --git a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py index 144d0184f3..24c6ac1a58 100644 --- a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py +++ b/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 diff --git a/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py index dc68dabad8..74a8ab150f 100644 --- a/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py +++ b/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 diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index ebe5ed86d1..30efc5da64 100644 --- a/src/megatron/bridge/training/setup.py +++ b/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: diff --git a/src/megatron/bridge/training/tokenizers/config.py b/src/megatron/bridge/training/tokenizers/config.py index 03fab40c7a..22508f0a85 100644 --- a/src/megatron/bridge/training/tokenizers/config.py +++ b/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. diff --git a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py index 2e21c8507f..24c1ae7768 100644 --- a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py +++ b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py @@ -31,6 +31,7 @@ MBRIDGE_CKPT = f"{BASE_DIR}/mbridge" MCORE_CKPT = f"{BASE_DIR}/mcore" TB_DIR = f"{BASE_DIR}/tb" +QWEN3_4B_VOCAB_SIZE = 151936 class TestQwen3Ckpt: @@ -42,6 +43,11 @@ def test_qwen3_4b_ckpt_mbridge(self): config = qwen3_4b_pretrain_config() + # Keep both checkpoint paths independent of the shared Hugging Face cache. + config.tokenizer.tokenizer_type = "NullTokenizer" + config.tokenizer.tokenizer_model = None + config.tokenizer.vocab_size = QWEN3_4B_VOCAB_SIZE + config.checkpoint.save = MBRIDGE_CKPT config.checkpoint.load = MCORE_CKPT if os.path.exists(MCORE_CKPT) else None config.checkpoint.load_optim = False @@ -127,7 +133,7 @@ def test_qwen3_4b_ckpt_mcore(self, monkeypatch): "--tokenizer-type", "NullTokenizer", "--vocab-size", - "151936", + str(QWEN3_4B_VOCAB_SIZE), "--train-iters", f"{train_iters}", "--save-interval", diff --git a/tests/unit_tests/models/gemma/test_gemma3_provider.py b/tests/unit_tests/models/gemma/test_gemma3_provider.py index cd727f4276..845d6bbcaf 100644 --- a/tests/unit_tests/models/gemma/test_gemma3_provider.py +++ b/tests/unit_tests/models/gemma/test_gemma3_provider.py @@ -89,6 +89,7 @@ def test_gemma3_provider_provide_method(self, mock_rotary_embedding, mock_langua # Mock the parent provide method mock_model = Mock() mock_model.embedding = Mock() + mock_model.vocab_size = 262272 mock_model.setup_embeddings_and_output_layer = Mock() provider = Gemma3ModelProvider( @@ -96,7 +97,7 @@ def test_gemma3_provider_provide_method(self, mock_rotary_embedding, mock_langua hidden_size=1152, num_attention_heads=4, kv_channels=256, - vocab_size=262144, + vocab_size=262145, seq_length=32768, ) @@ -109,7 +110,7 @@ def test_gemma3_provider_provide_method(self, mock_rotary_embedding, mock_langua # Verify that custom embedding was created mock_language_embedding.assert_called_once_with( config=provider, - vocab_size=provider.vocab_size, + vocab_size=mock_model.vocab_size, max_sequence_length=provider.seq_length, position_embedding_type=provider.position_embedding_type, scatter_to_sequence_parallel=provider.scatter_embedding_sequence_parallel, diff --git a/tests/unit_tests/models/gemma/test_gemma4_provider.py b/tests/unit_tests/models/gemma/test_gemma4_provider.py index df617211be..6ceafb164a 100644 --- a/tests/unit_tests/models/gemma/test_gemma4_provider.py +++ b/tests/unit_tests/models/gemma/test_gemma4_provider.py @@ -501,6 +501,29 @@ def test_provide_restores_dual_rotary_base(self, provider): mock_rotary.assert_called_once() mock_tied_kv.assert_called_once_with(mock_model, provider) + def test_provide_uses_padded_model_vocab_for_custom_embedding(self): + provider = Gemma4ModelProvider(vocab_size=262145) + mock_model = Mock(vocab_size=262272) + mock_model.embedding = Mock() + mock_model.setup_embeddings_and_output_layer = Mock() + + with ( + patch.object(GPTModelProvider, "provide", return_value=mock_model), + patch("megatron.bridge.models.gemma.gemma4_provider.Gemma3LanguageModelEmbedding") as mock_embedding, + patch("megatron.bridge.models.gemma.gemma4_provider.Gemma4RotaryEmbedding"), + patch("megatron.bridge.models.gemma.gemma4_provider._install_tied_kv"), + ): + provider.provide(pre_process=True, post_process=True) + + assert provider.vocab_size == 262145 + mock_embedding.assert_called_once_with( + config=provider, + vocab_size=262272, + max_sequence_length=provider.seq_length, + position_embedding_type=provider.position_embedding_type, + scatter_to_sequence_parallel=provider.scatter_embedding_sequence_parallel, + ) + def test_provide_restores_dual_rotary_base_on_error(self, provider): with patch.object(GPTModelProvider, "provide", side_effect=RuntimeError("boom")): with pytest.raises(RuntimeError, match="boom"): diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py index 344f5446dc..4407337145 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py @@ -23,6 +23,7 @@ import os from dataclasses import replace from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -411,6 +412,53 @@ def test_shared_embedding_or_output_weight(self, hf_config): weight_no_decoder = model_no_decoder.shared_embedding_or_output_weight() assert weight_no_decoder is None + @pytest.mark.parametrize( + ("vocab_size", "should_pad_vocab", "expected_vocab_size"), + [ + (151669, True, 152064), + (248077, True, 248320), + (151936, False, 151936), + ], + ) + def test_language_model_honors_vocab_padding_policy( + self, + hf_config, + monkeypatch, + vocab_size, + should_pad_vocab, + expected_vocab_size, + ): + """Apply tokenizer-derived padding before constructing the Qwen language model.""" + self._setup_parallel_state(tp_size=1, ep_size=1, pp_size=1) + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + language_transformer_config = self.get_language_transformer_config(hf_config) + language_transformer_config.vocab_size = vocab_size + language_transformer_config.should_pad_vocab = should_pad_vocab + language_transformer_config.make_vocab_size_divisible_by = 128 + language_transformer_config.tensor_model_parallel_size = 4 + + language_model = Mock() + language_model.config.cuda_graph_impl = "none" + language_model.share_embeddings_and_output_weights = False + language_model_constructor = Mock(return_value=language_model) + monkeypatch.setattr( + "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.Qwen3VLGPTModel", + language_model_constructor, + ) + + Qwen3VLModel( + vision_transformer_config=self.get_vision_transformer_config(hf_config), + language_transformer_config=language_transformer_config, + language_transformer_layer_spec=self.get_language_model_layer_spec(), + pre_process=False, + post_process=True, + add_encoder=False, + add_decoder=True, + pg_collection=pg_collection, + ) + + assert language_model_constructor.call_args.kwargs["vocab_size"] == expected_vocab_size + @pytest.mark.timeout(50) def test_set_input_tensor(self, hf_config): """Test set_input_tensor method.""" diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py index 33116c86b1..a20fde30dd 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py +++ b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py @@ -715,7 +715,9 @@ def test_each_qwen35_vl_pretrain_mock_recipe_builds_config(recipe_func: Callable _assert_basic_config(cfg) - assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert cfg.tokenizer.tokenizer_model == cfg.dataset.hf_processor_path + assert cfg.tokenizer.use_tokenizer_vocab_size is True assert getattr(cfg.model, "tensor_model_parallel_size", 1) >= 1 assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py index e9a295d54f..84850aa733 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py +++ b/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py @@ -600,7 +600,9 @@ def test_each_qwen3_vl_pretrain_mock_recipe_builds_config(recipe_func: Callable, _assert_basic_config(cfg) - assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert cfg.tokenizer.tokenizer_model == cfg.dataset.hf_processor_path + assert cfg.tokenizer.use_tokenizer_vocab_size is True assert getattr(cfg.model, "tensor_model_parallel_size", 1) >= 1 assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 diff --git a/tests/unit_tests/recipes/test_all_perf_recipe_factories.py b/tests/unit_tests/recipes/test_all_perf_recipe_factories.py index 4589795a51..79e8b84cf6 100644 --- a/tests/unit_tests/recipes/test_all_perf_recipe_factories.py +++ b/tests/unit_tests/recipes/test_all_perf_recipe_factories.py @@ -21,7 +21,7 @@ import pytest -from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.config import ConfigContainer, MockVLMSFTDatasetConfig from tests.unit_tests.recipes.recipe_test_utils import ( discover_recipe_factories, exported_recipe_factory_keys, @@ -84,3 +84,9 @@ def test_perf_recipe_factory_builds_config(recipe_factory: Callable[..., object] "dist", ): assert getattr(cfg, section) is not None + + assert cfg.tokenizer.use_tokenizer_vocab_size is False + + if "pretrain" in recipe_factory.__name__ and isinstance(cfg.dataset, MockVLMSFTDatasetConfig): + assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.vocab_size == cfg.model.vocab_size diff --git a/tests/unit_tests/recipes/test_all_recipe_factories.py b/tests/unit_tests/recipes/test_all_recipe_factories.py index 64ef8a4a4b..c3d32ff88e 100644 --- a/tests/unit_tests/recipes/test_all_recipe_factories.py +++ b/tests/unit_tests/recipes/test_all_recipe_factories.py @@ -20,6 +20,7 @@ import pytest +from megatron.bridge.recipes.common import _pretrain_common from megatron.bridge.training.config import ConfigContainer from tests.unit_tests.recipes.recipe_test_utils import ( discover_recipe_factories, @@ -65,6 +66,11 @@ def test_all_recipe_factories_are_exported() -> None: assert getattr(_RECIPES_PACKAGE, factory.__name__) is factory +def test_common_pretrain_uses_runtime_tokenizer_vocabulary() -> None: + """From-scratch pretraining derives the model vocabulary from its tokenizer.""" + assert _pretrain_common().tokenizer.use_tokenizer_vocab_size is True + + @pytest.mark.parametrize("recipe_factory", _RUNNABLE_RECIPE_FACTORIES, ids=recipe_factory_id) def test_recipe_factory_builds_config(recipe_factory: Callable[..., object]) -> None: """Every supported recipe can be called with defaults without GPU or network access.""" diff --git a/tests/unit_tests/recipes/test_gemma4_recipe.py b/tests/unit_tests/recipes/test_gemma4_recipe.py index f4ba2c557b..54439cbbb9 100644 --- a/tests/unit_tests/recipes/test_gemma4_recipe.py +++ b/tests/unit_tests/recipes/test_gemma4_recipe.py @@ -57,26 +57,28 @@ def to_megatron_provider(self, load_weights=True): return self.provider -def _load_gemma4_recipe_module(): +def _load_gemma4_recipe_module(monkeypatch): """Load the Gemma4 recipe without importing the umbrella recipes package.""" bridge_root = Path(__file__).resolve().parents[3] recipes_root = bridge_root / "src" / "megatron" / "bridge" / "recipes" recipes_pkg = types.ModuleType("megatron.bridge.recipes") recipes_pkg.__path__ = [str(recipes_root)] - sys.modules.setdefault("megatron.bridge.recipes", recipes_pkg) + if "megatron.bridge.recipes" not in sys.modules: + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes", recipes_pkg) common_mod = types.ModuleType("megatron.bridge.recipes.common") common_mod._pretrain_common = _minimal_pretrain_common - sys.modules["megatron.bridge.recipes.common"] = common_mod + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes.common", common_mod) utils_pkg = types.ModuleType("megatron.bridge.recipes.utils") utils_pkg.__path__ = [str(recipes_root / "utils")] - sys.modules.setdefault("megatron.bridge.recipes.utils", utils_pkg) + if "megatron.bridge.recipes.utils" not in sys.modules: + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes.utils", utils_pkg) tokenizer_mod = types.ModuleType("megatron.bridge.recipes.utils.tokenizer_utils") tokenizer_mod.DEFAULT_NULL_TOKENIZER_VOCAB_SIZE = 32000 - sys.modules["megatron.bridge.recipes.utils.tokenizer_utils"] = tokenizer_mod + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes.utils.tokenizer_utils", tokenizer_mod) recipe_path = recipes_root / "gemma" / "gemma4.py" spec = importlib.util.spec_from_file_location("_gemma4_recipe_under_test", recipe_path) @@ -124,8 +126,8 @@ def bridge_provider(hf_config_e4b): @pytest.fixture -def recipe_module(): - return _load_gemma4_recipe_module() +def recipe_module(monkeypatch): + return _load_gemma4_recipe_module(monkeypatch) @pytest.fixture diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index d644761000..fad2b67021 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -98,9 +98,10 @@ def test_common_environment_defaults_are_small_and_universal(): assert COMMON_PERF_ENV_VARS == {"TORCH_NCCL_HIGH_PRIORITY": 1} -def test_benchmark_common_preserves_legacy_manual_gc_defaults(): +def test_benchmark_common_preserves_legacy_manual_gc_and_model_shape_defaults(): cfg = SimpleNamespace( train=SimpleNamespace(train_iters=0, eval_iters=1, manual_gc=False, manual_gc_interval=0), + tokenizer=SimpleNamespace(use_tokenizer_vocab_size=True), checkpoint=SimpleNamespace(save="checkpoint"), logger=SimpleNamespace(log_interval=10, tensorboard_dir="tensorboard"), ddp=SimpleNamespace(check_for_nan_in_grad=True, check_for_large_grads=True, grad_reduce_in_fp32=True), @@ -121,6 +122,7 @@ def test_benchmark_common_preserves_legacy_manual_gc_defaults(): assert cfg.train.manual_gc is True assert cfg.train.manual_gc_interval == 100 + assert cfg.tokenizer.use_tokenizer_vocab_size is False def test_every_flat_recipe_builder_declares_its_environment_inline(): diff --git a/tests/unit_tests/scripts/performance/test_recipe_environment.py b/tests/unit_tests/scripts/performance/test_recipe_environment.py index b0564b54de..dd6083a580 100644 --- a/tests/unit_tests/scripts/performance/test_recipe_environment.py +++ b/tests/unit_tests/scripts/performance/test_recipe_environment.py @@ -948,6 +948,44 @@ def apply_hydra(config, _overrides): assert effective_recipe.model.moe_flex_dispatcher_backend is None +def test_training_tokenizer_override_preserves_recipe_vocab_policy(): + """Selecting a runtime tokenizer keeps the canonical pretraining vocabulary policy.""" + from argument_parser import parse_cli_args + + from megatron.bridge.recipes.common import _pretrain_common + + parser = parse_cli_args() + args, unknown = parser.parse_known_args( + [ + "--model_family_name", + "gpt", + "--model_recipe_name", + "vanilla_gpt", + "--num_gpus", + "1", + "--gpu", + "h100", + "--use_recipes", + "--max_steps", + "1000", + "--tokenizer_type", + "HuggingFaceTokenizer", + "--tokenizer_model", + "test-tokenizer", + ] + ) + assert unknown == [] + + recipe = _pretrain_common() + assert recipe.tokenizer.use_tokenizer_vocab_size is True + + updated = run_recipe._apply_training_argparse_overrides(recipe, args) + + assert updated.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert updated.tokenizer.tokenizer_model == "test-tokenizer" + assert updated.tokenizer.use_tokenizer_vocab_size is True + + def test_prepare_recipe_runs_base_override_and_finalize_stages(monkeypatch): """Recipe preparation has one visible path with exactly three config stages.""" args = SimpleNamespace( diff --git a/tests/unit_tests/training/test_modelopt_setup.py b/tests/unit_tests/training/test_modelopt_setup.py index 4c7358ee42..c8d0fc694f 100644 --- a/tests/unit_tests/training/test_modelopt_setup.py +++ b/tests/unit_tests/training/test_modelopt_setup.py @@ -46,7 +46,7 @@ class StopAfterHooksRegistered(Exception): model=SimpleNamespace( fine_grained_activation_offloading=False, restore_modelopt_state=True, - vocab_size=32, + vocab_size=64, ), peft=None, profiling=SimpleNamespace(), @@ -96,6 +96,8 @@ class StopAfterHooksRegistered(Exception): with pytest.raises(StopAfterHooksRegistered): training_setup.setup(state, Mock()) + assert cfg.model.vocab_size == 64 + assert cfg.model.should_pad_vocab is False assert len(hooks) == 1 if resume_has_modelopt_state: hooks[0]([]) diff --git a/tests/unit_tests/training/test_setup.py b/tests/unit_tests/training/test_setup.py index 6f099271ff..ad91729521 100644 --- a/tests/unit_tests/training/test_setup.py +++ b/tests/unit_tests/training/test_setup.py @@ -195,6 +195,28 @@ def test_vocab_size_larger_than_tokenizer_returns_same_value(self): assert vocab_size == 40960 assert should_pad_vocab is False + def test_pretraining_uses_tokenizer_vocab_over_larger_model_vocab(self): + """Tokenizer-derived pretraining vocab ignores the source model vocabulary.""" + vocab_size, should_pad_vocab = _validate_and_set_vocab_size( + model_vocab_size=248320, + tokenizer_vocab_size=32000, + use_tokenizer_vocab_size=True, + ) + + assert vocab_size == 32000 + assert should_pad_vocab is True + + def test_checkpoint_compatibility_override_preserves_model_vocab(self): + """Disabling the policy preserves the explicit vocabulary used by an existing checkpoint.""" + vocab_size, should_pad_vocab = _validate_and_set_vocab_size( + model_vocab_size=248320, + tokenizer_vocab_size=32000, + use_tokenizer_vocab_size=False, + ) + + assert vocab_size == 248320 + assert should_pad_vocab is False + def test_vocab_size_equal_to_tokenizer_returns_same_value(self): """Test that vocab_size equal to tokenizer returns the same value and disables padding.""" vocab_size, should_pad_vocab = _validate_and_set_vocab_size(