feat: add option to choose a specific tokenizer

README.md

@@ -329,7 +329,7 @@ You must specify exactly one determining factor for training duration in the `tr
 - **Offline**: `offline: true`
   Disables Hugging Face Hub / Weights & Biases network calls (essential for air-gapped nodes).
 - **Prefetch assets**: `prefetch_assets: true` (default)
-  When submitting via `scripts/submit.py`, pre-fetches models and datasets on the login node before the job is submitted to warm the shared Hugging Face caches (useful for faster startup and for `offline: true` air-gapped runs). Runs whenever `offline` or `prefetch_assets` is true — set `prefetch_assets: false` to skip this step (e.g., assets are already warm in the shared cache).
+  When submitting via `scripts/submit.py`, pre-fetches models and datasets on the login node before the job is submitted to warm the shared Hugging Face caches (useful for faster startup and for `offline: true` air-gapped runs). Runs whenever `offline` or `prefetch_assets` is true — set `prefetch_assets: false` to skip this step (e.g., assets are already warm in the shared cache). A tokenizer from a different repo or revision (`model.tokenizer_name_or_path` / `model.tokenizer_revision`) is fetched too, without its model weights.
 - **Debug**: `debug.enabled: true`
   Forces `report_to: none`, uses a separate output directory, and allows overwriting existing runs.
 - **Tokenize only**: `--tokenize-only` (CLI flag on `train.py` / `submit.py`)
@@ -444,6 +444,9 @@ model:
   name_or_path: "allenai/Olmo-3-1025-7B"
   attn_implementation: "flash_attention_3"
   dtype: "bfloat16"
+  revision: null                             # pin the model repo to a branch, tag, or commit
+  tokenizer_name_or_path: null               # null = take the tokenizer from the model repo
+  tokenizer_revision: null                   # pins the tokenizer (overrides revision above)
 
 # -- Training hyper-parameters -----------------------------------------------
 training:

configs/trl/dpo.yaml

@@ -26,6 +26,9 @@ model:
   name_or_path: "allenai/Olmo-3-1025-7B"
   attn_implementation: "flash_attention_3"
   dtype: "bfloat16"
+  revision: null                             # pin the model repo to a branch, tag, or commit
+  tokenizer_name_or_path: null               # null = take the tokenizer from the model repo
+  tokenizer_revision: null                   # pins the tokenizer (overrides revision above)
 
 # -- Training hyper-parameters -----------------------------------------------
 training:

configs/trl/sft.yaml

@@ -35,6 +35,9 @@ model:
   name_or_path: "allenai/Olmo-3-1025-7B"
   attn_implementation: "flash_attention_3"
   dtype: "bfloat16"
+  revision: null                             # pin the model repo to a branch, tag, or commit
+  tokenizer_name_or_path: null               # null = take the tokenizer from the model repo
+  tokenizer_revision: null                   # pins the tokenizer (overrides revision above)
 
 # -- Training hyper-parameters -----------------------------------------------
 training:

scripts/data.py

@@ -110,10 +110,9 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None:
 
     # ── Lazy imports (heavy) ────────────────────────────────────────
     from datasets import load_dataset
-    from transformers import AutoTokenizer
 
-    from post_training.chat_templates.registry import get_chat_template
     from post_training.data.transforms import get_transform
+    from post_training.methods.common import build_tokenizer
 
     n = args.num_samples
 
@@ -151,9 +150,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None:
         # ── Formatted via chat template ─────────────────────────────
         if args.show_formatted:
             _print_header(f"FORMATTED (chat_template='{config.data.chat_template}')")
-            tokenizer = AutoTokenizer.from_pretrained(config.model.name_or_path)
-            template_str = get_chat_template(config.data.chat_template)
-            tokenizer.chat_template = template_str
+            tokenizer = build_tokenizer(config)
 
             for i in range(min(n, len(ds))):
                 row = ds[i]
@@ -174,9 +171,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None:
         # ── Tokens ──────────────────────────────────────────────────
         if args.show_tokens:
             _print_header("TOKENS")
-            tokenizer = AutoTokenizer.from_pretrained(config.model.name_or_path)
-            template_str = get_chat_template(config.data.chat_template)
-            tokenizer.chat_template = template_str
+            tokenizer = build_tokenizer(config)
 
             for i in range(min(n, len(ds))):
                 row = ds[i]
@@ -208,9 +203,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None:
             print(f"  Total rows : {len(ds)}")
             print(f"  Columns    : {ds.column_names}")
 
-            tokenizer = AutoTokenizer.from_pretrained(config.model.name_or_path)
-            template_str = get_chat_template(config.data.chat_template)
-            tokenizer.chat_template = template_str
+            tokenizer = build_tokenizer(config)
 
             # Sample up to 200 rows for token-length stats.
             sample_size = min(200, len(ds))

scripts/submit.py

@@ -146,6 +146,7 @@ def main() -> None:
     # entirely — avoiding contention when every rank loads it at once.
     if prefetched is not None:
         config.model.name_or_path = prefetched.model
+        config.model.tokenizer_name_or_path = prefetched.tokenizer
         if prefetched.ref_model is not None:
             config.dpo.ref_model_name_or_path = prefetched.ref_model
 

src/post_training/config.py

@@ -29,6 +29,27 @@ class ModelConfig:
     dtype: str = "bfloat16"
     revision: str | None = None
 
+    # Tokenizer source. If ``tokenizer_name_or_path`` is None the tokenizer
+    # comes from the model repo, and ``tokenizer_revision`` then pins it to a
+    # revision other than ``revision``.
+    tokenizer_name_or_path: str | None = None
+    tokenizer_revision: str | None = None
+
+    def resolve_tokenizer(self) -> tuple[str, str | None]:
+        """Return the ``(name_or_path, revision)`` pair to load the tokenizer from.
+
+        ``tokenizer_name_or_path`` points at a different repo, so ``revision``
+        does not apply to it: such a tokenizer is pinned by
+        ``tokenizer_revision`` alone.  Without ``tokenizer_name_or_path`` the
+        tokenizer comes from the model repo, where ``tokenizer_revision``
+        overrides ``revision``.
+        """
+        if self.tokenizer_name_or_path is not None:
+            return self.tokenizer_name_or_path, self.tokenizer_revision
+        if self.tokenizer_revision is not None:
+            return self.name_or_path, self.tokenizer_revision
+        return self.name_or_path, self.revision
+
 
 @dataclass
 class LRSchedulerKwargs:

src/post_training/methods/common.py

@@ -42,9 +42,10 @@ def resolve_torch_dtype(name: str) -> torch.dtype:
 
 def build_tokenizer(config: PostTrainingConfig) -> AutoTokenizer:
     """Load the tokenizer and apply the configured chat template."""
-    tokenizer = AutoTokenizer.from_pretrained(
-        config.model.name_or_path, revision=config.model.revision
-    )
+    tokenizer_name_or_path, revision = config.model.resolve_tokenizer()
+    logger.info("Loading tokenizer '%s' (revision=%s)", tokenizer_name_or_path, revision or "main")
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=revision)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 

src/post_training/utils/guardrails.py

@@ -179,6 +179,13 @@ def run_guardrails(config: PostTrainingConfig, run_dir: Path, tokenize_only: boo
     _row("Attention impl", config.model.attn_implementation)
     _row("Dtype", config.model.dtype)
     _row("Revision", config.model.revision or "main")
+    tokenizer_name_or_path, tokenizer_revision = config.model.resolve_tokenizer()
+    _row(
+        "Tokenizer",
+        tokenizer_name_or_path,
+        warn=tokenizer_name_or_path != config.model.name_or_path,
+    )
+    _row("Tokenizer revision", tokenizer_revision or "main")
     _row("Chat template", config.data.chat_template)
     _row("Data seed", str(config.data.seed))
     for i, entry in enumerate(config.data.datasets):

src/post_training/utils/prefetch.py

@@ -31,16 +31,21 @@
 class PrefetchedPaths:
     """Resolved local snapshot directories for the models in *config*.
 
-    Substituting these back into ``config.model.name_or_path`` (and, for DPO,
+    Substituting these back into ``config.model.name_or_path``,
+    ``config.model.tokenizer_name_or_path`` (and, for DPO,
     ``config.dpo.ref_model_name_or_path``) before training starts makes
     ``from_pretrained`` treat the model as a plain local directory, skipping
     huggingface_hub's cache-resolution/filelock path entirely — the source of
     spurious "does not appear to have a file named ..." errors when many
     ranks call ``from_pretrained`` on the same shared (e.g. Lustre) cache at
     once.
+
+    ``tokenizer`` equals ``model`` when the tokenizer comes from the model
+    repo, because the model snapshot already contains the tokenizer files.
     """
 
     model: str
+    tokenizer: str
     ref_model: str | None = None
 
 
@@ -62,6 +67,48 @@ def _prefetch_model(name_or_path: str, revision: str | None = None) -> str:
     return local_path
 
 
+# Weight formats, skipped when a repo is fetched for its tokenizer alone.
+# Denying the large blobs is safer than allowing a list of tokenizer files:
+# every tokenizer format arrives, including the sentencepiece and legacy names
+# (``spiece.model``, ``sentencepiece.bpe.model``, ``bpe.codes``, ``*.spm``) and
+# the ``additional_chat_templates/`` directory.  The worst case is a few extra
+# small files, not a missing file on an offline compute node.
+_WEIGHT_PATTERNS = [
+    "*.safetensors",
+    "*.bin",
+    "*.pt",
+    "*.pth",
+    "*.ckpt",
+    "*.h5",
+    "*.msgpack",
+    "*.gguf",
+    "*.onnx",
+    "*.onnx_data",
+    "*.ot",  # rust_model.ot
+    "*.tflite",
+    "*.mlmodel",
+    "*.npz",
+]
+
+
+def _prefetch_tokenizer(name_or_path: str, revision: str | None = None) -> str:
+    """Cache *name_or_path* without its weights; return the local directory."""
+    if _is_local(name_or_path):
+        logger.info("Tokenizer '%s' is a local path, skipping download.", name_or_path)
+        return name_or_path
+    logger.info(
+        "Downloading tokenizer '%s' (revision=%s) to HF cache...", name_or_path, revision or "main"
+    )
+    local_path = snapshot_download(
+        repo_id=name_or_path,
+        repo_type="model",
+        revision=revision,
+        ignore_patterns=_WEIGHT_PATTERNS,
+    )
+    logger.info("Tokenizer '%s' cached at '%s'.", name_or_path, local_path)
+    return local_path
+
+
 def _prefetch_dataset(entry: DatasetEntry) -> None:
     if _is_local(entry.path):
         logger.info("Dataset '%s' is a local path, skipping download.", entry.name)
@@ -87,6 +134,16 @@ def prefetch_assets(config: PostTrainingConfig) -> PrefetchedPaths:
 
     model_path = _prefetch_model(config.model.name_or_path, revision=config.model.revision)
 
+    tokenizer_name_or_path, tokenizer_revision = config.model.resolve_tokenizer()
+    if (tokenizer_name_or_path, tokenizer_revision) == (
+        config.model.name_or_path,
+        config.model.revision,
+    ):
+        # The tokenizer files are part of the model snapshot fetched above.
+        tokenizer_path = model_path
+    else:
+        tokenizer_path = _prefetch_tokenizer(tokenizer_name_or_path, revision=tokenizer_revision)
+
     ref_model_path = None
     if config.method == "dpo" and config.dpo.ref_model_name_or_path is not None:
         ref_model_path = _prefetch_model(config.dpo.ref_model_name_or_path)
@@ -95,4 +152,4 @@ def prefetch_assets(config: PostTrainingConfig) -> PrefetchedPaths:
         _prefetch_dataset(entry)
 
     logger.info("All assets pre-fetched successfully.")
-    return PrefetchedPaths(model=model_path, ref_model=ref_model_path)
+    return PrefetchedPaths(model=model_path, tokenizer=tokenizer_path, ref_model=ref_model_path)

tests/test_config.py

@@ -3,7 +3,7 @@
 import pytest
 import yaml
 
-from post_training.config import PostTrainingConfig
+from post_training.config import ModelConfig, PostTrainingConfig
 from post_training.methods.common import build_common_training_kwargs
 
 
@@ -107,3 +107,30 @@ def test_deepspeed_old_style_config_path_rejected(tmp_path):
 
     with pytest.raises(ValueError, match="deepspeed.config_path is no longer supported"):
         PostTrainingConfig.load(config_path)
+
+
+@pytest.mark.parametrize(
+    ("tokenizer_name_or_path", "tokenizer_revision", "expected"),
+    [
+        # No tokenizer override: the tokenizer follows the model.
+        (None, None, ("org/model", "r1")),
+        # A tokenizer revision alone pins the tokenizer inside the model repo.
+        (None, "r2", ("org/model", "r2")),
+        # A separate repo ignores model.revision, which pins a different repo.
+        ("org/tokenizer", None, ("org/tokenizer", None)),
+        ("org/tokenizer", "r2", ("org/tokenizer", "r2")),
+    ],
+)
+def test_resolve_tokenizer(tokenizer_name_or_path, tokenizer_revision, expected):
+    model = ModelConfig(
+        name_or_path="org/model",
+        revision="r1",
+        tokenizer_name_or_path=tokenizer_name_or_path,
+        tokenizer_revision=tokenizer_revision,
+    )
+
+    assert model.resolve_tokenizer() == expected
+
+
+def test_resolve_tokenizer_without_any_revision():
+    assert ModelConfig(name_or_path="org/model").resolve_tokenizer() == ("org/model", None)

tests/test_prefetch.py

@@ -0,0 +1,80 @@
+"""Tests for which repos prefetch_assets downloads."""
+
+import pytest
+
+from post_training.config import PostTrainingConfig
+from post_training.utils import prefetch
+
+
+@pytest.fixture
+def downloads(monkeypatch):
+    """Record snapshot_download calls instead of contacting the Hub."""
+    calls = []
+
+    def fake_snapshot_download(repo_id, repo_type=None, revision=None, ignore_patterns=None):
+        calls.append((repo_id, revision, ignore_patterns))
+        return f"/cache/{repo_id}@{revision or 'main'}"
+
+    monkeypatch.setattr(prefetch, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(prefetch, "_prefetch_dataset", lambda entry: None)
+    return calls
+
+
+def _config(**model_overrides) -> PostTrainingConfig:
+    config = PostTrainingConfig()
+    config.model.name_or_path = "org/model"
+    for key, value in model_overrides.items():
+        setattr(config.model, key, value)
+    return config
+
+
+def test_tokenizer_reuses_the_model_snapshot(downloads):
+    """The model snapshot already holds the tokenizer, so nothing extra is fetched."""
+    paths = prefetch.prefetch_assets(_config(revision="r1"))
+
+    assert downloads == [("org/model", "r1", None)]
+    assert paths.tokenizer == paths.model
+
+
+def test_tokenizer_revision_is_fetched_separately(downloads):
+    """A tokenizer pinned to another revision of the model repo needs its own snapshot."""
+    paths = prefetch.prefetch_assets(_config(revision="r1", tokenizer_revision="r2"))
+
+    assert downloads == [
+        ("org/model", "r1", None),
+        ("org/model", "r2", prefetch._WEIGHT_PATTERNS),
+    ]
+    assert paths.model == "/cache/org/model@r1"
+    assert paths.tokenizer == "/cache/org/model@r2"
+
+
+def test_model_revision_is_ignored_when_tokenizer_name_is_specified(downloads):
+    """The model revision is ignored when a tokenizer revision is specified."""
+    paths = prefetch.prefetch_assets(_config(revision="r1", tokenizer_name_or_path="org/tokenizer"))
+
+    assert downloads == [
+        ("org/model", "r1", None),
+        ("org/tokenizer", None, prefetch._WEIGHT_PATTERNS),
+    ]
+    assert paths.model == "/cache/org/model@r1"
+    assert paths.tokenizer == "/cache/org/tokenizer@main"
+
+
+def test_tokenizer_repo_is_fetched_without_weights(downloads):
+    """A tokenizer from another repo is fetched without the weight files."""
+    paths = prefetch.prefetch_assets(_config(tokenizer_name_or_path="org/tokenizer"))
+
+    assert downloads == [
+        ("org/model", None, None),
+        ("org/tokenizer", None, prefetch._WEIGHT_PATTERNS),
+    ]
+    assert "*.safetensors" in prefetch._WEIGHT_PATTERNS
+    assert paths.tokenizer == "/cache/org/tokenizer@main"
+
+
+def test_local_tokenizer_path_is_not_downloaded(downloads, tmp_path):
+    """An existing local tokenizer directory is used as-is."""
+    paths = prefetch.prefetch_assets(_config(tokenizer_name_or_path=str(tmp_path)))
+
+    assert downloads == [("org/model", None, None)]
+    assert paths.tokenizer == str(tmp_path)