diff --git a/modelopt/torch/export/__init__.py b/modelopt/torch/export/__init__.py index b62dd041a18..874fa71b2b9 100644 --- a/modelopt/torch/export/__init__.py +++ b/modelopt/torch/export/__init__.py @@ -21,8 +21,10 @@ # names from modelopt.torch.export directly. from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_hf_ckpt_remote_code, + copy_non_model_files, copy_non_safetensor_files_from_ckpt, copy_off_index_safetensors, + ensure_local_checkpoint, indexed_weight_map, load_multimodal_components, locate_source_keys, diff --git a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py index 033716a9b6f..de8532e5297 100644 --- a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py @@ -15,11 +15,12 @@ """Hugging Face checkpoint utility. -General-purpose logic about the on-disk shape of an HF checkpoint (index, shards, sidecars) -- -not export-specific, so it lives under ``modelopt.torch.utils.plugins`` alongside -``model_load_utils`` rather than under ``modelopt.torch.export``. Deliberately independent of -that module's ``transformers``/``accelerate`` module-scope imports; only needs -``huggingface_hub`` + ``safetensors``. +General-purpose logic about the on-disk shape of an HF checkpoint (index, shards, off-index +weight files, non-model files) -- not export-specific, so it lives under +``modelopt.torch.utils.plugins`` alongside ``model_load_utils`` rather than under +``modelopt.torch.export``. Deliberately independent of that module's +``transformers``/``accelerate`` module-scope imports: beyond ``torch``, it needs only +``huggingface_hub`` and ``safetensors``. """ import contextlib @@ -30,17 +31,101 @@ import shutil import warnings from collections.abc import Callable, Iterable +from datetime import timedelta from pathlib import Path from typing import Any import torch -from huggingface_hub import snapshot_download -from huggingface_hub.errors import LocalEntryNotFoundError +from huggingface_hub import HfApi, snapshot_download +from huggingface_hub.errors import ( + GatedRepoError, + HFValidationError, + LocalEntryNotFoundError, + RepositoryNotFoundError, + RevisionNotFoundError, +) from safetensors.torch import safe_open from tqdm import tqdm +from modelopt.torch.utils.mlflow import EXPERIMENT_JSON + _HF_HUB_OFFLINE_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +# Standard HF weight-file names: ``model.safetensors`` or ``model-00001-of-00005.safetensors``. +_IS_MAIN_WEIGHT_SHARD = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors") + +# Off-index files that re-ship weights rather than add new ones. ``consolidated*.safetensors`` +# is Mistral's second full copy of the model (vLLM's mistral load-format looks for it BY NAME, +# so copying it into an export is not inert -- it can be served in place of the quantized +# weights). ``adapter_model.safetensors`` is a PEFT adapter, whose tensor names do not overlap +# the index, so only a name rule catches it. +_IS_WEIGHT_DUPLICATE = re.compile(r"(consolidated[^/]*|adapter_model)\.safetensors") + +# How long the other ranks wait for rank 0 to make a checkpoint local: long enough for a large +# model's download, which the default process-group timeout would not survive. +_CHECKPOINT_DOWNLOAD_TIMEOUT = timedelta(hours=24) + +# Default size cap for a file read through ``resolve_checkpoint_file``: it bounds metadata reads +# (indexes, configs); callers copying weight files pass ``max_bytes=None``. +_MAX_CHECKPOINT_METADATA_BYTES = 128 * 1024 * 1024 + +# --- Model files vs non-model files ------------------------------------------------------------ +# An export splits a checkpoint's files in two: +# +# * MODEL files are what the export writes itself: the weights, in any format +# (``_WEIGHT_FILE_PATTERNS``), and the metadata describing them (``_EXPORT_OWNED_FILES``). They +# are never copied from the source, since the source's versions describe the unquantized model. +# Off-index weight files such as GLM-4.7's ``mtp.safetensors`` are model files too; the exporter +# carries them over itself with :func:`copy_off_index_safetensors`. +# +# * NON-MODEL files are everything else the checkpoint ships: tokenizer and processor files, +# remote-code ``*.py`` and parser plugins, chat templates, README, LICENSE, and assets in +# subdirectories. The export does not produce them, so :func:`copy_non_model_files` carries them +# over verbatim. ``generation_config.json`` falls in between: some exports write their own, which +# then wins, and the source's is copied only when they do not. + +# Weight files in any format, at any depth in the checkpoint. +_WEIGHT_FILE_PATTERNS = ( + "*.safetensors", + "*.safetensors.index.json", + "*.bin", + "*.bin.index.json", + "*.ckpt", + "*.gguf", + "*.h5", + "*.msgpack", + "*.npy", + "*.npz", + "*.onnx", + "*.pb", + "*.pickle", + "*.pkl", + "*.pt", + "*.pth", + "*.tar", + "*.tar.bz2", + "*.tar.gz", + "*.tar.xz", + "*.tflite", + "*.tgz", + "*.zip", +) + +# Top-level metadata an export writes itself, or that describes the source's own quantization and +# would mislabel the export if copied. +_EXPORT_OWNED_FILES = frozenset( + { + EXPERIMENT_JSON, + "config.json", + "hf_quant_config.json", + "quant_config.json", + "quantization_config.json", + "quantize_config.json", + "recipe.yaml", + "recipe.yml", + } +) + def _as_nonnegative_int(value: Any) -> int | None: """Return ``value`` as an int when it is a non-negative integer.""" @@ -156,6 +241,9 @@ def copy_hf_ckpt_remote_code( ``HF_HUB_OFFLINE`` is set, the snapshot must already be available in the local Hugging Face cache. + To carry over every non-model file of the checkpoint rather than only its code, use + :func:`copy_non_model_files`. + Args: pretrained_model_path: Local path to the pretrained model or HuggingFace Hub model ID. save_directory: Path to the save directory. @@ -298,17 +386,6 @@ def _matches_any_pattern(file_name: str, patterns: tuple[str, ...]) -> bool: return any(fnmatch.fnmatchcase(file_name, pattern) for pattern in patterns) -# Standard HF weight-file names: ``model.safetensors`` or ``model-00001-of-00005.safetensors``. -_IS_MAIN_WEIGHT_SHARD = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors") - -# Off-index files that re-ship weights rather than add new ones. ``consolidated*.safetensors`` -# is Mistral's second full copy of the model (vLLM's mistral load-format looks for it BY NAME, -# so copying it into an export is not inert -- it can be served in place of the quantized -# weights). ``adapter_model.safetensors`` is a PEFT adapter, whose tensor names do not overlap -# the index, so only a name rule catches it. -_IS_WEIGHT_DUPLICATE = re.compile(r"(consolidated[^/]*|adapter_model)\.safetensors") - - # --- What reaches the export without passing through quantization -------------------------- # Two disjoint sets, distinguished by what the LOADER did with the file. That difference decides # both how we find them and how we move them, so it is worth keeping straight: @@ -319,10 +396,10 @@ def _matches_any_pattern(file_name: str, patterns: tuple[str, ...]) -> bool: # holds them and merged into the exporter's ``extra_state_dict``. # Found by: ``read_unplaced_weights`` / ``carryable_unplaced_keys`` / ``locate_source_keys``. # -# 2. OFF-INDEX sidecars. The index never names the file, so the loader never opened it and never -# had the chance to call anything unexpected -- GLM-4.7 keeps its MTP head in a standalone -# ``mtp.safetensors`` exactly this way. Moved as FILES: copied byte for byte, so no host -# memory is spent re-serialising tensors the export does not otherwise touch. +# 2. OFF-INDEX weight files. The index never names the file, so the loader never opened it and +# never had the chance to call anything unexpected -- GLM-4.7 keeps its MTP head in a +# standalone ``mtp.safetensors`` exactly this way. Moved as FILES: copied byte for byte, so +# no host memory is spent re-serialising tensors the export does not otherwise touch. # Found by: ``off_index_safetensors_files`` / ``off_index_tensor_names``. # # The index is what separates them, and it is NOT an inventory of the checkpoint: a tensor missing @@ -339,7 +416,6 @@ def _matches_any_pattern(file_name: str, patterns: tuple[str, ...]) -> bool: # "inside the checkpoint" has to mean the snapshot dir OR that blob root. Getting this wrong # in either direction is costly: reject links and no hub checkpoint works, follow them blindly # and a checkpoint can name any file on the host. -_MAX_CHECKPOINT_METADATA_BYTES = 128 * 1024 * 1024 def _is_relative_to(path: Path, root: Path) -> bool: @@ -382,7 +458,7 @@ def resolve_checkpoint_file( return resolved_src -def off_index_safetensors_files(src: "str | os.PathLike") -> list[str]: +def off_index_safetensors_files(src: "str | os.PathLike | None") -> list[str]: """Safetensors files in a checkpoint that model loading never opens. Transformers reads the shards named in ``model.safetensors.index.json`` -- or the single @@ -390,7 +466,7 @@ def off_index_safetensors_files(src: "str | os.PathLike") -> list[str]: GLM-4.7 keeps its MTP head in a standalone ``mtp.safetensors``. Those tensors are never loaded, never quantized, and never reported as ``unexpected_keys`` (the loader did not see them to call them unexpected), so they are not "the unquantized source weights" the export must avoid - re-emitting -- they are untouched sidecars that happen to be in safetensors format. See + re-emitting -- they are untouched weight files the export carries over as they are. See :func:`~modelopt.torch.utils.plugins.model_load_utils.record_unplaced_source_keys` for why "never opened" is a property of the FILE rather than of the individual tensor: a tensor the index omits from a file it DOES open is reported, and is carried rather than copied. @@ -403,6 +479,8 @@ def off_index_safetensors_files(src: "str | os.PathLike") -> list[str]: (``consolidated.safetensors``, ``adapter_model.safetensors``) and by tensor-name overlap for the ones we do not. See :func:`_without_reshipped_weights`. """ + if src is None: + return [] d = Path(src) if not d.is_dir(): return [] @@ -441,12 +519,13 @@ def _without_reshipped_weights( The name rules above only catch conventions we know. A checkpoint free to invent its own filename can still carry a second copy of the indexed weights, and copying that into an export puts unquantized tensors beside the quantized ones. Overlapping tensor names are the - general signal: a genuine sidecar (an MTP head) holds names the index does not have, which - is exactly why the loader never placed them. + general signal: a genuine off-index file (an MTP head) holds names the index does not have, + which is exactly why the loader never placed them. Best-effort. Reads safetensors headers, never tensor data, and keeps any candidate whose - header cannot be read: refusing to copy a real sidecar because of an unreadable header would - silently drop weights from the export, which is the failure this whole path exists to avoid. + header cannot be read: refusing to copy a real off-index file because of an unreadable header + would silently drop weights from the export, which is the failure this whole path exists to + avoid. """ if not candidates or not indexed_tensors: return candidates @@ -582,14 +661,19 @@ def locate_source_keys(ckpt: "str | Path", keys: list[str]) -> dict[str, str]: return located -def copy_off_index_safetensors(src: "str | os.PathLike", dst: "str | os.PathLike") -> list[str]: +def copy_off_index_safetensors( + src: "str | os.PathLike | None", dst: "str | os.PathLike" +) -> list[str]: """Copy the safetensors files model loading never reads, verbatim. Copying beats reading them into a state dict and re-serialising: no host memory is spent on tensors the export does not touch, the bytes and the file layout are preserved exactly, and a - consumer that finds them by filename (vLLM looks for the MTP sidecar) sees what it saw in the - source. See :func:`off_index_safetensors_files` for how "never reads" is decided. + consumer that finds them by filename (vLLM looks for the MTP file) sees what it saw in the + source. See :func:`off_index_safetensors_files` for how "never reads" is decided. A ``src`` + that is ``None`` or not a local directory (e.g. a Hub model ID) has nothing to copy. """ + if src is None: + return [] names = off_index_safetensors_files(src) copied = [] for name in names: @@ -600,8 +684,8 @@ def copy_off_index_safetensors(src: "str | os.PathLike", dst: "str | os.PathLike # copy whatever that names into the export under an approved-looking name. The question # is where the link LANDS, not whether it is a link: a Hugging Face snapshot stores every # file as a symlink into ``../../blobs/``, so refusing links outright drops the - # sidecar of every hub-downloaded checkpoint -- the GLM-4.7 ``mtp.safetensors`` this path - # exists to carry included. + # off-index weights of every hub-downloaded checkpoint -- the GLM-4.7 ``mtp.safetensors`` + # this path exists to carry included. # # resolve_checkpoint_file already draws that line and is tested against both shapes: it # resolves strictly, demands a regular file, and demands the target sit under the @@ -614,6 +698,8 @@ def copy_off_index_safetensors(src: "str | os.PathLike", dst: "str | os.PathLike continue shutil.copy2(source, target) copied.append(name) + if copied: + print(f"Copied off-index safetensors into {dst}: {copied}") return copied @@ -660,3 +746,204 @@ def copy_non_safetensor_files_from_ckpt( continue copied_files.append(entry) return copied_files + + +def _snapshot_download(hub_model_id: str) -> str: + """``snapshot_download``, but warn when it unexpectedly falls back to the cache. + + ``snapshot_download`` silently returns whatever snapshot is cached when the Hub cannot be + reached; ask the Hub first so that case is visible. A missing repo or revision still raises. + With ``HF_HUB_OFFLINE`` set the cache is the user's choice, so it is used without a warning. + """ + if _is_hf_hub_offline(): + try: + return snapshot_download(hub_model_id, local_files_only=True) + except LocalEntryNotFoundError as error: + raise RuntimeError( + f"HF_HUB_OFFLINE is set and {hub_model_id!r} is not in the local cache." + ) from error + try: + HfApi().repo_info(hub_model_id) + except (RepositoryNotFoundError, RevisionNotFoundError): + raise + except Exception as error: # offline, unreachable, or refused: only the cache is left + hub_error = f"{type(error).__name__}: {error}" + try: + local_checkpoint_path = snapshot_download(hub_model_id, local_files_only=True) + except LocalEntryNotFoundError: + raise RuntimeError( + f"Could not reach the Hugging Face Hub for {hub_model_id!r} ({hub_error}), and it " + "is not in the local cache." + ) from error + warnings.warn( + f"Could not reach the Hugging Face Hub for {hub_model_id!r} ({hub_error}); using its " + "cached snapshot as it is. It may lack files, e.g. non-model files an earlier " + "from_pretrained did not need." + ) + return local_checkpoint_path + return snapshot_download(hub_model_id) + + +def _make_local(model_name_or_path: str | os.PathLike) -> tuple[str | None, str]: + """``(hub_model_id, local_checkpoint_path)`` on this rank alone. + + See :func:`ensure_local_checkpoint`. + """ + if os.path.isdir(model_name_or_path): + return None, str(model_name_or_path) + hub_model_id = str(model_name_or_path) + try: + return hub_model_id, _snapshot_download(hub_model_id) + except GatedRepoError: + raise + except (HFValidationError, RepositoryNotFoundError) as error: + # A mistyped local path lands here as a Hub ID; say both lookups failed. + raise ValueError( + f"{hub_model_id!r} is not a local directory, and is not a model ID the Hugging Face " + "Hub can access." + ) from error + + +def ensure_local_checkpoint( + model_name_or_path: str | os.PathLike, + group: "torch.distributed.ProcessGroup | None" = None, +) -> tuple[str | None, str]: + """Ensure the whole checkpoint is on local storage. + + Returns ``(hub_model_id, local_checkpoint_path)``: the Hub model ID, or ``None`` when + ``model_name_or_path`` is already a local directory, and the local checkpoint directory. Read + checkpoint files from ``local_checkpoint_path``; use ``hub_model_id`` where a portable name is + wanted (a record that outlives this host's cache), falling back to ``local_checkpoint_path`` + for a local checkpoint. + + A local directory is returned as-is. A Hub model ID resolves to its snapshot in the Hugging + Face cache, downloading only what is not cached yet. The whole repo is made local, not a + subset: guessing which files are redundant would risk dropping ones the checkpoint needs. + Exporters read the source checkpoint only from local disk (see :func:`copy_non_model_files`), + so call this before loading a Hub model that will be exported. + + With ``HF_HUB_OFFLINE`` set, the cached snapshot is used as it is. With the Hub unexpectedly + unreachable it is too, but with a warning: nothing can tell whether it is complete, and one + filled by an earlier ``from_pretrained`` holds only the files loading needed, so non-model files + may be missing. A path that is neither a local directory nor a model ID the Hub can access (a + mistyped local path, say) raises ``ValueError``; a gated repo raises ``GatedRepoError``. + + Under ``torch.distributed`` rank 0 of ``group`` resolves the checkpoint -- whether it is a local + directory, and if not, downloading it once -- and the other ranks of ``group`` wait for its + result, so every rank of ``group`` (and only those) must call this. The checkpoint, or the + Hugging Face cache it lands in, must then be on storage all of them can see. + + Args: + model_name_or_path: Hub model ID or local checkpoint directory. + group: Process group sharing the download. Its timeout must outlast the download, since + the other ranks wait in it. ``None`` means every rank of the default group, waiting on + a temporary gloo group whose timeout does. + + Returns: + ``(hub_model_id, local_checkpoint_path)``. + """ + if group is None and not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + return _make_local(model_name_or_path) + + # Rank 0 decides even for a local path: ranks that each checked for themselves could disagree + # (a directory on one node's local disk) and part of them would wait here forever. Without a + # caller's group, the others wait on a gloo group of their own, whose timeout outlasts a large + # download; the default group's would not. Rank 0 sends its error too, so a failure does not + # strand them. + own_group = group is None + if own_group: + group = torch.distributed.new_group(backend="gloo", timeout=_CHECKPOINT_DOWNLOAD_TIMEOUT) + try: + # (hub_model_id, local_checkpoint_path, error): local_checkpoint_path or error is set. + result: list[tuple[str | None, str | None, str | None]] = [(None, None, None)] + if torch.distributed.get_rank(group) == 0: + try: + result = [(*_make_local(model_name_or_path), None)] + except Exception as error: + result = [(None, None, f"{type(error).__name__}: {error}")] + torch.distributed.broadcast_object_list(result, group=group, group_src=0) + finally: + if own_group: + torch.distributed.destroy_process_group(group) + hub_model_id, local_checkpoint_path, error_message = result[0] + if local_checkpoint_path is None: + raise RuntimeError( + f"Rank 0 could not make {str(model_name_or_path)!r} local: {error_message}" + ) + if not os.path.isdir(local_checkpoint_path): + raise RuntimeError( + f"Rank 0 has {str(model_name_or_path)!r} at {local_checkpoint_path}, which this rank " + "cannot see. Put the checkpoint, or the Hugging Face cache (HF_HOME) it downloads " + "into, on storage shared by every node." + ) + return hub_model_id, local_checkpoint_path + + +def copy_non_model_files(source: str | os.PathLike, export_dir: str | os.PathLike) -> list[str]: + """Copy the source checkpoint's non-model files into an export, verbatim. + + See "Model files vs non-model files" at the top of this module: the export writes the model + files (weights and the metadata describing them); this carries over everything else -- + tokenizer and processor files, remote code, chat templates, README, subdirectory assets. It + copies every file under ``source`` except: + + * model files: weight files in any format (``_WEIGHT_FILE_PATTERNS``) and top-level export + metadata (``_EXPORT_OWNED_FILES``); + * files the export already wrote (e.g. its own ``generation_config.json``), so call this after + the export has written its files, and paths that are links in the export; + * files under ``export_dir`` when it lies inside ``source``, e.g. an earlier export; + * files under hidden directories (``.git``, ``.cache``), which are version-control or download + state rather than part of the checkpoint; + * links that resolve outside the checkpoint (or, for a Hub snapshot, its ``blobs/``), and + dangling links, each with a warning. + + Args: + source: Local checkpoint directory. Download a Hub model first with + :func:`ensure_local_checkpoint`. + export_dir: Export directory; created if missing. + + Returns: + The copied paths, relative to the checkpoint root and ``/``-separated on every platform. + """ + source_dir = Path(source) + if not source_dir.is_dir(): + raise ValueError( + f"{source!s} is not a local checkpoint directory; get one with " + "ensure_local_checkpoint()." + ) + export_dir = Path(export_dir) + source_root, export_root = source_dir.resolve(), export_dir.resolve() + copied = [] + for src in sorted(source_dir.rglob("*")): + rel = src.relative_to(source_dir) + if ( + src.is_dir() + or (source_root / rel).is_relative_to(export_root) + or any(part.startswith(".") for part in rel.parts[:-1]) + or (len(rel.parts) == 1 and rel.name in _EXPORT_OWNED_FILES) + or _matches_any_pattern(rel.name, _WEIGHT_FILE_PATTERNS) + ): + continue + dst = export_dir / rel + # A dangling link fails exists(); copying through it would write outside the export. + if dst.exists() or dst.is_symlink(): + continue + # A link is followed only to where a checkpoint file may live -- the checkpoint itself or, + # for a Hub snapshot, its blobs/ directory -- or the export could copy any file on the host. + try: + resolved = resolve_checkpoint_file(source_dir, rel, max_bytes=None) + except ValueError as error: + warnings.warn(f"Skipping {rel}: {error}") + continue + try: + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(resolved, dst) + except OSError as error: + dst.unlink(missing_ok=True) # or a rerun would keep the partial file + warnings.warn(f"Failed to copy checkpoint file {rel}: {error}") + continue + copied.append(rel.as_posix()) + print(f"Copied {len(copied)} source checkpoint file(s) to {export_dir}: {copied}") + return copied diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py index 44bc3fa3544..490ca79d7ed 100644 --- a/tests/unit/torch/utils/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py @@ -16,12 +16,18 @@ """Tests for modelopt/torch/utils/plugins/hf_checkpoint_utils.py""" import json +import os import sys +import warnings +from functools import partial +from pathlib import Path from types import SimpleNamespace from unittest.mock import patch import pytest import torch +import torch.distributed as dist +from _test_utils.torch.distributed.utils import spawn_multiprocess_job from safetensors.torch import save_file from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -35,6 +41,7 @@ hf_hub_errors = pytest.importorskip("huggingface_hub.errors") LocalEntryNotFoundError = hf_hub_errors.LocalEntryNotFoundError +from modelopt.torch.utils.mlflow import EXPERIMENT_JSON from modelopt.torch.utils.plugins import hf_checkpoint_utils from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_hf_ckpt_remote_code, @@ -90,6 +97,26 @@ def copy2(source, *args, **kwargs): assert copied_files == ["good.py"] +def test_copy_non_model_files_continues_after_copy_failure(tmp_path, monkeypatch): + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "bad.py").write_text("bad") + (src_dir / "good.py").write_text("good") + + original_copy2 = hf_checkpoint_utils.shutil.copy2 + + def copy2(source, *args, **kwargs): + if str(source).endswith("bad.py"): + raise PermissionError("unreadable") + return original_copy2(source, *args, **kwargs) + + monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", copy2) + with pytest.warns(UserWarning, match="bad.py"): + copied_files = hf_checkpoint_utils.copy_non_model_files(src_dir, tmp_path / "dst") + + assert copied_files == ["good.py"] + + def test_copy_hf_ckpt_remote_code_local_dir(tmp_path): """copy_hf_ckpt_remote_code copies top-level .py files from a local directory.""" src_dir = tmp_path / "src" @@ -607,3 +634,345 @@ def test_read_safetensors_subset(tmp_path): assert set(result.keys()) == {"a.weight", "a.bias"} assert torch.equal(result["a.weight"], torch.tensor([1.0, 2.0])) assert torch.equal(result["a.bias"], torch.tensor([3.0])) + + +def test_copy_non_model_files_copies_everything_the_export_does_not_own(tmp_path): + source_dir = tmp_path / "source" + export_dir = tmp_path / "export" + source_dir.mkdir() + export_dir.mkdir() + + source_files = { + "super_v3_reasoning_parser.py": "class Parser: pass\n", + "modeling_custom.py": "class Model: pass\n", + "README.md": "# Source model\n", + "LICENSE": "license text\n", + ".gitattributes": "*.safetensors filter=lfs\n", + "chat_template.jinja": "{{ messages }}\n", + "tokenizer_config.json": '{"chat_template": "source"}\n', + "assets/audio/config.json": '{"nested": "config"}\n', + "generation_config.json": '{"source": "generation"}\n', + "config.json": '{"source": "config"}\n', + "hf_quant_config.json": '{"source": "quant"}\n', + "quant_config.json": '{"source": "stale quant"}\n', + "quantize_config.json": '{"source": "stale quant"}\n', + "recipe.yaml": "quantize: {}\n", + ".experiment.json": '{"run_id": "source"}\n', + "model.safetensors.index.json": '{"weight_map": {}}\n', + "model-00001-of-00001.safetensors": "source weights\n", + "mtp.safetensors": "off-index weights\n", + "pytorch_model.bin": "source weights\n", + "original/consolidated.00.pth": "source weights\n", + ".git/lfs/objects/blob": "source weights\n", + ".cache/huggingface/download/README.md.metadata": "download state\n", + } + for file_name, contents in source_files.items(): + (source_dir / file_name).parent.mkdir(parents=True, exist_ok=True) + (source_dir / file_name).write_text(contents) + + export_files = { + "config.json": '{"export": "config"}\n', + "generation_config.json": '{"export": "generation"}\n', + } + for file_name, contents in export_files.items(): + (export_dir / file_name).write_text(contents) + + copied = hf_checkpoint_utils.copy_non_model_files(str(source_dir), str(export_dir)) + + expected = [ + ".gitattributes", + "LICENSE", + "README.md", + "assets/audio/config.json", + "chat_template.jinja", + "modeling_custom.py", + "super_v3_reasoning_parser.py", + "tokenizer_config.json", + ] + assert sorted(copied) == expected + for file_name in expected: + assert (export_dir / file_name).read_text() == source_files[file_name] + for file_name, contents in export_files.items(): + assert (export_dir / file_name).read_text() == contents + exported = {p.relative_to(export_dir).as_posix() for p in export_dir.rglob("*") if p.is_file()} + assert exported == {*expected, *export_files} + + +def test_copy_non_model_files_skips_an_export_inside_the_source(tmp_path): + """An earlier export saved inside the checkpoint is not part of it.""" + source_dir = tmp_path / "source" + (source_dir / "quantized").mkdir(parents=True) + (source_dir / "tokenizer.json").write_text("{}") + (source_dir / "quantized" / "tokenizer.json").write_text("{}") + + copied = hf_checkpoint_utils.copy_non_model_files(source_dir, source_dir / "quantized") + + assert copied == [] + assert not (source_dir / "quantized" / "quantized").exists() + + +@pytest.mark.skipif(sys.platform == "win32", reason="symlinks need privileges on Windows") +def test_copy_non_model_files_does_not_write_through_an_export_link(tmp_path): + source_dir, export_dir, outside = tmp_path / "source", tmp_path / "export", tmp_path / "outside" + source_dir.mkdir() + export_dir.mkdir() + (source_dir / "tokenizer.json").write_text("{}") + (export_dir / "tokenizer.json").symlink_to(outside) # dangling: exists() is False + + assert hf_checkpoint_utils.copy_non_model_files(source_dir, export_dir) == [] + assert not outside.exists() + + +def test_copy_non_model_files_removes_a_partial_copy(monkeypatch, tmp_path): + source_dir, export_dir = tmp_path / "source", tmp_path / "export" + source_dir.mkdir() + (source_dir / "tokenizer.json").write_text("{}") + + def failing_copy(src, dst): + Path(dst).write_text("{") + raise OSError("disk full") + + monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", failing_copy) + with pytest.warns(UserWarning, match="disk full"): + assert hf_checkpoint_utils.copy_non_model_files(source_dir, export_dir) == [] + assert not (export_dir / "tokenizer.json").exists() + + +def test_ensure_local_checkpoint_returns_a_local_dir_untouched(monkeypatch, tmp_path): + def fail(*args, **kwargs): + raise AssertionError("a local checkpoint must not hit the hub") + + monkeypatch.setattr(hf_checkpoint_utils, "snapshot_download", fail) + assert hf_checkpoint_utils.ensure_local_checkpoint(str(tmp_path)) == (None, str(tmp_path)) + + +class _FakeHfApi: + """``HfApi`` whose ``repo_info`` raises ``error``, or succeeds when it is None.""" + + error: Exception | None = None + + def repo_info(self, repo_id): + if self.error is not None: + raise self.error + + +def _fake_hub(monkeypatch, tmp_path, error=None): + # tests/unit/conftest.py sets HF_HUB_OFFLINE for the session; these tests model a live Hub. + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + calls = [] + + def fake_snapshot_download(*args, **kwargs): + calls.append((args, kwargs)) + return str(tmp_path) + + monkeypatch.setattr(hf_checkpoint_utils, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(hf_checkpoint_utils, "HfApi", type("Api", (_FakeHfApi,), {"error": error})) + return calls + + +def test_ensure_local_checkpoint_fetches_a_hub_id_in_full(monkeypatch, tmp_path): + calls = _fake_hub(monkeypatch, tmp_path) + + assert hf_checkpoint_utils.ensure_local_checkpoint("org/model") == ("org/model", str(tmp_path)) + assert calls == [(("org/model",), {})] + + +def test_ensure_local_checkpoint_warns_when_only_the_cache_is_left(monkeypatch, tmp_path): + """snapshot_download falls back to the cache silently; the fallback must be visible.""" + calls = _fake_hub(monkeypatch, tmp_path, error=ConnectionError("no route to host")) + + with pytest.warns(UserWarning, match="cached snapshot"): + result = hf_checkpoint_utils.ensure_local_checkpoint("org/model") + + assert result == ("org/model", str(tmp_path)) + assert calls == [(("org/model",), {"local_files_only": True})] + + +def test_ensure_local_checkpoint_uses_the_cache_quietly_when_offline(monkeypatch, tmp_path): + """HF_HUB_OFFLINE is the user's choice: do not ask the Hub, and do not warn on every run.""" + calls = _fake_hub(monkeypatch, tmp_path, error=AssertionError("the Hub must not be asked")) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + + with warnings.catch_warnings(): + warnings.simplefilter("error") + result = hf_checkpoint_utils.ensure_local_checkpoint("org/model") + + assert result == ("org/model", str(tmp_path)) + assert calls == [(("org/model",), {"local_files_only": True})] + + +def test_ensure_local_checkpoint_offline_and_uncached_raises(monkeypatch, tmp_path): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + + def not_cached(*args, **kwargs): + raise hf_hub_errors.LocalEntryNotFoundError("not in cache") + + monkeypatch.setattr(hf_checkpoint_utils, "snapshot_download", not_cached) + + with pytest.raises(RuntimeError, match="HF_HUB_OFFLINE is set"): + hf_checkpoint_utils.ensure_local_checkpoint("org/model") + + +def test_ensure_local_checkpoint_names_the_hub_error_when_nothing_is_cached(monkeypatch, tmp_path): + _fake_hub(monkeypatch, tmp_path, error=ConnectionError("no route to host")) + + def not_cached(*args, **kwargs): + raise hf_hub_errors.LocalEntryNotFoundError("not in cache") + + monkeypatch.setattr(hf_checkpoint_utils, "snapshot_download", not_cached) + + with pytest.raises(RuntimeError, match=r"no route to host.*not in the local cache"): + hf_checkpoint_utils.ensure_local_checkpoint("org/model") + + +def test_ensure_local_checkpoint_raises_for_a_missing_repo(monkeypatch, tmp_path): + # Its constructor wants an HTTP response; only the type matters here. + missing = hf_hub_errors.RepositoryNotFoundError.__new__(hf_hub_errors.RepositoryNotFoundError) + calls = _fake_hub(monkeypatch, tmp_path, error=missing) + + with pytest.raises(ValueError, match="not a local directory") as excinfo: + hf_checkpoint_utils.ensure_local_checkpoint("org/typo") + assert isinstance(excinfo.value.__cause__, hf_hub_errors.RepositoryNotFoundError) + assert calls == [] + + +def test_ensure_local_checkpoint_raises_for_a_gated_repo_unchanged(monkeypatch, tmp_path): + gated = hf_hub_errors.GatedRepoError.__new__(hf_hub_errors.GatedRepoError) + _fake_hub(monkeypatch, tmp_path, error=gated) + + with pytest.raises(hf_hub_errors.GatedRepoError): + hf_checkpoint_utils.ensure_local_checkpoint("org/gated") + + +def test_ensure_local_checkpoint_names_a_missing_local_path(tmp_path): + """A mistyped path is not a valid Hub ID either; the error must say it is not a directory.""" + missing = tmp_path / "ckpt-typo" + + with pytest.raises(ValueError, match="not a local directory") as excinfo: + hf_checkpoint_utils.ensure_local_checkpoint(str(missing)) + assert isinstance(excinfo.value.__cause__, hf_hub_errors.HFValidationError) + + +def test_copy_non_model_files_never_fetches_from_the_hub(monkeypatch, tmp_path): + """Exporters read the source from local disk only; a Hub ID is the caller's to download.""" + + def fail(*args, **kwargs): + raise AssertionError("copy_non_model_files must not download") + + monkeypatch.setattr(hf_checkpoint_utils, "snapshot_download", fail) + + with pytest.raises(ValueError, match="ensure_local_checkpoint"): + hf_checkpoint_utils.copy_non_model_files("org/model", tmp_path / "export") + + +def test_experiment_json_is_export_owned(): + """copy_non_model_files copies dotfiles, so without this the source checkpoint's MLflow pointer + would follow it into every derived checkpoint.""" + assert EXPERIMENT_JSON in hf_checkpoint_utils._EXPORT_OWNED_FILES + + +def test_copy_non_model_files_follows_links_only_within_the_checkpoint(tmp_path): + src_dir, export = tmp_path / "src", tmp_path / "export" + src_dir.mkdir() + outside = tmp_path / "secret.txt" + outside.write_text("host file") + (src_dir / "README.md").write_text("readme") + (src_dir / "LICENSE").symlink_to(src_dir / "README.md") + (src_dir / "notes.txt").symlink_to(outside) + (src_dir / "gone.txt").symlink_to(tmp_path / "missing.txt") + + with pytest.warns(UserWarning) as record: + copied = hf_checkpoint_utils.copy_non_model_files(src_dir, export) + + assert sorted(copied) == ["LICENSE", "README.md"] + assert (export / "LICENSE").read_text() == "readme" + assert not (export / "notes.txt").exists() + assert not (export / "gone.txt").exists() + warned = " ".join(str(w.message) for w in record) + assert "notes.txt" in warned and "gone.txt" in warned + + +@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX symlink support") +def test_copy_non_model_files_follows_hub_snapshot_links_into_blobs(tmp_path): + repo = tmp_path / "models--org--model" + snapshot, blobs = repo / "snapshots" / "abc123", repo / "blobs" + snapshot.mkdir(parents=True) + blobs.mkdir() + (blobs / "sha-tokenizer").write_text("{}") + (snapshot / "tokenizer.json").symlink_to("../../blobs/sha-tokenizer") + + copied = hf_checkpoint_utils.copy_non_model_files(snapshot, tmp_path / "export") + + assert copied == ["tokenizer.json"] + assert (tmp_path / "export" / "tokenizer.json").read_text() == "{}" + + +def _ensure_local_checkpoint_job(tmp_path, fail, group_ranks, rank, size): + utils = hf_checkpoint_utils + snapshot = tmp_path / "snapshot" + + def fake_snapshot_download(repo_id, **kwargs): + (tmp_path / f"downloaded-by-rank{rank}").write_text(repo_id) + if fail: + raise OSError("hub unreachable") + snapshot.mkdir(exist_ok=True) + return str(snapshot) + + utils.snapshot_download = fake_snapshot_download + utils.HfApi = _FakeHfApi + os.environ.pop("HF_HUB_OFFLINE", None) # inherited from tests/unit/conftest.py + # new_group is collective over every rank, members or not. + group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") + if group_ranks is not None and rank not in group_ranks: + return + if fail: + with pytest.raises(RuntimeError, match="hub unreachable"): + utils.ensure_local_checkpoint("org/model", group=group) + else: + result = utils.ensure_local_checkpoint("org/model", group=group) + assert result == ("org/model", str(snapshot)) + + +def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, size): + utils = hf_checkpoint_utils + + def fail(*args, **kwargs): + raise AssertionError("a local path must not be downloaded") + + utils.snapshot_download = fail + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir(exist_ok=True) + if hidden_from_rank1 and rank == 1: + # As if the directory sat on rank 0's node-local disk only. + isdir = os.path.isdir + utils.os.path.isdir = lambda path: path != str(checkpoint) and isdir(path) + with pytest.raises(RuntimeError, match="this rank cannot see"): + utils.ensure_local_checkpoint(str(checkpoint)) + else: + assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) + + +@pytest.mark.parametrize("hidden_from_rank1", [False, True]) +def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidden_from_rank1): + """Ranks that each checked a local path could disagree about it and strand one another in the + collective; rank 0 decides, and a rank that cannot see its answer raises instead.""" + job = partial(_ensure_local_checkpoint_local_path_job, tmp_path, hidden_from_rank1) + spawn_multiprocess_job(2, job, backend="gloo") + + +@pytest.mark.parametrize( + ("fail", "group_ranks", "downloader"), + [ + (False, None, "rank0"), + (True, None, "rank0"), + (False, [0, 1], "rank0"), + # Rank 0 of the group, not of the world: global rank 1. + (False, [1], "rank1"), + ], +) +def test_ensure_local_checkpoint_downloads_once_per_group(tmp_path, fail, group_ranks, downloader): + """Rank 0 of the group downloads once and every member gets its result -- including its + failure, so the waiting ranks raise instead of hanging.""" + job = partial(_ensure_local_checkpoint_job, tmp_path, fail, group_ranks) + spawn_multiprocess_job(2, job, backend="gloo") + assert [p.name for p in tmp_path.glob("downloaded-by-*")] == [f"downloaded-by-{downloader}"]