diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 3007b1402c4a..e573d2f8cb71 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -20,11 +20,30 @@ to the provided Hugging Face checkpoint. ## Current Support Scope -The post-transform MX receive path currently supports only -`LlamaForCausalLM` with transform protocol version 1. TensorRT LLM publishes -post-transform weights together with source-identity and layout metadata. A -receiver whose model family is not allow-listed does not consume those bytes; -it falls back to the standard Hugging Face checkpoint path. +The post-transform MX receive path currently supports one exact qualification +profile: + +| Profile | Root class | Config identity | Scope | Protocol | Transform-layout ABI | Constraints | +|---------|------------|-----------------|-------|----------|----------------------|-------------| +| `llama-for-causal-lm-target-v1` | `LlamaForCausalLM` | `LlamaForCausalLM` / `llama` | Target model | 1 | `trtllm-llama-target-layout-v1` | No speculative mode or separately loaded draft model | + +The registry matches the exact root class and the architecture/model type +captured from the resolved config before model construction. An unregistered +subclass or config alias does not inherit support. It falls back to the +standard Hugging Face checkpoint path before any P2P transfer starts. + +TensorRT LLM applies two independent compatibility gates: + +- The qualification profile records that a model/config/lifecycle combination + has passed full-load versus staged-load equivalence testing. +- `SourceIdentity` format version 3 binds two concrete runs to the same + checkpoint artifact, runtime layout choices, local shard layout, and + transform-layout ABI. + +The transfer protocol version identifies the staged receiver protocol. The +transform-layout ABI identifies the meaning of the transferred tensor names, +layouts, aliases, and receiver finalization. A pre-version-3 identity, a +missing ABI, or a different ABI is rejected rather than treated as compatible. Loads that require a separately loaded draft model also fall back to the standard checkpoint path. Target-plus-draft post-transform transfer remains @@ -42,16 +61,40 @@ Support for another model family requires a focused qualification change: 2. Verify that every one-time transform is guarded by `_weights_transformed` and that the staged receiver can skip `transform_weights()` without changing aliases, derived state, tensor layout, or outputs. -3. Add the model class and transform protocol version to the MX staged-receiver - allow-list only after full-load and staged-load equivalence tests pass. +3. Add an exact qualification profile only after the reusable harness in + `tests/unittest/utils/post_transform_qualification.py` proves tensor, + alias, transform-guard, derived-state, and deterministic output + equivalence. Include an unregistered-root negative control. 4. Cover compatible transfer, source-identity mismatch, unsupported layout or - protocol, and non-allow-listed fallback. Keep target-plus-draft loading - disabled unless that combination has its own mixed-layout tests. + protocol/ABI, no-disk staged reception, and unqualified-profile fallback. + Keep target-plus-draft loading disabled unless that combination has its own + mixed-layout tests. 5. Run a real ModelExpress donor/receiver test with the model configurations being claimed, including the supported quantization and TP/PP/EP layouts. Compare deterministic output token IDs with the standard Hugging Face load path before documenting the family as supported. +### Transform-Layout ABI Rules + +An existing transform-layout ABI ID is immutable. Introduce a new ID when a +change affects any transferred tensor name, shape, dtype, packing, sharding, +alias relationship, one-shot transform result, or receiver-side +`setup_aliases()`/`cache_derived_state()` interpretation. Keep the existing ID +for implementation-only changes that preserve all of those observable +semantics. + +When adding an ABI ID: + +1. Give the qualified profile the new ID and propagate it through + `SourceIdentity` and MX source metadata. +2. Add matching, missing, and mismatched producer/receiver compatibility + tests. ABI mismatches remain incompatible even under the `ENFORCE` identity + policy. +3. Re-run the qualification harness and the real donor/receiver GPU test for + every profile that adopts the ID. +4. Never reinterpret an already published ID. Supporting two ABIs requires an + explicit compatibility decision and tests for each producer/receiver pair. + ## Installation The official TensorRT LLM release container includes the MX Python client. No @@ -140,7 +183,7 @@ path. - Post-transform MX reception is currently limited to the Llama model family. Other model families safely fall back to Hugging Face loading until they are - explicitly qualified and added to the staged-receiver allow-list. + explicitly qualified and added as exact capability profiles. - The MX server and Redis lifecycle is external to TensorRT LLM. Every TensorRT LLM instance must be able to reach the configured MX server URL. - The MX server coordinates source discovery but does not store model weights. diff --git a/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py index 91dbe496976b..05db808f1d5e 100644 --- a/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py @@ -66,6 +66,7 @@ _MX_SOURCE_IDENTITY_METADATA_KEY = "trtllm_source_identity" _MX_WEIGHT_LAYOUT_METADATA_KEY = "trtllm_weight_layout" _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY = "trtllm_transform_protocol_version" +_MX_TRANSFORM_ABI_ID_METADATA_KEY = "trtllm_transform_abi_id" _MX_WEIGHT_LAYOUT_POST_TRANSFORM = "post_transform" _MX_STAGED_TRANSFORM_PROTOCOL_VERSION = 1 @@ -425,13 +426,24 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[ **kwargs, ) - layout_status = _metadata_weight_layout_status(source_metadata) + expected_transform_abi_id = ( + self._local_source_identity.transform_abi_id + if self._local_source_identity is not None + else None + ) + layout_status = _metadata_weight_layout_status( + source_metadata, + expected_transform_abi_id=expected_transform_abi_id, + ) if layout_status is _MxWeightLayoutStatus.UNSUPPORTED: self._source_identity_compatible_for_last_load = False return self._fallback_to_disk( checkpoint_dir, mapping, - reason=_metadata_unsupported_layout_reason(source_metadata), + reason=_metadata_unsupported_layout_reason( + source_metadata, + expected_transform_abi_id=expected_transform_abi_id, + ), **kwargs, ) @@ -721,6 +733,12 @@ def publish_as_source( "unavailable; receivers cannot safely verify transformed weights." ) return + if source_identity.transform_abi_id is None: + logger.warning( + "Skipping MX post-transform publish because SourceIdentity has " + "no qualified transform-layout ABI." + ) + return try: from modelexpress import ( @@ -884,6 +902,8 @@ def _build_mx_source_metadata(source_identity: Optional[SourceIdentity]) -> dict } if source_identity is not None: metadata[_MX_SOURCE_IDENTITY_METADATA_KEY] = _serialize_source_identity(source_identity) + if source_identity.transform_abi_id is not None: + metadata[_MX_TRANSFORM_ABI_ID_METADATA_KEY] = source_identity.transform_abi_id return metadata @@ -941,6 +961,7 @@ def _metadata_has_trtllm_key(metadata: dict[str, Any]) -> bool: _MX_SOURCE_IDENTITY_METADATA_KEY, _MX_WEIGHT_LAYOUT_METADATA_KEY, _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY, + _MX_TRANSFORM_ABI_ID_METADATA_KEY, ) ) @@ -970,13 +991,25 @@ def _source_identity_from_metadata(metadata: Optional[dict[str, Any]]) -> Option return None -def _metadata_is_post_transform(metadata: Optional[dict[str, Any]]) -> bool: +def _metadata_is_post_transform( + metadata: Optional[dict[str, Any]], + *, + expected_transform_abi_id: Optional[str], +) -> bool: return ( - _metadata_weight_layout_status(metadata) is _MxWeightLayoutStatus.POST_TRANSFORM_SUPPORTED + _metadata_weight_layout_status( + metadata, + expected_transform_abi_id=expected_transform_abi_id, + ) + is _MxWeightLayoutStatus.POST_TRANSFORM_SUPPORTED ) -def _metadata_weight_layout_status(metadata: Optional[dict[str, Any]]) -> _MxWeightLayoutStatus: +def _metadata_weight_layout_status( + metadata: Optional[dict[str, Any]], + *, + expected_transform_abi_id: Optional[str], +) -> _MxWeightLayoutStatus: layout = _metadata_get(metadata, _MX_WEIGHT_LAYOUT_METADATA_KEY) if layout is None: return _MxWeightLayoutStatus.PRE_TRANSFORM @@ -994,16 +1027,41 @@ def _metadata_weight_layout_status(metadata: Optional[dict[str, Any]]) -> _MxWei return _MxWeightLayoutStatus.UNSUPPORTED if protocol_version != _MX_STAGED_TRANSFORM_PROTOCOL_VERSION: return _MxWeightLayoutStatus.UNSUPPORTED + + source_transform_abi_id = _metadata_get(metadata, _MX_TRANSFORM_ABI_ID_METADATA_KEY) + if not isinstance(source_transform_abi_id, str) or not source_transform_abi_id: + return _MxWeightLayoutStatus.UNSUPPORTED + if expected_transform_abi_id is None or source_transform_abi_id != expected_transform_abi_id: + return _MxWeightLayoutStatus.UNSUPPORTED return _MxWeightLayoutStatus.POST_TRANSFORM_SUPPORTED -def _metadata_unsupported_layout_reason(metadata: Optional[dict[str, Any]]) -> str: +def _metadata_unsupported_layout_reason( + metadata: Optional[dict[str, Any]], + *, + expected_transform_abi_id: Optional[str], +) -> str: layout = _metadata_get(metadata, _MX_WEIGHT_LAYOUT_METADATA_KEY) if str(layout).lower() == _MX_WEIGHT_LAYOUT_POST_TRANSFORM: version = _metadata_get(metadata, _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY) + try: + protocol_version = int(version) + except (TypeError, ValueError): + protocol_version = None + if protocol_version != _MX_STAGED_TRANSFORM_PROTOCOL_VERSION: + return ( + "source publishes post-transform weights with unsupported " + f"transform protocol {version!r}" + ) + + source_transform_abi_id = _metadata_get(metadata, _MX_TRANSFORM_ABI_ID_METADATA_KEY) + if not isinstance(source_transform_abi_id, str) or not source_transform_abi_id: + return "source publishes post-transform weights without a transform-layout ABI" + if expected_transform_abi_id is None: + return "receiver has no qualified transform-layout ABI for post-transform weights" return ( - "source publishes post-transform weights with unsupported " - f"transform protocol {version!r}" + "source publishes post-transform weights with transform-layout ABI " + f"{source_transform_abi_id!r}; receiver requires {expected_transform_abi_id!r}" ) return f"source publishes unsupported MX weight layout {layout!r}" diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f3218f2e192..74fba60a06ae 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -13,10 +13,11 @@ from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import ( AutoCheckpointMapper, BaseCheckpointLoader) from tensorrt_llm._torch.weight_sharing import ( - ArtifactIdentity, IdentityCheckPolicy, PostTransformFeature, - PostTransformProfile, PostTransformProfileRegistry, - PostTransformQualificationDecision, PostTransformTransferScope, - SourceIdentity, check_weight_sharing_compatibility) + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity, IdentityCheckPolicy, + PostTransformConfigIdentity, 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, @@ -310,6 +311,7 @@ class ModelLoader: model_type="llama", speculative_mode=None, protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope=PostTransformTransferScope.TARGET_MODEL, ), )) @@ -439,6 +441,7 @@ def _build_source_identity( *, checkpoint_dir: str, model_name: str, + transform_abi_id: Optional[str], fallback_on_artifact_error: bool, ) -> Optional[SourceIdentity]: """Build the local identity without weakening artifact validation. @@ -463,6 +466,7 @@ def _build_source_identity( model, artifact_identity=artifact_identity, model_name=model_name, + transform_abi_id=transform_abi_id, ) def load( @@ -482,6 +486,12 @@ def load( """ config = self._load_and_validate_config(checkpoint_dir, checkpoint_loader) + # Some model constructors normalize or rewrite config fields. Capture + # the registry identity from the resolved input before construction so + # publication and reception qualify the architecture the user asked + # to load, not a post-construction alias. + post_transform_config_identity = PostTransformConfigIdentity.from_model_config( + config) load_format = self.llm_args.load_format with timing("Model init total"), maybe_create_moe_load_balancer( @@ -500,6 +510,16 @@ def load( model = AutoModelForCausalLM.from_config(config) is_meta_init = False + loads_draft_weights = ( + self.spec_config is not None + and self.spec_config.spec_dec_mode.need_load_draft_weights()) + speculative_mode = self._speculative_mode_name(self.spec_config) + post_transform_qualification = self._qualify_post_transform_profile( + model, + config_identity=post_transform_config_identity, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) + self._source_identity: Optional[SourceIdentity] = None if self._needs_source_identity(checkpoint_loader, load_format): # Receiver's local SourceIdentity, built once from the final @@ -515,6 +535,8 @@ def load( model_name=str( getattr(self.llm_args, "model", None) or checkpoint_dir), + transform_abi_id=( + post_transform_qualification.transform_abi_id), fallback_on_artifact_error=( load_format != LoadFormat.GMS and checkpoint_loader.checkpoint_format == "MX"), @@ -593,10 +615,6 @@ def init_meta_tensor(t: torch.Tensor): f"Use {rank_model_storage / (1024**3):.2f} GB for model weights." ) weights_preloaded = False - loads_draft_weights = ( - self.spec_config is not None - and self.spec_config.spec_dec_mode.need_load_draft_weights()) - speculative_mode = self._speculative_mode_name(self.spec_config) # Set when either GMS RW or GMS RO branch has already run the # post_load_* hooks itself, so the shared post-load block below # must skip them. RW handles them inside `mem_pool_scope` so the @@ -621,13 +639,9 @@ def init_meta_tensor(t: torch.Tensor): # do not accept post-transform bytes for only the target # model. Enable this only after target and draft subgraphs # have an explicit mixed-layout policy. - qualification = self._qualify_post_transform_profile( - model, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) load_weights_kwargs[ - "allow_post_transform_weights"] = qualification.qualified - if qualification.qualified: + "allow_post_transform_weights"] = post_transform_qualification.qualified + if post_transform_qualification.qualified: load_weights_kwargs[ "prepare_post_transform_receiver"] = self._setup_aliases @@ -748,13 +762,9 @@ def init_meta_tensor_in_pool(t: torch.Tensor): "source_identity": self._source_identity, } if checkpoint_loader.checkpoint_format == "MX": - qualification = self._qualify_post_transform_profile( - model, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) load_weights_kwargs[ - "allow_post_transform_weights"] = qualification.qualified - if qualification.qualified: + "allow_post_transform_weights"] = post_transform_qualification.qualified + if post_transform_qualification.qualified: load_weights_kwargs[ "prepare_post_transform_receiver"] = self._setup_aliases weights = checkpoint_loader.load_weights( @@ -820,8 +830,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor): checkpoint_loader, model, weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + qualification=post_transform_qualification) if mx_staged_receiver_path: self._setup_aliases(model) self._mark_weights_transformed(model) @@ -847,8 +856,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor): model, checkpoint_dir=checkpoint_dir, weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + qualification=post_transform_qualification) # Pool closed. Commit the post-post_load layout. gms_backend.finalize_write(model) @@ -901,8 +909,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor): model, checkpoint_dir=checkpoint_dir, weights_preloaded=True, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + qualification=post_transform_qualification) gms_post_load_handled = True logger.info("LoadFormat.GMS (RO): materialized weights") else: @@ -941,20 +948,19 @@ def init_meta_tensor_in_pool(t: torch.Tensor): checkpoint_loader, model, weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + qualification=post_transform_qualification) if mx_staged_receiver_path: self._setup_aliases(model) self._mark_weights_transformed(model) self._walk_cache_state(model) else: self._walk_full_post_load(model) - self._post_load_publish(checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + self._post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + qualification=post_transform_qualification) # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, # `register_weight_slots_after_to_cuda` and `finalize_model` @@ -1009,13 +1015,9 @@ def _check_gms_source_identity(self, gms_backend) -> None: @classmethod def _should_run_mx_staged_receiver_path( - cls, - checkpoint_loader: BaseCheckpointLoader, - model: DecoderModelForCausalLM, - *, - weights_preloaded: bool, - speculative_mode: Optional[str] = None, - loads_draft_weights: bool = False) -> bool: + cls, checkpoint_loader: BaseCheckpointLoader, + model: DecoderModelForCausalLM, *, weights_preloaded: bool, + qualification: PostTransformQualificationDecision) -> bool: """Whether an MX receiver can skip one-shot weight transforms. MXCheckpointLoader only accepts post-transform P2P bytes when this same @@ -1033,18 +1035,15 @@ def _should_run_mx_staged_receiver_path( ): return False - qualification = cls._qualify_post_transform_profile( - model, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) profile = qualification.profile if qualification.qualified and profile is not None: logger.info( "MX receiver using staged post-load profile %s for %s " - "(transform protocol v%d).", + "(transform protocol v%d, layout ABI %s).", profile.profile_id, type(model).__name__, cls._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + profile.transform_abi_id, ) return True @@ -1068,49 +1067,49 @@ def _speculative_mode_name( return None spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) mode_name = getattr(spec_dec_mode, "name", None) - return mode_name.lower() if isinstance(mode_name, str) else "unknown" + if not isinstance(mode_name, str): + logger.warning( + "Unable to identify the speculative decoding mode from %s; " + "post-transform sharing is disabled for this load.", + type(spec_dec_mode).__name__, + ) + return "unknown" + return mode_name.lower() @classmethod def _qualify_post_transform_profile( - cls, model: DecoderModelForCausalLM, *, + cls, + model: DecoderModelForCausalLM, + *, + config_identity: Optional[PostTransformConfigIdentity] = None, speculative_mode: Optional[str], loads_draft_weights: bool) -> PostTransformQualificationDecision: - pretrained_config = model.model_config.pretrained_config - architectures = getattr(pretrained_config, "architectures", None) - architecture = (architectures[0] - if isinstance(architectures, - (list, tuple)) and architectures - and isinstance(architectures[0], str) else None) - configured_model_type = getattr(pretrained_config, "model_type", None) - model_type = (configured_model_type if isinstance( - configured_model_type, str) else None) + if config_identity is None: + config_identity = PostTransformConfigIdentity.from_model_config( + model.model_config) enabled_features = set() if loads_draft_weights: enabled_features.add(PostTransformFeature.SEPARATE_DRAFT_MODEL) return cls._POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=type(model), - architecture=architecture, - model_type=model_type, + architecture=config_identity.architecture, + model_type=config_identity.model_type, speculative_mode=speculative_mode, protocol_version=cls._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, enabled_features=frozenset(enabled_features), ) - def _post_load_publish(self, checkpoint_loader: BaseCheckpointLoader, - model: DecoderModelForCausalLM, *, - checkpoint_dir: str, weights_preloaded: bool, - speculative_mode: Optional[str], - loads_draft_weights: bool) -> None: + def _post_load_publish( + self, checkpoint_loader: BaseCheckpointLoader, + model: DecoderModelForCausalLM, *, checkpoint_dir: str, + weights_preloaded: bool, + qualification: PostTransformQualificationDecision) -> None: kwargs = { "checkpoint_dir": checkpoint_dir, "weights_preloaded": weights_preloaded, } if checkpoint_loader.checkpoint_format == "MX": - qualification = self._qualify_post_transform_profile( - model, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) if not qualification.qualified: if not weights_preloaded: logger.info( diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index 38ec2471c41c..e48de4784b31 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -19,6 +19,8 @@ ArtifactIdentity, ) from tensorrt_llm._torch.weight_sharing.post_transform_profiles import ( + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + PostTransformConfigIdentity, PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry, @@ -39,7 +41,9 @@ __all__ = [ "ARTIFACT_IDENTITY_FORMAT_VERSION", "ArtifactIdentity", + "LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1", "SOURCE_IDENTITY_FORMAT_VERSION", + "PostTransformConfigIdentity", "PostTransformFeature", "PostTransformProfile", "PostTransformProfileRegistry", diff --git a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py index e6b4c82b3d2c..10df784a63d1 100644 --- a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py +++ b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py @@ -34,6 +34,39 @@ from torch import nn +# Stable identifier for the tensor names, layouts, aliases, and receiver-side +# finalization contract produced by the currently qualified Llama target path. +# Never change the meaning of an existing ABI ID. Introduce a new ID whenever +# a transform changes transferred tensor semantics or the receiver must +# interpret/finalize the transferred state differently. +LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-llama-target-layout-v1" + + +@dataclass(frozen=True) +class PostTransformConfigIdentity: + """Canonical model identity captured before constructor normalization.""" + + architecture: str | None + model_type: str | None + + @classmethod + def from_model_config(cls, model_config: object) -> "PostTransformConfigIdentity": + """Capture registry dimensions from a resolved pre-construction config.""" + + pretrained_config = getattr(model_config, "pretrained_config", None) + architectures = getattr(pretrained_config, "architectures", None) + architecture = ( + architectures[0] + if isinstance(architectures, (list, tuple)) + and architectures + and isinstance(architectures[0], str) + else None + ) + configured_model_type = getattr(pretrained_config, "model_type", None) + model_type = configured_model_type if isinstance(configured_model_type, str) else None + return cls(architecture=architecture, model_type=model_type) + + class PostTransformTransferScope(str, Enum): """The model component represented by a post-transform transfer.""" @@ -72,6 +105,7 @@ class PostTransformProfile: speculative_mode: str | None protocol_version: int transfer_scope: PostTransformTransferScope + transform_abi_id: str supported_features: frozenset[PostTransformFeature] = field(default_factory=frozenset) def __post_init__(self) -> None: @@ -85,6 +119,8 @@ def __post_init__(self) -> None: raise ValueError("Post-transform speculative_mode must not be empty") if self.protocol_version < 1: raise ValueError("Post-transform protocol_version must be positive") + if not isinstance(self.transform_abi_id, str) or not self.transform_abi_id: + raise ValueError("Post-transform transform_abi_id must be a non-empty string") object.__setattr__(self, "supported_features", frozenset(self.supported_features)) @@ -116,6 +152,14 @@ def qualified(self) -> bool: self.reason is PostTransformQualificationReason.QUALIFIED and self.profile is not None ) + @property + def transform_abi_id(self) -> str | None: + """The qualified layout ABI, or `None` for a rejected profile.""" + + return ( + self.profile.transform_abi_id if self.qualified and self.profile is not None else None + ) + @dataclass(frozen=True) class PostTransformProfileRegistry: diff --git a/tensorrt_llm/_torch/weight_sharing/source_identity.py b/tensorrt_llm/_torch/weight_sharing/source_identity.py index c6b010634760..4fa8305d5f92 100644 --- a/tensorrt_llm/_torch/weight_sharing/source_identity.py +++ b/tensorrt_llm/_torch/weight_sharing/source_identity.py @@ -42,15 +42,16 @@ -------------- The fingerprint is split so comparison can be selective: -* **global fingerprint** -- immutable checkpoint artifact, rank-invariant - model identity, quantization, backend selection, fusion flags, and parallel - *sizes* (TP/PP/EP/CP). +* **global fingerprint** -- immutable checkpoint artifact, transform-layout + ABI, 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`. A caller that wants *enforced* sharing across otherwise-divergent runs can skip -the global comparison (`compare_global=False`) and trust the source. +the global comparison (`compare_global=False`) and trust the source. Identity +format and transform-layout ABI compatibility remain mandatory. Adding fields ------------- @@ -82,7 +83,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 = 2 +SOURCE_IDENTITY_FORMAT_VERSION = 3 _PRETRAINED_METADATA_FIELDS = frozenset( { @@ -198,10 +199,10 @@ class IdentityCheckPolicy(Enum): * `WARN_FALLBACK` (default): log a warning and fall back to non-shared loading. Never raises. * `STRICT`: raise :class:`SourceIdentityMismatchError` on mismatch. - * `ENFORCE`: always share regardless of concrete-identity mismatch (the + * `ENFORCE`: share despite concrete artifact/config/shard differences (the caller explicitly trusts the source, e.g. enforced cross-run sharing). - Still requires both local and source identities to be present. Logs at - debug. + Identity format and transform-layout ABI must still match, and both + identities must be present. """ WARN_FALLBACK = "warn_fallback" @@ -242,6 +243,17 @@ class SourceIdentity: pp_size: int = 1 ep_size: int = -1 dtype: Optional[str] = None + # Layout semantics are safety-critical and always compared, including + # under ENFORCE. `None` denotes a source that does not use a qualified + # post-transform layout contract. Keep this last to preserve existing + # positional construction of the discovery fields above. + transform_abi_id: Optional[str] = None + + def __post_init__(self) -> None: + if self.transform_abi_id is not None and ( + not isinstance(self.transform_abi_id, str) or not self.transform_abi_id + ): + raise ValueError("SourceIdentity transform_abi_id must be a non-empty string") # ---- construction -------------------------------------------------- @@ -254,6 +266,7 @@ def from_model_config( checkpoint_dir: Optional[str] = None, artifact_identity: Optional[ArtifactIdentity] = None, model_name: Optional[str] = None, + transform_abi_id: Optional[str] = None, ) -> "SourceIdentity": """Build an identity from a torch-backend :class:`ModelConfig`. @@ -274,6 +287,9 @@ def from_model_config( model_name: Human-readable model identity used by discovery layers (e.g. the MX server's source catalog). Does not affect the compatibility fingerprints. + transform_abi_id: Stable identifier for the post-transform tensor + layout and receiver-finalization contract. `None` when the + source does not use a qualified post-transform layout. Returns: A fully populated :class:`SourceIdentity` for @@ -305,6 +321,7 @@ def from_model_config( parallel_fingerprint=cls._build_parallel_fingerprint(mapping), rank=rank, shard_fingerprint=cls._build_shard_fingerprint(mapping, model), + transform_abi_id=transform_abi_id, model_name=model_name, tp_size=getattr(mapping, "tp_size", 1), pp_size=getattr(mapping, "pp_size", 1), @@ -443,6 +460,7 @@ def global_fingerprint(self) -> str: { "format_version": self.format_version, "artifact": self.artifact_identity.to_dict(), + "transform_abi_id": self.transform_abi_id, "model": self.model_fingerprint, "quant": self.quant_fingerprint, "backend": self.backend_fingerprint, @@ -471,6 +489,11 @@ def matches( if self.format_version != other.format_version: mismatched.append("format_version") + # A transform ABI mismatch changes the meaning of transferred tensors, + # so no identity policy may bypass it. + if self.transform_abi_id != other.transform_abi_id: + mismatched.append("transform_abi_id") + if compare_global: if self.artifact_identity != other.artifact_identity: mismatched.append("artifact_identity") @@ -507,6 +530,7 @@ def to_dict(self) -> dict: "parallel_fingerprint": self.parallel_fingerprint, "rank": self.rank, "shard_fingerprint": self.shard_fingerprint, + "transform_abi_id": self.transform_abi_id, "model_name": self.model_name, "tp_size": self.tp_size, "pp_size": self.pp_size, @@ -537,6 +561,7 @@ def from_dict(cls, data: dict) -> "SourceIdentity": parallel_fingerprint=data["parallel_fingerprint"], rank=data["rank"], shard_fingerprint=data["shard_fingerprint"], + transform_abi_id=data["transform_abi_id"], model_name=data.get("model_name"), tp_size=data.get("tp_size", 1), pp_size=data.get("pp_size", 1), @@ -596,10 +621,15 @@ def check_weight_sharing_compatibility( if policy is IdentityCheckPolicy.ENFORCE: result = local.matches(source, compare_global=False, compare_shard=False) if not result.matched: - logger.debug( - f"SourceIdentity ENFORCE: sharing despite mismatch in {result.mismatched_fields}." + logger.warning( + "SourceIdentity ENFORCE cannot bypass identity format or " + f"transform-layout ABI mismatch in {result.mismatched_fields}." ) - return IdentityCheckDecision(should_share=True, match_result=result, policy=policy) + return IdentityCheckDecision( + should_share=result.matched, + match_result=result, + policy=policy, + ) result = local.matches(source, compare_global=compare_global, compare_shard=compare_shard) if result.matched: diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index b56935da6378..5b20f4f81653 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -2,7 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Unit tests for GMS-specific branches in ``ModelLoader``.""" +from collections.abc import Callable from contextlib import contextmanager, nullcontext +from dataclasses import replace from types import SimpleNamespace from unittest.mock import MagicMock @@ -15,6 +17,7 @@ from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( ARTIFACT_IDENTITY_FORMAT_VERSION, + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, SOURCE_IDENTITY_FORMAT_VERSION, ArtifactIdentity, PostTransformProfile, @@ -100,7 +103,14 @@ def _make_loader(monkeypatch, *, events, spec_config=None): side_effect=lambda fn, weights, mapper, **kwargs: fn(weights, mapper) ) loader._load_and_validate_config = MagicMock( - return_value=SimpleNamespace(name="config", mapping=SimpleNamespace()) + return_value=SimpleNamespace( + name="config", + mapping=SimpleNamespace(), + pretrained_config=SimpleNamespace( + architectures=["TinyForCausalLM"], + model_type="tiny", + ), + ) ) monkeypatch.setattr(model_loader_mod, "timing", lambda *_args, **_kwargs: nullcontext()) @@ -121,7 +131,10 @@ def _build_artifact_identity(_cls, checkpoint_dir): def _build_source_identity(_cls, *_args, **kwargs): assert kwargs["artifact_identity"] is _SOURCE_IDENTITY.artifact_identity - return _SOURCE_IDENTITY + return replace( + _SOURCE_IDENTITY, + transform_abi_id=kwargs["transform_abi_id"], + ) monkeypatch.setattr( model_loader_mod.SourceIdentity, @@ -173,8 +186,13 @@ def __init__(self) -> None: self.post_load_publish = MagicMock() @staticmethod - def _load_weights(*_args, **kwargs): - kwargs["prepare_post_transform_receiver"](kwargs["model"]) + def _load_weights( + *_args: object, + prepare_post_transform_receiver: Callable[[nn.Module], None], + model: nn.Module, + **_kwargs: object, + ) -> dict[str, torch.Tensor]: + prepare_post_transform_receiver(model) return {} def is_post_transform_weights_preloaded(self) -> bool: @@ -198,6 +216,7 @@ def _tiny_profile_registry() -> PostTransformProfileRegistry: model_type="tiny", speculative_mode=None, protocol_version=(ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION), + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope=PostTransformTransferScope.TARGET_MODEL, ), ) @@ -482,6 +501,7 @@ def test_gms_rw_mx_post_transform_preload_uses_staged_path(monkeypatch): _args, kwargs = checkpoint_loader.load_weights.call_args assert kwargs["allow_post_transform_weights"] is True assert callable(kwargs["prepare_post_transform_receiver"]) + assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 loader._call_load_weights.assert_not_called() checkpoint_loader.post_load_publish.assert_called_once_with( model, diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index ccdf960ea786..907d754cfae5 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -2,8 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 """Unit tests for MX-specific ModelLoader branches.""" +from collections.abc import Callable from contextlib import contextmanager, nullcontext +from dataclasses import replace +from pathlib import Path from types import SimpleNamespace +from typing import cast from unittest.mock import MagicMock import pytest @@ -25,6 +29,7 @@ from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( ARTIFACT_IDENTITY_FORMAT_VERSION, + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, SOURCE_IDENTITY_FORMAT_VERSION, ArtifactIdentity, PostTransformFeature, @@ -119,7 +124,15 @@ def _moe_context(config, mapping): yield None -def _tiny_llama_model(monkeypatch): +class _UnqualifiedLlamaForCausalLM(model_loader_mod.LlamaForCausalLM): + pass + + +def _tiny_llama_model( + monkeypatch: pytest.MonkeyPatch, + *, + model_class: type[nn.Module] = model_loader_mod.LlamaForCausalLM, +) -> nn.Module: monkeypatch.setattr(modeling_llama_mod, "get_sm_version", lambda: 90) llama_config = LlamaConfig( architectures=["LlamaForCausalLM"], @@ -137,13 +150,22 @@ def _tiny_llama_model(monkeypatch): torch_dtype=torch.float32, vocab_size=32, ) - return model_loader_mod.LlamaForCausalLM( + model = model_class( ModelConfig( pretrained_config=llama_config, max_num_tokens=16, max_seq_len=16, ) ) + with torch.no_grad(): + for index, parameter in enumerate(model.parameters()): + values = torch.arange( + parameter.numel(), + dtype=torch.float32, + device=parameter.device, + ).reshape(parameter.shape) + parameter.copy_(((values + index) % 17).to(parameter.dtype) / 17) + return model def _llama_alias_state(model): @@ -158,6 +180,15 @@ def _llama_alias_state(model): } +def _llama_embedding_logits(model: nn.Module) -> torch.Tensor: + input_ids = torch.tensor( + [0, 1, 2], + dtype=torch.long, + device=model.model.embed_tokens.weight.device, + ) + return model.lm_head(model.model.embed_tokens(input_ids)) + + def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransformProfileRegistry: return PostTransformProfileRegistry( profiles=( @@ -168,6 +199,7 @@ def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransf model_type="tiny", speculative_mode=speculative_mode, protocol_version=(ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION), + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope=PostTransformTransferScope.TARGET_MODEL, ), ) @@ -188,7 +220,14 @@ def _make_loader(monkeypatch, *, events, spec_config=None): side_effect=lambda fn, weights, mapper, **kwargs: fn(weights, mapper) ) loader._load_and_validate_config = MagicMock( - return_value=SimpleNamespace(name="config", mapping=SimpleNamespace()) + return_value=SimpleNamespace( + name="config", + mapping=SimpleNamespace(), + pretrained_config=SimpleNamespace( + architectures=["TinyForCausalLM"], + model_type="tiny", + ), + ) ) monkeypatch.setattr(model_loader_mod, "timing", lambda *_args, **_kwargs: nullcontext()) @@ -209,7 +248,10 @@ def _build_artifact_identity(_cls, checkpoint_dir): def _build_source_identity(_cls, *_args, **kwargs): assert kwargs["artifact_identity"] is _SOURCE_IDENTITY.artifact_identity - return _SOURCE_IDENTITY + return replace( + _SOURCE_IDENTITY, + transform_abi_id=kwargs["transform_abi_id"], + ) monkeypatch.setattr( model_loader_mod.SourceIdentity, @@ -250,9 +292,44 @@ def test_construct_checkpoint_loader_passes_mx_config(): assert checkpoint_loader.model_name == "Qwen/Qwen2.5-7B-Instruct" -def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works(monkeypatch): +def test_public_support_table_matches_qualified_profile_registry() -> None: + profiles = ModelLoader._POST_TRANSFORM_PROFILE_REGISTRY.profiles + documentation = (Path(__file__).parents[4] / "docs/source/features/model-express.md").read_text( + encoding="utf-8" + ) + lines = documentation.splitlines() + table_header = ( + "| Profile | Root class | Config identity | Scope | Protocol | " + "Transform-layout ABI | Constraints |" + ) + table_header_index = lines.index(table_header) + table_rows = [] + for line in lines[table_header_index + 2 :]: + if not line.startswith("|"): + break + table_rows.append(line) + + assert len(table_rows) == len(profiles) + for profile in profiles: + scope = profile.transfer_scope.value.replace("_", " ").capitalize() + expected_row_prefix = ( + f"| `{profile.profile_id}` | `{profile.root_model_class.__name__}` | " + f"`{profile.architecture}` / `{profile.model_type}` | {scope} | " + f"{profile.protocol_version} | `{profile.transform_abi_id}` |" + ) + assert any(row.startswith(expected_row_prefix) for row in table_rows) + + +def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works( + monkeypatch: pytest.MonkeyPatch, +) -> None: events = [] loader = _make_loader(monkeypatch, events=events) + monkeypatch.setattr( + ModelLoader, + "_POST_TRANSFORM_PROFILE_REGISTRY", + _tiny_profile_registry(), + ) checkpoint_loader = MagicMock(name="checkpoint_loader") checkpoint_loader.checkpoint_format = "MX" checkpoint_loader.is_weights_preloaded.return_value = True @@ -265,11 +342,17 @@ def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works(mon assert kwargs["mapping"] is loader.mapping assert kwargs["model"] is model assert kwargs["source_identity"] is loader._source_identity - assert kwargs["allow_post_transform_weights"] is False + assert kwargs["allow_post_transform_weights"] is True + assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 assert loader._call_load_weights.call_count == 0 checkpoint_loader.get_initialized_weight_mapper.assert_called_once() assert loader.weight_mapper is checkpoint_loader.get_initialized_weight_mapper.return_value - checkpoint_loader.post_load_publish.assert_not_called() + checkpoint_loader.post_load_publish.assert_called_once_with( + model, + checkpoint_dir="/ckpt", + weights_preloaded=True, + source_identity=loader._source_identity, + ) # reload() uses self.weight_mapper unconditionally; MX success must # initialize it even though the initial load skipped _call_load_weights. @@ -299,9 +382,16 @@ def test_reload_partial_loading_preserves_weights_transformed_flags(monkeypatch) assert events == ["load_weights"] -def test_mx_partial_fallback_merges_returned_weights(monkeypatch): +def test_mx_partial_fallback_merges_returned_weights( + monkeypatch: pytest.MonkeyPatch, +) -> None: events = [] loader = _make_loader(monkeypatch, events=events) + monkeypatch.setattr( + ModelLoader, + "_POST_TRANSFORM_PROFILE_REGISTRY", + _tiny_profile_registry(), + ) checkpoint_loader = MagicMock(name="checkpoint_loader") checkpoint_loader.checkpoint_format = "MX" checkpoint_loader.is_weights_preloaded.return_value = True @@ -315,7 +405,12 @@ def test_mx_partial_fallback_merges_returned_weights(monkeypatch): assert load_fn == model.load_weights assert weights is fallback_weights assert mapper is loader.weight_mapper - checkpoint_loader.post_load_publish.assert_not_called() + checkpoint_loader.post_load_publish.assert_called_once_with( + model, + checkpoint_dir="/ckpt", + weights_preloaded=True, + source_identity=loader._source_identity, + ) class _PostTransformMxLoader: @@ -331,13 +426,16 @@ def __init__(self, *, post_transform: bool) -> None: self.post_load_apply = MagicMock() self.post_load_publish = MagicMock() - def _load_weights(self, *_args, **kwargs): + def _load_weights(self, *_args: object, **kwargs: object) -> dict[str, object]: if self._post_transform and kwargs.get("allow_post_transform_weights") is False: self._post_transform = False self._weights_preloaded = False return {"disk.weight": self._disk_weight} if self._post_transform: - kwargs["prepare_post_transform_receiver"](kwargs["model"]) + prepare_receiver = cast( + Callable[[nn.Module], None], kwargs["prepare_post_transform_receiver"] + ) + prepare_receiver(cast(nn.Module, kwargs["model"])) return {} def is_post_transform_weights_preloaded(self) -> bool: @@ -345,11 +443,13 @@ def is_post_transform_weights_preloaded(self) -> bool: class _UnsafePostTransformMxLoader(_PostTransformMxLoader): - def _load_weights(self, *_args, **_kwargs): + def _load_weights(self, *_args: object, **_kwargs: object) -> dict[str, object]: return {} -def test_mx_post_transform_receiver_uses_staged_path_when_qualified(monkeypatch): +def test_mx_post_transform_receiver_uses_staged_path_when_qualified( + monkeypatch: pytest.MonkeyPatch, +) -> None: events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( @@ -380,22 +480,31 @@ def test_mx_post_transform_receiver_uses_staged_path_when_qualified(monkeypatch) assert events == ["setup_aliases", "setup_aliases", "cache_derived_state"] -def test_default_profile_qualifies_real_tiny_llama_lifecycle(monkeypatch): +def test_default_profile_qualifies_real_tiny_llama_lifecycle( + monkeypatch: pytest.MonkeyPatch, +) -> None: case = PostTransformQualificationCase( profile_id="llama-for-causal-lm-target-v1", model_factory=lambda: _tiny_llama_model(monkeypatch), + unqualified_model_factory=lambda: _tiny_llama_model( + monkeypatch, + model_class=_UnqualifiedLlamaForCausalLM, + ), qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, ), state_probes=(("aliases", _llama_alias_state),), + output_probes=(("embedding-logits", _llama_embedding_logits),), ) assert_post_transform_lifecycle_equivalent(case) -def test_separate_draft_model_is_not_qualified_by_target_only_profile(monkeypatch): +def test_separate_draft_model_is_not_qualified_by_target_only_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr( ModelLoader, "_POST_TRANSFORM_PROFILE_REGISTRY", @@ -413,7 +522,9 @@ def test_separate_draft_model_is_not_qualified_by_target_only_profile(monkeypatc assert decision.unsupported_features == frozenset({PostTransformFeature.SEPARATE_DRAFT_MODEL}) -def test_one_engine_speculative_mode_is_not_qualified_by_target_only_profile(monkeypatch): +def test_one_engine_speculative_mode_is_not_qualified_by_target_only_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr( ModelLoader, "_POST_TRANSFORM_PROFILE_REGISTRY", @@ -431,21 +542,35 @@ def test_one_engine_speculative_mode_is_not_qualified_by_target_only_profile(mon assert decision.unsupported_features == frozenset() -def test_speculative_mode_name_is_canonical_and_fails_closed() -> None: +def test_speculative_mode_name_is_canonical_and_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + warning = MagicMock() + monkeypatch.setattr(model_loader_mod.logger, "warning", warning) + assert ModelLoader._speculative_mode_name(None) is None + warning.assert_not_called() assert ( ModelLoader._speculative_mode_name( SimpleNamespace(spec_dec_mode=SimpleNamespace(name="MTP")) ) == "mtp" ) + warning.assert_not_called() assert ( ModelLoader._speculative_mode_name(SimpleNamespace(spec_dec_mode=SimpleNamespace())) == "unknown" ) + warning.assert_called_once_with( + "Unable to identify the speculative decoding mode from %s; " + "post-transform sharing is disabled for this load.", + "SimpleNamespace", + ) -def test_mx_post_transform_receiver_falls_back_for_unqualified_model(monkeypatch): +def test_mx_post_transform_receiver_falls_back_for_unqualified_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: events = [] loader = _make_loader(monkeypatch, events=events) checkpoint_loader = _PostTransformMxLoader(post_transform=True) @@ -461,11 +586,46 @@ def test_mx_post_transform_receiver_falls_back_for_unqualified_model(monkeypatch assert load_fn == model.load_weights assert weights == {"disk.weight": checkpoint_loader._disk_weight} assert mapper is loader.weight_mapper + assert loader._source_identity.transform_abi_id is None assert events == ["load_weights", "post_load_weights"] checkpoint_loader.post_load_publish.assert_not_called() -def test_mx_rejects_post_transform_preload_after_failed_qualification(monkeypatch): +def test_load_qualifies_with_preconstruction_identity_after_model_normalization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[str] = [] + loader = _make_loader(monkeypatch, events=events) + registry = MagicMock(wraps=_tiny_profile_registry()) + monkeypatch.setattr( + ModelLoader, + "_POST_TRANSFORM_PROFILE_REGISTRY", + registry, + ) + normalized_model = _TinyModel(events) + normalized_model.model_config.pretrained_config.architectures = ["NormalizedForCausalLM"] + normalized_model.model_config.pretrained_config.model_type = "normalized" + monkeypatch.setattr( + model_loader_mod.AutoModelForCausalLM, + "from_config", + MagicMock(return_value=normalized_model), + ) + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} + checkpoint_loader.is_weights_preloaded.return_value = False + + loader.load("/ckpt", checkpoint_loader) + + _args, kwargs = checkpoint_loader.load_weights.call_args + registry.qualify.assert_called_once() + assert kwargs["allow_post_transform_weights"] is True + assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + + +def test_mx_rejects_post_transform_preload_after_failed_qualification( + monkeypatch: pytest.MonkeyPatch, +) -> None: loader = _make_loader(monkeypatch, events=[]) checkpoint_loader = _UnsafePostTransformMxLoader(post_transform=True) @@ -480,7 +640,9 @@ def test_mx_rejects_post_transform_preload_after_failed_qualification(monkeypatc checkpoint_loader.post_load_publish.assert_not_called() -def test_mx_fallback_runs_standard_weight_mapping(monkeypatch): +def test_mx_fallback_runs_standard_weight_mapping( + monkeypatch: pytest.MonkeyPatch, +) -> None: events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( 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 e80026d31114..06e7a85091a8 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 @@ -33,6 +33,7 @@ from tensorrt_llm._torch.models.checkpoints.mx.checkpoint_loader import ( _MX_SOURCE_IDENTITY_METADATA_KEY, _MX_STAGED_TRANSFORM_PROTOCOL_VERSION, + _MX_TRANSFORM_ABI_ID_METADATA_KEY, _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY, _MX_WEIGHT_LAYOUT_METADATA_KEY, _MX_WEIGHT_LAYOUT_POST_TRANSFORM, @@ -44,6 +45,7 @@ ) from tensorrt_llm._torch.weight_sharing import ( ARTIFACT_IDENTITY_FORMAT_VERSION, + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, SOURCE_IDENTITY_FORMAT_VERSION, ArtifactIdentity, SourceIdentity, @@ -52,7 +54,12 @@ _MISSING = object() -def _identity(rank: int = 0, suffix: str = "same") -> SourceIdentity: +def _identity( + rank: int = 0, + suffix: str = "same", + *, + transform_abi_id: str | None = LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, +) -> SourceIdentity: return SourceIdentity( format_version=SOURCE_IDENTITY_FORMAT_VERSION, artifact_identity=ArtifactIdentity( @@ -66,6 +73,7 @@ def _identity(rank: int = 0, suffix: str = "same") -> SourceIdentity: parallel_fingerprint=f"parallel-{suffix}", rank=rank, shard_fingerprint=f"shard-{rank}-{suffix}", + transform_abi_id=transform_abi_id, model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0", ) @@ -352,7 +360,7 @@ def test_mixed_success_returns_fallback_weights(self): assert result is fallback mock_super_load.assert_not_called() - def test_post_transform_full_success_prepares_receiver_before_p2p(self): + def test_post_transform_full_success_prepares_receiver_before_p2p(self) -> None: identity = _identity() loader = MXCheckpointLoader(mx_server_url="http://mx:8001") model = MagicMock(name="model") @@ -379,7 +387,9 @@ def test_post_transform_full_success_prepares_receiver_before_p2p(self): prepare_receiver.assert_called_once_with(model) assert events == ["prepare_receiver", "p2p"] - def test_post_transform_source_without_receiver_preparer_falls_back_before_p2p(self): + def test_post_transform_source_without_receiver_preparer_falls_back_before_p2p( + self, + ) -> None: identity = _identity() loader = MXCheckpointLoader(mx_server_url="http://mx:8001") disk_weights = {"disk.weight": MagicMock()} @@ -409,7 +419,9 @@ def test_post_transform_source_without_receiver_preparer_falls_back_before_p2p(s mx_loader.load_weights.assert_not_called() mock_super_load.assert_called_once() - def test_post_transform_source_falls_back_before_p2p_when_not_allowlisted(self): + def test_post_transform_source_falls_back_before_p2p_when_profile_is_not_qualified( + self, + ) -> None: identity = _identity() loader = MXCheckpointLoader(mx_server_url="http://mx:8001") disk_weights = {"disk.weight": MagicMock()} @@ -449,7 +461,7 @@ def test_post_transform_source_falls_back_before_p2p_when_not_allowlisted(self): ) def test_post_transform_source_with_unsupported_protocol_falls_back_before_p2p( self, protocol_value - ): + ) -> None: identity = _identity() loader = MXCheckpointLoader(mx_server_url="http://mx:8001") disk_weights = {"disk.weight": MagicMock()} @@ -487,7 +499,52 @@ def test_post_transform_source_with_unsupported_protocol_falls_back_before_p2p( prepare_receiver.assert_not_called() mock_super_load.assert_called_once() - def test_selects_matching_source_metadata_from_multiple_instances(self): + @pytest.mark.parametrize( + "transform_abi_id", + ["trtllm-llama-target-layout-v2", _MISSING], + ids=["mismatched-abi", "missing-abi"], + ) + def test_post_transform_source_with_unsupported_abi_falls_back_before_p2p( + self, transform_abi_id: object + ) -> None: + identity = _identity() + loader = MXCheckpointLoader(mx_server_url="http://mx:8001") + disk_weights = {"disk.weight": MagicMock()} + prepare_receiver = MagicMock() + source_instance = _source_instance(identity) + if transform_abi_id is _MISSING: + source_instance.metadata.pop(_MX_TRANSFORM_ABI_ID_METADATA_KEY) + else: + source_instance.metadata[_MX_TRANSFORM_ABI_ID_METADATA_KEY] = transform_abi_id + fake_mx = _build_fake_modelexpress( + load_weights_return={}, + source_instances=[source_instance], + ) + + with ( + _install_fake_modelexpress(fake_mx), + patch.object( + HfCheckpointLoader, "load_weights", return_value=disk_weights + ) as mock_super_load, + ): + result = loader.load_weights( + "/nonexistent", + mapping=MagicMock(), + model=MagicMock(), + source_identity=identity, + allow_post_transform_weights=True, + prepare_post_transform_receiver=prepare_receiver, + ) + + assert result is disk_weights + assert loader.is_weights_preloaded() is False + assert loader.is_post_transform_weights_preloaded() is False + mx_loader = fake_mx.trtllm_live_transfer.MxLiveWeightLoader.return_value + mx_loader.load_weights.assert_not_called() + prepare_receiver.assert_not_called() + mock_super_load.assert_called_once() + + def test_selects_matching_source_metadata_from_multiple_instances(self) -> None: rank0_identity = _identity(rank=0) rank1_identity = _identity(rank=1) loader = MXCheckpointLoader(mx_server_url="http://mx:8001") @@ -628,6 +685,7 @@ def test_publish_called_with_model(self): assert metadata[_MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY] == str( _MX_STAGED_TRANSFORM_PROTOCOL_VERSION ) + assert metadata[_MX_TRANSFORM_ABI_ID_METADATA_KEY] == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 assert _MX_SOURCE_IDENTITY_METADATA_KEY in metadata def test_publish_synchronizes_cuda_before_exposing_source(self, monkeypatch): @@ -657,6 +715,18 @@ def test_source_identity_required_for_post_transform_publish(self): fake_mx.trtllm_live_transfer.publish_model_params.assert_not_called() + def test_transform_abi_required_for_post_transform_publish(self): + loader = MXCheckpointLoader(mx_server_url="http://mx:8001") + fake_mx = _build_fake_modelexpress() + + with _install_fake_modelexpress(fake_mx): + loader.publish_as_source( + MagicMock(), + source_identity=_identity(transform_abi_id=None), + ) + + fake_mx.trtllm_live_transfer.publish_model_params.assert_not_called() + def test_publish_without_metadata_kwarg_uses_identity_metadata(self): loader = MXCheckpointLoader(mx_server_url="http://mx:8001") calls = [] @@ -683,6 +753,7 @@ def _publish_without_metadata(model): assert metadata[_MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY] == str( _MX_STAGED_TRANSFORM_PROTOCOL_VERSION ) + assert metadata[_MX_TRANSFORM_ABI_ID_METADATA_KEY] == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 def test_env_var_set_during_publish_then_restored(self): loader = MXCheckpointLoader(mx_server_url="http://mx-instance:9999") @@ -760,6 +831,10 @@ def _publish_side_effect(model, **_kwargs): assert captured["identity"].extra_parameters[ _MX_TRANSFORM_PROTOCOL_VERSION_METADATA_KEY ] == str(_MX_STAGED_TRANSFORM_PROTOCOL_VERSION) + assert ( + captured["identity"].extra_parameters[_MX_TRANSFORM_ABI_ID_METADATA_KEY] + == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + ) def test_serialized_identity_ignores_local_checkpoint_path(self): donor_identity = _identity() diff --git a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py index ff567b742cec..e3c31c33712f 100644 --- a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py +++ b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py @@ -223,6 +223,7 @@ def identity_from( *, model_name: Optional[str] = None, artifact_key: str = "same", + transform_abi_id: Optional[str] = None, ) -> SourceIdentity: """Build a :class:`SourceIdentity` from a fake config and derived model.""" return SourceIdentity.from_model_config( @@ -230,6 +231,7 @@ def identity_from( FakeModel(config.pretrained_config), artifact_identity=make_artifact_identity(artifact_key), model_name=model_name, + transform_abi_id=transform_abi_id, ) diff --git a/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py b/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py index 981b963e81c4..e697076b32fb 100644 --- a/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py +++ b/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace + import pytest from torch import nn from tensorrt_llm._torch.weight_sharing import ( + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + PostTransformConfigIdentity, PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry, @@ -41,6 +45,7 @@ def _profile( model_type: str = "model", speculative_mode: str | None = None, protocol_version: int = 1, + transform_abi_id: str = LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope: PostTransformTransferScope = PostTransformTransferScope.TARGET_MODEL, supported_features: frozenset[PostTransformFeature] = frozenset(), ) -> PostTransformProfile: @@ -51,6 +56,7 @@ def _profile( model_type=model_type, speculative_mode=speculative_mode, protocol_version=protocol_version, + transform_abi_id=transform_abi_id, transfer_scope=transfer_scope, supported_features=supported_features, ) @@ -72,6 +78,7 @@ def test_exact_profile_is_qualified() -> None: assert decision.qualified assert decision.reason is PostTransformQualificationReason.QUALIFIED assert decision.profile is profile + assert decision.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 assert decision.unsupported_features == frozenset() @@ -232,6 +239,8 @@ def test_registry_rejects_duplicate_match_key() -> None: id="speculative-mode", ), pytest.param({"protocol_version": 0}, "protocol_version", id="protocol"), + pytest.param({"transform_abi_id": ""}, "transform_abi_id", id="transform-abi"), + pytest.param({"transform_abi_id": 1}, "transform_abi_id", id="transform-abi-type"), ], ) def test_profile_rejects_invalid_required_fields( @@ -239,3 +248,45 @@ def test_profile_rejects_invalid_required_fields( ) -> None: with pytest.raises(ValueError, match=expected_message): _profile(**kwargs) + + +def test_config_identity_is_captured_before_later_normalization() -> None: + pretrained_config = SimpleNamespace( + architectures=["ModelForCausalLM"], + model_type="model", + ) + model_config = SimpleNamespace(pretrained_config=pretrained_config) + + identity = PostTransformConfigIdentity.from_model_config(model_config) + pretrained_config.architectures[0] = "NormalizedForCausalLM" + pretrained_config.model_type = "normalized" + + assert identity == PostTransformConfigIdentity( + architecture="ModelForCausalLM", + model_type="model", + ) + + +@pytest.mark.parametrize( + "architectures, model_type, expected", + [ + pytest.param([], "model", (None, "model"), id="missing-architecture"), + pytest.param([1], "model", (None, "model"), id="non-string-architecture"), + pytest.param(["ModelForCausalLM"], 1, ("ModelForCausalLM", None), id="model-type"), + ], +) +def test_config_identity_fails_closed_for_noncanonical_dimensions( + architectures: list[object], + model_type: object, + expected: tuple[str | None, str | None], +) -> None: + identity = PostTransformConfigIdentity.from_model_config( + SimpleNamespace( + pretrained_config=SimpleNamespace( + architectures=architectures, + model_type=model_type, + ) + ) + ) + + assert (identity.architecture, identity.model_type) == expected diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index 779dc8c355bd..6d574f351a61 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -34,6 +34,7 @@ ) from tensorrt_llm._torch.weight_sharing import ( + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, IdentityCheckPolicy, SourceIdentity, SourceIdentityMismatchError, @@ -68,6 +69,26 @@ def test_from_model_config_requires_one_artifact_source() -> None: ) +def test_from_model_config_binds_transform_abi() -> None: + identity = SourceIdentity.from_model_config( + FakeModelConfig(), + artifact_identity=make_artifact_identity(), + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + ) + + assert identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + + +@pytest.mark.parametrize("transform_abi_id", ["", 1]) +def test_source_identity_rejects_invalid_transform_abi(transform_abi_id: object) -> None: + with pytest.raises(ValueError, match="transform_abi_id"): + SourceIdentity.from_model_config( + FakeModelConfig(), + artifact_identity=make_artifact_identity(), + transform_abi_id=transform_abi_id, + ) + + def test_rank_defaults_from_mapping(): cfg = FakeModelConfig(mapping=FakeMapping(rank=3, tp_rank=3)) identity = identity_from(cfg) @@ -201,7 +222,10 @@ def test_enforced_sharing_skips_global(): def test_serialization_roundtrip(): - a = identity_from(FakeModelConfig()) + a = identity_from( + FakeModelConfig(), + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + ) restored = SourceIdentity.from_dict(a.to_dict()) assert restored == a assert a.matches(restored).matched @@ -215,6 +239,13 @@ def test_deserialization_rejects_missing_artifact_identity(): SourceIdentity.from_dict(payload) +def test_deserialization_rejects_missing_transform_abi_field(): + payload = identity_from(FakeModelConfig()).to_dict() + payload.pop("transform_abi_id") + 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 @@ -230,6 +261,14 @@ def test_deserialization_rejects_v1_identity_without_artifact_binding(): SourceIdentity.from_dict(payload) +def test_deserialization_rejects_v2_identity_without_transform_abi_binding(): + payload = identity_from(FakeModelConfig()).to_dict() + payload["format_version"] = 2 + payload.pop("transform_abi_id") + with pytest.raises(ValueError, match="Unsupported SourceIdentity format version"): + SourceIdentity.from_dict(payload) + + def test_check_warn_fallback_on_mismatch(): local = identity_from(FakeModelConfig(attn_backend="TRTLLM")) source = identity_from(FakeModelConfig(attn_backend="FLASHINFER")) @@ -272,6 +311,41 @@ def test_check_enforce_shares_despite_mismatch(): assert decision.should_share is True +def test_transform_abi_mismatch_is_never_bypassed(): + local = identity_from( + FakeModelConfig(), + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + ) + source = identity_from( + FakeModelConfig(), + transform_abi_id="trtllm-llama-target-layout-v2", + ) + + result = local.matches(source, compare_global=False, compare_shard=False) + decision = check_weight_sharing_compatibility( + local, + source, + IdentityCheckPolicy.ENFORCE, + ) + + assert not result.matched + assert result.mismatched_fields == ["transform_abi_id"] + assert decision.should_share is False + + +def test_transform_abi_participates_in_global_fingerprint(): + without_abi = identity_from(FakeModelConfig()) + with_abi = identity_from( + FakeModelConfig(), + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + ) + + result = without_abi.matches(with_abi) + + assert result.mismatched_fields == ["transform_abi_id"] + assert without_abi.global_fingerprint != with_abi.global_fingerprint + + def test_format_version_mismatch_never_matches(): a = identity_from(FakeModelConfig()) b = ( @@ -289,5 +363,11 @@ def test_format_version_mismatch_never_matches(): ) ) result = a.matches(b) + enforce_decision = check_weight_sharing_compatibility( + a, + b, + IdentityCheckPolicy.ENFORCE, + ) assert not result.matched assert "format_version" in result.mismatched_fields + assert enforce_decision.should_share is False diff --git a/tests/unittest/utils/post_transform_qualification.py b/tests/unittest/utils/post_transform_qualification.py index 4583b1623924..81f0f8dacdbf 100644 --- a/tests/unittest/utils/post_transform_qualification.py +++ b/tests/unittest/utils/post_transform_qualification.py @@ -23,12 +23,16 @@ from torch import nn from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader -from tensorrt_llm._torch.weight_sharing import PostTransformQualificationDecision +from tensorrt_llm._torch.weight_sharing import ( + PostTransformQualificationDecision, + PostTransformQualificationReason, +) ModelFactory = Callable[[], nn.Module] ModelQualifier = Callable[[nn.Module], PostTransformQualificationDecision] ReceiverPreparer = Callable[[nn.Module, nn.Module], None] StateProbe = Callable[[nn.Module], object] +OutputProbe = Callable[[nn.Module], object] def copy_post_transform_parameters(producer: nn.Module, receiver: nn.Module) -> None: @@ -73,6 +77,11 @@ def copy_post_transform_parameters(producer: nn.Module, receiver: nn.Module) -> class PostTransformQualificationCase: """Inputs needed to prove full and staged post-load lifecycle equivalence. + `unqualified_model_factory` must preserve the candidate config dimensions + while returning an unregistered exact root type. `output_probes` should be + deterministic and exercise the deepest practical output path for the unit + fixture; real donor/receiver GPU output coverage remains a separate gate. + `prepare_receiver` simulates installation of a producer's post-transform tensors when a family needs value-level coverage. A real MX donor/receiver test is still required before enabling a production profile. @@ -80,8 +89,10 @@ class PostTransformQualificationCase: profile_id: str model_factory: ModelFactory + unqualified_model_factory: ModelFactory qualify_model: ModelQualifier state_probes: tuple[tuple[str, StateProbe], ...] + output_probes: tuple[tuple[str, OutputProbe], ...] prepare_receiver: ReceiverPreparer = copy_post_transform_parameters def __post_init__(self) -> None: @@ -89,9 +100,14 @@ def __post_init__(self) -> None: raise ValueError("Qualification case profile_id must not be empty") if not self.state_probes: raise ValueError("Qualification case must define at least one state probe") - probe_names = tuple(name for name, _probe in self.state_probes) - if len(probe_names) != len(set(probe_names)): + if not self.output_probes: + raise ValueError("Qualification case must define at least one output probe") + state_probe_names = tuple(name for name, _probe in self.state_probes) + if len(state_probe_names) != len(set(state_probe_names)): raise ValueError("Qualification case state probe names must be unique") + output_probe_names = tuple(name for name, _probe in self.output_probes) + if len(output_probe_names) != len(set(output_probe_names)): + raise ValueError("Qualification case output probe names must be unique") def _transform_guard_state(model: nn.Module) -> dict[str, bool]: @@ -190,6 +206,20 @@ def assert_post_transform_lifecycle_equivalent( assert decision.profile is not None assert decision.profile.profile_id == case.profile_id + unqualified_model = case.unqualified_model_factory() + unqualified_decision = case.qualify_model(unqualified_model) + assert not unqualified_decision.qualified, ( + f"Negative-control root {type(unqualified_model).__name__!r} unexpectedly " + f"inherited profile {case.profile_id!r}" + ) + assert ( + unqualified_decision.reason + is PostTransformQualificationReason.ROOT_MODEL_CLASS_NOT_REGISTERED + ), ( + f"Negative-control root {type(unqualified_model).__name__!r} was rejected for " + f"{unqualified_decision.reason.value!r}, not exact root-class registration" + ) + transform_calls = _install_transform_call_recorders(receiver) ModelLoader._walk_full_post_load(producer) ModelLoader._setup_aliases(receiver) @@ -221,4 +251,22 @@ def assert_post_transform_lifecycle_equivalent( f"for profile {case.profile_id!r}" ) + for probe_name, probe in case.output_probes: + with torch.no_grad(): + producer_output = probe(producer) + receiver_output = probe(receiver) + try: + torch.testing.assert_close( + receiver_output, + producer_output, + rtol=0, + atol=0, + equal_nan=True, + ) + except AssertionError as error: + raise AssertionError( + f"Post-transform output probe {probe_name!r} differs " + f"for profile {case.profile_id!r}" + ) from error + return producer, receiver