diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 48c72c95a..d7c47e22a 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -143,6 +143,35 @@ is no-AllGather-only, and it selects the backend transfer protocol. This is coordination through typed configuration and `HostWeightPlan`, not through model-specific global state. +### Runtime host-layout identity (v1) + +Host weight cache v1 identifies the loader-produced host representation, not +the topology used later to move weights to devices. These are separate +contracts: + +- `RuntimeHostLayout` describes tensor ownership and representation after the + ordinary loader has applied its parallel-aware transforms. This is the + layout stored in the host weight cache. +- `DLOTransferLayout` describes whether each worker copies its local weights + directly or participates in the existing DLO AllGather protocol. It is not + part of the host weight cache in v1; AllGather continues to use its existing + process-owned host shards. + +The v1 physical-sharing invariant is one cache entry per TP coordinate, +shared by equivalent DP and SP consumers on the same node. The identity +therefore includes TP world size and TP rank, because different TP coordinates +own different tensors. It excludes DP world size/rank and SP rank, because +those dimensions do not change the final tensor ownership supported by v1. +SP world size and implementation details remain conservative compatibility +guards until equivalent layouts across those configurations are proven. + +PP, HSDP/DTensor, expert parallelism, and CFG parallelism are rejected rather +than assigned speculative cache identities. Supporting one of these layouts +requires first defining its loader-produced ownership, representation, and +sharing domain; it does not follow automatically from supporting its runtime +collectives. Consequently, this phase implements a TP-coordinate runtime host +layout, not a general parallel-layout cache. + ### Shared startup and lifecycle The shared control path selects host backing and transfers its ownership to diff --git a/tests/diffusion/model_loader/test_host_weight_cache.py b/tests/diffusion/model_loader/test_host_weight_cache.py index 634299dcf..74f34e1fb 100644 --- a/tests/diffusion/model_loader/test_host_weight_cache.py +++ b/tests/diffusion/model_loader/test_host_weight_cache.py @@ -162,6 +162,14 @@ def test_expert_parallel_ownership_fails_closed(tmp_path): assert list(tmp_path.rglob("*.safetensors")) == [] +def test_pipeline_parallel_ownership_fails_closed(tmp_path): + result = _build(_Pipeline(), tmp_path, pipeline_parallel_size=2) + + assert result.plan is None + assert result.fallback_code == "unsupported_pp" + assert list(tmp_path.rglob("*.safetensors")) == [] + + def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path): first = _build(_Pipeline(), tmp_path) assert first.plan is not None diff --git a/vllm_omni/diffusion/model_loader/host_weight_cache.py b/vllm_omni/diffusion/model_loader/host_weight_cache.py index 1b6d4be78..e0e4ec2ab 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_cache.py +++ b/vllm_omni/diffusion/model_loader/host_weight_cache.py @@ -480,6 +480,9 @@ def build_host_weight_cache_plan( try: records = _collect_runtime_tensors(pipeline, dit_modules) content_digest = runtime_content_digest(records) + # The cache stores loader-produced ownership: TP coordinates differ, + # while equivalent DP replicas and SP ranks may share one entry. The + # conservative SP guard prevents reuse across unproven implementations. layout_identity = { "schema_version": HOST_WEIGHT_CACHE_SCHEMA_VERSION, "model": (canonicalize_existing_local_path(model_identity) if model_identity is not None else None),