cp: Enable overriding TokenizerConfig when instantiating from checkpoint (1030) into r0.2.0

src/megatron/bridge/training/model_load_save.py

@@ -118,7 +118,7 @@ def temporary_distributed_context(backend: str = "gloo") -> Generator[None, None
         dist.destroy_process_group()
 
 
-def load_tokenizer(checkpoint_path: str) -> MegatronTokenizer:
+def load_tokenizer(checkpoint_path: str, **kwargs) -> MegatronTokenizer:
     """Create a tokenizer from a training checkpoint.
 
     Obtains tokenizer configuration from the checkpoint and builds the tokenizer.
@@ -127,6 +127,9 @@ def load_tokenizer(checkpoint_path: str) -> MegatronTokenizer:
     Args:
         checkpoint_path: path to an MCore distributed checkpoint directory
                           (e.g., /path/to/model/checkpoints/iter_0000001).
+        **kwargs: Overrides to the TokenizerConfig used to build the tokenizer.
+                Useful if the tokenizer assets have moved since training.
+
     """
     from megatron.bridge.training.checkpointing import (
         get_checkpoint_run_config_filename,
@@ -152,6 +155,14 @@ def load_tokenizer(checkpoint_path: str) -> MegatronTokenizer:
     else:
         cfg = _tokenizer_config_from_args(mlm_args)
 
+    for key, val in kwargs.items():
+        if hasattr(cfg, key):
+            setattr(cfg, key, val)
+        else:
+            raise AttributeError(
+                f"Attempting to set a non-existent attribute '{key}' on TokenizerConfig.\nState of TokenizerConfig before attempting this override: {cfg}"
+            )
+
     return build_tokenizer(cfg)
 
 

tests/unit_tests/training/test_model_load_save.py

@@ -20,6 +20,7 @@
 import torch
 
 from megatron.bridge.models.model_provider import ModelProviderMixin
+from megatron.bridge.training.config import TokenizerConfig
 from megatron.bridge.training.model_load_save import (
     dtype_from_hf,
     dtype_from_str,
@@ -905,3 +906,37 @@ def test_load_mlm_saved_tokenizer(self, mock_load_args, mock_cfg_from_args, mock
         mock_load_args.assert_called_once_with(ckpt_path)
         mock_cfg_from_args.assert_called_once_with(mock_args)
         mock_build_tokenizer.assert_called_once_with(mock_tokenizer_cfg)
+
+    @patch("megatron.bridge.training.model_load_save.build_tokenizer")
+    @patch("megatron.bridge.utils.instantiate_utils.instantiate")
+    @patch("megatron.bridge.training.checkpointing.read_run_config")
+    def test_load_tokenizer_with_kwargs(self, mock_read_cfg, mock_instantiate, mock_build_tokenizer, mock_tokenizer):
+        """Test loading tokenizer config and overriding."""
+        # Setup mocks
+        mock_run_cfg_dict = {
+            "model": {"tensor_model_parallel_size": 1, "make_vocab_size_divisible_by": 128},
+            "tokenizer": {},
+        }
+        mock_read_cfg.return_value = mock_run_cfg_dict
+
+        mock_tokenizer_cfg = Mock(spec=TokenizerConfig)
+        mock_tokenizer_cfg.vocab_size = 32000
+        mock_tokenizer_cfg.tokenizer_model = "/path/to/tokenizer.model"
+        mock_instantiate.return_value = mock_tokenizer_cfg
+
+        mock_build_tokenizer.return_value = mock_tokenizer
+
+        # test changing asset filepath
+        new_asset_path = "/path/to/different/tokenizer.model"
+        with tempfile.TemporaryDirectory() as ckpt_path:
+            config_file = Path(ckpt_path) / "run_config.yaml"
+            config_file.touch()
+            _ = load_tokenizer(ckpt_path, tokenizer_model=new_asset_path)
+
+            assert mock_tokenizer_cfg.tokenizer_model == new_asset_path
+
+            # test setting attribute that doesn't exist
+            with pytest.raises(
+                AttributeError, match="Attempting to set a non-existent attribute 'tensor_model_parallel_size'"
+            ):
+                load_tokenizer(ckpt_path, tensor_model_parallel_size=1)