diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index 08d4f58e5..5a202e3c3 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -15,8 +15,11 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti ## Status The first implementation provides contracts and a CPU local-filesystem store. -It does not enable cached loading for a model by itself. A consumer must still -provide an exact identity, a representation producer, and a model restorer. +The initial diffusion consumer contract additionally defines typed, +representation-independent final-layout identity/restoration mechanics plus a +concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains +library-only: no loader or DLO path selects, publishes, restores, or transports +that artifact yet. V1 includes: @@ -30,8 +33,9 @@ V1 includes: V1 does not include: -- a public CLI or default-on model integration; -- a generic BF16, FP8, quantized, or model-specific producer; +- a public CLI, loader activation, or default-on model integration; +- lease handoff to DLO or another transport consumer; +- FP8, quantized, merged-adaptation, or additional model producers; - CUDA registration, pinned staging, H2D scheduling, or GPU kernels; - a remote artifact provider or cross-node coordination; - automatic eviction; or @@ -216,6 +220,56 @@ Dynamic LoRA overlays are not part of a reusable base-weight artifact. A statically merged adapter is cacheable only as a separate identity containing the adapter fingerprint and merge semantics. +### Initial diffusion final-layout contract + +The shared diffusion contract covers complete final-layout DiT parameters and +persistent buffers. Text encoders, VAEs, non-persistent derived state, and other +pipeline components remain outside the artifact. One explicit representation +policy selects allowed dtypes, tensor roles, physical layout identity, producer +ABI, manifest schema, and restoration schema. + +The contract is intentionally separate from loader activation: + +- `FinalLayoutRequest` contains typed loader identity/configuration fingerprints, + TP coordinate, and conservative SP semantics. It has no open metadata bag, + DP coordinate, SP rank, device identity, DLO transfer mode, registration + policy, or store path. +- `FinalLayoutArtifactSpec` binds one `WeightRepresentation` and runtime-layout + name to explicit producer/restorer schemas and a canonical, versioned + implementation ABI descriptor. Compatibility never depends on reflective + source inspection. +- `PreparedWeightSource` snapshots immutable revisions or exact local file + content plus a typed checkpoint-adapter identity before ordinary + materialization. Source replacement before or during production fails + publication. A hash-looking symlink basename is trusted only for an explicit + Hugging Face Hub source whose repository ID and + `models--.../snapshots/ -> blobs/` topology validate; every + local or otherwise unverified symlink target is content-hashed. +- the tensor ownership digest records exact runtime names, kinds, shapes, + semantic roles, dtypes, and strides from a CPU or meta model skeleton; +- `FinalLayoutTensorRestorer` accepts only an exact lease identity, validates + complete policy-defined coverage without mutation, and returns a one-shot + commit plan; +- each model declares one dtype-neutral `FinalLayoutModelContract` with an + explicit implementation version and a post-commit validator; and +- `FinalLayoutBF16Producer` accepts only the matching identity context and a + finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP + coordinate. Its BF16 policy preserves model-declared FP32 parameters and + buffers and revalidates MiniMax H3 mixed-precision invariants. + +Other representations reuse source identity, typed parallel identity, tensor +ownership, and exact restoration only when their policy proves those semantics. +For example, runtime FP8 needs a separate policy/producer for generated scales, +quantization metadata, and Cutlass physical layouts; it is not enabled by +changing a dtype string on the BF16 producer. + +This stage makes no startup, sharing, or DLO performance claim. A following +consumer PR owns disabled/preferred/required precedence, mixed-component loader +transactions, warm-hit restoration, and transactional lease handoff. A TP2 +prewarm deployment will require a matching TP2 producer cohort to populate both +TP-coordinate identities even though the store coordinates each artifact +independently. + ## Host sharing and GPU transport Every process receives its own virtual mappings and `HostWeightLease`, but diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md index 7268d6a3a..ef6845965 100644 --- a/docs/design/module/host_weight_runtime.md +++ b/docs/design/module/host_weight_runtime.md @@ -56,7 +56,8 @@ The contracts have deliberately narrow ownership: - `HostWeightRuntime` applies resolution and fallback policy and emits one terminal resolution report. Explicit post-load publication emits a separate - report and may return a lease, but does not revise the completed resolution. + report and closes its store-returned lease before returning; it does not + revise the completed resolution. The runtime does not own a model or call the canonical loader. - `HostWeightStore` owns artifact lookup, coordinated construction, publication, validation, quarantine, and lifecycle. diff --git a/tests/diffusion/model_loader/test_final_layout_host_weights.py b/tests/diffusion/model_loader/test_final_layout_host_weights.py new file mode 100644 index 000000000..90cdd4998 --- /dev/null +++ b/tests/diffusion/model_loader/test_final_layout_host_weights.py @@ -0,0 +1,1041 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Contract tests for final-layout BF16 Host Weight Runtime artifacts.""" + +from __future__ import annotations + +import dataclasses +import gc +from collections.abc import Sequence +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.model_loader.host_weights import ( + FINAL_LAYOUT_BF16_POLICY, + FINAL_LAYOUT_BF16_SPEC, + FinalLayoutBF16Producer, + FinalLayoutIdentityContext, + FinalLayoutLoaderIdentity, + FinalLayoutParallelIdentity, + FinalLayoutRequest, + FinalLayoutTensorRestorer, + ImplementationIdentity, + PreparedWeightSource, + WeightSourceKind, + build_final_layout_identity, +) +from vllm_omni.diffusion.model_loader.host_weights.contracts import ( + FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + FinalLayoutArtifactSpec, + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutTensorPolicy, + implementation_abi_fingerprint, +) +from vllm_omni.diffusion.model_loader.host_weights.tensor_layout import ( + RuntimeTensorTarget, + collect_final_layout_targets, +) +from vllm_omni.diffusion.models.host_weight_contract import ( + FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + FinalLayoutModelContract, +) +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3DiTModel +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + CoordinationScope, + FailureCode, + HostWeightRuntime, + HostWeightRuntimeConfig, + LookupPhase, + PostLoadPublicationOutcome, + ProducerIdentity, + ProductionPolicy, + ProductionSourceMode, + ResolutionOutcome, + ResolutionStage, + RuntimeMode, + StorageDomainPolicy, + StorageScope, + TensorKind, + WaitPolicy, + WeightRepresentation, +) +from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _TinyDiT(nn.Module): + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="test-tiny-dit", + version="1", + ) + + def __init__(self, *, device: torch.device | str = "cpu") -> None: + super().__init__() + target = torch.device(device) + self.restore_validations = 0 + self.proj = nn.Linear(3, 2, dtype=torch.bfloat16, device=target) + self.fp32_gain = nn.Parameter(torch.empty(1, dtype=torch.float32, device=target)) + self.register_buffer("scale", torch.empty(1, dtype=torch.float32, device=target)) + self.register_buffer( + "derived", + torch.tensor([7.0], dtype=torch.float32), + persistent=False, + ) + + def validate_restored_host_weights(self) -> None: + self.restore_validations += 1 + assert not self.proj.weight.is_meta + assert self.proj.weight.dtype is torch.bfloat16 + assert self.fp32_gain.dtype is torch.float32 + assert self.scale.dtype is torch.float32 + + +class _TinyPipeline(nn.Module): + def __init__(self, *, dit_device: torch.device | str = "cpu") -> None: + super().__init__() + self.transformer = _TinyDiT(device=dit_device) + self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16) + self.vae = nn.Module() + self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32)) + + +class _AlternateTinyPipeline(_TinyPipeline): + pass + + +class _MultiTinyPipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer_a = _TinyDiT() + self.transformer_b = _TinyDiT() + + +class _Float32DiT(nn.Module): + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="test-float32-dit", + version="1", + ) + + def __init__(self) -> None: + super().__init__() + self.proj = nn.Linear(2, 2, dtype=torch.float32) + + def validate_restored_host_weights(self) -> None: + assert self.proj.weight.dtype is torch.float32 + + +class _Float32Pipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = _Float32DiT() + + +_SYNTHETIC_FP32_ABI = CanonicalJson.from_value( + { + "artifact_identity": "diffusion-final-layout-identity-v1", + "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + "producer": "test-fp32-producer-v1", + "representation_policy": "test-all-fp32-v1", + "restorer": "exact-final-layout-tensor-rebind-v1", + "source_identity": "prepared-diffusion-weight-source-v2", + "tensor_contract": "complete-strided-tensor-ownership-v1", + } +) +_SYNTHETIC_FP32_SPEC = FinalLayoutArtifactSpec( + representation=WeightRepresentation( + name="test-final-layout-fp32", + dtype=str(torch.float32), + metadata=CanonicalJson.from_value({"policy": "all-fp32"}), + ), + producer=ProducerIdentity( + producer_id="test.final-layout-fp32", + version="1", + implementation_fingerprint=implementation_abi_fingerprint(_SYNTHETIC_FP32_ABI), + manifest_schema="test-final-layout-fp32-manifest-v1", + restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + ), + implementation_abi=_SYNTHETIC_FP32_ABI, +) + + +class _SyntheticFP32Policy: + @property + def spec(self) -> FinalLayoutArtifactSpec: + return _SYNTHETIC_FP32_SPEC + + def validate_request(self, request: FinalLayoutRequest) -> None: + if request.load_format != "default": + raise ValueError("synthetic policy requires default loading") + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: + del name, tensor + return "weight" if kind is TensorKind.PARAMETER else "persistent_buffer" + + def validate_target(self, target: RuntimeTensorTarget) -> None: + if target.tensor.dtype is not torch.float32: + raise ValueError("synthetic policy requires FP32 tensors") + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: + if not targets: + raise ValueError("synthetic policy requires tensors") + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: + return CanonicalJson.from_value( + { + "component_names": list(component_names), + "format": "test-final-layout-fp32", + "tensor_contract_sha256": tensor_contract_digest, + "tensor_count": tensor_count, + } + ) + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: + assert metadata == self.build_format_metadata( + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + + +_SYNTHETIC_FP32_POLICY = _SyntheticFP32Policy() + + +class _PolicyWithSpec: + def __init__(self, base: FinalLayoutTensorPolicy, spec: FinalLayoutArtifactSpec) -> None: + self._base = base + self._spec = spec + + @property + def spec(self) -> FinalLayoutArtifactSpec: + return self._spec + + def validate_request(self, request: FinalLayoutRequest) -> None: + self._base.validate_request(request) + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: + return self._base.tensor_role(name, tensor, kind) + + def validate_target(self, target: RuntimeTensorTarget) -> None: + self._base.validate_target(target) + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: + self._base.validate_collection(targets) + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: + return self._base.build_format_metadata( + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: + self._base.validate_format_metadata( + metadata, + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + + +class _FakeH3Linear(nn.Module): + def __init__( + self, + *args: object, + bias: bool = False, + params_dtype: torch.dtype | None = None, + total_num_heads: int | None = None, + total_num_kv_heads: int | None = None, + **kwargs: object, + ) -> None: + del args, kwargs + super().__init__() + dtype = params_dtype or torch.get_default_dtype() + self.weight = nn.Parameter(torch.empty(1, dtype=dtype)) + if bias: + self.bias = nn.Parameter(torch.empty(1, dtype=dtype)) + else: + self.register_parameter("bias", None) + self.num_heads = total_num_heads + self.num_kv_heads = total_num_kv_heads + + +class _FakeH3Attention(nn.Module): + def __init__(self, **kwargs: object) -> None: + del kwargs + super().__init__() + + +class _H3Pipeline(nn.Module): + def __init__(self, transformer: nn.Module) -> None: + super().__init__() + self.transformer = transformer + + +def _small_h3_config() -> SimpleNamespace: + return SimpleNamespace( + tf_model_config={ + "num_layers": 1, + "token_refiner_num_layers": 1, + "hidden_size": 8, + "num_attention_heads": 2, + "attention_head_dim": 4, + "ffn_hidden_size": 16, + "latents_dim": 2, + "audio_latents_dim": 2, + "patch_size": (1, 2, 2), + "text_dim": 6, + "timestep_input_dim": 4, + "time_embed_hidden_size": 8, + "time_embed_dim": 4, + "adaln_out_features": 18 * 8, + "final_adaln_out_features": 2 * 8, + "rope_inv_freq_len": 2, + }, + parallel_config=SimpleNamespace(ulysses_degree=1), + ) + + +def _prepared_source( + tmp_path: Path, + *, + content: bytes = b"canonical-source-for-identity", + requested_revision: str | None = None, + resolved_revision: str | None = None, + directory: str = "canonical", + prefix: str = "transformer.", +) -> PreparedWeightSource: + source_root = tmp_path / directory + if resolved_revision is not None: + source_root = source_root / "snapshots" / resolved_revision + source_root.mkdir(parents=True, exist_ok=True) + weight_file = source_root / "model.safetensors" + if not weight_file.exists(): + weight_file.write_bytes(content) + elif weight_file.read_bytes() != content: + weight_file.write_bytes(content) + return PreparedWeightSource( + model_or_path="test-org/tiny-diffusion", + subfolder=None, + requested_revision=requested_revision, + prefix=prefix, + resolved_root=source_root, + weight_files=(weight_file,), + use_safetensors=True, + ) + + +def _request(**changes: object) -> FinalLayoutRequest: + request = FinalLayoutRequest( + model_id="test-org/tiny-diffusion", + loader=FinalLayoutLoaderIdentity( + implementation=ImplementationIdentity( + implementation_id="test-diffusion-loader", + version="1", + fingerprint="test-loader-implementation-v1", + ), + model_config_fingerprint="test-model-config-v1", + weight_transform_fingerprint="test-weight-transform-v1", + ), + ) + return dataclasses.replace(request, **changes) + + +def _identity( + model: nn.Module, + source: PreparedWeightSource, + *, + request: FinalLayoutRequest | None = None, + policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY, +) -> FinalLayoutIdentityContext: + transformer = model.get_submodule("transformer") + return build_final_layout_identity( + model, + dit_modules=(("transformer", transformer),), + prepared_sources=(source,), + request=request or _request(), + policy=policy, + ) + + +def _multi_identity( + model: _MultiTinyPipeline, + sources: tuple[PreparedWeightSource, ...], +) -> FinalLayoutIdentityContext: + return build_final_layout_identity( + model, + dit_modules=( + ("transformer_a", model.transformer_a), + ("transformer_b", model.transformer_b), + ), + prepared_sources=sources, + request=_request(), + policy=FINAL_LAYOUT_BF16_POLICY, + ) + + +def _runtime(root: Path) -> HostWeightRuntime: + return HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=StorageDomainPolicy( + root=root, + scope=StorageScope.NODE, + domain_id="node", + storage_class=detect_storage_class(root.parent), + ), + production=ProductionPolicy( + allow_local_build=False, + allow_post_load_publish=True, + ), + wait=WaitPolicy(coordination_timeout_seconds=5.0), + ) + ) + + +def _dit_modules(model: _TinyPipeline) -> tuple[tuple[str, nn.Module], ...]: + return (("transformer", model.transformer),) + + +def _fill_final_weights(model: _TinyPipeline) -> None: + with torch.no_grad(): + model.transformer.proj.weight.copy_(torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3)) + model.transformer.proj.bias.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16)) + model.transformer.fp32_gain.copy_(torch.tensor([6.5])) + model.transformer.scale.copy_(torch.tensor([2.5])) + model.text_encoder.weight.fill_(11) + model.text_encoder.bias.fill_(12) + + +def _pointer_snapshot(model: nn.Module) -> dict[str, tuple[int, int, str]]: + snapshot: dict[str, tuple[int, int, str]] = {} + for name, tensor in list(model.named_parameters()) + list(model.named_buffers()): + pointer = 0 if tensor.is_meta else tensor.untyped_storage().data_ptr() + snapshot[name] = (id(tensor), pointer, tensor.device.type) + return snapshot + + +def test_preload_identity_is_stable_for_cpu_and_meta_skeletons(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + cpu_model = _TinyPipeline() + meta_model = _TinyPipeline(dit_device="meta") + + cpu = _identity(cpu_model, source) + meta = _identity(meta_model, source) + + assert cpu.identity == meta.identity + assert cpu.tensor_names == { + "transformer.fp32_gain", + "transformer.proj.bias", + "transformer.proj.weight", + "transformer.scale", + } + assert "transformer.derived" not in cpu.tensor_names + assert all("text_encoder" not in name and "vae" not in name for name in cpu.tensor_names) + + request_fields = {field.name for field in dataclasses.fields(FinalLayoutRequest)} + assert not any("data_parallel" in name for name in request_fields) + assert not any("metadata" in name for name in request_fields) + identity_bytes = cpu.identity.canonical_bytes + assert b"data_parallel" not in identity_bytes + assert b"dlo_use_allgather" not in identity_bytes + assert b"registration" not in identity_bytes + with pytest.raises(TypeError, match="unexpected keyword argument"): + dataclasses.replace(_request(), loader_metadata=CanonicalJson.empty()) + + +def test_generic_identity_and_tensor_discovery_accept_a_second_policy(tmp_path: Path) -> None: + model = _Float32Pipeline() + source = _prepared_source(tmp_path, directory="fp32-source") + + context = _identity(model, source, policy=_SYNTHETIC_FP32_POLICY) + targets = collect_final_layout_targets( + model, + (("transformer", model.transformer),), + policy=_SYNTHETIC_FP32_POLICY, + require_materialized=True, + ) + + assert context.spec is _SYNTHETIC_FP32_SPEC + assert context.identity.representation.name == "test-final-layout-fp32" + assert context.identity.representation.dtype == str(torch.float32) + assert all(target.tensor.dtype is torch.float32 for target in targets) + assert context.identity.producer.implementation_fingerprint != ( + FINAL_LAYOUT_BF16_SPEC.producer.implementation_fingerprint + ) + + +def test_source_identity_requires_complete_multi_component_coverage(tmp_path: Path) -> None: + model = _MultiTinyPipeline() + source_a = _prepared_source( + tmp_path, + directory="source-a", + prefix="transformer_a.", + ) + + with pytest.raises(FinalLayoutContractError, match="transformer_b") as exc_info: + _multi_identity(model, (source_a,)) + assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID + + source_b = _prepared_source( + tmp_path, + directory="source-b", + prefix="transformer_b.", + ) + context = _multi_identity(model, (source_a, source_b)) + metadata = context.identity.source.metadata.to_value() + assert isinstance(metadata, dict) + bindings = metadata["target_bindings"] + assert isinstance(bindings, list) + assert {binding["source_prefix"] for binding in bindings} == { + "transformer_a.", + "transformer_b.", + } + assert {binding["target_name"].split(".", 1)[0] for binding in bindings} == { + "transformer_a", + "transformer_b", + } + + +def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path) -> None: + model = _MultiTinyPipeline() + root_source = _prepared_source( + tmp_path, + directory="root-source", + prefix="", + ) + source_a = _prepared_source( + tmp_path, + directory="specific-a", + prefix="transformer_a.", + ) + + context = _multi_identity(model, (root_source, source_a)) + reordered = _multi_identity(model, (source_a, root_source)) + assert context.identity == reordered.identity + metadata = context.identity.source.metadata.to_value() + assert isinstance(metadata, dict) + bindings = metadata["target_bindings"] + assert isinstance(bindings, list) + by_target = {binding["target_name"]: binding["source_prefix"] for binding in bindings} + assert all(prefix == "transformer_a." for name, prefix in by_target.items() if name.startswith("transformer_a.")) + assert all(prefix == "" for name, prefix in by_target.items() if name.startswith("transformer_b.")) + + conflicting_a = _prepared_source( + tmp_path, + directory="conflicting-a", + prefix="transformer_a.", + ) + with pytest.raises(FinalLayoutContractError, match="equally specific") as exc_info: + _multi_identity(model, (root_source, source_a, conflicting_a)) + assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID + + +def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None: + model = _TinyPipeline() + source_root = tmp_path / "local-symlink-source" + blobs_root = tmp_path / "local-blobs" + source_root.mkdir() + blobs_root.mkdir() + blob_path = blobs_root / ("a" * 64) + blob_path.write_bytes(b"a" * 32) + weight_path = source_root / "model.safetensors" + weight_path.symlink_to(blob_path) + source = PreparedWeightSource( + model_or_path=str(source_root), + subfolder=None, + requested_revision=None, + prefix="transformer.", + resolved_root=source_root, + weight_files=(weight_path,), + use_safetensors=True, + ) + + first = _identity(model, source) + first_metadata = first.identity.source.metadata.to_value() + assert isinstance(first_metadata, dict) + first_sources = first_metadata["sources"] + assert isinstance(first_sources, list) + assert first_sources[0]["files"][0]["content_id"].startswith("sha256:") + + blob_path.write_bytes(b"b" * 32) + second = _identity(model, source) + + assert second.identity != first.identity + + +def test_hf_blob_shortcut_requires_validated_snapshot_topology(tmp_path: Path) -> None: + model = _TinyPipeline() + model_id = "test-org/tiny-diffusion" + repo_root = tmp_path / "models--test-org--tiny-diffusion" + blobs_root = repo_root / "blobs" + revision = "1" * 40 + snapshot_root = repo_root / "snapshots" / revision + blobs_root.mkdir(parents=True) + snapshot_root.mkdir(parents=True) + blob_name = "2" * 64 + blob_path = blobs_root / blob_name + blob_path.write_bytes(b"trusted-hf-blob") + weight_path = snapshot_root / "model.safetensors" + weight_path.symlink_to(Path("..") / ".." / "blobs" / blob_name) + source = PreparedWeightSource( + model_or_path=model_id, + subfolder=None, + requested_revision="main", + prefix="transformer.", + resolved_root=snapshot_root, + weight_files=(weight_path,), + use_safetensors=True, + source_kind=WeightSourceKind.HUGGING_FACE_HUB, + ) + + context = _identity(model, source) + metadata = context.identity.source.metadata.to_value() + assert isinstance(metadata, dict) + sources = metadata["sources"] + assert isinstance(sources, list) + assert sources[0]["files"][0]["content_id"] == f"immutable-blob:{blob_name}" + + +def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) -> None: + model = _TinyPipeline() + commit = "0123456789abcdef0123456789abcdef01234567" + source = _prepared_source(tmp_path, requested_revision="main", resolved_revision=commit) + by_alias = _identity(model, source) + by_commit = _identity(model, dataclasses.replace(source, requested_revision=commit)) + + assert by_alias.identity == by_commit.identity + + base = _request() + variants = ( + dataclasses.replace( + base, + loader=dataclasses.replace(base.loader, model_config_fingerprint="test-model-config-v2"), + ), + dataclasses.replace( + base, + loader=dataclasses.replace( + base.loader, + implementation=dataclasses.replace(base.loader.implementation, version="2"), + ), + ), + dataclasses.replace( + base, + parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0), + ), + dataclasses.replace( + base, + parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1), + ), + dataclasses.replace( + base, + parallel=FinalLayoutParallelIdentity( + sequence_parallel_size=2, + ulysses_degree=2, + ulysses_mode="strict", + ), + ), + ) + base_identity = _identity(model, source, request=base).identity + assert all(_identity(model, source, request=variant).identity != base_identity for variant in variants) + + model_v2 = _TinyPipeline() + model_v2.transformer.host_weight_restore_contract = dataclasses.replace( + model_v2.transformer.host_weight_restore_contract, + version="2", + ) + assert _identity(model_v2, source).identity != base_identity + + changed_source = _prepared_source( + tmp_path, + directory="changed", + content=b"different-canonical-source", + ) + assert _identity(model, changed_source).identity != base_identity + source_with_adapter = dataclasses.replace( + source, + checkpoint_adapter=ImplementationIdentity( + implementation_id="test-checkpoint-adapter", + version="1", + fingerprint="test-checkpoint-adapter-v1", + ), + ) + assert _identity(model, source_with_adapter).identity != base_identity + + abi_v2 = CanonicalJson.from_value( + { + "artifact_identity": "diffusion-final-layout-identity-v1", + "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + "producer": "finalized-bf16-tensor-writer-v1", + "representation_policy": "bf16-with-preserved-fp32-v1", + "restorer": "exact-final-layout-tensor-rebind-v2", + "source_identity": "prepared-diffusion-weight-source-v2", + "tensor_contract": "complete-strided-tensor-ownership-v1", + } + ) + with pytest.raises(ValueError, match="fingerprint differs"): + dataclasses.replace(FINAL_LAYOUT_BF16_SPEC, implementation_abi=abi_v2) + spec_v2 = dataclasses.replace( + FINAL_LAYOUT_BF16_SPEC, + producer=dataclasses.replace( + FINAL_LAYOUT_BF16_SPEC.producer, + version="2", + implementation_fingerprint=implementation_abi_fingerprint(abi_v2), + ), + implementation_abi=abi_v2, + ) + identity_v2 = _identity( + model, + source, + policy=_PolicyWithSpec(FINAL_LAYOUT_BF16_POLICY, spec_v2), + ).identity + assert identity_v2.representation == base_identity.representation + assert identity_v2.key != base_identity.key + + +@pytest.mark.parametrize( + ("semantic_request", "message"), + [ + (_request(load_format="diffusers"), "load_format='default'"), + ( + _request(parallel=FinalLayoutParallelIdentity(pipeline_parallel_size=2)), + "pipeline-parallel", + ), + ( + _request(parallel=FinalLayoutParallelIdentity(cfg_parallel_size=2)), + "CFG-parallel", + ), + (_request(parallel=FinalLayoutParallelIdentity(use_hsdp=True)), "HSDP"), + ( + _request(parallel=FinalLayoutParallelIdentity(enable_expert_parallel=True)), + "expert-parallel", + ), + ( + _request(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256")), + "LoRA", + ), + ], +) +def test_bf16_policy_fails_closed_for_unsupported_semantics( + tmp_path: Path, + semantic_request: FinalLayoutRequest, + message: str, +) -> None: + model = _TinyPipeline() + with pytest.raises(ValueError, match=message): + _identity(model, _prepared_source(tmp_path), request=semantic_request) + + +def test_tensor_ownership_is_complete_mixed_precision_and_alias_free(tmp_path: Path) -> None: + model = _TinyPipeline() + records = collect_final_layout_targets( + model, + _dit_modules(model), + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + by_name = {record.name: record for record in records} + + assert set(by_name) == { + "transformer.fp32_gain", + "transformer.proj.bias", + "transformer.proj.weight", + "transformer.scale", + } + assert by_name["transformer.proj.weight"].tensor.dtype is torch.bfloat16 + assert by_name["transformer.fp32_gain"].tensor.dtype is torch.float32 + assert by_name["transformer.scale"].role == "persistent_buffer" + + model.transformer.register_parameter("alias", model.transformer.proj.weight) + with pytest.raises(FinalLayoutContractError, match="aliases tensor object"): + _identity(model, _prepared_source(tmp_path)) + + +def test_identity_requires_explicit_model_restore_contract(tmp_path: Path) -> None: + model = _TinyPipeline() + model.transformer.host_weight_restore_contract = "unsupported-contract" + + with pytest.raises(FinalLayoutContractError, match="does not declare"): + _identity(model, _prepared_source(tmp_path)) + + +def test_publication_is_warm_only_and_restore_is_transactional( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = _prepared_source(tmp_path) + cold_model = _TinyPipeline() + _fill_final_weights(cold_model) + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "store") + producer = FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model), max_shard_bytes=16) + assert producer.spec.lookup_phase is LookupPhase.POST_LOAD_ONLY + assert producer.spec.source_mode is ProductionSourceMode.FINALIZED_MODEL + assert producer.spec.coordination_scope is CoordinationScope.SINGLE_PROCESS + + miss = runtime.resolve(context.identity) + assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + cold_pointers = _pointer_snapshot(cold_model) + cold_values = {name: tensor.detach().clone() for name, tensor in cold_model.named_parameters()} + + publication = runtime.publish_after_load(context.identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED + assert publication.failure is None + assert _pointer_snapshot(cold_model) == cold_pointers + assert all(torch.equal(dict(cold_model.named_parameters())[name], value) for name, value in cold_values.items()) + assert cold_model.transformer.restore_validations == 1 + + def unexpected_produce(_writer: object) -> object: + raise AssertionError("already-present publication must not invoke the producer") + + monkeypatch.setattr(producer, "produce", unexpected_produce) + already_present = runtime.publish_after_load(context.identity, producer=producer) + assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT + + warm_model = _TinyPipeline(dit_device="meta") + warm_context = _identity(warm_model, source) + assert warm_context.identity == context.identity + hit = runtime.resolve(warm_context.identity) + assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT + assert hit.lease is not None + + before_plan = _pointer_snapshot(warm_model) + text_weight = warm_model.text_encoder.weight + vae_gain = warm_model.vae.gain + plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease) + assert _pointer_snapshot(warm_model) == before_plan + assert warm_model.transformer.restore_validations == 0 + + plan.commit() + with pytest.raises(RuntimeError, match="already committed"): + plan.commit() + + expected = { + "transformer.proj.weight": torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3), + "transformer.proj.bias": torch.tensor([3.0, 4.0], dtype=torch.bfloat16), + "transformer.fp32_gain": torch.tensor([6.5]), + "transformer.scale": torch.tensor([2.5]), + } + restored = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers()) + assert all(torch.equal(restored[name], value) for name, value in expected.items()) + assert all(not restored[name].is_meta for name in warm_context.tensor_names) + assert warm_model.transformer.restore_validations == 1 + assert warm_model.text_encoder.weight is text_weight + assert warm_model.vae.gain is vae_gain + assert warm_model.transformer.derived.item() == 7.0 + for name in warm_context.tensor_names: + assert restored[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr() + + del restored, warm_model, plan + gc.collect() + hit.lease.close() + assert isinstance(runtime.store, FilesystemHostWeightStore) + assert runtime.store.cleanup(context.identity) is None + + +def test_independent_publications_are_content_deterministic(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + model = _TinyPipeline() + _fill_final_weights(model) + context = _identity(model, source) + first_runtime = _runtime(tmp_path / "first-store") + second_runtime = _runtime(tmp_path / "second-store") + + first = first_runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=16), + ) + second = second_runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=1024), + ) + assert first.outcome is PostLoadPublicationOutcome.PUBLISHED + assert second.outcome is PostLoadPublicationOutcome.PUBLISHED + + first_hit = first_runtime.resolve(context.identity) + second_hit = second_runtime.resolve(context.identity) + assert first_hit.lease is not None and second_hit.lease is not None + assert first_hit.lease.manifest.artifact_content_sha256 == second_hit.lease.manifest.artifact_content_sha256 + assert first_hit.lease.manifest.format_metadata == second_hit.lease.manifest.format_metadata + assert {entry.name: entry.sha256 for entry in first_hit.lease.manifest.tensors} == { + entry.name: entry.sha256 for entry in second_hit.lease.manifest.tensors + } + first_hit.lease.close() + second_hit.lease.close() + + +def test_source_replacement_prevents_publication(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + model = _TinyPipeline() + _fill_final_weights(model) + context = _identity(model, source) + runtime = _runtime(tmp_path / "store") + source.weight_files[0].write_bytes(b"replacement-source") + + publication = runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, model, _dit_modules(model)), + ) + + assert publication.outcome is PostLoadPublicationOutcome.FAILED + assert publication.failure is not None + assert publication.failure.stage is ResolutionStage.CANONICAL_LOADING + assert publication.failure.code is FailureCode.CANONICAL_SOURCE_FAILED + assert publication.failure.details.to_value() == {"final_layout_code": "source_changed"} + assert not context.sources_unchanged() + assert runtime.resolve(context.identity).report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + + +def test_source_replacement_between_restore_plan_and_commit_causes_no_mutation(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + cold_model = _TinyPipeline() + _fill_final_weights(cold_model) + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "store") + assert ( + runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)), + ).outcome + is PostLoadPublicationOutcome.PUBLISHED + ) + hit = runtime.resolve(context.identity) + assert hit.lease is not None + + warm_model = _TinyPipeline() + warm_context = _identity(warm_model, source) + plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease) + before = _pointer_snapshot(warm_model) + source.weight_files[0].write_bytes(b"source-changed-before-commit") + + with pytest.raises(FinalLayoutContractError, match="canonical source changed"): + plan.commit() + assert _pointer_snapshot(warm_model) == before + del plan + gc.collect() + hit.lease.close() + + +def test_restore_rejects_wrong_identity_and_coverage_without_mutation(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + cold_model = _TinyPipeline() + _fill_final_weights(cold_model) + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "store") + assert ( + runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)), + ).outcome + is PostLoadPublicationOutcome.PUBLISHED + ) + hit = runtime.resolve(context.identity) + assert hit.lease is not None + + alternate = _AlternateTinyPipeline() + _fill_final_weights(alternate) + with pytest.raises(ValueError, match="producer model implementation differs"): + FinalLayoutBF16Producer(context, alternate, _dit_modules(alternate)) + with pytest.raises(ValueError, match="restore model implementation differs"): + FinalLayoutTensorRestorer(context).plan_restore(alternate, hit.lease) + + model = _TinyPipeline() + wrong_context = _identity( + model, + source, + request=_request( + parallel=FinalLayoutParallelIdentity( + tensor_parallel_size=2, + tensor_parallel_rank=1, + ) + ), + ) + with pytest.raises(ValueError, match="semantic identity differs"): + FinalLayoutTensorRestorer(wrong_context).plan_restore(model, hit.lease) + + exact_context = _identity(model, source) + before = _pointer_snapshot(model) + model.transformer.register_buffer("new_persistent_state", torch.tensor([1.0])) + with pytest.raises(FinalLayoutContractError, match="tensor contract differs"): + FinalLayoutTensorRestorer(exact_context).plan_restore(model, hit.lease) + after = _pointer_snapshot(model) + assert all(after[name] == value for name, value in before.items()) + hit.lease.close() + + +def test_reduced_minimax_h3_satisfies_real_tensor_ownership_contract( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "Attention", _FakeH3Attention) + monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1) + + transformer = h3.MiniMaxH3DiTModel(_small_h3_config(), quant_config=None) + pipeline = _H3Pipeline(transformer) + source = _prepared_source(tmp_path, directory="minimax-h3-source") + context = _identity(pipeline, source) + targets = collect_final_layout_targets( + pipeline, + (("transformer", transformer),), + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + by_name = {target.name: target for target in targets} + + assert context.model_contracts == (MiniMaxH3DiTModel.host_weight_restore_contract,) + assert "transformer.video_patch_proj.weight" in by_name + assert "transformer.blocks.0.attn.qkv_proj.weight" in by_name + assert "transformer.rope.inv_freq" in by_name + assert by_name["transformer.video_patch_proj.weight"].tensor.dtype is torch.float32 + assert by_name["transformer.blocks.0.attn.qkv_proj.weight"].tensor.dtype is torch.bfloat16 + assert by_name["transformer.rope.inv_freq"].role == "persistent_buffer" + transformer.validate_restored_host_weights() + + +def test_minimax_h3_declares_the_final_layout_restore_contract() -> None: + contract = MiniMaxH3DiTModel.host_weight_restore_contract + assert isinstance(contract, FinalLayoutModelContract) + assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + assert contract.implementation_id == "minimax-h3-dit" + assert callable(MiniMaxH3DiTModel.validate_restored_host_weights) + assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__ diff --git a/vllm_omni/diffusion/model_loader/host_weights/__init__.py b/vllm_omni/diffusion/model_loader/host_weights/__init__.py new file mode 100644 index 000000000..a7709353d --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed diffusion adapters for runtime-ready final-layout artifacts.""" + +from .contracts import ( + FinalLayoutArtifactSpec, + FinalLayoutLoaderIdentity, + FinalLayoutParallelIdentity, + FinalLayoutRequest, + ImplementationIdentity, +) +from .identity_adapter import ( + FinalLayoutIdentityContext, + build_final_layout_identity, +) +from .producers import ( + FINAL_LAYOUT_BF16_POLICY, + FINAL_LAYOUT_BF16_SPEC, + FinalLayoutBF16Producer, +) +from .restorer import FinalLayoutTensorRestorer +from .source_identity import PreparedWeightSource, WeightSourceKind + +__all__ = [ + "FINAL_LAYOUT_BF16_POLICY", + "FINAL_LAYOUT_BF16_SPEC", + "FinalLayoutArtifactSpec", + "FinalLayoutBF16Producer", + "FinalLayoutIdentityContext", + "FinalLayoutLoaderIdentity", + "FinalLayoutParallelIdentity", + "FinalLayoutRequest", + "FinalLayoutTensorRestorer", + "ImplementationIdentity", + "PreparedWeightSource", + "WeightSourceKind", + "build_final_layout_identity", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/contracts.py b/vllm_omni/diffusion/model_loader/host_weights/contracts.py new file mode 100644 index 000000000..f62bb7bb0 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/contracts.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed semantic contracts shared by final-layout artifact representations.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Protocol + +import torch + +from vllm_omni.diffusion.models.host_weight_contract import ( + FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, +) +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + FailureCode, + HostWeightError, + HostWeightFailure, + ProducerIdentity, + ResolutionStage, + TensorKind, + WeightRepresentation, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from .tensor_layout import RuntimeTensorTarget + +FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA = "diffusion-final-layout-tensor-restorer-v1" + + +def _require_text(name: str, value: object) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{name} must not be empty") + return value + + +def _require_positive_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + +class FinalLayoutContractCode(str, Enum): + SOURCE_CHANGED = "source_changed" + SOURCE_COVERAGE_INVALID = "source_coverage_invalid" + MODEL_CONTRACT_UNSUPPORTED = "model_contract_unsupported" + OWNERSHIP_AMBIGUOUS = "ownership_ambiguous" + TENSOR_UNSUPPORTED = "tensor_unsupported" + DTYPE_UNSUPPORTED = "dtype_unsupported" + IDENTITY_INCOMPATIBLE = "identity_incompatible" + TENSOR_CONTRACT_CHANGED = "tensor_contract_changed" + + +class FinalLayoutContractError(ValueError): + """Typed fail-closed error at a final-layout model integration boundary.""" + + def __init__(self, code: FinalLayoutContractCode, message: str) -> None: + if not isinstance(code, FinalLayoutContractCode): + raise ValueError("final-layout contract errors require FinalLayoutContractCode") + super().__init__(message) + self.code = code + + +def final_layout_producer_error(error: FinalLayoutContractError) -> HostWeightError: + """Translate representation failures into the runtime's stable taxonomy.""" + if error.code is FinalLayoutContractCode.SOURCE_CHANGED: + stage = ResolutionStage.CANONICAL_LOADING + code = FailureCode.CANONICAL_SOURCE_FAILED + elif error.code in { + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + }: + stage = ResolutionStage.PRODUCTION + code = FailureCode.PRODUCER_UNSUPPORTED + else: + stage = ResolutionStage.PRODUCTION + code = FailureCode.PRODUCER_FAILED + return HostWeightError( + HostWeightFailure( + stage=stage, + code=code, + retryable=False, + message=str(error), + details=CanonicalJson.from_value({"final_layout_code": error.code.value}), + ) + ) + + +@dataclass(frozen=True) +class ImplementationIdentity: + """Explicit stable identity for one semantic implementation boundary.""" + + implementation_id: str + version: str + fingerprint: str + + def __post_init__(self) -> None: + _require_text("implementation_id", self.implementation_id) + _require_text("implementation version", self.version) + _require_text("implementation fingerprint", self.fingerprint) + + def to_dict(self) -> dict[str, str]: + return { + "implementation_id": self.implementation_id, + "version": self.version, + "fingerprint": self.fingerprint, + } + + +@dataclass(frozen=True) +class FinalLayoutLoaderIdentity: + """Loader ABI and exact configuration facts that may change final bytes.""" + + implementation: ImplementationIdentity + model_config_fingerprint: str + weight_transform_fingerprint: str + + def __post_init__(self) -> None: + if not isinstance(self.implementation, ImplementationIdentity): + raise ValueError("loader implementation must use ImplementationIdentity") + _require_text("model config fingerprint", self.model_config_fingerprint) + _require_text("weight transform fingerprint", self.weight_transform_fingerprint) + + def to_dict(self) -> dict[str, object]: + return { + "implementation": self.implementation.to_dict(), + "model_config_fingerprint": self.model_config_fingerprint, + "weight_transform_fingerprint": self.weight_transform_fingerprint, + } + + +@dataclass(frozen=True) +class FinalLayoutParallelIdentity: + """Semantic parallel layout with DP and SP rank intentionally absent.""" + + tensor_parallel_size: int = 1 + tensor_parallel_rank: int = 0 + sequence_parallel_size: int = 1 + ulysses_degree: int = 1 + ring_degree: int = 1 + allgather_degree: int = 1 + ulysses_mode: str = "strict" + pipeline_parallel_size: int = 1 + cfg_parallel_size: int = 1 + use_hsdp: bool = False + enable_expert_parallel: bool = False + + def __post_init__(self) -> None: + tp_size = _require_positive_int("tensor_parallel_size", self.tensor_parallel_size) + if isinstance(self.tensor_parallel_rank, bool) or not isinstance(self.tensor_parallel_rank, int): + raise ValueError("tensor_parallel_rank must be an integer") + if not 0 <= self.tensor_parallel_rank < tp_size: + raise ValueError("tensor_parallel_rank must be within tensor_parallel_size") + for name, value in ( + ("sequence_parallel_size", self.sequence_parallel_size), + ("ulysses_degree", self.ulysses_degree), + ("ring_degree", self.ring_degree), + ("allgather_degree", self.allgather_degree), + ("pipeline_parallel_size", self.pipeline_parallel_size), + ("cfg_parallel_size", self.cfg_parallel_size), + ): + _require_positive_int(name, value) + _require_text("ulysses_mode", self.ulysses_mode) + if not isinstance(self.use_hsdp, bool) or not isinstance(self.enable_expert_parallel, bool): + raise ValueError("HSDP and expert-parallel flags must be booleans") + + @property + def sequence_parallel_backend(self) -> str: + if self.sequence_parallel_size == 1: + return "none" + active = [ + name + for name, degree in ( + ("ulysses", self.ulysses_degree), + ("ring", self.ring_degree), + ("allgather", self.allgather_degree), + ) + if degree > 1 + ] + return "+".join(active) if active else "replicated-sp" + + def to_dict(self) -> dict[str, object]: + return { + "tensor_parallel_size": self.tensor_parallel_size, + "tensor_parallel_rank": self.tensor_parallel_rank, + "sequence_parallel_size": self.sequence_parallel_size, + "sequence_parallel_backend": self.sequence_parallel_backend, + "ulysses_degree": self.ulysses_degree, + "ring_degree": self.ring_degree, + "allgather_degree": self.allgather_degree, + "ulysses_mode": self.ulysses_mode, + "pipeline_parallel_size": self.pipeline_parallel_size, + "cfg_parallel_size": self.cfg_parallel_size, + "use_hsdp": self.use_hsdp, + "enable_expert_parallel": self.enable_expert_parallel, + } + + +@dataclass(frozen=True) +class FinalLayoutRequest: + """Representation-independent semantic request assembled by the loader.""" + + model_id: str + loader: FinalLayoutLoaderIdentity + parallel: FinalLayoutParallelIdentity = field(default_factory=FinalLayoutParallelIdentity) + load_format: str = "default" + adaptation: AdaptationIdentity = field(default_factory=AdaptationIdentity) + + def __post_init__(self) -> None: + _require_text("canonical model ID or path", self.model_id) + if not isinstance(self.loader, FinalLayoutLoaderIdentity): + raise ValueError("loader must use FinalLayoutLoaderIdentity") + if not isinstance(self.parallel, FinalLayoutParallelIdentity): + raise ValueError("parallel must use FinalLayoutParallelIdentity") + _require_text("load_format", self.load_format) + if not isinstance(self.adaptation, AdaptationIdentity): + raise ValueError("adaptation must use AdaptationIdentity") + + +def implementation_abi_fingerprint(abi: CanonicalJson) -> str: + if not isinstance(abi, CanonicalJson): + raise ValueError("implementation ABI must use CanonicalJson") + return hashlib.sha256(abi.encoded).hexdigest() + + +@dataclass(frozen=True) +class FinalLayoutArtifactSpec: + """Concrete representation, layout, producer, and restoration ABI.""" + + representation: WeightRepresentation + producer: ProducerIdentity + implementation_abi: CanonicalJson + layout_name: str = "diffusion-final-module-layout-v1" + component_name: str = "diffusion-dit" + component_ownership: str = "complete-final-layout-tensors" + model_contract_schema: str = FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + + def __post_init__(self) -> None: + if not isinstance(self.representation, WeightRepresentation): + raise ValueError("artifact representation must use WeightRepresentation") + if not isinstance(self.producer, ProducerIdentity): + raise ValueError("artifact producer must use ProducerIdentity") + if not isinstance(self.implementation_abi, CanonicalJson): + raise ValueError("artifact implementation ABI must use CanonicalJson") + for name, value in ( + ("layout_name", self.layout_name), + ("component_name", self.component_name), + ("component_ownership", self.component_ownership), + ("model_contract_schema", self.model_contract_schema), + ): + _require_text(name, value) + abi = self.implementation_abi.to_value() + required_abi_fields = { + "artifact_identity", + "model_contract", + "producer", + "representation_policy", + "restorer", + "source_identity", + "tensor_contract", + } + if not isinstance(abi, dict) or not required_abi_fields <= set(abi): + raise ValueError( + "artifact implementation ABI must version identity, source, tensor, policy, producer, " + "restorer, and model contracts" + ) + for name in required_abi_fields: + _require_text(f"implementation ABI {name}", abi[name]) + if self.producer.implementation_fingerprint != implementation_abi_fingerprint(self.implementation_abi): + raise ValueError("producer implementation fingerprint differs from its explicit ABI descriptor") + + +class FinalLayoutTensorPolicy(Protocol): + """Representation-specific validation over shared final-layout mechanics.""" + + @property + def spec(self) -> FinalLayoutArtifactSpec: ... + + def validate_request(self, request: FinalLayoutRequest) -> None: ... + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: ... + + def validate_target(self, target: RuntimeTensorTarget) -> None: ... + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: ... + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: ... + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: ... + + +__all__ = [ + "FinalLayoutArtifactSpec", + "FinalLayoutContractCode", + "FinalLayoutContractError", + "FinalLayoutLoaderIdentity", + "FinalLayoutParallelIdentity", + "FinalLayoutRequest", + "FinalLayoutTensorPolicy", + "FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA", + "ImplementationIdentity", + "final_layout_producer_error", + "implementation_abi_fingerprint", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py b/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py new file mode 100644 index 000000000..416f03a45 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact pre-load identity assembly for policy-defined final-layout artifacts.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +from torch import nn + +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.host_weight_runtime import ( + CanonicalJson, + ComponentIdentity, + RuntimeWeightLayout, + WeightArtifactIdentity, + WeightRepresentation, +) + +from .contracts import ( + FinalLayoutArtifactSpec, + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutRequest, + FinalLayoutTensorPolicy, +) +from .source_identity import ( + FinalLayoutSourceContext, + PreparedWeightSource, + resolve_final_layout_source_identity, +) +from .tensor_layout import ( + RuntimeTensorTarget, + collect_final_layout_targets, + tensor_contract_sha256, + validate_final_layout_model_contract, +) + + +@dataclass(frozen=True) +class FinalLayoutIdentityContext: + """Exact semantic identity plus process-local transaction guards.""" + + identity: WeightArtifactIdentity + spec: FinalLayoutArtifactSpec + policy: FinalLayoutTensorPolicy + request: FinalLayoutRequest + tensor_names: frozenset[str] + dit_names: tuple[str, ...] + model_contracts: tuple[FinalLayoutModelContract, ...] + source: FinalLayoutSourceContext + pipeline_type: type[nn.Module] + dit_types: tuple[type[nn.Module], ...] + + def sources_unchanged(self) -> bool: + return self.source.sources_unchanged() + + def ensure_sources_unchanged(self) -> None: + self.source.ensure_sources_unchanged() + + +def _require_metadata_object(name: str, metadata: CanonicalJson) -> Mapping[str, object]: + value = metadata.to_value() + if not isinstance(value, dict): + raise ValueError(f"{name} must encode one JSON object") + return value + + +def _model_contract_values(contracts: Sequence[FinalLayoutModelContract]) -> list[dict[str, str]]: + return [contract.to_dict() for contract in contracts] + + +def build_final_layout_identity( + pipeline: nn.Module, + *, + dit_modules: Sequence[tuple[str, nn.Module]], + prepared_sources: Sequence[PreparedWeightSource], + request: FinalLayoutRequest, + policy: FinalLayoutTensorPolicy, +) -> FinalLayoutIdentityContext: + """Build exact source, representation, layout, and implementation identity.""" + if not isinstance(request, FinalLayoutRequest): + raise TypeError("final-layout identity requires FinalLayoutRequest") + policy.validate_request(request) + spec = policy.spec + bindings = validate_final_layout_model_contract( + dit_modules, + expected_schema=spec.model_contract_schema, + ) + structural_targets = collect_final_layout_targets( + pipeline, + dit_modules, + policy=policy, + require_materialized=False, + ) + target_names = frozenset(target.name for target in structural_targets) + source = resolve_final_layout_source_identity( + prepared_sources, + model_id=request.model_id, + target_names=target_names, + ) + dit_names = tuple(name for name, _ in dit_modules) + model_contracts = tuple(binding.contract for binding in bindings) + + representation = WeightRepresentation( + name=spec.representation.name, + dtype=spec.representation.dtype, + metadata=CanonicalJson.from_value( + { + "artifact_policy": _require_metadata_object( + "artifact representation metadata", + spec.representation.metadata, + ), + "load_format": request.load_format, + "loader": request.loader.to_dict(), + } + ), + ) + identity = WeightArtifactIdentity( + schema_version=1, + source=source.identity, + component=ComponentIdentity( + name=spec.component_name, + ownership=spec.component_ownership, + metadata=CanonicalJson.from_value( + { + "component_names": list(dit_names), + "model_contracts": _model_contract_values(model_contracts), + } + ), + ), + representation=representation, + layout=RuntimeWeightLayout( + name=spec.layout_name, + tensor_parallel_size=request.parallel.tensor_parallel_size, + tensor_parallel_rank=request.parallel.tensor_parallel_rank, + sequence_parallel_size=request.parallel.sequence_parallel_size, + sequence_parallel_backend=request.parallel.sequence_parallel_backend, + metadata=CanonicalJson.from_value( + { + "model_contract_schema": spec.model_contract_schema, + "parallel": request.parallel.to_dict(), + "tensor_contract_sha256": tensor_contract_sha256(structural_targets), + } + ), + ), + adaptation=request.adaptation, + producer=spec.producer, + ) + return FinalLayoutIdentityContext( + identity=identity, + spec=spec, + policy=policy, + request=request, + tensor_names=target_names, + dit_names=dit_names, + model_contracts=model_contracts, + source=source, + pipeline_type=type(pipeline), + dit_types=tuple(type(module) for _, module in dit_modules), + ) + + +def validate_final_layout_identity( + context: FinalLayoutIdentityContext, + targets: Sequence[RuntimeTensorTarget], +) -> str: + """Validate a context against current targets and return their digest.""" + identity = context.identity + spec = context.spec + request = context.request + digest = tensor_contract_sha256(targets) + if identity.source != context.source.identity or identity.adaptation != request.adaptation: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity source or adaptation differs from the exact context", + ) + if identity.producer != spec.producer: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity producer differs from the artifact specification", + ) + if identity.representation.name != spec.representation.name or identity.representation.dtype != ( + spec.representation.dtype + ): + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity representation differs from the artifact specification", + ) + if identity.component.name != spec.component_name or identity.component.ownership != spec.component_ownership: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity component ownership differs from the artifact specification", + ) + if identity.layout.name != spec.layout_name: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity runtime layout differs from the artifact specification", + ) + + component = _require_metadata_object("component metadata", identity.component.metadata) + representation = _require_metadata_object("representation metadata", identity.representation.metadata) + layout = _require_metadata_object("layout metadata", identity.layout.metadata) + if representation != { + "artifact_policy": _require_metadata_object( + "artifact representation metadata", + spec.representation.metadata, + ), + "load_format": request.load_format, + "loader": request.loader.to_dict(), + }: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity loader or representation policy differs from the exact context", + ) + if component.get("component_names") != list(context.dit_names) or component.get("model_contracts") != ( + _model_contract_values(context.model_contracts) + ): + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity model ownership differs from the exact context", + ) + if layout.get("model_contract_schema") != spec.model_contract_schema: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity model restore schema differs from the artifact specification", + ) + if ( + identity.layout.tensor_parallel_size != request.parallel.tensor_parallel_size + or identity.layout.tensor_parallel_rank != request.parallel.tensor_parallel_rank + or identity.layout.sequence_parallel_size != request.parallel.sequence_parallel_size + or identity.layout.sequence_parallel_backend != request.parallel.sequence_parallel_backend + or layout.get("parallel") != request.parallel.to_dict() + ): + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity parallel semantics differ from the exact context", + ) + if layout.get("tensor_contract_sha256") != digest: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_CONTRACT_CHANGED, + "identity tensor contract differs from the current model", + ) + if frozenset(target.name for target in targets) != context.tensor_names: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_CONTRACT_CHANGED, + "current tensor ownership differs from the exact identity context", + ) + return digest + + +__all__ = [ + "FinalLayoutIdentityContext", + "build_final_layout_identity", + "validate_final_layout_identity", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py b/vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py new file mode 100644 index 000000000..268e26981 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Diffusion Host Weight Runtime producers.""" + +from .final_layout_bf16 import ( + FINAL_LAYOUT_BF16_MANIFEST_SCHEMA, + FINAL_LAYOUT_BF16_POLICY, + FINAL_LAYOUT_BF16_PRODUCER_ID, + FINAL_LAYOUT_BF16_REPRESENTATION, + FINAL_LAYOUT_BF16_SPEC, + FINAL_LAYOUT_BF16_VERSION, + FinalLayoutBF16Policy, + FinalLayoutBF16Producer, +) + +__all__ = [ + "FINAL_LAYOUT_BF16_MANIFEST_SCHEMA", + "FINAL_LAYOUT_BF16_POLICY", + "FINAL_LAYOUT_BF16_PRODUCER_ID", + "FINAL_LAYOUT_BF16_REPRESENTATION", + "FINAL_LAYOUT_BF16_SPEC", + "FINAL_LAYOUT_BF16_VERSION", + "FinalLayoutBF16Policy", + "FinalLayoutBF16Producer", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py b/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py new file mode 100644 index 000000000..622256f24 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Concrete BF16-with-preserved-FP32 final-layout artifact policy and producer.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +from torch import nn + +from vllm_omni.diffusion.models.host_weight_contract import ( + FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, +) +from vllm_omni.host_weight_runtime import ( + ArtifactWriter, + CanonicalJson, + CoordinationScope, + LookupPhase, + ProducerIdentity, + ProductionMetadata, + ProductionSourceMode, + TensorKind, + TensorWriteSpec, + WeightProductionSpec, + WeightRepresentation, +) + +from ..contracts import ( + FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + FinalLayoutArtifactSpec, + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutRequest, + final_layout_producer_error, + implementation_abi_fingerprint, +) +from ..identity_adapter import ( + FinalLayoutIdentityContext, + validate_final_layout_identity, +) +from ..tensor_layout import ( + RuntimeTensorTarget, + collect_final_layout_targets, + validate_final_layout_model_contract, +) + +FINAL_LAYOUT_BF16_PRODUCER_ID = "vllm-omni.diffusion.final-layout-bf16" +FINAL_LAYOUT_BF16_VERSION = "1" +FINAL_LAYOUT_BF16_REPRESENTATION = "diffusion-final-layout-bf16" +FINAL_LAYOUT_BF16_MANIFEST_SCHEMA = "diffusion-final-layout-bf16-manifest-v1" +DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3 + +_BF16_PARAMETER_DTYPES = {torch.bfloat16, torch.float32} +_BF16_BUFFER_DTYPES = { + torch.bool, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, +} +for _dtype_name in ( + "float8_e4m3fn", + "float8_e5m2", + "float8_e4m3fnuz", + "float8_e5m2fnuz", +): + if (_dtype := getattr(torch, _dtype_name, None)) is not None: + _BF16_BUFFER_DTYPES.add(_dtype) + +_BF16_IMPLEMENTATION_ABI = CanonicalJson.from_value( + { + "artifact_identity": "diffusion-final-layout-identity-v1", + "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + "producer": "finalized-bf16-tensor-writer-v1", + "representation_policy": "bf16-with-preserved-fp32-v1", + "restorer": "exact-final-layout-tensor-rebind-v1", + "source_identity": "prepared-diffusion-weight-source-v2", + "tensor_contract": "complete-strided-tensor-ownership-v1", + } +) + +FINAL_LAYOUT_BF16_SPEC = FinalLayoutArtifactSpec( + representation=WeightRepresentation( + name=FINAL_LAYOUT_BF16_REPRESENTATION, + dtype=str(torch.bfloat16), + metadata=CanonicalJson.from_value( + { + "format": FINAL_LAYOUT_BF16_REPRESENTATION, + "mixed_precision_policy": "bf16-with-preserved-fp32", + } + ), + ), + producer=ProducerIdentity( + producer_id=FINAL_LAYOUT_BF16_PRODUCER_ID, + version=FINAL_LAYOUT_BF16_VERSION, + implementation_fingerprint=implementation_abi_fingerprint(_BF16_IMPLEMENTATION_ABI), + manifest_schema=FINAL_LAYOUT_BF16_MANIFEST_SCHEMA, + restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + ), + implementation_abi=_BF16_IMPLEMENTATION_ABI, +) + + +class FinalLayoutBF16Policy: + """Concrete compatibility and metadata policy for finalized BF16 tensors.""" + + @property + def spec(self) -> FinalLayoutArtifactSpec: + return FINAL_LAYOUT_BF16_SPEC + + def validate_request(self, request: FinalLayoutRequest) -> None: + if request.load_format != "default": + raise ValueError(f"final-layout BF16 artifacts require load_format='default', got {request.load_format!r}") + if request.adaptation.kind != "base" or request.adaptation.fingerprint is not None: + raise ValueError("merged or static LoRA requires a representation-specific producer") + parallel = request.parallel + if parallel.pipeline_parallel_size != 1: + raise ValueError("pipeline-parallel ownership is not supported by the final-layout BF16 producer") + if parallel.cfg_parallel_size != 1: + raise ValueError("CFG-parallel ownership is not supported by the final-layout BF16 producer") + if parallel.use_hsdp: + raise ValueError("HSDP/DTensor layouts are not supported by the final-layout BF16 producer") + if parallel.enable_expert_parallel: + raise ValueError("expert-parallel ownership is not supported by the final-layout BF16 producer") + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: + del name, tensor + if kind is TensorKind.PARAMETER: + return "weight" + return "persistent_buffer" + + def validate_target(self, target: RuntimeTensorTarget) -> None: + if target.kind is TensorKind.PARAMETER and target.tensor.dtype not in _BF16_PARAMETER_DTYPES: + raise FinalLayoutContractError( + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + f"{target.name!r} must be BF16 or an explicitly preserved FP32 parameter, got {target.tensor.dtype}", + ) + if target.kind is TensorKind.BUFFER and target.tensor.dtype not in _BF16_BUFFER_DTYPES: + raise FinalLayoutContractError( + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + f"{target.name!r} uses unsupported buffer dtype {target.tensor.dtype}", + ) + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: + if not any(target.kind is TensorKind.PARAMETER and target.tensor.dtype is torch.bfloat16 for target in targets): + raise FinalLayoutContractError( + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + "final-layout BF16 representation requires at least one BF16 parameter", + ) + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: + return CanonicalJson.from_value( + { + "component_names": list(component_names), + "format": FINAL_LAYOUT_BF16_REPRESENTATION, + "mixed_precision_policy": "bf16-with-preserved-fp32", + "tensor_contract_sha256": tensor_contract_digest, + "tensor_count": tensor_count, + "tensor_layout": "contiguous", + } + ) + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: + expected = self.build_format_metadata( + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + if metadata != expected: + raise ValueError("lease format metadata differs from the BF16 artifact policy") + + +FINAL_LAYOUT_BF16_POLICY = FinalLayoutBF16Policy() + + +def _split_shards( + records: Sequence[RuntimeTensorTarget], + max_shard_bytes: int, +) -> tuple[tuple[RuntimeTensorTarget, ...], ...]: + if max_shard_bytes <= 0: + raise ValueError("final-layout BF16 shard size must be positive") + shards: list[tuple[RuntimeTensorTarget, ...]] = [] + current: list[RuntimeTensorTarget] = [] + current_bytes = 0 + for record in records: + if current and current_bytes + record.nbytes > max_shard_bytes: + shards.append(tuple(current)) + current = [] + current_bytes = 0 + current.append(record) + current_bytes += record.nbytes + if current: + shards.append(tuple(current)) + return tuple(shards) + + +class FinalLayoutBF16Producer: + """Publish finalized BF16-policy tensors through a store-scoped writer.""" + + def __init__( + self, + context: FinalLayoutIdentityContext, + pipeline: nn.Module, + dit_modules: Sequence[tuple[str, nn.Module]], + *, + max_shard_bytes: int = DEFAULT_SHARD_SIZE_BYTES, + ) -> None: + if max_shard_bytes <= 0: + raise ValueError("final-layout BF16 shard size must be positive") + if context.spec != FINAL_LAYOUT_BF16_SPEC: + raise ValueError("producer requires the exact BF16 artifact specification") + self._context = context + self._pipeline = pipeline + self._dit_modules = tuple(dit_modules) + self._max_shard_bytes = max_shard_bytes + if tuple(name for name, _ in self._dit_modules) != context.dit_names: + raise ValueError("producer DiT components differ from the exact identity context") + if type(pipeline) is not context.pipeline_type or tuple(type(module) for _, module in self._dit_modules) != ( + context.dit_types + ): + raise ValueError("producer model implementation differs from the exact identity context") + self._spec = WeightProductionSpec( + producer_id=FINAL_LAYOUT_BF16_PRODUCER_ID, + outputs=(context.identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + coordination_scope=CoordinationScope.SINGLE_PROCESS, + lookup_phase=LookupPhase.POST_LOAD_ONLY, + ) + + @property + def spec(self) -> WeightProductionSpec: + return self._spec + + def produce(self, writer: ArtifactWriter) -> ProductionMetadata: + try: + return self._produce(writer) + except FinalLayoutContractError as exc: + raise final_layout_producer_error(exc) from exc + + def _produce(self, writer: ArtifactWriter) -> ProductionMetadata: + self._context.ensure_sources_unchanged() + bindings = validate_final_layout_model_contract( + self._dit_modules, + expected_schema=self._context.spec.model_contract_schema, + ) + records = collect_final_layout_targets( + self._pipeline, + self._dit_modules, + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + contract_digest = validate_final_layout_identity(self._context, records) + for binding in bindings: + binding.validator() + + shards = _split_shards(records, self._max_shard_bytes) + shard_count = len(shards) + for index, shard in enumerate(shards, start=1): + specs = tuple( + TensorWriteSpec( + name=record.name, + shape=tuple(record.tensor.shape), + dtype=record.tensor.dtype, + kind=record.kind, + role=record.role, + ) + for record in shard + ) + file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors" + with writer.open_tensor_file(file_name, specs) as output: + for record in shard: + output.write_tensor(record.name, record.tensor.detach()) + + self._context.ensure_sources_unchanged() + return ProductionMetadata( + producer_schema=FINAL_LAYOUT_BF16_MANIFEST_SCHEMA, + restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + format_metadata=FINAL_LAYOUT_BF16_POLICY.build_format_metadata( + component_names=self._context.dit_names, + tensor_contract_digest=contract_digest, + tensor_count=len(records), + ), + ) + + +__all__ = [ + "DEFAULT_SHARD_SIZE_BYTES", + "FINAL_LAYOUT_BF16_MANIFEST_SCHEMA", + "FINAL_LAYOUT_BF16_POLICY", + "FINAL_LAYOUT_BF16_PRODUCER_ID", + "FINAL_LAYOUT_BF16_REPRESENTATION", + "FINAL_LAYOUT_BF16_SPEC", + "FINAL_LAYOUT_BF16_VERSION", + "FinalLayoutBF16Policy", + "FinalLayoutBF16Producer", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/restorer.py b/vllm_omni/diffusion/model_loader/host_weights/restorer.py new file mode 100644 index 000000000..04581172f --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/restorer.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generic validation-first restorer for exact final-layout tensor artifacts.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn + +from vllm_omni.host_weight_runtime import HostWeightLease, TensorKind, WeightRestorePlan + +from .identity_adapter import ( + FinalLayoutIdentityContext, + validate_final_layout_identity, +) +from .tensor_layout import collect_final_layout_targets, validate_final_layout_model_contract + + +@dataclass(frozen=True) +class _Replacement: + name: str + parent: nn.Module + leaf_name: str + target: torch.Tensor + source: torch.Tensor + kind: TensorKind + + def current_target(self) -> torch.Tensor | None: + if self.kind is TensorKind.PARAMETER: + return self.parent._parameters.get(self.leaf_name) + return self.parent._buffers.get(self.leaf_name) + + +class FinalLayoutTensorRestorePlan(WeightRestorePlan): + """One-shot exact tensor rebinding validated before the first mutation.""" + + def __init__( + self, + lease: HostWeightLease, + replacements: tuple[_Replacement, ...], + validators: tuple[Callable[[], None], ...], + source_guard: Callable[[], None], + ) -> None: + self._lease = lease + self._replacements = replacements + self._validators = validators + self._source_guard = source_guard + self._committed = False + + def commit(self) -> None: + if self._committed: + raise RuntimeError("final-layout tensor restore plan was already committed") + if self._lease.closed: + raise RuntimeError("cannot commit a final-layout tensor restore from a closed lease") + self._source_guard() + + for replacement in self._replacements: + current = replacement.current_target() + if current is not replacement.target: + raise RuntimeError(f"restore target {replacement.name!r} changed after preflight") + if ( + tuple(current.shape) != tuple(replacement.source.shape) + or current.dtype != replacement.source.dtype + or tuple(current.stride()) != tuple(replacement.source.stride()) + or not replacement.source.is_contiguous() + ): + raise RuntimeError(f"restore target {replacement.name!r} changed layout after preflight") + + # Mark first: assignment or model-validator failure makes the plan + # permanently non-retryable and the partially restored model disposable. + self._committed = True + for replacement in self._replacements: + if replacement.target.is_meta: + if replacement.kind is TensorKind.PARAMETER: + replacement.parent._parameters[replacement.leaf_name] = nn.Parameter( + replacement.source, + requires_grad=replacement.target.requires_grad, + ) + else: + replacement.parent._buffers[replacement.leaf_name] = replacement.source + else: + replacement.target.data = replacement.source + for validator in self._validators: + validator() + + +class FinalLayoutTensorRestorer: + """Validate and plan restoration for one exact policy-defined lease.""" + + def __init__(self, context: FinalLayoutIdentityContext) -> None: + if not context.dit_names or any(not name for name in context.dit_names): + raise ValueError("final-layout restorer requires DiT component names") + self._context = context + + @property + def schema(self) -> str: + return self._context.spec.producer.restorer_schema + + def plan_restore(self, model: object, lease: HostWeightLease) -> FinalLayoutTensorRestorePlan: + if not isinstance(model, nn.Module): + raise TypeError("final-layout restoration requires an nn.Module pipeline") + if type(model) is not self._context.pipeline_type: + raise ValueError("restore model implementation differs from the exact identity context") + if lease.closed: + raise ValueError("cannot restore from a closed HostWeightLease") + if lease.identity.canonical_bytes != self._context.identity.canonical_bytes: + raise ValueError("lease semantic identity differs from the exact restore request") + if lease.manifest.identity.canonical_bytes != self._context.identity.canonical_bytes: + raise ValueError("lease manifest identity differs from the exact restore request") + if lease.identity.producer.restorer_schema != self.schema or lease.manifest.restorer_schema != self.schema: + raise ValueError("lease restorer schema is incompatible with the exact restore request") + self._context.ensure_sources_unchanged() + + try: + dit_modules = tuple((name, model.get_submodule(name)) for name in self._context.dit_names) + except AttributeError as exc: + raise ValueError("one or more lease-owned DiT modules do not exist in the pipeline") from exc + if tuple(type(module) for _, module in dit_modules) != self._context.dit_types: + raise ValueError("restore DiT implementation differs from the exact identity context") + bindings = validate_final_layout_model_contract( + dit_modules, + expected_schema=self._context.spec.model_contract_schema, + ) + targets = collect_final_layout_targets( + model, + dit_modules, + policy=self._context.policy, + require_materialized=False, + ) + contract_digest = validate_final_layout_identity(self._context, targets) + + manifest_entries = {entry.name: entry for entry in lease.manifest.tensors} + target_names = {target.name for target in targets} + if set(manifest_entries) != target_names or target_names != self._context.tensor_names: + missing = sorted(target_names - set(manifest_entries)) + unexpected = sorted(set(manifest_entries) - target_names) + raise ValueError( + "lease tensor coverage differs from the structurally initialized DiT: " + f"missing={missing[:5]}, unexpected={unexpected[:5]}" + ) + self._context.policy.validate_format_metadata( + lease.manifest.format_metadata, + component_names=self._context.dit_names, + tensor_contract_digest=contract_digest, + tensor_count=len(manifest_entries), + ) + + replacements: list[_Replacement] = [] + source_storages: dict[tuple[int, int], str] = {} + for target in targets: + entry = manifest_entries[target.name] + source = lease.tensors[target.name] + if entry.kind is not target.kind or entry.role != target.role: + raise ValueError(f"lease tensor kind or role differs for {target.name!r}") + if ( + tuple(source.shape) != tuple(target.tensor.shape) + or source.dtype != target.tensor.dtype + or tuple(source.stride()) != tuple(target.tensor.stride()) + or not source.is_contiguous() + or source.device.type != "cpu" + ): + raise ValueError(f"lease tensor layout differs for {target.name!r}") + if source.numel(): + storage = source.untyped_storage() + storage_id = (storage.data_ptr(), storage.nbytes()) + owner = source_storages.setdefault(storage_id, target.name) + if owner != target.name: + raise ValueError(f"lease tensor {target.name!r} aliases {owner!r}") + parent_path, _, leaf_name = target.name.rpartition(".") + replacements.append( + _Replacement( + name=target.name, + parent=model.get_submodule(parent_path), + leaf_name=leaf_name, + target=target.tensor, + source=source, + kind=target.kind, + ) + ) + + self._context.ensure_sources_unchanged() + return FinalLayoutTensorRestorePlan( + lease, + tuple(replacements), + tuple(binding.validator for binding in bindings), + self._context.ensure_sources_unchanged, + ) + + +__all__ = [ + "FinalLayoutTensorRestorePlan", + "FinalLayoutTensorRestorer", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/source_identity.py b/vllm_omni/diffusion/model_loader/host_weights/source_identity.py new file mode 100644 index 000000000..800bdf288 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/source_identity.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Canonical diffusion weight-source identity and one-load replacement guard.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from pathlib import Path + +import regex as re + +from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity +from vllm_omni.host_weight_runtime.identity import canonical_json + +from .contracts import ( + FinalLayoutContractCode, + FinalLayoutContractError, + ImplementationIdentity, +) + +_HASH_CHUNK_BYTES = 8 * 1024**2 +_IMMUTABLE_REVISION_RE = re.compile(r"(?:[0-9a-fA-F]{40}|[0-9a-fA-F]{64})\Z") + + +class WeightSourceKind(str, Enum): + LOCAL_PATH = "local_path" + HUGGING_FACE_HUB = "hugging_face_hub" + + +@dataclass(frozen=True) +class PreparedWeightSource: + """Canonical source files and adapter ABI selected by the ordinary loader.""" + + model_or_path: str + subfolder: str | None + requested_revision: str | None + prefix: str + resolved_root: Path + weight_files: tuple[Path, ...] + use_safetensors: bool + checkpoint_adapter: ImplementationIdentity | None = None + source_kind: WeightSourceKind = WeightSourceKind.LOCAL_PATH + + def __post_init__(self) -> None: + if not self.model_or_path or not isinstance(self.prefix, str): + raise ValueError("prepared source model and prefix must be valid strings") + if not isinstance(self.resolved_root, Path): + raise ValueError("prepared source root must use pathlib.Path") + if not isinstance(self.weight_files, tuple) or not self.weight_files: + raise ValueError("prepared source must contain an immutable tuple of weight files") + if any(not isinstance(path, Path) for path in self.weight_files): + raise ValueError("prepared source weight files must use pathlib.Path") + if not isinstance(self.use_safetensors, bool): + raise ValueError("prepared source safetensors mode must be a boolean") + if self.checkpoint_adapter is not None and not isinstance( + self.checkpoint_adapter, + ImplementationIdentity, + ): + raise ValueError("checkpoint adapter must use ImplementationIdentity") + if not isinstance(self.source_kind, WeightSourceKind): + raise ValueError("prepared source kind must use WeightSourceKind") + + +@dataclass(frozen=True) +class _FileSnapshot: + path: Path + relative_name: str + size: int + device: int + inode: int + mtime_ns: int + ctime_ns: int + symlink_target: str | None + content_id: str + + def semantic_dict(self) -> dict[str, object]: + return { + "relative_name": self.relative_name, + "size": self.size, + "content_id": self.content_id, + } + + def unchanged(self) -> bool: + try: + current = self.path.stat() + symlink_target = os.readlink(self.path) if self.path.is_symlink() else None + except OSError: + return False + return ( + current.st_size == self.size + and current.st_dev == self.device + and current.st_ino == self.inode + and current.st_mtime_ns == self.mtime_ns + and current.st_ctime_ns == self.ctime_ns + and symlink_target == self.symlink_target + ) + + +@dataclass(frozen=True) +class _SourceSnapshot: + semantic: CanonicalJson + revision: str + files: tuple[_FileSnapshot, ...] + + def unchanged(self) -> bool: + return all(file.unchanged() for file in self.files) + + +@dataclass(frozen=True) +class FinalLayoutSourceContext: + """Resolved source identity plus process-local replacement observations.""" + + identity: WeightSourceIdentity + snapshots: tuple[_SourceSnapshot, ...] + + def sources_unchanged(self) -> bool: + return all(snapshot.unchanged() for snapshot in self.snapshots) + + def ensure_sources_unchanged(self) -> None: + if not self.sources_unchanged(): + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_CHANGED, + "canonical source changed after final-layout identity resolution", + ) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(_HASH_CHUNK_BYTES): + digest.update(chunk) + return digest.hexdigest() + + +def _hf_blob_content_id( + path: Path, + *, + source: PreparedWeightSource, + source_root: Path, +) -> str | None: + if source.source_kind is not WeightSourceKind.HUGGING_FACE_HUB or not path.is_symlink(): + return None + if Path(source.model_or_path).expanduser().exists(): + return None + + parts = source_root.parts + snapshot_indexes = [index for index, part in enumerate(parts[:-1]) if part == "snapshots"] + if not snapshot_indexes: + return None + snapshot_index = snapshot_indexes[-1] + revision = parts[snapshot_index + 1] + if _IMMUTABLE_REVISION_RE.fullmatch(revision) is None: + return None + repo_root = Path(*parts[:snapshot_index]) + expected_repo_name = f"models--{source.model_or_path.replace('/', '--')}" + if repo_root.name != expected_repo_name: + return None + snapshot_root = repo_root / "snapshots" / revision + try: + path.relative_to(snapshot_root) + except ValueError: + return None + blobs_root = repo_root / "blobs" + try: + resolved_target = path.resolve(strict=True) + resolved_blobs_root = blobs_root.resolve(strict=True) + except OSError: + return None + if resolved_target.parent != resolved_blobs_root: + return None + blob_name = resolved_target.name + if _IMMUTABLE_REVISION_RE.fullmatch(blob_name) is None: + return None + return f"immutable-blob:{blob_name.lower()}" + + +def _snapshot_revision(path: Path) -> str | None: + parts = path.absolute().parts + for index, part in enumerate(parts[:-1]): + if part != "snapshots": + continue + candidate = parts[index + 1] + if _IMMUTABLE_REVISION_RE.fullmatch(candidate) is not None: + return candidate.lower() + return None + + +def _logical_model_id(value: str) -> str: + candidate = Path(value).expanduser() + try: + if candidate.exists(): + return str(candidate.resolve()) + except OSError: + pass + return value + + +def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot: + root = source.resolved_root.expanduser().resolve() + if not root.is_dir(): + raise ValueError(f"canonical source root is not a directory: {root}") + immutable_revision = _snapshot_revision(root) + if immutable_revision is None and source.requested_revision is not None: + requested = source.requested_revision.strip() + if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None: + immutable_revision = requested.lower() + + files: list[_FileSnapshot] = [] + for path in sorted(source.weight_files, key=lambda item: str(item)): + candidate = path.expanduser() + if not candidate.is_absolute(): + candidate = root / candidate + candidate = candidate.absolute() + try: + relative_name = candidate.relative_to(root).as_posix() + except ValueError as exc: + raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc + try: + current = candidate.stat() + except OSError as exc: + raise ValueError(f"cannot stat canonical weight file {candidate}") from exc + if not stat.S_ISREG(current.st_mode): + raise ValueError(f"canonical weight source is not a regular file: {candidate}") + symlink_target = os.readlink(candidate) if candidate.is_symlink() else None + content_id = _hf_blob_content_id( + candidate, + source=source, + source_root=root, + ) + if content_id is None: + content_id = f"sha256:{_sha256_file(candidate)}" + files.append( + _FileSnapshot( + path=candidate, + relative_name=relative_name, + size=current.st_size, + device=current.st_dev, + inode=current.st_ino, + mtime_ns=current.st_mtime_ns, + ctime_ns=current.st_ctime_ns, + symlink_target=symlink_target, + content_id=content_id, + ) + ) + + file_semantics = [file.semantic_dict() for file in files] + file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest() + revision = immutable_revision or f"content-{file_fingerprint}" + semantic = CanonicalJson.from_value( + { + "checkpoint_adapter": ( + source.checkpoint_adapter.to_dict() if source.checkpoint_adapter is not None else None + ), + "files": file_semantics, + "model_or_path": _logical_model_id(source.model_or_path), + "prefix": source.prefix, + "resolved_revision": revision, + "source_kind": source.source_kind.value, + "subfolder": source.subfolder, + "use_safetensors": source.use_safetensors, + } + ) + return _SourceSnapshot(semantic=semantic, revision=revision, files=tuple(files)) + + +def _resolve_target_sources( + prepared_sources: Sequence[PreparedWeightSource], + target_names: frozenset[str], +) -> dict[str, PreparedWeightSource]: + """Bind every target to one source using deterministic longest-prefix wins.""" + if not target_names: + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_COVERAGE_INVALID, + "final-layout source identity requires at least one owned tensor", + ) + + bindings: dict[str, PreparedWeightSource] = {} + for target_name in sorted(target_names): + matches = [source for source in prepared_sources if target_name.startswith(source.prefix)] + if not matches: + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_COVERAGE_INVALID, + f"no canonical weight source covers final-layout tensor {target_name!r}", + ) + longest_prefix = max(len(source.prefix) for source in matches) + winners = [source for source in matches if len(source.prefix) == longest_prefix] + if len(winners) != 1: + candidates = sorted( + f"{source.model_or_path}:{source.subfolder or ''}:{source.prefix}" for source in winners + ) + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_COVERAGE_INVALID, + f"multiple equally specific canonical sources cover {target_name!r}: {candidates}", + ) + bindings[target_name] = winners[0] + return bindings + + +def resolve_final_layout_source_identity( + prepared_sources: Sequence[PreparedWeightSource], + *, + model_id: str, + target_names: frozenset[str], +) -> FinalLayoutSourceContext: + """Resolve exact source identity for files covering the owned tensors.""" + bindings = _resolve_target_sources(prepared_sources, target_names) + selected_sources = frozenset(bindings.values()) + snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources} + snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded)) + source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots] + target_bindings = [ + { + "source_fingerprint": hashlib.sha256(snapshot_by_source[source].semantic.encoded).hexdigest(), + "source_prefix": source.prefix, + "target_name": target_name, + } + for target_name, source in sorted(bindings.items()) + ] + source_document = { + "sources": source_semantics, + "target_bindings": target_bindings, + } + source_fingerprint = hashlib.sha256(canonical_json(source_document)).hexdigest() + unique_revisions = tuple(sorted({snapshot.revision for snapshot in snapshots})) + aggregate_revision = ( + unique_revisions[0] + if len(unique_revisions) == 1 + else f"aggregate-{hashlib.sha256(canonical_json(unique_revisions)).hexdigest()}" + ) + return FinalLayoutSourceContext( + identity=WeightSourceIdentity( + model_id=_logical_model_id(model_id), + revision=aggregate_revision, + fingerprint=source_fingerprint, + metadata=CanonicalJson.from_value(source_document), + ), + snapshots=snapshots, + ) + + +__all__ = [ + "FinalLayoutSourceContext", + "PreparedWeightSource", + "WeightSourceKind", + "resolve_final_layout_source_identity", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py b/vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py new file mode 100644 index 000000000..d8759417c --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Representation-neutral ownership discovery for final-layout tensors.""" + +from __future__ import annotations + +import hashlib +from collections.abc import Callable, Iterator, Sequence +from dataclasses import dataclass +from itertools import chain + +import torch +from torch import nn + +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.host_weight_runtime import TensorKind +from vllm_omni.host_weight_runtime.identity import canonical_json + +from .contracts import ( + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutTensorPolicy, +) + + +@dataclass(frozen=True) +class ModelRestoreBinding: + """Model-declared restore ABI and its post-commit validator.""" + + name: str + contract: FinalLayoutModelContract + validator: Callable[[], None] + + +@dataclass(frozen=True) +class RuntimeTensorTarget: + """One complete final-layout tensor owned by a DiT component.""" + + name: str + tensor: torch.Tensor + kind: TensorKind + role: str + + @property + def nbytes(self) -> int: + return self.tensor.numel() * self.tensor.element_size() + + +def validate_final_layout_model_contract( + dit_modules: Sequence[tuple[str, nn.Module]], + *, + expected_schema: str, +) -> tuple[ModelRestoreBinding, ...]: + """Require an explicit tensor-complete restore ABI from every DiT.""" + if not dit_modules: + raise FinalLayoutContractError( + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + "final-layout artifacts require at least one DiT component", + ) + + bindings: list[ModelRestoreBinding] = [] + seen_names: set[str] = set() + for name, module in dit_modules: + if not name or name in seen_names: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"DiT component name {name!r} is empty or duplicated", + ) + seen_names.add(name) + contract = getattr(module, "host_weight_restore_contract", None) + if not isinstance(contract, FinalLayoutModelContract) or contract.schema != expected_schema: + raise FinalLayoutContractError( + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + f"DiT component {name!r} does not declare restore schema {expected_schema!r}", + ) + validator = getattr(module, "validate_restored_host_weights", None) + if not callable(validator): + raise FinalLayoutContractError( + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + f"DiT component {name!r} does not implement validate_restored_host_weights()", + ) + bindings.append(ModelRestoreBinding(name, contract, validator)) + return tuple(bindings) + + +def _named_parameters_with_duplicates(module: nn.Module) -> Iterator[tuple[str, nn.Parameter]]: + try: + return module.named_parameters(remove_duplicate=False) + except TypeError: # pragma: no cover - compatibility with old torch + return module.named_parameters() + + +def _named_buffers_with_duplicates(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: + try: + return module.named_buffers(remove_duplicate=False) + except TypeError: # pragma: no cover - compatibility with old torch + return module.named_buffers() + + +def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool: + parent_path, _, leaf_name = local_name.rpartition(".") + owner = module.get_submodule(parent_path) + return leaf_name not in owner._non_persistent_buffers_set + + +def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None: + parent_path, _, leaf_name = runtime_name.rpartition(".") + parent = pipeline.get_submodule(parent_path) + tensor = parent._parameters.get(leaf_name) + if tensor is None: + tensor = parent._buffers.get(leaf_name) + return tensor + + +def collect_final_layout_targets( + pipeline: nn.Module, + dit_modules: Sequence[tuple[str, nn.Module]], + *, + policy: FinalLayoutTensorPolicy, + require_materialized: bool, +) -> tuple[RuntimeTensorTarget, ...]: + """Return complete, alias-free targets validated by one representation policy.""" + validate_final_layout_model_contract( + dit_modules, + expected_schema=policy.spec.model_contract_schema, + ) + records: dict[str, RuntimeTensorTarget] = {} + for dit_name, dit_module in dit_modules: + for local_name, tensor in _named_parameters_with_duplicates(dit_module): + runtime_name = f"{dit_name}.{local_name}" + candidate = RuntimeTensorTarget( + runtime_name, + tensor, + TensorKind.PARAMETER, + policy.tensor_role(runtime_name, tensor, TensorKind.PARAMETER), + ) + existing = records.get(runtime_name) + if existing is not None and existing.tensor is not tensor: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"multiple DiT tensors resolve to {runtime_name!r}", + ) + records[runtime_name] = candidate + for local_name, tensor in _named_buffers_with_duplicates(dit_module): + if not _is_persistent_buffer(dit_module, local_name): + continue + runtime_name = f"{dit_name}.{local_name}" + candidate = RuntimeTensorTarget( + runtime_name, + tensor, + TensorKind.BUFFER, + policy.tensor_role(runtime_name, tensor, TensorKind.BUFFER), + ) + existing = records.get(runtime_name) + if existing is not None and existing.tensor is not tensor: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"multiple DiT tensors resolve to {runtime_name!r}", + ) + records[runtime_name] = candidate + + if not records: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + "no DiT parameters or persistent buffers were discovered", + ) + + object_owners: dict[int, str] = {} + storage_owners: dict[tuple[int, int], str] = {} + for record in records.values(): + tensor = record.tensor + if not record.role: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} has an empty semantic role", + ) + try: + pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name) + except AttributeError as exc: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"DiT tensor {record.name!r} is not owned by the pipeline", + ) from exc + if pipeline_tensor is not tensor: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"DiT tensor {record.name!r} does not resolve to the discovered object", + ) + + object_owner = object_owners.setdefault(id(tensor), record.name) + if object_owner != record.name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} aliases tensor object {object_owner!r}", + ) + if hasattr(tensor, "to_local"): + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} is a distributed tensor", + ) + if tensor.layout != torch.strided: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} uses unsupported layout {tensor.layout}", + ) + if tensor.device.type not in {"cpu", "meta"}: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} must be a CPU or meta tensor, got {tensor.device}", + ) + if not tensor.is_contiguous(): + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} is non-contiguous with stride {tensor.stride()}", + ) + policy.validate_target(record) + + if not require_materialized: + continue + if tensor.device.type != "cpu" or tensor.is_meta: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}", + ) + storage = tensor.untyped_storage() + if tensor.storage_offset() != 0 or storage.nbytes() != record.nbytes: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} is a view into a larger storage", + ) + if record.nbytes: + storage_id = (storage.data_ptr(), storage.nbytes()) + storage_owner = storage_owners.setdefault(storage_id, record.name) + if storage_owner != record.name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} shares storage with {storage_owner!r}", + ) + + for pipeline_name, tensor in chain( + _named_parameters_with_duplicates(pipeline), + _named_buffers_with_duplicates(pipeline), + ): + owner = object_owners.get(id(tensor)) + if owner is not None and owner != pipeline_name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"cached tensor {owner!r} aliases pipeline tensor {pipeline_name!r}", + ) + if not require_materialized or tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0: + continue + storage = tensor.untyped_storage() + pipeline_storage_owner = storage_owners.get((storage.data_ptr(), storage.nbytes())) + if pipeline_storage_owner is not None and pipeline_storage_owner != pipeline_name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"cached tensor {pipeline_storage_owner!r} shares storage with pipeline tensor {pipeline_name!r}", + ) + + ordered = tuple(records[name] for name in sorted(records)) + policy.validate_collection(ordered) + return ordered + + +def tensor_contract_sha256(records: Sequence[RuntimeTensorTarget]) -> str: + """Hash exact structural ownership and representation-specific tensor roles.""" + contract = [ + { + "dtype": str(record.tensor.dtype), + "kind": record.kind.value, + "name": record.name, + "role": record.role, + "shape": list(record.tensor.shape), + "stride": list(record.tensor.stride()), + } + for record in records + ] + return hashlib.sha256(canonical_json(contract)).hexdigest() + + +__all__ = [ + "ModelRestoreBinding", + "RuntimeTensorTarget", + "collect_final_layout_targets", + "tensor_contract_sha256", + "validate_final_layout_model_contract", +] diff --git a/vllm_omni/diffusion/models/host_weight_contract.py b/vllm_omni/diffusion/models/host_weight_contract.py new file mode 100644 index 000000000..939c9af2e --- /dev/null +++ b/vllm_omni/diffusion/models/host_weight_contract.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Model-owned declaration for complete final-layout tensor restoration.""" + +from __future__ import annotations + +from dataclasses import dataclass + +FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA = "vllm-omni.diffusion.final-layout-tensors-v1" + + +@dataclass(frozen=True) +class FinalLayoutModelContract: + """Versioned model ABI for reconstruction from exact final-layout tensors.""" + + implementation_id: str + version: str + schema: str = FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + + def __post_init__(self) -> None: + for name, value in ( + ("schema", self.schema), + ("implementation_id", self.implementation_id), + ("version", self.version), + ): + if not isinstance(value, str) or not value: + raise ValueError(f"final-layout model contract {name} must not be empty") + + def to_dict(self) -> dict[str, str]: + return { + "schema": self.schema, + "implementation_id": self.implementation_id, + "version": self.version, + } + + +__all__ = [ + "FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA", + "FinalLayoutModelContract", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index af5f9b1dd..91e03c865 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 packed-token audio/video DiT for vLLM-Omni. vLLM tensor parallel linears and the unified attention layer provide TP and @@ -43,6 +44,7 @@ from vllm_omni.diffusion.layers.activation import SiluAndMul from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -889,6 +891,15 @@ class MiniMaxH3SPGather(nn.Module): class MiniMaxH3DiTModel(nn.Module): + # Loading is tensor-complete: constructor state plus final-layout + # parameters and persistent buffers is sufficient to reconstruct a ready + # inference model. The model-specific validator below checks the preserved + # FP32 portion after a lease-backed restore commits. + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="minimax-h3-dit", + version="1", + ) + _cache_dit_adapter_config = CacheDiTAdapterConfig( block_forward_patterns={"blocks": ForwardPattern.Pattern_3}, # H3 is CFG-distilled and performs one transformer forward per step. @@ -1064,6 +1075,10 @@ class MiniMaxH3DiTModel(nn.Module): if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE: raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.") + def validate_restored_host_weights(self) -> None: + """Validate mixed-precision invariants after lease-backed restore.""" + self.post_load_weights() + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]],