docs(training): clarify vocab checkpoint migration

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

Author: cuichenx

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:

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

@@ -39,7 +39,9 @@ class TokenizerConfig(MTrainTokenizerConfig):
     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.
+    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)

tests/unit_tests/training/test_setup.py

@@ -206,6 +206,17 @@ def test_pretraining_uses_tokenizer_vocab_over_larger_model_vocab(self):
         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(