diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 688debc80692..9f3218f2e192 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -13,10 +13,10 @@ from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import ( AutoCheckpointMapper, BaseCheckpointLoader) from tensorrt_llm._torch.weight_sharing import ( - IdentityCheckPolicy, PostTransformFeature, PostTransformProfile, - PostTransformProfileRegistry, PostTransformQualificationDecision, - PostTransformTransferScope, SourceIdentity, - check_weight_sharing_compatibility) + ArtifactIdentity, IdentityCheckPolicy, PostTransformFeature, + PostTransformProfile, PostTransformProfileRegistry, + PostTransformQualificationDecision, PostTransformTransferScope, + SourceIdentity, check_weight_sharing_compatibility) from tensorrt_llm._utils import str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, ExecutorMemoryType, @@ -432,6 +432,39 @@ def _needs_source_identity(checkpoint_loader: BaseCheckpointLoader, """ return load_format == LoadFormat.GMS or checkpoint_loader.checkpoint_format == "MX" + @staticmethod + def _build_source_identity( + config: ModelConfig, + model: DecoderModelForCausalLM, + *, + checkpoint_dir: str, + model_name: str, + fallback_on_artifact_error: bool, + ) -> Optional[SourceIdentity]: + """Build the local identity without weakening artifact validation. + + Artifact construction remains fail-closed. MX may convert an artifact + error into an unavailable local identity so its compatibility gate + falls back to disk; GMS propagates the error because it has no fallback. + """ + try: + artifact_identity = ArtifactIdentity.from_checkpoint(checkpoint_dir) + except (OSError, RuntimeError, ValueError) as error: + if not fallback_on_artifact_error: + raise + logger.warning( + "Unable to build checkpoint artifact identity for MX checkpoint " + f"{checkpoint_dir}; falling back to regular checkpoint loading: {error}" + ) + return None + + return SourceIdentity.from_model_config( + config, + model, + artifact_identity=artifact_identity, + model_name=model_name, + ) + def load( self, checkpoint_dir: str, @@ -475,12 +508,16 @@ def load( # ground truth; building it here (post-construction, # pre-weight-load) gives producer and consumer a common, # comparable lifecycle point. - self._source_identity = SourceIdentity.from_model_config( + self._source_identity = self._build_source_identity( config, model, + checkpoint_dir=checkpoint_dir, model_name=str( getattr(self.llm_args, "model", None) or checkpoint_dir), + fallback_on_artifact_error=( + load_format != LoadFormat.GMS + and checkpoint_loader.checkpoint_format == "MX"), ) memo: dict[torch.Tensor, torch.Tensor] = {} diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index 80b70c6c515d..38ec2471c41c 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -14,6 +14,10 @@ # limitations under the License. """Backend-agnostic weight-sharing utilities (MX, GMS, ...).""" +from tensorrt_llm._torch.weight_sharing.artifact_identity import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, +) from tensorrt_llm._torch.weight_sharing.post_transform_profiles import ( PostTransformFeature, PostTransformProfile, @@ -33,6 +37,8 @@ ) __all__ = [ + "ARTIFACT_IDENTITY_FORMAT_VERSION", + "ArtifactIdentity", "SOURCE_IDENTITY_FORMAT_VERSION", "PostTransformFeature", "PostTransformProfile", diff --git a/tensorrt_llm/_torch/weight_sharing/artifact_identity.py b/tensorrt_llm/_torch/weight_sharing/artifact_identity.py new file mode 100644 index 000000000000..d1b8908e241a --- /dev/null +++ b/tensorrt_llm/_torch/weight_sharing/artifact_identity.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Immutable checkpoint identity for shared-weight compatibility checks.""" + +from __future__ import annotations + +import hashlib +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +ARTIFACT_IDENTITY_FORMAT_VERSION = 1 + +_HF_SNAPSHOT_SCHEME = "hf_snapshot_revision" +_CHECKPOINT_MANIFEST_SCHEME = "checkpoint_manifest_sha256" +_SUPPORTED_SCHEMES = frozenset({_HF_SNAPSHOT_SCHEME, _CHECKPOINT_MANIFEST_SCHEME}) +_IGNORED_DIRECTORY_NAMES = frozenset({".cache", ".git", "__pycache__"}) +_IGNORED_FILE_NAMES = frozenset({".DS_Store"}) +_HASH_CHUNK_SIZE = 1024 * 1024 + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps(value, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _is_hex(value: str, lengths: tuple[int, ...]) -> bool: + return len(value) in lengths and all(char in "0123456789abcdef" for char in value) + + +def _hf_snapshot_descriptor(path: Path) -> tuple[str, str] | None: + """Return an immutable HF revision and repository-relative subpath.""" + parts = path.resolve().parts + for index, part in enumerate(parts[:-1]): + if part != "snapshots" or index == 0: + continue + if not parts[index - 1].startswith("models--"): + continue + + revision = parts[index + 1].lower() + if not _is_hex(revision, (40, 64)): + continue + subpath = "/".join(parts[index + 2 :]) + return revision, subpath + return None + + +def _raise_walk_error(error: OSError) -> None: + raise error + + +def _checkpoint_files(path: Path) -> tuple[Path, list[Path]]: + if path.is_file(): + return path.parent, [path] + + files = [] + for directory, directory_names, file_names in os.walk(path, onerror=_raise_walk_error): + retained_directories = [] + for directory_name in directory_names: + if directory_name in _IGNORED_DIRECTORY_NAMES: + continue + nested_directory = Path(directory) / directory_name + if nested_directory.is_symlink(): + raise ValueError( + "Checkpoint manifests do not support nested symlinked directories: " + f"{nested_directory}" + ) + retained_directories.append(directory_name) + directory_names[:] = retained_directories + for file_name in file_names: + if file_name in _IGNORED_FILE_NAMES: + continue + candidate = Path(directory) / file_name + if candidate.is_file(): + files.append(candidate) + files.sort(key=lambda candidate: candidate.relative_to(path).as_posix()) + if not files: + raise ValueError(f"Checkpoint path contains no files: {path}") + return path, files + + +def _sha256_file(path: Path) -> tuple[int, str]: + before = path.stat() + digest = hashlib.sha256() + with path.open("rb") as checkpoint_file: + for chunk in iter(lambda: checkpoint_file.read(_HASH_CHUNK_SIZE), b""): + digest.update(chunk) + after = path.stat() + if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns): + raise RuntimeError(f"Checkpoint file changed while being fingerprinted: {path}") + return after.st_size, digest.hexdigest() + + +def _checkpoint_manifest_digest(path: Path) -> str: + root, files = _checkpoint_files(path) + manifest = [] + for checkpoint_file in files: + size, digest = _sha256_file(checkpoint_file) + manifest.append( + { + "path": checkpoint_file.relative_to(root).as_posix(), + "size": size, + "sha256": digest, + } + ) + return _canonical_hash( + { + "format_version": ARTIFACT_IDENTITY_FORMAT_VERSION, + "files": manifest, + } + ) + + +@dataclass(frozen=True) +class ArtifactIdentity: + """Versioned identity of the immutable checkpoint artifact being loaded. + + `SourceIdentity` embeds this value as a global compatibility component. + Hugging Face cache snapshots use their immutable commit revision; local + checkpoints use a canonical manifest of relative paths, sizes, and file + content digests. Absolute paths are intentionally excluded. + """ + + format_version: int + scheme: str + digest: str + + def __post_init__(self) -> None: + if not isinstance(self.format_version, int) or isinstance(self.format_version, bool): + raise ValueError("ArtifactIdentity format version must be an integer") + if self.format_version != ARTIFACT_IDENTITY_FORMAT_VERSION: + raise ValueError(f"Unsupported ArtifactIdentity format version: {self.format_version}") + if not isinstance(self.scheme, str): + raise ValueError("ArtifactIdentity scheme must be a string") + if self.scheme not in _SUPPORTED_SCHEMES: + raise ValueError(f"Unsupported ArtifactIdentity scheme: {self.scheme}") + if not isinstance(self.digest, str): + raise ValueError("ArtifactIdentity digest must be a string") + + normalized_digest = self.digest.lower() + if not _is_hex(normalized_digest, (64,)): + raise ValueError("ArtifactIdentity digest must be a 64-character hex value") + object.__setattr__(self, "digest", normalized_digest) + + @classmethod + def from_checkpoint(cls, checkpoint_path: str | os.PathLike[str]) -> "ArtifactIdentity": + """Build an identity from an immutable snapshot or local checkpoint. + + Args: + checkpoint_path: A model checkpoint file or directory. + + Returns: + The path-independent checkpoint identity. + + Raises: + FileNotFoundError: If `checkpoint_path` does not exist. + ValueError: If a local checkpoint directory contains no files. + RuntimeError: If a local checkpoint changes while it is hashed. + + Note: + Local checkpoints have no authoritative immutable revision, so + their regular files are read in full to derive a content-bound + manifest. Hugging Face cache snapshots use the resolved immutable + revision without rereading model shards. + """ + path = Path(checkpoint_path).expanduser() + if not path.exists(): + raise FileNotFoundError(f"Checkpoint path does not exist: {path}") + + snapshot_descriptor = _hf_snapshot_descriptor(path) + if snapshot_descriptor is not None: + revision, subpath = snapshot_descriptor + digest = _canonical_hash( + { + "scheme": _HF_SNAPSHOT_SCHEME, + "revision": revision, + "subpath": subpath, + } + ) + return cls( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme=_HF_SNAPSHOT_SCHEME, + digest=digest, + ) + + return cls( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme=_CHECKPOINT_MANIFEST_SCHEME, + digest=_checkpoint_manifest_digest(path), + ) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return { + "format_version": self.format_version, + "scheme": self.scheme, + "digest": self.digest, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ArtifactIdentity": + """Reconstruct and validate a serialized artifact identity.""" + return cls( + format_version=data["format_version"], + scheme=data["scheme"], + digest=data["digest"], + ) diff --git a/tensorrt_llm/_torch/weight_sharing/source_identity.py b/tensorrt_llm/_torch/weight_sharing/source_identity.py index 18168703c868..c6b010634760 100644 --- a/tensorrt_llm/_torch/weight_sharing/source_identity.py +++ b/tensorrt_llm/_torch/weight_sharing/source_identity.py @@ -14,17 +14,19 @@ # limitations under the License. """Backend-agnostic source identity for weight-sharing receivers. -A :class:`SourceIdentity` is a serializable fingerprint of configuration choices -that affect how a model's weights are laid out in memory. It exists so that a -*receiver* of pre-laid-out weights (e.g. MX peer-to-peer transfer, or a GMS -read-only materialize) can verify that both the producer ("source") and the -consumer built identities and agree on every layout-affecting choice before the -receiver consumes shared weights. +A :class:`SourceIdentity` is a serializable fingerprint of an immutable +checkpoint artifact and the configuration choices that affect how its weights +are laid out in memory. It exists so that a *receiver* of pre-laid-out weights +(e.g. MX peer-to-peer transfer, or a GMS read-only materialize) can verify that +both the producer ("source") and the consumer built identities and agree on the +artifact and every layout-affecting choice before consuming shared weights. The identity is intentionally decoupled from any specific weight-sharing technology (neither MX nor GMS appears here). Both consume it identically: - local = SourceIdentity.from_model_config(model_config) + local = SourceIdentity.from_model_config( + model_config, checkpoint_dir="/path/to/checkpoint" + ) decision = check_weight_sharing_compatibility(local, source_identity, policy) if decision.should_share: ... # pull / materialize shared weights @@ -40,8 +42,9 @@ -------------- The fingerprint is split so comparison can be selective: -* **global fingerprint** -- rank-invariant model identity, quantization, - backend selection, fusion flags, and parallel *sizes* (TP/PP/EP/CP). +* **global fingerprint** -- immutable checkpoint artifact, rank-invariant + model identity, quantization, backend selection, fusion flags, and parallel + *sizes* (TP/PP/EP/CP). * **shard fingerprint** -- this rank's TP/PP/EP/CP *rank* slice plus the realized local parameter/buffer `(shape, dtype)` layout. Receiver rank `N` must align with the source rank that produced shard `N`. @@ -67,6 +70,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any, List, Optional +from tensorrt_llm._torch.weight_sharing.artifact_identity import ArtifactIdentity from tensorrt_llm.logger import logger if TYPE_CHECKING: @@ -78,7 +82,7 @@ # Bump when the fingerprint projection changes in a way that makes previously # stored identities incomparable. Two identities with different format versions # never match. -SOURCE_IDENTITY_FORMAT_VERSION = 1 +SOURCE_IDENTITY_FORMAT_VERSION = 2 _PRETRAINED_METADATA_FIELDS = frozenset( { @@ -224,6 +228,7 @@ class SourceIdentity: format_version: int # --- global parts (must match across all ranks) --- + artifact_identity: ArtifactIdentity model_fingerprint: str quant_fingerprint: str backend_fingerprint: str @@ -246,6 +251,8 @@ def from_model_config( model_config: "ModelConfig", model: Optional["nn.Module"] = None, *, + checkpoint_dir: Optional[str] = None, + artifact_identity: Optional[ArtifactIdentity] = None, model_name: Optional[str] = None, ) -> "SourceIdentity": """Build an identity from a torch-backend :class:`ModelConfig`. @@ -258,6 +265,12 @@ def from_model_config( Producer and consumer must build the identity at the same lifecycle point (model construction, before weight load). When `None`, the shard fingerprint contains no tensor-layout data. + checkpoint_dir: Checkpoint file or directory used to derive the + nested artifact identity. Required unless `artifact_identity` + is supplied explicitly. + artifact_identity: Precomputed immutable checkpoint identity, + primarily for callers that resolve provenance outside this + method. Mutually exclusive with `checkpoint_dir`. model_name: Human-readable model identity used by discovery layers (e.g. the MX server's source catalog). Does not affect the compatibility fingerprints. @@ -265,7 +278,17 @@ def from_model_config( Returns: A fully populated :class:`SourceIdentity` for `model_config.mapping.rank`. + + Raises: + ValueError: If neither or both artifact identity inputs are given. """ + if checkpoint_dir is None and artifact_identity is None: + raise ValueError("Exactly one of checkpoint_dir or artifact_identity must be provided") + if checkpoint_dir is not None and artifact_identity is not None: + raise ValueError("Exactly one of checkpoint_dir or artifact_identity must be provided") + if artifact_identity is None: + artifact_identity = ArtifactIdentity.from_checkpoint(checkpoint_dir) + mapping = model_config.mapping rank = getattr(mapping, "rank", 0) @@ -275,6 +298,7 @@ def from_model_config( return cls( format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=artifact_identity, model_fingerprint=cls._build_model_fingerprint(model_config), quant_fingerprint=cls._build_quant_fingerprint(model_config), backend_fingerprint=cls._build_backend_fingerprint(model_config), @@ -418,6 +442,7 @@ def global_fingerprint(self) -> str: return _canonical_hash( { "format_version": self.format_version, + "artifact": self.artifact_identity.to_dict(), "model": self.model_fingerprint, "quant": self.quant_fingerprint, "backend": self.backend_fingerprint, @@ -447,6 +472,8 @@ def matches( mismatched.append("format_version") if compare_global: + if self.artifact_identity != other.artifact_identity: + mismatched.append("artifact_identity") for name in ( "model_fingerprint", "quant_fingerprint", @@ -473,6 +500,7 @@ def to_dict(self) -> dict: """ return { "format_version": self.format_version, + "artifact_identity": self.artifact_identity.to_dict(), "model_fingerprint": self.model_fingerprint, "quant_fingerprint": self.quant_fingerprint, "backend_fingerprint": self.backend_fingerprint, @@ -496,8 +524,13 @@ def from_dict(cls, data: dict) -> "SourceIdentity": Returns: The reconstructed :class:`SourceIdentity`. """ + format_version = data["format_version"] + if format_version != SOURCE_IDENTITY_FORMAT_VERSION: + raise ValueError(f"Unsupported SourceIdentity format version: {format_version}") + return cls( - format_version=data["format_version"], + format_version=format_version, + artifact_identity=ArtifactIdentity.from_dict(data["artifact_identity"]), model_fingerprint=data["model_fingerprint"], quant_fingerprint=data["quant_fingerprint"], backend_fingerprint=data["backend_fingerprint"], @@ -548,7 +581,8 @@ def check_weight_sharing_compatibility( result = IdentityMatchResult(matched=False, mismatched_fields=missing_fields) message = ( "SourceIdentity unavailable for fields " - f"{missing_fields}; receiver cannot verify source weight layout." + f"{missing_fields}; receiver cannot verify the source checkpoint " + "artifact and weight layout." ) if policy is IdentityCheckPolicy.STRICT: raise SourceIdentityMismatchError(message) @@ -574,7 +608,7 @@ def check_weight_sharing_compatibility( message = ( "SourceIdentity mismatch on fields " f"{result.mismatched_fields}; receiver and source disagree on " - "weight layout." + "the checkpoint artifact or weight layout." ) if policy is IdentityCheckPolicy.STRICT: raise SourceIdentityMismatchError(message) diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 34dd8223a6a5..b56935da6378 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -14,6 +14,9 @@ from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + SOURCE_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, PostTransformProfile, PostTransformProfileRegistry, PostTransformTransferScope, @@ -21,7 +24,12 @@ from tensorrt_llm.llmapi.llm_args import LoadFormat _SOURCE_IDENTITY = model_loader_mod.SourceIdentity( - format_version=1, + format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest="0" * 64, + ), model_fingerprint="model", quant_fingerprint="quant", backend_fingerprint="backend", @@ -100,10 +108,25 @@ def _make_loader(monkeypatch, *, events, spec_config=None): monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) # These tests stub ModelConfig, while SourceIdentity has dedicated # coverage. Keep this file focused on ModelLoader GMS branch behavior. + + def _build_artifact_identity(_cls, checkpoint_dir): + assert checkpoint_dir == "/ckpt" + return _SOURCE_IDENTITY.artifact_identity + + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + classmethod(_build_artifact_identity), + ) + + def _build_source_identity(_cls, *_args, **kwargs): + assert kwargs["artifact_identity"] is _SOURCE_IDENTITY.artifact_identity + return _SOURCE_IDENTITY + monkeypatch.setattr( model_loader_mod.SourceIdentity, "from_model_config", - classmethod(lambda cls, *_args, **_kwargs: _SOURCE_IDENTITY), + classmethod(_build_source_identity), ) monkeypatch.setattr( model_loader_mod.AutoModelForCausalLM, @@ -181,6 +204,31 @@ def _tiny_profile_registry() -> PostTransformProfileRegistry: ) +def test_gms_artifact_identity_failure_remains_fatal(monkeypatch): + loader = _make_loader(monkeypatch, events=[]) + artifact_error = ValueError( + "Checkpoint manifests do not support nested symlinked directories: /ckpt/shards" + ) + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + MagicMock(side_effect=artifact_error), + ) + source_identity_factory = MagicMock() + monkeypatch.setattr( + model_loader_mod.SourceIdentity, + "from_model_config", + source_identity_factory, + ) + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "MX" + + with pytest.raises(ValueError, match="nested symlinked directories"): + loader.load("/ckpt", checkpoint_loader) + + source_identity_factory.assert_not_called() + + @pytest.mark.parametrize( "is_rw, expected_events", [ diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 82a92622fd30..ccdf960ea786 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -24,6 +24,9 @@ from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + SOURCE_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry, @@ -33,7 +36,12 @@ from tensorrt_llm.llmapi.llm_args import LoadFormat _SOURCE_IDENTITY = model_loader_mod.SourceIdentity( - format_version=1, + format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest="0" * 64, + ), model_fingerprint="model", quant_fingerprint="quant", backend_fingerprint="backend", @@ -188,10 +196,25 @@ def _make_loader(monkeypatch, *, events, spec_config=None): monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) # These tests stub ModelConfig, while SourceIdentity has dedicated # coverage. Keep this file focused on ModelLoader MX branch behavior. + + def _build_artifact_identity(_cls, checkpoint_dir): + assert checkpoint_dir == "/ckpt" + return _SOURCE_IDENTITY.artifact_identity + + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + classmethod(_build_artifact_identity), + ) + + def _build_source_identity(_cls, *_args, **kwargs): + assert kwargs["artifact_identity"] is _SOURCE_IDENTITY.artifact_identity + return _SOURCE_IDENTITY + monkeypatch.setattr( model_loader_mod.SourceIdentity, "from_model_config", - classmethod(lambda cls, *_args, **_kwargs: _SOURCE_IDENTITY), + classmethod(_build_source_identity), ) monkeypatch.setattr( model_loader_mod.AutoModelForCausalLM, @@ -484,6 +507,53 @@ def test_mx_fallback_runs_standard_weight_mapping(monkeypatch): ) +def test_mx_artifact_identity_failure_falls_back_to_disk(monkeypatch): + events = [] + loader = _make_loader(monkeypatch, events=events) + monkeypatch.setattr( + ModelLoader, + "_POST_TRANSFORM_PROFILE_REGISTRY", + _tiny_profile_registry(), + ) + artifact_error = ValueError( + "Checkpoint manifests do not support nested symlinked directories: /ckpt/shards" + ) + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + MagicMock(side_effect=artifact_error), + ) + source_identity_factory = MagicMock() + monkeypatch.setattr( + model_loader_mod.SourceIdentity, + "from_model_config", + source_identity_factory, + ) + warning = MagicMock() + monkeypatch.setattr(model_loader_mod.logger, "warning", warning) + + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader.is_weights_preloaded.return_value = False + checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} + checkpoint_loader.get_initialized_weight_mapper.return_value = MagicMock() + + model, _ = loader.load("/ckpt", checkpoint_loader) + + assert loader._source_identity is None + source_identity_factory.assert_not_called() + warning.assert_called_once() + assert "falling back to regular checkpoint loading" in warning.call_args.args[0] + _args, kwargs = checkpoint_loader.load_weights.call_args + assert kwargs["source_identity"] is None + checkpoint_loader.post_load_publish.assert_called_once_with( + model, + checkpoint_dir="/ckpt", + weights_preloaded=False, + source_identity=None, + ) + + class _HookRecorder(nn.Module): def __init__( self, diff --git a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py index c8e845c5d3b0..e80026d31114 100644 --- a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py +++ b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py @@ -42,14 +42,24 @@ _resolve_mx_model_name, _serialize_source_identity, ) -from tensorrt_llm._torch.weight_sharing import SourceIdentity +from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + SOURCE_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, + SourceIdentity, +) _MISSING = object() def _identity(rank: int = 0, suffix: str = "same") -> SourceIdentity: return SourceIdentity( - format_version=1, + format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest="0" * 64, + ), model_fingerprint=f"model-{suffix}", quant_fingerprint=f"quant-{suffix}", backend_fingerprint=f"backend-{suffix}", @@ -753,12 +763,9 @@ def _publish_side_effect(model, **_kwargs): def test_serialized_identity_ignores_local_checkpoint_path(self): donor_identity = _identity() - receiver_identity = SourceIdentity( - **{ - **donor_identity.to_dict(), - "model_name": "/tmp/no-shards/TinyLlama", - } - ) + receiver_payload = donor_identity.to_dict() + receiver_payload["model_name"] = "/tmp/no-shards/TinyLlama" + receiver_identity = SourceIdentity.from_dict(receiver_payload) assert donor_identity.model_name != receiver_identity.model_name assert _serialize_source_identity(donor_identity) == _serialize_source_identity( diff --git a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py index c68b10506b7a..ff567b742cec 100644 --- a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py +++ b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py @@ -19,9 +19,14 @@ `test_source_identity.py` and `test_mx_source_identity_gate.py`. """ +import hashlib from typing import Optional, Sequence -from tensorrt_llm._torch.weight_sharing import SourceIdentity +from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, + SourceIdentity, +) _UNSET = object() @@ -204,16 +209,37 @@ def named_buffers(self): return list(self._buffers.items()) -def identity_from(config: FakeModelConfig, *, model_name: Optional[str] = None) -> SourceIdentity: +def make_artifact_identity(key: str = "same") -> ArtifactIdentity: + """Build a deterministic local-checkpoint identity for tests.""" + return ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest=hashlib.sha256(key.encode("utf-8")).hexdigest(), + ) + + +def identity_from( + config: FakeModelConfig, + *, + model_name: Optional[str] = None, + artifact_key: str = "same", +) -> SourceIdentity: """Build a :class:`SourceIdentity` from a fake config and derived model.""" return SourceIdentity.from_model_config( - config, FakeModel(config.pretrained_config), model_name=model_name + config, + FakeModel(config.pretrained_config), + artifact_identity=make_artifact_identity(artifact_key), + model_name=model_name, ) def make_identity( - *, attn_backend: str = "TRTLLM", rank: int = 0, model_name: str = "m" + *, + attn_backend: str = "TRTLLM", + rank: int = 0, + model_name: str = "m", + artifact_key: str = "same", ) -> SourceIdentity: """Build a :class:`SourceIdentity` from a fake config for `rank`.""" cfg = FakeModelConfig(mapping=FakeMapping(rank=rank, tp_rank=rank), attn_backend=attn_backend) - return identity_from(cfg, model_name=model_name) + return identity_from(cfg, model_name=model_name, artifact_key=artifact_key) diff --git a/tests/unittest/_torch/weight_sharing/test_artifact_identity.py b/tests/unittest/_torch/weight_sharing/test_artifact_identity.py new file mode 100644 index 000000000000..30526ce77c66 --- /dev/null +++ b/tests/unittest/_torch/weight_sharing/test_artifact_identity.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for immutable checkpoint artifact identities.""" + +from pathlib import Path + +import pytest + +from tensorrt_llm._torch.weight_sharing import ArtifactIdentity + + +def _write_checkpoint(path: Path, weights: bytes = b"weights") -> None: + path.mkdir(parents=True) + (path / "config.json").write_text('{"architectures":["LlamaForCausalLM"]}') + (path / "model.safetensors").write_bytes(weights) + + +def test_local_checkpoint_identity_is_path_independent(tmp_path: Path) -> None: + left = tmp_path / "left" / "checkpoint" + right = tmp_path / "right" / "checkpoint" + _write_checkpoint(left) + _write_checkpoint(right) + + assert ArtifactIdentity.from_checkpoint(left) == ArtifactIdentity.from_checkpoint(right) + + +def test_local_checkpoint_identity_binds_file_contents(tmp_path: Path) -> None: + left = tmp_path / "left" + right = tmp_path / "right" + _write_checkpoint(left, weights=b"fine-tune-a") + _write_checkpoint(right, weights=b"fine-tune-b") + + assert ArtifactIdentity.from_checkpoint(left) != ArtifactIdentity.from_checkpoint(right) + + +def test_local_checkpoint_identity_ignores_cache_and_scm_metadata(tmp_path: Path) -> None: + left = tmp_path / "left" + right = tmp_path / "right" + _write_checkpoint(left) + _write_checkpoint(right) + (left / ".cache").mkdir() + (left / ".cache" / "download.lock").write_text("transient") + (right / ".git").mkdir() + (right / ".git" / "HEAD").write_text("ref: refs/heads/main") + + assert ArtifactIdentity.from_checkpoint(left) == ArtifactIdentity.from_checkpoint(right) + + +def test_local_checkpoint_identity_rejects_nested_directory_symlink(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + external_weights = tmp_path / "external-weights" + external_weights.mkdir() + (external_weights / "model.safetensors").write_bytes(b"weights") + (checkpoint / "weights").symlink_to(external_weights, target_is_directory=True) + + with pytest.raises(ValueError, match="nested symlinked directories"): + ArtifactIdentity.from_checkpoint(checkpoint) + + +def test_hf_snapshot_identity_binds_revision_across_cache_roots(tmp_path: Path) -> None: + revision = "a" * 40 + left = tmp_path / "cache-a" / "models--org--model" / "snapshots" / revision + right = tmp_path / "cache-b" / "models--org--model" / "snapshots" / revision + left.mkdir(parents=True) + right.mkdir(parents=True) + + left_identity = ArtifactIdentity.from_checkpoint(left) + right_identity = ArtifactIdentity.from_checkpoint(right) + assert left_identity == right_identity + assert left_identity.scheme == "hf_snapshot_revision" + + +def test_hf_snapshot_identity_binds_revision_and_subpath(tmp_path: Path) -> None: + snapshot = tmp_path / "models--org--model" / "snapshots" + revision_a = snapshot / ("a" * 40) + revision_b = snapshot / ("b" * 40) + (revision_a / "variant-a").mkdir(parents=True) + (revision_a / "variant-b").mkdir() + revision_b.mkdir(parents=True) + + root_identity = ArtifactIdentity.from_checkpoint(revision_a) + assert root_identity != ArtifactIdentity.from_checkpoint(revision_b) + assert root_identity != ArtifactIdentity.from_checkpoint(revision_a / "variant-a") + assert ArtifactIdentity.from_checkpoint( + revision_a / "variant-a" + ) != ArtifactIdentity.from_checkpoint(revision_a / "variant-b") + + +def test_serialization_roundtrip(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + identity = ArtifactIdentity.from_checkpoint(checkpoint) + + assert ArtifactIdentity.from_dict(identity.to_dict()) == identity + + +def test_rejects_unknown_format_version(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + payload = ArtifactIdentity.from_checkpoint(checkpoint).to_dict() + payload["format_version"] += 1 + + with pytest.raises(ValueError, match="Unsupported ArtifactIdentity format version"): + ArtifactIdentity.from_dict(payload) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("scheme", "unknown", "Unsupported ArtifactIdentity scheme"), + ("scheme", [], "scheme must be a string"), + ("digest", "not-a-digest", "64-character hex value"), + ("digest", 1, "digest must be a string"), + ], +) +def test_rejects_invalid_serialized_fields( + tmp_path: Path, field: str, value: object, message: str +) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + payload = ArtifactIdentity.from_checkpoint(checkpoint).to_dict() + payload[field] = value + + with pytest.raises(ValueError, match=message): + ArtifactIdentity.from_dict(payload) + + +@pytest.mark.parametrize("version", [True, "1"]) +def test_rejects_non_integer_format_version(tmp_path: Path, version: object) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + payload = ArtifactIdentity.from_checkpoint(checkpoint).to_dict() + payload["format_version"] = version + + with pytest.raises(ValueError, match="format version must be an integer"): + ArtifactIdentity.from_dict(payload) + + +def test_rejects_missing_or_empty_checkpoint(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError): + ArtifactIdentity.from_checkpoint(tmp_path / "missing") + + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(ValueError, match="contains no files"): + ArtifactIdentity.from_checkpoint(empty) diff --git a/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py b/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py index f6689bfd0caa..99f21b4bd9b9 100644 --- a/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py +++ b/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py @@ -79,6 +79,14 @@ def test_gate_raises_on_mismatch(): loader._check_gms_source_identity(_FakeGMSBackend(writer)) +def test_gate_raises_on_checkpoint_artifact_mismatch(): + local = _identity(artifact_key="fine-tune-a") + writer = _identity(artifact_key="fine-tune-b") + loader = _new_loader(local) + with pytest.raises(SourceIdentityMismatchError): + loader._check_gms_source_identity(_FakeGMSBackend(writer)) + + def test_gate_raises_when_writer_identity_unavailable(): # Publisher metadata not wired yet (get_source_identity returns None); # GMS has no disk fallback, so unverified sharing must raise. diff --git a/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py b/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py index 2a3d3af7d179..e2658ca86d21 100644 --- a/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py +++ b/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py @@ -51,6 +51,13 @@ def test_gate_falls_back_on_mismatch(): assert loader._source_metadata_identity_compatible(_build_mx_source_metadata(source)) is False +def test_gate_falls_back_on_checkpoint_artifact_mismatch(): + local = _identity(artifact_key="fine-tune-a") + source = _identity(artifact_key="fine-tune-b") + loader = _new_loader(local) + assert loader._source_metadata_identity_compatible(_build_mx_source_metadata(source)) is False + + def test_gate_falls_back_when_no_local_identity(): # MX must not consume shared weights unless the receiver identity exists. loader = _new_loader(None) diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index a02995c4f520..779dc8c355bd 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -19,6 +19,7 @@ """ import copy +from pathlib import Path import pytest from _source_identity_fakes import ( @@ -29,6 +30,7 @@ FakeQuantConfig, FakeQuantConfigWithPythonOnlyField, identity_from, + make_artifact_identity, ) from tensorrt_llm._torch.weight_sharing import ( @@ -48,6 +50,24 @@ def test_identical_configs_match(): assert bool(result) is True +def test_from_model_config_derives_artifact_identity(tmp_path: Path) -> None: + (tmp_path / "model.safetensors").write_bytes(b"checkpoint") + identity = SourceIdentity.from_model_config(FakeModelConfig(), checkpoint_dir=str(tmp_path)) + assert identity.artifact_identity.scheme == "checkpoint_manifest_sha256" + + +def test_from_model_config_requires_one_artifact_source() -> None: + config = FakeModelConfig() + with pytest.raises(ValueError, match="Exactly one"): + SourceIdentity.from_model_config(config) + with pytest.raises(ValueError, match="Exactly one"): + SourceIdentity.from_model_config( + config, + checkpoint_dir="/checkpoint", + artifact_identity=make_artifact_identity(), + ) + + def test_rank_defaults_from_mapping(): cfg = FakeModelConfig(mapping=FakeMapping(rank=3, tp_rank=3)) identity = identity_from(cfg) @@ -83,10 +103,14 @@ def test_param_dtype_override_flags_shard(): # the realized-layout fingerprint catches it. cfg = FakeModelConfig() a = SourceIdentity.from_model_config( - cfg, FakeModel(cfg.pretrained_config, dtype="torch.bfloat16") + cfg, + FakeModel(cfg.pretrained_config, dtype="torch.bfloat16"), + artifact_identity=make_artifact_identity(), ) b = SourceIdentity.from_model_config( - cfg, FakeModel(cfg.pretrained_config, dtype="torch.float16") + cfg, + FakeModel(cfg.pretrained_config, dtype="torch.float16"), + artifact_identity=make_artifact_identity(), ) result = a.matches(b) assert not result.matched @@ -109,11 +133,24 @@ def test_cross_architecture_same_shapes_flags_global(): assert "model_fingerprint" in result.mismatched_fields +def test_different_checkpoint_artifacts_flag_global(): + a = identity_from(FakeModelConfig(), artifact_key="fine-tune-a") + b = identity_from(FakeModelConfig(), artifact_key="fine-tune-b") + result = a.matches(b) + assert not result.matched + assert result.mismatched_fields == ["artifact_identity"] + assert a.global_fingerprint != b.global_fingerprint + + def test_no_model_degrades_to_architecture_only(): - # Without a module, the fingerprint still builds (architecture-only) and - # two identical configs still match. - a = SourceIdentity.from_model_config(FakeModelConfig(), None) - b = SourceIdentity.from_model_config(FakeModelConfig(), None) + # Without a module, the shard fingerprint has no realized tensor layout, + # while matching artifacts and configurations remain comparable. + a = SourceIdentity.from_model_config( + FakeModelConfig(), None, artifact_identity=make_artifact_identity() + ) + b = SourceIdentity.from_model_config( + FakeModelConfig(), None, artifact_identity=make_artifact_identity() + ) assert a.matches(b).matched @@ -168,6 +205,29 @@ def test_serialization_roundtrip(): restored = SourceIdentity.from_dict(a.to_dict()) assert restored == a assert a.matches(restored).matched + assert restored.artifact_identity == a.artifact_identity + + +def test_deserialization_rejects_missing_artifact_identity(): + payload = identity_from(FakeModelConfig()).to_dict() + payload.pop("artifact_identity") + with pytest.raises(KeyError): + SourceIdentity.from_dict(payload) + + +def test_deserialization_rejects_unknown_format_version(): + payload = identity_from(FakeModelConfig()).to_dict() + payload["format_version"] += 1 + with pytest.raises(ValueError, match="Unsupported SourceIdentity format version"): + SourceIdentity.from_dict(payload) + + +def test_deserialization_rejects_v1_identity_without_artifact_binding(): + payload = identity_from(FakeModelConfig()).to_dict() + payload["format_version"] = 1 + payload.pop("artifact_identity") + with pytest.raises(ValueError, match="Unsupported SourceIdentity format version"): + SourceIdentity.from_dict(payload) def test_check_warn_fallback_on_mismatch(): @@ -219,6 +279,7 @@ def test_format_version_mismatch_never_matches(): if hasattr(copy, "replace") else SourceIdentity( format_version=a.format_version + 1, + artifact_identity=a.artifact_identity, model_fingerprint=a.model_fingerprint, quant_fingerprint=a.quant_fingerprint, backend_fingerprint=a.backend_fingerprint,