diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index 83fcda787e..1b50666dc5 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/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) diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index d1fb5c287c..958594610e 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/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)