diff --git a/README.md b/README.md index 98016a8..76c092d 100644 --- a/README.md +++ b/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: diff --git a/configs/trl/dpo.yaml b/configs/trl/dpo.yaml index d24a611..d5d908d 100644 --- a/configs/trl/dpo.yaml +++ b/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: diff --git a/configs/trl/sft.yaml b/configs/trl/sft.yaml index 5699cf2..a361aee 100644 --- a/configs/trl/sft.yaml +++ b/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: diff --git a/scripts/data.py b/scripts/data.py index 1f0f49f..9c4933e 100755 --- a/scripts/data.py +++ b/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)) diff --git a/scripts/submit.py b/scripts/submit.py index f00137e..555241b 100644 --- a/scripts/submit.py +++ b/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 diff --git a/src/post_training/config.py b/src/post_training/config.py index fb65058..bd3be4e 100644 --- a/src/post_training/config.py +++ b/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: diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index 60540d8..213260b 100644 --- a/src/post_training/methods/common.py +++ b/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 diff --git a/src/post_training/utils/guardrails.py b/src/post_training/utils/guardrails.py index cd6bf96..006348d 100644 --- a/src/post_training/utils/guardrails.py +++ b/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): diff --git a/src/post_training/utils/prefetch.py b/src/post_training/utils/prefetch.py index 391bca4..3097114 100644 --- a/src/post_training/utils/prefetch.py +++ b/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) diff --git a/tests/test_config.py b/tests/test_config.py index 8857f00..3033faf 100644 --- a/tests/test_config.py +++ b/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) diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py new file mode 100644 index 0000000..a957c63 --- /dev/null +++ b/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)