feat(tokenizers): support loading a chat template from a file

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

@@ -61,6 +61,12 @@ class TokenizerConfig(MTrainTokenizerConfig):
         }
     """
 
+    chat_template_path: Optional[str] = None
+    """Path to a jinja chat template file, loaded at build time as ``chat_template``. Supports local
+    paths and ``msc://`` URLs. Mutually exclusive with ``chat_template``. Useful for supplying a
+    template from an external/process caller (e.g. CLI overrides) where inlining the jinja is
+    impractical."""
+
     tokenizer_prompt_format: Optional[str] = None
     """Prompt format for the tokenizer."""
 

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

@@ -2,12 +2,46 @@
 
 """Megatron tokenizers."""
 
+import dataclasses
+from typing import Optional
+
+from megatron.core.msc_utils import MultiStorageClientFeature
 from megatron.core.tokenizers import MegatronTokenizer
 from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer as build_mcore_tokenizer
 
 from megatron.bridge.training.tokenizers.config import TokenizerConfig
 
 
+def _resolve_chat_template(config: TokenizerConfig) -> Optional[str]:
+    """Resolve the effective chat template from ``config``.
+
+    Returns ``chat_template`` verbatim when set, or the contents of the file at
+    ``chat_template_path`` (local path or ``msc://`` URL) otherwise.
+
+    Args:
+        config: Tokenizer configuration.
+
+    Returns:
+        The chat template string, or ``None`` when neither field is set.
+
+    Raises:
+        ValueError: If both ``chat_template`` and ``chat_template_path`` are set.
+    """
+    if config.chat_template_path is None:
+        return config.chat_template
+    if config.chat_template is not None:
+        raise ValueError("Set only one of `chat_template` or `chat_template_path`, not both.")
+
+    path = config.chat_template_path
+    if MultiStorageClientFeature.is_enabled():
+        msc = MultiStorageClientFeature.import_package()
+        with msc.open(str(path), "r") as f:
+            return f.read()
+    else:
+        with open(path, "r", encoding="utf-8") as f:
+            return f.read()
+
+
 def build_tokenizer(config: TokenizerConfig, **kwargs) -> MegatronTokenizer:
     """Initialize tokenizer from megatron.core.tokenizers based on the provided configuration.
 
@@ -26,4 +60,7 @@ def build_tokenizer(config: TokenizerConfig, **kwargs) -> MegatronTokenizer:
         "Please, use `megatron.core.tokenizers.utils.build_tokenizer` instead."
     )
 
+    if config.chat_template_path is not None:
+        config = dataclasses.replace(config, chat_template=_resolve_chat_template(config), chat_template_path=None)
+
     return build_mcore_tokenizer(config, **kwargs)

tests/unit_tests/training/test_tokenizer.py

@@ -166,3 +166,80 @@ def test_hf_tokenizer_as_local_path_object(self, tmp_path):
         # verify that the directory actually contains files (sanity check)
         assert (local_model_path / "tokenizer_config.json").exists()
         assert (local_model_path / "tokenizer.json").exists()
+
+
+CHAT_TEMPLATE = "{% generation %}{{ messages }}{% endgeneration %}"
+
+
+class TestChatTemplatePathOverride:
+    """`TokenizerConfig.chat_template_path` loads a jinja template from a local or msc:// path."""
+
+    def test_resolve_from_local_file(self, tmp_path):
+        from megatron.bridge.training.tokenizers import tokenizer as tok_mod
+
+        template_file = tmp_path / "template.jinja"
+        template_file.write_text(CHAT_TEMPLATE)
+        config = TokenizerConfig(chat_template_path=str(template_file))
+
+        with patch.object(tok_mod.MultiStorageClientFeature, "is_enabled", return_value=False):
+            assert tok_mod._resolve_chat_template(config) == CHAT_TEMPLATE
+
+    def test_inline_chat_template_passthrough(self):
+        from megatron.bridge.training.tokenizers.tokenizer import _resolve_chat_template
+
+        config = TokenizerConfig(chat_template=CHAT_TEMPLATE)
+
+        assert _resolve_chat_template(config) == CHAT_TEMPLATE
+
+    def test_none_when_neither_set(self):
+        from megatron.bridge.training.tokenizers.tokenizer import _resolve_chat_template
+
+        assert _resolve_chat_template(TokenizerConfig()) is None
+
+    def test_inline_and_path_are_mutually_exclusive(self, tmp_path):
+        from megatron.bridge.training.tokenizers.tokenizer import _resolve_chat_template
+
+        template_file = tmp_path / "template.jinja"
+        template_file.write_text(CHAT_TEMPLATE)
+        config = TokenizerConfig(chat_template=CHAT_TEMPLATE, chat_template_path=str(template_file))
+
+        with pytest.raises(ValueError):
+            _resolve_chat_template(config)
+
+    def test_resolve_via_msc_when_enabled(self):
+        from unittest.mock import MagicMock
+
+        from megatron.bridge.training.tokenizers import tokenizer as tok_mod
+
+        fake_handle = MagicMock()
+        fake_handle.read.return_value = CHAT_TEMPLATE
+        fake_msc = MagicMock()
+        fake_msc.open.return_value.__enter__.return_value = fake_handle
+
+        config = TokenizerConfig(chat_template_path="msc://bucket/template.jinja")
+        with patch.object(tok_mod, "MultiStorageClientFeature") as msc_feat:
+            msc_feat.is_enabled.return_value = True
+            msc_feat.import_package.return_value = fake_msc
+            assert tok_mod._resolve_chat_template(config) == CHAT_TEMPLATE
+        fake_msc.open.assert_called_once_with("msc://bucket/template.jinja", "r")
+
+    def test_build_hf_tokenizer_passes_resolved_template(self, tmp_path):
+        from megatron.bridge.training.tokenizers import tokenizer as tok_mod
+
+        template_file = tmp_path / "template.jinja"
+        template_file.write_text(CHAT_TEMPLATE)
+        config = TokenizerConfig(
+            tokenizer_type="HuggingFaceTokenizer",
+            tokenizer_model="meta-llama/Llama-2-7b-chat-hf",
+            chat_template_path=str(template_file),
+        )
+
+        with (
+            patch.object(tok_mod.MultiStorageClientFeature, "is_enabled", return_value=False),
+            patch.object(tok_mod, "build_mcore_tokenizer") as mock_build_mcore,
+        ):
+            build_tokenizer(config)
+
+        (called_config,), _ = mock_build_mcore.call_args
+        assert called_config.chat_template == CHAT_TEMPLATE
+        assert called_config.chat_template_path is None