diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 24dcbb753e..f069bdb3f6 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -12,6 +12,7 @@ /mooncake-integration/transfer_engine @ShangmingCai @alogfans /mooncake-integration/store @ykwd @stmatengss @zxpdemonio /mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys +/mooncake-reshard @ShangmingCai @stmatengss @Bo-Vincent @zxpdemonio /mooncake-store @ykwd @stmatengss @XucSh @YiXR /mooncake-store/*/ha/ @Libotry @YiXR @00fish0 @Icedcoco /mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q @staryxchen diff --git a/.github/labeler.yml b/.github/labeler.yml index aa76d0be22..92650848b7 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -33,7 +33,9 @@ Integration: Common: - changed-files: - - any-glob-to-any-file: 'mooncake-common/**/*' + - any-glob-to-any-file: + - 'mooncake-common/**/*' + - 'mooncake-reshard/**/*' CI/Build: - changed-files: @@ -54,6 +56,7 @@ Tests: - any-glob-to-any-file: - 'scripts/test_*' - 'mooncake-wheel/tests/**/*' + - 'mooncake-reshard/tests/**/*' - 'scripts/tone_tests/**/*' Ascend/NPU: diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 2e44ed37cf..0ca038db46 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -6,6 +6,7 @@ - [ ] Transfer Engine (`mooncake-transfer-engine`) - [ ] Mooncake Store (`mooncake-store`) +- [ ] Reshard (`mooncake-reshard`) - [ ] Mooncake EP (`mooncake-ep`) - [ ] Mooncake PG (`mooncake-pg`) - [ ] Integration (`mooncake-integration`) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9aeabbb2d7..e3195b0b83 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -573,6 +573,26 @@ jobs: - name: Spell Check Repo uses: crate-ci/typos@v1.30.2 + reshard-type-check: + name: Check reshard manifest types + if: *run-ci + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Check canonical reshard contract types + run: | + python -m pip install --disable-pip-version-check pyright==1.1.411 + bash scripts/check_reshard_types.sh + shell: bash + clang-format: name: Check code format if: *run-ci diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 46607d9270..06600a0e65 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -62,7 +62,7 @@ repos: hooks: - id: codespell exclude: '^(extern/|FAST25-release/)' - args: ['--ignore-words-list=te,mooncake,KVCache,cann,hsa'] + args: ['--ignore-words-list=te,mooncake,KVCache,cann,hsa,crate'] - repo: https://github.com/cheshirekow/cmake-format-precommit rev: v0.6.13 diff --git a/docs/source/design/reshard-manifest.md b/docs/source/design/reshard-manifest.md new file mode 100644 index 0000000000..1b310234c9 --- /dev/null +++ b/docs/source/design/reshard-manifest.md @@ -0,0 +1,181 @@ +# Resource Reshard Manifest Contract + +This document defines the framework-neutral resource contract used by +Mooncake resharding and its model-weight specialization. The contract separates +complete logical placement from live physical addresses so planning can finish +before a runtime binding is available. + +The implementation is owned by the top-level `mooncake-reshard` module. Common +contracts are exposed through `mooncake.reshard.contracts`; the public weight API +is `mooncake.reshard.weight`. + +Framework-owned adapters outside Mooncake inspect framework runtime objects, +normalize framework-specific values, and construct the typed canonical +manifests. Mooncake core accepts only those typed values; it does not import or +inspect framework objects or accept alternate field names or duck-typed +records. + +## Contract Split + +| Contract | Contents | Lifetime | +|----------|----------|----------| +| `ResourceManifest` | structural protocol for resource identity and kind | shared public contract | +| `PlacementManifest` | structural protocol for address-free placement identity and digest | serializable and reusable | +| `RuntimeBindingManifest` | structural protocol for placement attestation, runtime instance, generation, and lease | one live runtime snapshot | +| `ParallelTopology` | TP/PP/EP/DP sizes and the explicit participant-to-rank mapping | one logical placement | +| `SplitAxis` | a parallel kind that shards one explicit logical dimension | one tensor descriptor | +| `ReplicatedAxis` | a parallel kind whose ranks each hold a complete replica | one tensor descriptor | +| `OwnershipAxis` | a parallel kind that assigns tensor or object ownership without splitting a dimension | one tensor descriptor | +| `WeightPlacementPart` | one participant's address-free tensors and logical fragments | framework-local contribution | +| `WeightPlacementManifest` | one complete global logical placement of a weight generation | serializable and reusable | +| `WeightRuntimeBindingManifest` | one participant's physical fragments for that global placement | one live runtime snapshot | + +The three common manifest contracts are public structural `Protocol` types. +Consumers depend on their fields and behavior, not inheritance from a Mooncake +base class. + +Weight revision, tensor geometry, model semantics, parallel ownership, and +weight generation belong to the weight specialization. GPU addresses, +endpoints, owners, generations, and leases never appear in +`WeightPlacementManifest`. + +`model_weight` is the serialized resource discriminator. Typed manifests carry +their `ResourceKind` explicitly; Mooncake does not infer a resource or model +type from parameter names. + +## Global Placement Assembly + +`ParallelTopology` declares the runtime's TP, PP, EP, and DP sizes and the exact +participants selected for this placement. Its `world_size` is the number of +declared participants, not `tp_size * pp_size * ep_size * dp_size`. Frameworks +may map axes such as TP and EP onto the same workers, and a placement may select +one DP replica while retaining the runtime's declared `dp_size`. + +For each participant, the framework adapter constructs one typed +`WeightPlacementPart`. A part carries the common resource ID, revision, weight +generation, placement-set ID, topology ID, participant ID, parallel rank, +tensor descriptors, and logical fragments. It contains no physical address. A +part declares exactly the tensor descriptors referenced by its fragments; an +empty part declares neither. + +A collection barrier assembles all declared parts into one +`WeightPlacementManifest`. Assembly fails when a participant is missing or +duplicated, when a part belongs to a different resource, generation, placement +set, or topology, or when its rank disagrees with the topology. Only after the +complete placement validates are its canonical `placement_id` and digest +available. + +For each live participant, the framework adapter then constructs a typed +`WeightRuntimeBindingManifest` that names its `participant_id` and attests the +same global `placement_id` and digest. Binding-set validation requires every +participant that owns fragments exactly once, and exact logical-fragment +membership for each such participant. Empty participants require no runtime +binding. + +## Logical Semantics + +Each tensor has a stable `tensor_id`, full `global_shape`, dtype, item size, +layout fingerprint, and optional layer or expert identity. Each fragment is an +N-D logical box described by `global_offset` and `local_shape`. + +`TensorDescriptor.shard_dims` is the only canonical shard representation. +`SplitAxis(kind, dim)` explicitly shards one logical dimension; +`ReplicatedAxis(kind)` requires each selected rank to provide a complete copy; +and `OwnershipAxis(kind)` assigns tensor or object ownership without splitting +a logical dimension. Axis size comes from `ParallelTopology`, and a fragment's +axis rank comes from `ParallelRank`. The dimensions named by all `SplitAxis` +values must match `shard_dims` exactly. + +The global manifest validates complete logical coverage. Every selected DP +replica must provide a gap-free cover of every tensor. `OwnershipAxis` and +`ReplicatedAxis` values form independent covers. Fragments across a `SplitAxis` +instead form one non-overlapping cover, and every split-axis rank declared by +the topology must participate. The explicit participant mapping defines the +selected workers and may be non-Cartesian overall. Within one tensor's owner and +replica cover, however, coordinates for multiple declared `SplitAxis` values +must form their Cartesian product so that each rank-to-dimension assignment is +provable. A physical coordinate coupled to another split rank but not +independently sharding the tensor is left out of that tensor's `parallel_axes`. +DP may therefore select one complete replica for transfer while the topology +retains the original `dp_size`. + +PP is layer or tensor ownership. A logical tensor may have complete replicas on +multiple PP owners, but every owner must independently provide a gap-free +cover; fragments from different PP owners cannot be combined to satisfy +coverage. For grouped expert tensors, EP uses `SplitAxis` on the leading logical +expert dimension rather than only an EP rank label. Independently allocated +experts use `OwnershipAxis` and remain independent tensors with an explicit +expert identity. + +Mooncake does not infer layer, expert, layout, or partition semantics from +model parameter names. Framework adapters must provide those facts. + +`placement_fragment_id` defaults to a canonical hash of tensor identity, +logical box, parallel rank, byte size, and alias group. Frameworks may supply +an explicit stable ID when they intentionally need a different identity. An +alias group is valid only when it contains the fragment's own `tensor_id`; two +fragments may share one runtime range only when both tensor IDs belong to the +same compatible alias group. + +Because an alias group can cross placement participants, a local +`WeightPlacementPart` validates only its own fragments. Complete +`WeightPlacementManifest` assembly is the authorization boundary: every alias +member must be in the global tensor catalog and every fragment of every member +must declare the same alias group before any runtime binding is accepted. + +## Identity And Fencing + +Canonical placement identity covers the resource, revision, weight generation, +placement-set ID, topology, global tensor descriptors, participant ownership, +and logical fragments. Runtime addresses, workers, endpoints, owners, and +leases do not affect placement identity. + +Every runtime binding carries the global placement ID and digest. Validation +rejects a binding when the logical placement changes, a participant is unknown, +a fragment is missing or unexpected, or its byte range differs. Generation and +lease fences remain live-runtime state and must be checked before transfer. + +Every runtime `address` points to the first transferable byte of a contiguous +tensor view. A runtime fragment preserves `itemsize`, `local_shape`, byte +strides normalized on singleton dimensions, storage base address, normalized +storage byte offset, and storage allocation size. Binding validation compares +item size, shape, and contiguous +row-major byte strides with the logical placement; singleton dimensions do not +constrain their corresponding stride. It also verifies +`address = storage_address + storage_offset_bytes`, and requires the complete +view range to remain inside the allocation. An optional +framework `is_contiguous` flag may reject a view early but is never accepted as +the sole proof of contiguity. Address zero is reserved as a null sentinel, and +all address ranges must have representable unsigned 64-bit exclusive ends. +Owner objects may keep framework allocations alive but are never serialized. + +## Integration Flow + +1. A framework-owned adapter reads framework state and constructs one typed + `ParallelTopology` plus the shared resource, revision, weight generation, + and placement-set ID. +2. The adapter constructs one typed `WeightPlacementPart` for every selected + participant. +3. A barrier collects the exact part set and constructs one complete + `WeightPlacementManifest`. +4. The adapter constructs a typed `WeightRuntimeBindingManifest` for each live + participant against the resulting placement ID and digest. +5. Planning consumes one source and one target `WeightPlacementManifest`. +6. Binding and execution use only the participant bindings referenced by the + logical plan, while preserving their generation and lease fences. + +`weight_placement_to_json` and `weight_placement_from_json` are the explicit +public JSON APIs for the canonical wire schema. Deserialization accepts exactly +the canonical fields and values; it does not accept aliases, attribute-based +records, or other framework-shaped inputs. Integer-valued contract fields +require Python `int` values and reject `bool`; framework adapters must normalize +framework-specific scalar types before constructing a manifest. + +## Boundaries + +The manifest contract does not inspect framework objects, infer model semantics, +synthesize framework placements, execute transfers, or define discovery, +activation, rollback, and other control-plane policies. Framework adapters own +object inspection and normalization. Planner, Store, and Transfer Engine +adapters consume the resulting canonical manifests without changing their +logical identity rules. diff --git a/docs/source/index.md b/docs/source/index.md index 9bf417894c..cb2faa8e2f 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -125,6 +125,7 @@ performance/vllm/index design/architecture design/transfer-engine/index +design/reshard-manifest design/tent/overview design/store/mooncake-store design/mooncake-backend-pg diff --git a/mooncake-reshard/README.md b/mooncake-reshard/README.md new file mode 100644 index 0000000000..25a22c86dc --- /dev/null +++ b/mooncake-reshard/README.md @@ -0,0 +1,90 @@ +# Mooncake Reshard + +`mooncake-reshard` defines framework-neutral contracts for reusable runtime +resources. This change adds the model-weight manifest contract; planning, +storage, and transfer execution are added separately. + +Framework-owned adapters outside Mooncake inspect framework runtime objects, +normalize framework-specific values, and construct the typed canonical +manifests. Mooncake core accepts only those typed values; it does not import or +inspect framework objects or accept alternate field names or duck-typed +records. + +The public Python API is split by responsibility: + +- `mooncake.reshard.contracts` exposes `ResourceManifest`, + `PlacementManifest`, and `RuntimeBindingManifest` as structural `Protocol` + contracts for resource-neutral identity and lifecycle; +- `mooncake.reshard.weight` defines model-weight placement and runtime binding. + +## Weight Placement Model + +`WeightPlacementManifest` describes one complete, address-free global logical +placement of a model-weight generation. It contains: + +- a `ParallelTopology` with TP, PP, EP, and DP sizes plus the exact selected + participants; +- per-tensor `SplitAxis(kind, dim)`, `ReplicatedAxis(kind)`, and + `OwnershipAxis(kind)` entries that distinguish logical sharding, complete + replicas, and ownership without overloading an optional dimension; +- canonical `TensorDescriptor` values whose only shard representation is + `shard_dims`; +- one `WeightPlacementPart` for every selected participant; +- canonical global tensor descriptors and N-D logical fragments; +- a placement ID and digest computed after the full part set validates. + +`ParallelTopology.world_size` is the selected participant count. It is not +inferred from `tp_size * pp_size * ep_size * dp_size`: parallel axes may share +workers, and a placement may select one complete DP replica while retaining the +runtime's declared `dp_size`. The overall participant map may be non-Cartesian, +but a tensor that declares multiple independent `SplitAxis` values must provide +the Cartesian rank combinations needed to prove each axis-to-dimension split. + +For each framework participant, the framework adapter first constructs an +address-free `WeightPlacementPart`. A collection barrier assembles the exact +participant set and validates complete logical tensor coverage. Each part +declares exactly the tensor descriptors referenced by its fragments. For each +live participant that owns fragments, the adapter then constructs one +`WeightRuntimeBindingManifest` with physical fragments, generation, and lease, +attesting the global placement ID and digest. A physical fragment preserves its +item size, view shape, byte strides, storage base, byte offset, and allocation +size so binding validation can prove canonical contiguity and address bounds. + +An alias group may span placement parts, so an individual part checks only its +local fragment invariants. `WeightPlacementManifest` performs the global check +after collection: every alias member must be in the complete tensor catalog and +every fragment for every member must declare the identical group. Runtime paths +consume only this globally validated placement. + +Empty participants need no runtime binding; any participant referenced by +execution must provide one. + +The weight implementation is split by responsibility: + +- `types.py` defines tensor and logical-fragment contracts; +- `topology.py` defines parallel sizes and selected participants; +- `part.py` defines one participant's address-free contribution; +- `placement.py` assembles and identifies the complete global placement; +- `runtime.py` defines typed physical bindings; +- `validation.py` checks logical geometry, coverage, declared storage alias + groups, and addresses; +- `binding.py` validates placement and binding attestation; +- `manifest.py` preserves the public import surface. + +`kv_cache` is reserved as a resource discriminator, but this change does not +define a KVCache manifest. Framework adapters must provide tensor semantics; +Mooncake does not infer them from parameter names. + +`weight_placement_to_json` and `weight_placement_from_json` are the explicit +public JSON APIs. Their wire format contains only canonical fields, and +deserialization rejects alternate field names rather than translating +framework-specific input. + +Run the contract and static type checks from the repository root: + +```bash +PYTHONPATH=mooncake-wheel:mooncake-reshard/python \ +python -m pytest -q mooncake-reshard/tests + +bash scripts/check_reshard_types.sh +``` diff --git a/mooncake-reshard/conftest.py b/mooncake-reshard/conftest.py new file mode 100644 index 0000000000..872c9f374b --- /dev/null +++ b/mooncake-reshard/conftest.py @@ -0,0 +1,8 @@ +from pathlib import Path + +import mooncake + + +RESHARD_PACKAGE = str(Path(__file__).parent / "python" / "mooncake") +if RESHARD_PACKAGE not in mooncake.__path__: + mooncake.__path__.append(RESHARD_PACKAGE) diff --git a/mooncake-reshard/pyrightconfig.json b/mooncake-reshard/pyrightconfig.json new file mode 100644 index 0000000000..e59a8aae8f --- /dev/null +++ b/mooncake-reshard/pyrightconfig.json @@ -0,0 +1,19 @@ +{ + "include": [ + "python/mooncake/reshard/contracts", + "python/mooncake/reshard/weight/types.py", + "python/mooncake/reshard/weight/topology.py", + "python/mooncake/reshard/weight/part.py", + "python/mooncake/reshard/weight/placement.py", + "python/mooncake/reshard/weight/runtime.py", + "python/mooncake/reshard/weight/validation.py", + "python/mooncake/reshard/weight/binding.py", + "python/mooncake/reshard/weight/serde.py" + ], + "extraPaths": ["python"], + "pythonVersion": "3.10", + "typeCheckingMode": "strict", + "reportPrivateUsage": "none", + "reportUnnecessaryIsInstance": "none", + "reportUnusedFunction": "none" +} diff --git a/mooncake-reshard/python/mooncake/__init__.py b/mooncake-reshard/python/mooncake/__init__.py new file mode 100644 index 0000000000..3de8145264 --- /dev/null +++ b/mooncake-reshard/python/mooncake/__init__.py @@ -0,0 +1,5 @@ +"""Mooncake split-package namespace for source-tree development.""" + +from pkgutil import extend_path + +__path__ = extend_path(__path__, __name__) diff --git a/mooncake-reshard/python/mooncake/reshard/__init__.py b/mooncake-reshard/python/mooncake/reshard/__init__.py new file mode 100644 index 0000000000..665868acb0 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/__init__.py @@ -0,0 +1,15 @@ +"""Framework-neutral contracts for reusable model runtime resources.""" + +from .contracts import ( + ResourceKind, + ResourceManifest, + PlacementManifest, + RuntimeBindingManifest, +) + +__all__ = [ + "ResourceKind", + "ResourceManifest", + "PlacementManifest", + "RuntimeBindingManifest", +] diff --git a/mooncake-reshard/python/mooncake/reshard/contracts/__init__.py b/mooncake-reshard/python/mooncake/reshard/contracts/__init__.py new file mode 100644 index 0000000000..9e2410725e --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/contracts/__init__.py @@ -0,0 +1,41 @@ +"""Public resource-neutral contracts for Mooncake resharding.""" + +from .ids import ( + LeaseId, + ParticipantId, + PlacementFragmentId, + PlacementId, + PlacementSetId, + ResourceId, + RevisionId, + RuntimeFragmentId, + RuntimeInstanceId, + TensorId, + TopologyId, +) +from .manifest import ( + PlacementManifest, + ResourceKind, + ResourceManifest, + RuntimeBindingManifest, + validate_resource_binding_identity, +) + +__all__ = [ + "LeaseId", + "ParticipantId", + "PlacementFragmentId", + "PlacementId", + "PlacementManifest", + "PlacementSetId", + "ResourceId", + "ResourceKind", + "ResourceManifest", + "RevisionId", + "RuntimeBindingManifest", + "RuntimeFragmentId", + "RuntimeInstanceId", + "TensorId", + "TopologyId", + "validate_resource_binding_identity", +] diff --git a/mooncake-reshard/python/mooncake/reshard/contracts/ids.py b/mooncake-reshard/python/mooncake/reshard/contracts/ids.py new file mode 100644 index 0000000000..eb48962b06 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/contracts/ids.py @@ -0,0 +1,37 @@ +"""Static identity categories shared by canonical reshard contracts. + +``NewType`` preserves the string wire representation while allowing checked +Python callers to distinguish unrelated resource, placement, and runtime IDs. +Wire decoders and framework adapters construct these values at their boundary. +""" + +from __future__ import annotations + +from typing import NewType + +ResourceId = NewType("ResourceId", str) +PlacementId = NewType("PlacementId", str) +ParticipantId = NewType("ParticipantId", str) +PlacementFragmentId = NewType("PlacementFragmentId", str) +RuntimeFragmentId = NewType("RuntimeFragmentId", str) +TensorId = NewType("TensorId", str) +TopologyId = NewType("TopologyId", str) +PlacementSetId = NewType("PlacementSetId", str) +RevisionId = NewType("RevisionId", str) +RuntimeInstanceId = NewType("RuntimeInstanceId", str) +LeaseId = NewType("LeaseId", str) + + +__all__ = [ + "LeaseId", + "ParticipantId", + "PlacementFragmentId", + "PlacementId", + "PlacementSetId", + "ResourceId", + "RevisionId", + "RuntimeFragmentId", + "RuntimeInstanceId", + "TensorId", + "TopologyId", +] diff --git a/mooncake-reshard/python/mooncake/reshard/contracts/manifest.py b/mooncake-reshard/python/mooncake/reshard/contracts/manifest.py new file mode 100644 index 0000000000..bea4798ca2 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/contracts/manifest.py @@ -0,0 +1,88 @@ +"""Structural manifest identity and lifecycle contracts.""" + +from __future__ import annotations + +from enum import Enum +from typing import Protocol + +from .ids import LeaseId, PlacementId, ResourceId, RuntimeInstanceId + + +class ResourceKind(str, Enum): + """Stable discriminator for a reusable runtime resource.""" + + MODEL_WEIGHT = "model_weight" + KV_CACHE = "kv_cache" + + +class ResourceManifest(Protocol): + """Structural contract for a typed reusable resource.""" + + @property + def resource_id(self) -> ResourceId: + """Return the stable canonical resource identity.""" + ... + + @property + def resource_kind(self) -> ResourceKind: + """Return the stable discriminator for the concrete resource.""" + ... + + +class PlacementManifest(ResourceManifest, Protocol): + """Address-free logical placement shared by reusable resources.""" + + @property + def placement_id(self) -> PlacementId: + """Return the stable logical placement identity.""" + ... + + @property + def digest(self) -> str: + """Return the digest that attests the serialized placement.""" + ... + + +class RuntimeBindingManifest(ResourceManifest, Protocol): + """Physical locations and lifetime fences for one placement.""" + + @property + def placement_id(self) -> PlacementId: + """Return the logical placement this runtime state attests.""" + ... + + @property + def placement_digest(self) -> str: + """Return the digest of the attested logical placement.""" + ... + + @property + def instance_id(self) -> RuntimeInstanceId: + """Return the runtime instance that owns this binding.""" + ... + + @property + def generation(self) -> int: + """Return the runtime generation fence.""" + ... + + @property + def lease_id(self) -> LeaseId: + """Return the live lease fence.""" + ... + + +def validate_resource_binding_identity( + placement: PlacementManifest, + binding: RuntimeBindingManifest, +) -> None: + """Fence a physical binding to the exact typed logical placement.""" + + if placement.resource_kind != binding.resource_kind: + raise ValueError("placement and runtime binding resource_kind differ") + if placement.resource_id != binding.resource_id: + raise ValueError("placement and runtime binding resource_id differ") + if placement.placement_id != binding.placement_id: + raise ValueError("placement_id and runtime binding placement_id differ") + if placement.digest != binding.placement_digest: + raise ValueError("placement digest and runtime binding placement digest differ") diff --git a/mooncake-reshard/python/mooncake/reshard/weight/__init__.py b/mooncake-reshard/python/mooncake/reshard/weight/__init__.py new file mode 100644 index 0000000000..2b21170f54 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/__init__.py @@ -0,0 +1,38 @@ +"""Public contracts for framework-neutral model-weight resharding.""" + +from .manifest import ( + OwnershipAxis, + ParallelRank, + ParallelTopology, + PlacementFragment, + ReplicatedAxis, + RuntimeBindingFragment, + SplitAxis, + TensorDescriptor, + TopologyParticipant, + WeightPlacementManifest, + WeightPlacementPart, + WeightRuntimeBindingManifest, + validate_runtime_binding, + validate_runtime_bindings, +) +from .serde import weight_placement_from_json, weight_placement_to_json + +__all__ = [ + "ParallelRank", + "ParallelTopology", + "PlacementFragment", + "WeightPlacementManifest", + "WeightPlacementPart", + "RuntimeBindingFragment", + "WeightRuntimeBindingManifest", + "SplitAxis", + "ReplicatedAxis", + "OwnershipAxis", + "TensorDescriptor", + "TopologyParticipant", + "validate_runtime_binding", + "validate_runtime_bindings", + "weight_placement_from_json", + "weight_placement_to_json", +] diff --git a/mooncake-reshard/python/mooncake/reshard/weight/binding.py b/mooncake-reshard/python/mooncake/reshard/weight/binding.py new file mode 100644 index 0000000000..e06e733354 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/binding.py @@ -0,0 +1,142 @@ +"""Validation between logical weight placement and runtime locations.""" + +from __future__ import annotations + +from ..contracts import RuntimeInstanceId, validate_resource_binding_identity +from .placement import WeightPlacementManifest +from .runtime import WeightRuntimeBindingManifest +from .types import canonical_strides_bytes, require_manifest_items +from .validation import _validate_runtime_binding_address_ranges + + +def validate_runtime_binding( + placement: WeightPlacementManifest, + binding: WeightRuntimeBindingManifest, +) -> None: + """Validate an ephemeral binding against one exact logical placement.""" + + if not isinstance(placement, WeightPlacementManifest): + raise ValueError("placement must be a WeightPlacementManifest") # noqa: TRY004 + if not isinstance(binding, WeightRuntimeBindingManifest): + raise ValueError("binding must be a WeightRuntimeBindingManifest") # noqa: TRY004 + validate_resource_binding_identity(placement, binding) + if placement.revision != binding.revision: + raise ValueError("placement and runtime binding revision differ") + + try: + placement_part = next( + part + for part in placement.parts + if part.participant_id == binding.participant_id + ) + except StopIteration as error: + raise ValueError( + f"unknown runtime binding participant: {binding.participant_id}" + ) from error + + placement_by_id = { + fragment.placement_fragment_id: fragment + for fragment in placement_part.fragments + } + binding_by_id = { + fragment.placement_fragment_id: fragment for fragment in binding.fragments + } + tensor_by_id = {tensor.tensor_id: tensor for tensor in placement.tensors} + unknown = sorted(binding_by_id.keys() - placement_by_id.keys()) + if unknown: + raise ValueError(f"unknown placement fragment in runtime binding: {unknown[0]}") + missing = sorted(placement_by_id.keys() - binding_by_id.keys()) + if missing: + raise ValueError(f"missing placement fragment in runtime binding: {missing[0]}") + + for placement_fragment in placement_part.fragments: + runtime_fragment = binding_by_id[placement_fragment.placement_fragment_id] + if runtime_fragment.nbytes != placement_fragment.nbytes: + raise ValueError( + "runtime binding byte size does not match placement: " + f"{placement_fragment.placement_fragment_id}" + ) + if runtime_fragment.local_shape != placement_fragment.local_shape: + raise ValueError( + "runtime binding local_shape does not match placement: " + f"{placement_fragment.placement_fragment_id}" + ) + tensor = tensor_by_id[placement_fragment.tensor_id] + if runtime_fragment.itemsize != tensor.itemsize: + raise ValueError( + "runtime binding itemsize does not match placement: " + f"{placement_fragment.placement_fragment_id}" + ) + expected_strides_bytes = canonical_strides_bytes( + placement_fragment.local_shape, + tensor.itemsize, + ) + if runtime_fragment.strides_bytes != expected_strides_bytes: + raise ValueError( + "runtime binding stride does not describe a canonical " + "contiguous view: " + f"{placement_fragment.placement_fragment_id}" + ) + + _validate_runtime_binding_address_ranges( + instance_id=binding.instance_id, + tensors=placement.tensors, + placements=placement_part.fragments, + bindings=binding.fragments, + ) + + +def _validate_runtime_binding_subset( + placement: WeightPlacementManifest, + bindings: object, +) -> tuple[WeightRuntimeBindingManifest, ...]: + """Validate one or more participants and their shared address spaces.""" + + items = require_manifest_items( + bindings, + "runtime bindings", + WeightRuntimeBindingManifest, + ) + participant_ids = [item.participant_id for item in items] + if len(participant_ids) != len(set(participant_ids)): + raise ValueError("duplicate runtime binding participant") + for binding in items: + validate_runtime_binding(placement, binding) + + part_by_participant = {part.participant_id: part for part in placement.parts} + by_instance: dict[RuntimeInstanceId, list[WeightRuntimeBindingManifest]] = {} + for binding in items: + by_instance.setdefault(binding.instance_id, []).append(binding) + for instance_id, instance_bindings in by_instance.items(): + _validate_runtime_binding_address_ranges( + instance_id=instance_id, + tensors=placement.tensors, + placements=tuple( + fragment + for binding in instance_bindings + for fragment in part_by_participant[binding.participant_id].fragments + ), + bindings=tuple( + fragment + for binding in instance_bindings + for fragment in binding.fragments + ), + ) + return items + + +def validate_runtime_bindings( + placement: WeightPlacementManifest, + bindings: object, +) -> None: + """Validate the exact runtime binding set for a complete placement.""" + + items = _validate_runtime_binding_subset(placement, bindings) + expected = {part.participant_id for part in placement.parts if part.fragments} + actual = {item.participant_id for item in items} + missing = sorted(expected - actual) + if missing: + raise ValueError(f"missing runtime binding participant: {missing[0]}") + unknown = sorted(actual - expected) + if unknown: + raise ValueError(f"unknown runtime binding participant: {unknown[0]}") diff --git a/mooncake-reshard/python/mooncake/reshard/weight/manifest.py b/mooncake-reshard/python/mooncake/reshard/weight/manifest.py new file mode 100644 index 0000000000..673e2b23a1 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/manifest.py @@ -0,0 +1,34 @@ +"""Stable public facade for model-weight manifest contracts.""" + +from .binding import validate_runtime_binding, validate_runtime_bindings +from .part import WeightPlacementPart +from .placement import WeightPlacementManifest +from .runtime import WeightRuntimeBindingManifest +from .topology import ParallelTopology, TopologyParticipant +from .types import ( + OwnershipAxis, + ParallelRank, + PlacementFragment, + ReplicatedAxis, + RuntimeBindingFragment, + SplitAxis, + TensorDescriptor, +) + + +__all__ = [ + "ParallelRank", + "ParallelTopology", + "PlacementFragment", + "TopologyParticipant", + "WeightPlacementManifest", + "WeightPlacementPart", + "RuntimeBindingFragment", + "WeightRuntimeBindingManifest", + "SplitAxis", + "ReplicatedAxis", + "OwnershipAxis", + "TensorDescriptor", + "validate_runtime_binding", + "validate_runtime_bindings", +] diff --git a/mooncake-reshard/python/mooncake/reshard/weight/part.py b/mooncake-reshard/python/mooncake/reshard/weight/part.py new file mode 100644 index 0000000000..c7cb5947dc --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/part.py @@ -0,0 +1,86 @@ +"""Per-participant logical weight placement contracts.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ..contracts import ( + ParticipantId, + PlacementSetId, + ResourceId, + RevisionId, + TopologyId, +) +from .types import ( + ParallelRank, + PlacementFragment, + TensorDescriptor, + _require_nonempty_string, + _require_u64, + require_manifest_items, +) +from .validation import _validate_fragments + + +@dataclass(frozen=True) +class WeightPlacementPart: + """One participant's address-free contribution to a global placement.""" + + resource_id: ResourceId + revision: RevisionId + weight_generation: int + placement_set_id: PlacementSetId + topology_id: TopologyId + participant_id: ParticipantId + rank: ParallelRank + tensors: tuple[TensorDescriptor, ...] + fragments: tuple[PlacementFragment, ...] + + def __post_init__(self) -> None: + for name in ( + "resource_id", + "revision", + "placement_set_id", + "topology_id", + "participant_id", + ): + _require_nonempty_string(getattr(self, name), name) + _require_u64(self.weight_generation, "weight_generation") + if not isinstance(self.rank, ParallelRank): + raise ValueError("placement part rank must be a ParallelRank") # noqa: TRY004 + + tensors = require_manifest_items( + self.tensors, + "WeightPlacementPart tensors", + TensorDescriptor, + ) + fragments = require_manifest_items( + self.fragments, + "WeightPlacementPart fragments", + PlacementFragment, + ) + if any(fragment.rank != self.rank for fragment in fragments): + raise ValueError("placement part fragment rank differs from part rank") + referenced_tensor_ids = {fragment.tensor_id for fragment in fragments} + unreferenced_tensor_ids = sorted( + {tensor.tensor_id for tensor in tensors} - referenced_tensor_ids + ) + if unreferenced_tensor_ids: + raise ValueError( + "placement part contains an unreferenced tensor: " + f"{unreferenced_tensor_ids[0]}" + ) + object.__setattr__( + self, + "tensors", + tuple(sorted(tensors, key=lambda item: item.tensor_id)), + ) + object.__setattr__( + self, + "fragments", + tuple(sorted(fragments, key=lambda item: item.placement_fragment_id)), + ) + _validate_fragments(self.tensors, self.fragments) + + +__all__ = ["WeightPlacementPart"] diff --git a/mooncake-reshard/python/mooncake/reshard/weight/placement.py b/mooncake-reshard/python/mooncake/reshard/weight/placement.py new file mode 100644 index 0000000000..babd79ebcf --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/placement.py @@ -0,0 +1,341 @@ +"""Complete logical weight placement and canonical identity.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Sequence +from dataclasses import asdict, dataclass, field + +from ..contracts import ( + ParticipantId, + PlacementId, + PlacementSetId, + ResourceId, + ResourceKind, + RevisionId, + TensorId, +) +from .part import WeightPlacementPart +from .topology import ParallelTopology +from .types import ( + OwnershipAxis, + PlacementFragment, + ReplicatedAxis, + SplitAxis, + TensorDescriptor, + _require_nonempty_string, + _require_u64, + require_manifest_items, +) +from .validation import ( + _validate_complete_weight_placement, + _validate_fragments, +) + + +@dataclass(frozen=True, init=False) +class WeightPlacementManifest: + """One complete address-free placement of a model-weight generation.""" + + resource_id: ResourceId + placement_id: PlacementId + revision: RevisionId + weight_generation: int + placement_set_id: PlacementSetId + topology: ParallelTopology + parts: tuple[WeightPlacementPart, ...] + tensors: tuple[TensorDescriptor, ...] + fragments: tuple[PlacementFragment, ...] + _digest_cache: str | None = field(init=False, repr=False, compare=False) + + def __init__( + self, + *, + resource_id: ResourceId, + revision: RevisionId, + weight_generation: int, + placement_set_id: PlacementSetId, + topology: ParallelTopology, + parts: tuple[WeightPlacementPart, ...], + placement_id: PlacementId | None = None, + ) -> None: + _require_nonempty_string(resource_id, "resource_id") + _require_nonempty_string(revision, "revision") + _require_u64(weight_generation, "weight_generation") + _require_nonempty_string(placement_set_id, "placement_set_id") + if not isinstance(topology, ParallelTopology): + raise ValueError("topology must be a ParallelTopology") # noqa: TRY004 + if placement_id is not None: + _require_nonempty_string(placement_id, "placement_id") + + normalized_parts = tuple( + sorted( + require_manifest_items( + parts, + "WeightPlacementManifest parts", + WeightPlacementPart, + ), + key=lambda item: item.participant_id, + ) + ) + _validate_parts( + normalized_parts, + topology=topology, + resource_id=resource_id, + revision=revision, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + ) + + tensors_by_id: dict[TensorId, TensorDescriptor] = {} + collected_fragments: list[PlacementFragment] = [] + for part in normalized_parts: + for tensor in part.tensors: + previous = tensors_by_id.setdefault(tensor.tensor_id, tensor) + if previous != tensor: + raise ValueError( + f"placement part tensor descriptor mismatch: {tensor.tensor_id}" + ) + collected_fragments.extend(part.fragments) + tensors = tuple(sorted(tensors_by_id.values(), key=lambda item: item.tensor_id)) + fragments = tuple( + sorted(collected_fragments, key=lambda item: item.placement_fragment_id) + ) + _validate_fragments( + tensors, + fragments, + require_complete_alias_groups=True, + ) + _validate_complete_weight_placement( + tensors, + fragments, + topology=topology, + ) + canonical_placement_id = _logical_placement_id( + resource_id=resource_id, + revision=revision, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + topology=topology, + tensors=tensors, + parts=normalized_parts, + ) + if placement_id is not None and placement_id != canonical_placement_id: + raise ValueError("placement_id does not match canonical logical content") + + object.__setattr__(self, "resource_id", resource_id) + object.__setattr__(self, "placement_id", canonical_placement_id) + object.__setattr__(self, "revision", revision) + object.__setattr__(self, "weight_generation", weight_generation) + object.__setattr__(self, "placement_set_id", placement_set_id) + object.__setattr__(self, "topology", topology) + object.__setattr__(self, "parts", normalized_parts) + object.__setattr__(self, "tensors", tensors) + object.__setattr__(self, "fragments", fragments) + object.__setattr__(self, "_digest_cache", None) + + @property + def resource_kind(self) -> ResourceKind: + """Identify this placement as model weight data.""" + + return ResourceKind.MODEL_WEIGHT + + @property + def digest(self) -> str: + """Return the stable SHA-256 digest of the canonical JSON form.""" + + digest = self._digest_cache + if digest is None: + from .serde import weight_placement_to_json + + digest = hashlib.sha256(weight_placement_to_json(self).encode()).hexdigest() + object.__setattr__(self, "_digest_cache", digest) + return digest + + @classmethod + def from_fragments( + cls, + *, + resource_id: ResourceId, + revision: RevisionId, + weight_generation: int, + placement_set_id: PlacementSetId, + topology: ParallelTopology, + tensors: Sequence[TensorDescriptor], + fragments: Sequence[PlacementFragment], + placement_id: PlacementId | None = None, + ) -> WeightPlacementManifest: + """Group a complete flat fragment inventory by topology participant.""" + + tensor_items = require_manifest_items( + tensors, "placement tensors", TensorDescriptor + ) + fragment_items = require_manifest_items( + fragments, "placement fragments", PlacementFragment + ) + tensor_by_id = {tensor.tensor_id: tensor for tensor in tensor_items} + if len(tensor_by_id) != len(tensor_items): + raise ValueError("duplicate tensor_id in placement tensors") + referenced_tensor_ids = {fragment.tensor_id for fragment in fragment_items} + unknown_tensor_ids = sorted(referenced_tensor_ids - set(tensor_by_id)) + if unknown_tensor_ids: + raise ValueError(f"unknown tensor_id: {unknown_tensor_ids[0]}") + unreferenced_tensor_ids = sorted(set(tensor_by_id) - referenced_tensor_ids) + if unreferenced_tensor_ids: + raise ValueError( + "global placement contains an unreferenced tensor: " + f"{unreferenced_tensor_ids[0]}" + ) + participant_by_rank = { + participant.rank: participant for participant in topology.participants + } + fragments_by_participant: dict[ParticipantId, list[PlacementFragment]] = { + participant.participant_id: [] for participant in topology.participants + } + for fragment in fragment_items: + participant = participant_by_rank.get(fragment.rank) + if participant is None: + raise ValueError( + "placement fragment rank has no topology participant: " + f"{fragment.fragment_id}" + ) + fragments_by_participant[participant.participant_id].append(fragment) + + parts: list[WeightPlacementPart] = [] + for participant in topology.participants: + local_fragments = tuple( + fragments_by_participant[participant.participant_id] + ) + local_tensor_ids = {fragment.tensor_id for fragment in local_fragments} + parts.append( + WeightPlacementPart( + resource_id=resource_id, + revision=revision, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + topology_id=topology.topology_id, + participant_id=participant.participant_id, + rank=participant.rank, + tensors=tuple( + tensor_by_id[tensor_id] + for tensor_id in sorted(local_tensor_ids) + ), + fragments=local_fragments, + ) + ) + return cls( + resource_id=resource_id, + revision=revision, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + topology=topology, + parts=tuple(parts), + placement_id=placement_id, + ) + + +def _validate_parts( + parts: tuple[WeightPlacementPart, ...], + *, + topology: ParallelTopology, + resource_id: ResourceId, + revision: RevisionId, + weight_generation: int, + placement_set_id: PlacementSetId, +) -> None: + participant_ids = [part.participant_id for part in parts] + if len(participant_ids) != len(set(participant_ids)): + raise ValueError("duplicate placement participant") + + expected = { + participant.participant_id: participant for participant in topology.participants + } + actual = set(participant_ids) + missing = sorted(set(expected) - actual) + if missing: + raise ValueError(f"missing topology participant: {missing[0]}") + unknown = sorted(actual - set(expected)) + if unknown: + raise ValueError(f"unknown topology participant: {unknown[0]}") + + for part in parts: + if part.resource_id != resource_id: + raise ValueError("placement part resource_id differs") + if part.revision != revision: + raise ValueError("placement part revision differs") + if part.weight_generation != weight_generation: + raise ValueError("placement part weight_generation differs") + if part.placement_set_id != placement_set_id: + raise ValueError("placement part placement_set_id differs") + if part.topology_id != topology.topology_id: + raise ValueError("placement part topology_id differs") + if part.rank != expected[part.participant_id].rank: + raise ValueError("placement part rank differs from topology") + + +def _canonical_json_digest(value: object) -> str: + encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _logical_placement_id( + *, + resource_id: ResourceId, + revision: RevisionId, + weight_generation: int, + placement_set_id: PlacementSetId, + topology: ParallelTopology, + tensors: Sequence[TensorDescriptor], + parts: Sequence[WeightPlacementPart], +) -> PlacementId: + content = { + "schema": "complete-weight-placement", + "resource_id": resource_id, + "revision": revision, + "weight_generation": weight_generation, + "placement_set_id": placement_set_id, + "topology": { + "tp_size": topology.tp_size, + "pp_size": topology.pp_size, + "ep_size": topology.ep_size, + "dp_size": topology.dp_size, + "topology_id": topology.topology_id, + "participants": [asdict(item) for item in topology.participants], + }, + "tensors": [ + { + "tensor_id": tensor.tensor_id, + "global_shape": tensor.global_shape, + "dtype": tensor.dtype, + "itemsize": tensor.itemsize, + "shard_dims": tensor.shard_dims, + "layout_fingerprint": tensor.layout_fingerprint, + "parallel_axes": [ + _parallel_axis_identity(axis) for axis in tensor.parallel_axes + ], + "layer_id": tensor.layer_id, + "expert_id": tensor.expert_id, + } + for tensor in tensors + ], + "parts": [ + { + "participant_id": part.participant_id, + "rank": asdict(part.rank), + "fragments": [asdict(fragment) for fragment in part.fragments], + } + for part in parts + ], + } + return PlacementId(f"sha256:{_canonical_json_digest(content)}") + + +def _parallel_axis_identity(axis: object) -> dict[str, object]: + if isinstance(axis, SplitAxis): + return {"semantics": "split", "kind": axis.kind, "dim": axis.dim} + if isinstance(axis, ReplicatedAxis): + return {"semantics": "replicated", "kind": axis.kind} + if isinstance(axis, OwnershipAxis): + return {"semantics": "ownership", "kind": axis.kind} + raise ValueError("unsupported parallel axis value") diff --git a/mooncake-reshard/python/mooncake/reshard/weight/runtime.py b/mooncake-reshard/python/mooncake/reshard/weight/runtime.py new file mode 100644 index 0000000000..f972c8d4e2 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/runtime.py @@ -0,0 +1,80 @@ +"""Canonical ephemeral runtime binding contracts.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ..contracts import ( + LeaseId, + ParticipantId, + PlacementId, + ResourceId, + ResourceKind, + RevisionId, + RuntimeInstanceId, +) +from .types import ( + RuntimeBindingFragment, + _require_nonempty_string, + _require_u64, + require_manifest_items, + require_sha256_digest, +) + + +@dataclass(frozen=True) +class WeightRuntimeBindingManifest: + """Ephemeral physical locations and lifetime fence for one placement.""" + + resource_id: ResourceId + placement_id: PlacementId + placement_digest: str + instance_id: RuntimeInstanceId + generation: int + lease_id: LeaseId + revision: RevisionId + participant_id: ParticipantId + fragments: tuple[RuntimeBindingFragment, ...] + + @property + def resource_kind(self) -> ResourceKind: + """Identify this binding as model weight data.""" + + return ResourceKind.MODEL_WEIGHT + + def __post_init__(self) -> None: + for name in ( + "resource_id", + "placement_id", + "instance_id", + "lease_id", + "revision", + "participant_id", + ): + _require_nonempty_string(getattr(self, name), name) + require_sha256_digest(self.placement_digest, "placement_digest") + _require_u64(self.generation, "generation") + fragments = require_manifest_items( + self.fragments, + "WeightRuntimeBindingManifest fragments", + RuntimeBindingFragment, + ) + object.__setattr__( + self, + "fragments", + tuple( + sorted( + fragments, + key=lambda item: item.placement_fragment_id, + ) + ), + ) + placement_ids = [item.placement_fragment_id for item in self.fragments] + if len(placement_ids) != len(set(placement_ids)): + raise ValueError("duplicate placement fragment in runtime binding") + fragment_ids = [item.fragment_id for item in self.fragments] + if len(fragment_ids) != len(set(fragment_ids)): + raise ValueError("duplicate runtime fragment_id in runtime binding") + + +__all__ = ["WeightRuntimeBindingManifest"] diff --git a/mooncake-reshard/python/mooncake/reshard/weight/serde.py b/mooncake-reshard/python/mooncake/reshard/weight/serde.py new file mode 100644 index 0000000000..04433976fd --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/serde.py @@ -0,0 +1,480 @@ +"""Canonical JSON serialization for model-weight placement manifests. + +This is the only untyped wire boundary in the manifest package. It accepts JSON +values, rejects unknown or malformed fields, and returns only canonical typed +domain values. Framework runtime objects are deliberately not accepted here. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from collections.abc import Set as AbstractSet +from typing import TypeAlias, cast + +from ..contracts import ( + ParticipantId, + PlacementFragmentId, + PlacementId, + PlacementSetId, + ResourceId, + ResourceKind, + RevisionId, + TensorId, + TopologyId, +) +from .part import WeightPlacementPart +from .placement import WeightPlacementManifest +from .topology import ParallelTopology, TopologyParticipant +from .types import ( + OwnershipAxis, + ParallelAxis, + ParallelAxisKind, + ParallelRank, + PlacementFragment, + ReplicatedAxis, + SplitAxis, + SplitAxisKind, + TensorDescriptor, +) + +JsonScalar: TypeAlias = None | bool | int | float | str +JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"] +JsonObject: TypeAlias = Mapping[str, JsonValue] + + +def weight_placement_to_json(placement: WeightPlacementManifest) -> str: + """Serialize one typed placement using the canonical wire schema.""" + + if not isinstance(placement, WeightPlacementManifest): + raise ValueError("placement must be a WeightPlacementManifest") # noqa: TRY004 + payload: dict[str, JsonValue] = { + "resource_kind": placement.resource_kind.value, + "resource_id": placement.resource_id, + "revision": placement.revision, + "weight_generation": placement.weight_generation, + "placement_set_id": placement.placement_set_id, + "placement_id": placement.placement_id, + "topology": _topology_to_wire(placement.topology), + "tensors": [_tensor_to_wire(tensor) for tensor in placement.tensors], + "parts": [_part_to_wire(part) for part in placement.parts], + } + return json.dumps(payload, sort_keys=True, separators=(",", ":")) + + +def weight_placement_from_json(value: str) -> WeightPlacementManifest: + """Parse only the canonical model-weight placement wire schema.""" + + manifest = _require_exact_fields( + _load_json_object(value, "placement manifest"), + { + "resource_kind", + "resource_id", + "revision", + "weight_generation", + "placement_set_id", + "placement_id", + "topology", + "tensors", + "parts", + }, + "placement manifest", + ) + if _require_nonempty_string(manifest["resource_kind"], "resource_kind") != ( + ResourceKind.MODEL_WEIGHT.value + ): + raise ValueError("placement manifest resource_kind must be model_weight") + + resource_id = ResourceId( + _require_nonempty_string(manifest["resource_id"], "resource_id") + ) + revision = RevisionId(_require_nonempty_string(manifest["revision"], "revision")) + weight_generation = _require_integer( + manifest["weight_generation"], "weight_generation", minimum=0 + ) + placement_set_id = PlacementSetId( + _require_nonempty_string(manifest["placement_set_id"], "placement_set_id") + ) + placement_id = PlacementId( + _require_nonempty_string(manifest["placement_id"], "placement_id") + ) + topology = _topology_from_wire(manifest["topology"]) + tensors = tuple( + _tensor_from_wire(item, index) + for index, item in enumerate( + _require_sequence(manifest["tensors"], "placement tensors") + ) + ) + tensor_by_id = {tensor.tensor_id: tensor for tensor in tensors} + if len(tensor_by_id) != len(tensors): + raise ValueError("duplicate tensor_id in placement JSON") + + parts: list[WeightPlacementPart] = [] + referenced_tensor_ids: set[TensorId] = set() + for index, item in enumerate( + _require_sequence(manifest["parts"], "placement parts") + ): + part = _require_exact_fields( + item, + {"participant_id", "rank", "fragments"}, + f"placement part {index}", + ) + participant_id = ParticipantId( + _require_nonempty_string(part["participant_id"], "participant_id") + ) + rank = _rank_from_wire(part["rank"], f"placement part rank {index}") + fragments = tuple( + _fragment_from_wire(fragment, fragment_index) + for fragment_index, fragment in enumerate( + _require_sequence(part["fragments"], "placement part fragments") + ) + ) + local_tensor_ids = {fragment.tensor_id for fragment in fragments} + referenced_tensor_ids.update(local_tensor_ids) + unknown = sorted(local_tensor_ids - set(tensor_by_id)) + if unknown: + raise ValueError(f"unknown tensor_id in placement part: {unknown[0]}") + parts.append( + WeightPlacementPart( + resource_id=resource_id, + revision=revision, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + topology_id=topology.topology_id, + participant_id=participant_id, + rank=rank, + tensors=tuple( + tensor_by_id[tensor_id] for tensor_id in sorted(local_tensor_ids) + ), + fragments=fragments, + ) + ) + + unreferenced_tensor_ids = sorted(set(tensor_by_id) - referenced_tensor_ids) + if unreferenced_tensor_ids: + raise ValueError( + "placement JSON contains an unreferenced tensor: " + f"{unreferenced_tensor_ids[0]}" + ) + return WeightPlacementManifest( + resource_id=resource_id, + revision=revision, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + placement_id=placement_id, + topology=topology, + parts=tuple(parts), + ) + + +def _part_to_wire(part: WeightPlacementPart) -> dict[str, JsonValue]: + return { + "participant_id": part.participant_id, + "rank": _rank_to_wire(part.rank), + "fragments": [_fragment_to_wire(fragment) for fragment in part.fragments], + } + + +def _topology_to_wire(topology: ParallelTopology) -> dict[str, JsonValue]: + return { + "tp_size": topology.tp_size, + "pp_size": topology.pp_size, + "ep_size": topology.ep_size, + "dp_size": topology.dp_size, + "topology_id": topology.topology_id, + "participants": [ + { + "participant_id": participant.participant_id, + "rank": _rank_to_wire(participant.rank), + } + for participant in topology.participants + ], + } + + +def _topology_from_wire(value: object) -> ParallelTopology: + topology = _require_exact_fields( + value, + {"tp_size", "pp_size", "ep_size", "dp_size", "topology_id", "participants"}, + "parallel topology", + ) + participants: list[TopologyParticipant] = [] + for index, item in enumerate( + _require_sequence(topology["participants"], "parallel topology participants") + ): + participant = _require_exact_fields( + item, + {"participant_id", "rank"}, + f"parallel topology participant {index}", + ) + participants.append( + TopologyParticipant( + participant_id=ParticipantId( + _require_nonempty_string( + participant["participant_id"], "participant_id" + ) + ), + rank=_rank_from_wire( + participant["rank"], f"parallel topology rank {index}" + ), + ) + ) + return ParallelTopology( + tp_size=_require_integer(topology["tp_size"], "tp_size", minimum=1), + pp_size=_require_integer(topology["pp_size"], "pp_size", minimum=1), + ep_size=_require_integer(topology["ep_size"], "ep_size", minimum=1), + dp_size=_require_integer(topology["dp_size"], "dp_size", minimum=1), + participants=tuple(participants), + topology_id=TopologyId( + _require_nonempty_string(topology["topology_id"], "topology_id") + ), + ) + + +def _tensor_to_wire(tensor: TensorDescriptor) -> dict[str, JsonValue]: + return { + "tensor_id": tensor.tensor_id, + "global_shape": list(tensor.global_shape), + "dtype": tensor.dtype, + "itemsize": tensor.itemsize, + "shard_dims": list(tensor.shard_dims), + "layout_fingerprint": tensor.layout_fingerprint, + "parallel_axes": [_axis_to_wire(axis) for axis in tensor.parallel_axes], + "layer_id": tensor.layer_id, + "expert_id": tensor.expert_id, + } + + +def _tensor_from_wire(value: object, index: int) -> TensorDescriptor: + tensor = _require_exact_fields( + value, + { + "tensor_id", + "global_shape", + "dtype", + "itemsize", + "shard_dims", + "layout_fingerprint", + "parallel_axes", + "layer_id", + "expert_id", + }, + f"placement tensor {index}", + ) + return TensorDescriptor( + tensor_id=TensorId(_require_nonempty_string(tensor["tensor_id"], "tensor_id")), + global_shape=_integer_tuple(tensor["global_shape"], "global_shape", minimum=1), + dtype=_require_nonempty_string(tensor["dtype"], "dtype"), + itemsize=_require_integer(tensor["itemsize"], "itemsize", minimum=1), + shard_dims=_integer_tuple(tensor["shard_dims"], "shard_dims", minimum=0), + layout_fingerprint=_require_nonempty_string( + tensor["layout_fingerprint"], "layout_fingerprint" + ), + parallel_axes=tuple( + _axis_from_wire(axis, axis_index) + for axis_index, axis in enumerate( + _require_sequence( + tensor["parallel_axes"], "placement tensor parallel_axes" + ) + ) + ), + layer_id=_optional_integer(tensor["layer_id"], "layer_id", minimum=0), + expert_id=_optional_integer(tensor["expert_id"], "expert_id", minimum=0), + ) + + +def _axis_to_wire(axis: ParallelAxis) -> dict[str, JsonValue]: + if isinstance(axis, SplitAxis): + return {"semantics": "split", "kind": axis.kind, "dim": axis.dim} + if isinstance(axis, ReplicatedAxis): + return {"semantics": "replicated", "kind": axis.kind} + if isinstance(axis, OwnershipAxis): + return {"semantics": "ownership", "kind": axis.kind} + raise ValueError("unsupported parallel axis value") + + +def _axis_from_wire(value: object, index: int) -> ParallelAxis: + axis = _require_mapping(value, f"placement tensor parallel axis {index}") + semantics = _require_nonempty_string(axis.get("semantics"), "axis semantics") + if semantics == "split": + split = _require_exact_fields( + axis, + {"semantics", "kind", "dim"}, + f"placement tensor parallel axis {index}", + ) + return SplitAxis( + kind=_split_axis_kind(split["kind"]), + dim=_require_integer(split["dim"], "split axis dim", minimum=0), + ) + if semantics == "replicated": + replicated = _require_exact_fields( + axis, + {"semantics", "kind"}, + f"placement tensor parallel axis {index}", + ) + return ReplicatedAxis(kind=_parallel_axis_kind(replicated["kind"])) + if semantics == "ownership": + ownership = _require_exact_fields( + axis, + {"semantics", "kind"}, + f"placement tensor parallel axis {index}", + ) + return OwnershipAxis(kind=_parallel_axis_kind(ownership["kind"])) + raise ValueError(f"unsupported parallel axis semantics: {semantics}") + + +def _fragment_to_wire(fragment: PlacementFragment) -> dict[str, JsonValue]: + return { + "placement_fragment_id": fragment.placement_fragment_id, + "tensor_id": fragment.tensor_id, + "global_offset": list(fragment.global_offset), + "local_shape": list(fragment.local_shape), + "nbytes": fragment.nbytes, + "rank": _rank_to_wire(fragment.rank), + "aliases": list(fragment.aliases), + } + + +def _fragment_from_wire(value: object, index: int) -> PlacementFragment: + fragment = _require_exact_fields( + value, + { + "placement_fragment_id", + "tensor_id", + "global_offset", + "local_shape", + "nbytes", + "rank", + "aliases", + }, + f"placement fragment {index}", + ) + return PlacementFragment( + placement_fragment_id=PlacementFragmentId( + _require_nonempty_string( + fragment["placement_fragment_id"], "placement_fragment_id" + ) + ), + tensor_id=TensorId( + _require_nonempty_string(fragment["tensor_id"], "tensor_id") + ), + global_offset=_integer_tuple( + fragment["global_offset"], "global_offset", minimum=0 + ), + local_shape=_integer_tuple(fragment["local_shape"], "local_shape", minimum=1), + nbytes=_require_integer(fragment["nbytes"], "nbytes", minimum=1), + rank=_rank_from_wire(fragment["rank"], f"placement rank {index}"), + aliases=tuple( + TensorId(_require_nonempty_string(alias, "alias")) + for alias in _require_sequence(fragment["aliases"], "aliases") + ), + ) + + +def _rank_to_wire(rank: ParallelRank) -> dict[str, JsonValue]: + return {"dp": rank.dp, "tp": rank.tp, "pp": rank.pp, "ep": rank.ep} + + +def _rank_from_wire(value: object, label: str) -> ParallelRank: + rank = _require_exact_fields(value, {"dp", "tp", "pp", "ep"}, label) + return ParallelRank( + dp=_require_integer(rank["dp"], "rank dp", minimum=0), + tp=_require_integer(rank["tp"], "rank tp", minimum=0), + pp=_require_integer(rank["pp"], "rank pp", minimum=0), + ep=_require_integer(rank["ep"], "rank ep", minimum=0), + ) + + +def _load_json_object(value: str, label: str) -> JsonObject: + def reject_constant(constant: str) -> None: + raise ValueError(f"non-finite JSON number is unsupported: {constant}") + + def reject_duplicate_fields( + pairs: list[tuple[str, object]], + ) -> dict[str, JsonValue]: + result: dict[str, JsonValue] = {} + for key, item in pairs: + if key in result: + raise ValueError(f"duplicate JSON field: {key}") + result[key] = cast(JsonValue, item) + return result + + try: + raw = json.loads( + value, + parse_constant=reject_constant, + object_pairs_hook=reject_duplicate_fields, + ) + except (TypeError, json.JSONDecodeError) as error: + raise ValueError(f"{label} is not valid JSON") from error + return _require_mapping(cast(object, raw), label) + + +def _require_mapping(value: object, label: str) -> JsonObject: + if not isinstance(value, Mapping): + raise ValueError(f"{label} must be a JSON object") # noqa: TRY004 + mapping = cast(Mapping[object, object], value) + if any(type(key) is not str for key in mapping): + raise ValueError(f"{label} must be a JSON object") + return cast(JsonObject, mapping) + + +def _require_exact_fields( + value: object, + expected: AbstractSet[str], + label: str, +) -> JsonObject: + mapping = _require_mapping(value, label) + if set(mapping) != set(expected): + raise ValueError(f"{label} schema fields do not match contract") + return mapping + + +def _require_sequence(value: object, label: str) -> Sequence[JsonValue]: + if isinstance(value, (str, bytes, bytearray)) or not isinstance(value, Sequence): + raise ValueError(f"{label} must be a sequence") # noqa: TRY004 + return cast(Sequence[JsonValue], value) + + +def _require_nonempty_string(value: object, label: str) -> str: + if type(value) is not str or not value: + raise ValueError(f"{label} must be a non-empty string") + return value + + +def _require_integer(value: object, label: str, *, minimum: int) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{label} must be an integer at least {minimum}") + return value + + +def _optional_integer(value: object, label: str, *, minimum: int) -> int | None: + if value is None: + return None + return _require_integer(value, label, minimum=minimum) + + +def _integer_tuple( + value: object, + label: str, + *, + minimum: int, +) -> tuple[int, ...]: + return tuple( + _require_integer(item, label, minimum=minimum) + for item in _require_sequence(value, label) + ) + + +def _parallel_axis_kind(value: object) -> ParallelAxisKind: + kind = _require_nonempty_string(value, "parallel axis kind") + if kind not in {"dp", "pp", "ep", "tp"}: + raise ValueError(f"unsupported parallel axis kind: {kind}") + return cast(ParallelAxisKind, kind) + + +def _split_axis_kind(value: object) -> SplitAxisKind: + kind = _parallel_axis_kind(value) + if kind not in {"ep", "tp"}: + raise ValueError(f"{kind} cannot use split semantics") + return cast(SplitAxisKind, kind) diff --git a/mooncake-reshard/python/mooncake/reshard/weight/topology.py b/mooncake-reshard/python/mooncake/reshard/weight/topology.py new file mode 100644 index 0000000000..9db065895b --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/topology.py @@ -0,0 +1,151 @@ +"""Global parallel topology for one model-weight placement.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass + +from ..contracts import ParticipantId, TopologyId +from .types import ( + ParallelRank, + _require_integer, + _require_nonempty_string, + require_manifest_items, +) + + +@dataclass(frozen=True) +class TopologyParticipant: + """One runtime participant and its framework-defined parallel coordinates.""" + + participant_id: ParticipantId + rank: ParallelRank + + def __post_init__(self) -> None: + _require_nonempty_string(self.participant_id, "participant_id") + if not isinstance(self.rank, ParallelRank): + raise ValueError( # noqa: TRY004 + "topology participant rank must be a ParallelRank" + ) + + +@dataclass(frozen=True, init=False) +class ParallelTopology: + """Complete TP/PP/EP/DP dimensions and selected participant mapping. + + ``world_size`` is the number of explicit participants. It is deliberately + not inferred from the product of the four axis sizes because frameworks may + map TP and EP onto the same workers, and a placement may select one DP + replica while retaining the source runtime's declared ``dp_size``. + """ + + tp_size: int + pp_size: int + ep_size: int + dp_size: int + participants: tuple[TopologyParticipant, ...] + topology_id: TopologyId + + def __init__( + self, + *, + tp_size: int, + pp_size: int, + ep_size: int, + dp_size: int, + participants: tuple[TopologyParticipant, ...], + topology_id: TopologyId | None = None, + ) -> None: + for value, name in ( + (tp_size, "tp_size"), + (pp_size, "pp_size"), + (ep_size, "ep_size"), + (dp_size, "dp_size"), + ): + _require_integer(value, name, minimum=1) + normalized_participants = require_manifest_items( + participants, + "ParallelTopology participants", + TopologyParticipant, + ) + if not normalized_participants: + raise ValueError("parallel topology participants must not be empty") + normalized_participants = tuple( + sorted(normalized_participants, key=lambda item: item.participant_id) + ) + + participant_ids = [item.participant_id for item in normalized_participants] + if len(participant_ids) != len(set(participant_ids)): + raise ValueError("duplicate topology participant_id") + ranks = [item.rank for item in normalized_participants] + if len(ranks) != len(set(ranks)): + raise ValueError("duplicate topology parallel rank") + + axis_sizes = { + "tp": tp_size, + "pp": pp_size, + "ep": ep_size, + "dp": dp_size, + } + for participant in normalized_participants: + for axis, size in axis_sizes.items(): + value = getattr(participant.rank, axis) + if value >= size: + raise ValueError( + f"{axis} rank is out of range for participant " + f"{participant.participant_id}" + ) + + canonical_id = _topology_id( + tp_size=tp_size, + pp_size=pp_size, + ep_size=ep_size, + dp_size=dp_size, + participants=normalized_participants, + ) + if topology_id is not None and topology_id != canonical_id: + raise ValueError("topology_id does not match canonical topology") + object.__setattr__(self, "tp_size", tp_size) + object.__setattr__(self, "pp_size", pp_size) + object.__setattr__(self, "ep_size", ep_size) + object.__setattr__(self, "dp_size", dp_size) + object.__setattr__(self, "participants", normalized_participants) + object.__setattr__(self, "topology_id", canonical_id) + + @property + def world_size(self) -> int: + """Return the number of actual runtime participants.""" + + return len(self.participants) + + def participant(self, participant_id: ParticipantId) -> TopologyParticipant: + """Return one declared participant or reject an unknown identifier.""" + + for participant in self.participants: + if participant.participant_id == participant_id: + return participant + raise ValueError(f"unknown topology participant: {participant_id}") + + +def _topology_id( + *, + tp_size: int, + pp_size: int, + ep_size: int, + dp_size: int, + participants: tuple[TopologyParticipant, ...], +) -> TopologyId: + content = { + "schema": "weight-parallel-topology", + "tp_size": tp_size, + "pp_size": pp_size, + "ep_size": ep_size, + "dp_size": dp_size, + "participants": [asdict(item) for item in participants], + } + encoded = json.dumps(content, sort_keys=True, separators=(",", ":")).encode() + return TopologyId(f"sha256:{hashlib.sha256(encoded).hexdigest()}") + + +__all__ = ["ParallelTopology", "TopologyParticipant"] diff --git a/mooncake-reshard/python/mooncake/reshard/weight/types.py b/mooncake-reshard/python/mooncake/reshard/weight/types.py new file mode 100644 index 0000000000..43a679a352 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/types.py @@ -0,0 +1,432 @@ +"""Public model-weight tensor and fragment contracts.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Sequence +from dataclasses import asdict, dataclass, field +from typing import Literal, TypeVar, cast + +from ..contracts.ids import PlacementFragmentId, RuntimeFragmentId, TensorId + +_MAX_U64 = (1 << 64) - 1 +ParallelAxisKind = Literal["dp", "pp", "ep", "tp"] +SplitAxisKind = Literal["ep", "tp"] +_PARALLEL_AXIS_ORDER: dict[ParallelAxisKind, int] = { + "dp": 0, + "pp": 1, + "ep": 2, + "tp": 3, +} +_T = TypeVar("_T") + + +@dataclass(frozen=True) +class ParallelRank: + """Framework-provided owner coordinates used only for routing. + + Logical sharding is defined by ``global_offset``, ``local_shape``, and + ``shard_dims``. This coordinate is not a replacement for the topology and + axis metadata used to synthesize a target placement. + """ + + dp: int = 0 + tp: int = 0 + pp: int = 0 + ep: int = 0 + + def __post_init__(self) -> None: + for name in ("dp", "tp", "pp", "ep"): + _require_integer(getattr(self, name), f"parallel rank {name}", minimum=0) + + +@dataclass(frozen=True) +class SplitAxis: + """An axis whose ranks collectively partition one logical tensor.""" + + kind: SplitAxisKind + dim: int + + def __post_init__(self) -> None: + _validate_parallel_axis_kind(self.kind) + if self.kind not in {"ep", "tp"}: + raise ValueError(f"{self.kind} cannot use split semantics") + _require_integer(self.dim, "split axis dim", minimum=0) + if self.kind == "ep" and self.dim != 0: + raise ValueError("EP must split the leading logical expert dimension") + + +@dataclass(frozen=True) +class ReplicatedAxis: + """An axis whose ranks each contain an independent complete tensor.""" + + kind: ParallelAxisKind + + def __post_init__(self) -> None: + _validate_parallel_axis_kind(self.kind) + + +@dataclass(frozen=True) +class OwnershipAxis: + """An axis where only explicitly declared owner ranks contain the tensor.""" + + kind: ParallelAxisKind + + def __post_init__(self) -> None: + _validate_parallel_axis_kind(self.kind) + + +ParallelAxis = SplitAxis | ReplicatedAxis | OwnershipAxis + + +@dataclass(frozen=True) +class TensorDescriptor: + """Logical tensor identity, shape, dtype, and framework-supplied semantics.""" + + tensor_id: TensorId + global_shape: tuple[int, ...] + dtype: str + itemsize: int + shard_dims: tuple[int, ...] + layout_fingerprint: str + parallel_axes: tuple[ParallelAxis, ...] + layer_id: int | None = None + expert_id: int | None = None + + def __post_init__(self) -> None: + shape = _require_integer_tuple(self.global_shape, "global_shape", minimum=1) + if not shape: + raise ValueError("global_shape must not be empty") + object.__setattr__(self, "global_shape", shape) + _require_nonempty_string(self.tensor_id, "tensor_id") + _require_nonempty_string(self.dtype, "dtype") + _require_integer(self.itemsize, "itemsize", minimum=1) + shard_dims = _require_integer_tuple(self.shard_dims, "shard_dims", minimum=0) + if len(shard_dims) != len(set(shard_dims)): + raise ValueError("shard_dims must not contain duplicates") + if tuple(sorted(shard_dims)) != shard_dims: + raise ValueError("shard_dims must be sorted") + if any(dim >= len(shape) for dim in shard_dims): + raise ValueError("shard_dims contains an out-of-range dimension") + object.__setattr__(self, "shard_dims", shard_dims) + raw_axes = _require_sequence( + self.parallel_axes, + "TensorDescriptor parallel_axes", + ) + if not all( + isinstance(axis, (SplitAxis, ReplicatedAxis, OwnershipAxis)) + for axis in raw_axes + ): + raise ValueError( + "TensorDescriptor parallel_axes must contain explicit axis values" + ) + parallel_axes: tuple[ParallelAxis, ...] = tuple( + cast(ParallelAxis, axis) for axis in raw_axes + ) + parallel_axes = tuple( + sorted(parallel_axes, key=lambda axis: _PARALLEL_AXIS_ORDER[axis.kind]) + ) + axis_kinds = tuple(axis.kind for axis in parallel_axes) + if len(axis_kinds) != len(set(axis_kinds)): + raise ValueError("parallel_axes must not contain duplicate kinds") + split_dims = tuple( + axis.dim for axis in parallel_axes if isinstance(axis, SplitAxis) + ) + if len(split_dims) != len(set(split_dims)): + raise ValueError("split axes must not share a dimension") + if any(dim >= len(shape) for dim in split_dims): + raise ValueError("split axis dimension is out of range") + if tuple(sorted(split_dims)) != shard_dims: + raise ValueError("parallel axis split dimensions conflict with shard_dims") + object.__setattr__(self, "parallel_axes", parallel_axes) + for name in ("layer_id", "expert_id"): + value = getattr(self, name) + if value is not None: + _require_integer(value, name, minimum=0) + if self.expert_id is not None and any( + axis.kind == "ep" and isinstance(axis, SplitAxis) for axis in parallel_axes + ): + raise ValueError( + "an individually allocated expert requires explicit EP ownership" + ) + _require_nonempty_string(self.layout_fingerprint, "layout_fingerprint") + + +def canonical_strides_bytes( + shape: tuple[int, ...], + itemsize: int, +) -> tuple[int, ...]: + strides: list[int] = [] + running = itemsize + for extent in reversed(shape): + strides.append(running) + running *= extent + return tuple(reversed(strides)) + + +@dataclass(frozen=True, init=False) +class PlacementFragment: + """An address-free logical tensor box assigned to one parallel rank.""" + + tensor_id: TensorId + global_offset: tuple[int, ...] + local_shape: tuple[int, ...] + nbytes: int + rank: ParallelRank + aliases: tuple[TensorId, ...] + placement_fragment_id: PlacementFragmentId + + def __init__( + self, + tensor_id: TensorId, + global_offset: tuple[int, ...], + local_shape: tuple[int, ...], + nbytes: int, + rank: ParallelRank, + aliases: tuple[TensorId, ...] = (), + placement_fragment_id: PlacementFragmentId | None = None, + ) -> None: + normalized_offset = _require_integer_tuple( + global_offset, + "global_offset", + minimum=0, + ) + normalized_shape = _require_integer_tuple( + local_shape, + "local_shape", + minimum=1, + ) + _require_nonempty_string(tensor_id, "tensor_id") + _require_u64(nbytes, "nbytes", minimum=1) + if not isinstance(rank, ParallelRank): + raise ValueError("rank must be a ParallelRank") # noqa: TRY004 + normalized_aliases = _normalize_aliases(aliases) + if normalized_aliases: + if len(normalized_aliases) < 2: + raise ValueError("alias group must contain at least two tensor IDs") + if tensor_id not in normalized_aliases: + raise ValueError("alias group must contain the fragment tensor_id") + if placement_fragment_id is not None: + _require_nonempty_string( + placement_fragment_id, + "placement_fragment_id", + ) + resolved_fragment_id = placement_fragment_id + else: + resolved_fragment_id = _canonical_placement_fragment_id( + tensor_id=tensor_id, + global_offset=normalized_offset, + local_shape=normalized_shape, + nbytes=nbytes, + rank=rank, + aliases=normalized_aliases, + ) + object.__setattr__(self, "tensor_id", tensor_id) + object.__setattr__(self, "global_offset", normalized_offset) + object.__setattr__(self, "local_shape", normalized_shape) + object.__setattr__(self, "nbytes", nbytes) + object.__setattr__(self, "rank", rank) + object.__setattr__(self, "aliases", normalized_aliases) + object.__setattr__(self, "placement_fragment_id", resolved_fragment_id) + + @property + def fragment_id(self) -> PlacementFragmentId: + """Expose the common fragment identifier used by future planners.""" + + return self.placement_fragment_id + + +@dataclass(frozen=True) +class RuntimeBindingFragment: + """Contiguous physical runtime view for one placement fragment.""" + + placement_fragment_id: PlacementFragmentId + fragment_id: RuntimeFragmentId + address: int + nbytes: int + worker_id: str + endpoint: str + device: str + itemsize: int + local_shape: tuple[int, ...] + strides_bytes: tuple[int, ...] + storage_address: int + storage_nbytes: int + storage_offset_bytes: int + owner: object | None = field(default=None, compare=False, repr=False) + + def __post_init__(self) -> None: + for name in ( + "placement_fragment_id", + "fragment_id", + "worker_id", + "endpoint", + "device", + ): + _require_nonempty_string(getattr(self, name), name) + _require_address_range(self.address, self.nbytes) + _require_integer(self.itemsize, "runtime itemsize", minimum=1) + local_shape = _require_integer_tuple( + self.local_shape, + "runtime local_shape", + minimum=1, + ) + strides_bytes = _require_integer_tuple( + self.strides_bytes, + "runtime strides_bytes", + minimum=0, + ) + if len(strides_bytes) != len(local_shape): + raise ValueError("runtime stride rank differs from local_shape") + if any( + extent > 1 and stride == 0 + for extent, stride in zip(local_shape, strides_bytes) + ): + raise ValueError( + "runtime stride must be positive for non-singleton dimensions" + ) + strides_bytes = _normalize_singleton_strides_bytes( + local_shape, + strides_bytes, + self.itemsize, + ) + object.__setattr__(self, "local_shape", local_shape) + object.__setattr__(self, "strides_bytes", strides_bytes) + _require_address_range(self.storage_address, self.storage_nbytes) + _require_u64(self.storage_offset_bytes, "storage_offset_bytes") + if self.storage_offset_bytes > _MAX_U64 - self.storage_address: + raise ValueError("normalized runtime address must fit in 64 bits") + if self.address != self.storage_address + self.storage_offset_bytes: + raise ValueError( + "runtime address must equal storage_address plus storage_offset_bytes" + ) + if self.storage_offset_bytes > self.storage_nbytes - self.nbytes: + raise ValueError("runtime view exceeds storage allocation bounds") + + +def _require_nonempty_string(value: object, name: str) -> str: + if type(value) is not str or not value: + raise ValueError(f"{name} must be a non-empty string") + return value + + +def _require_integer( + value: object, + name: str, + *, + minimum: int | None = None, +) -> int: + if type(value) is not int: + raise ValueError(f"{name} must be an integer") + if minimum is not None and value < minimum: + raise ValueError(f"{name} must be at least {minimum}") + return value + + +def _require_u64(value: object, name: str, *, minimum: int = 0) -> int: + integer = _require_integer(value, name, minimum=minimum) + if integer > _MAX_U64: + raise ValueError(f"{name} must fit in an unsigned 64-bit integer") + return integer + + +def _require_address_range(address: object, nbytes: object) -> None: + normalized_address = _require_u64(address, "address", minimum=1) + normalized_nbytes = _require_u64(nbytes, "nbytes", minimum=1) + if normalized_nbytes > _MAX_U64 - normalized_address: + raise ValueError("address range must fit in an unsigned 64-bit integer") + + +def require_sha256_digest(value: object, name: str) -> None: + digest = _require_nonempty_string(value, name) + if len(digest) != 64 or any( + character not in "0123456789abcdef" for character in digest + ): + raise ValueError(f"{name} must be a lowercase SHA-256 digest") + + +def _require_integer_tuple( + value: object, + name: str, + *, + minimum: int, +) -> tuple[int, ...]: + if isinstance(value, (str, bytes, bytearray)) or not isinstance(value, Sequence): + raise ValueError(f"{name} must contain integers") # noqa: TRY004 + items = cast(Sequence[object], value) + return tuple(_require_integer(item, name, minimum=minimum) for item in items) + + +def _require_sequence(value: object, name: str) -> Sequence[object]: + if isinstance(value, (str, bytes, bytearray)) or not isinstance(value, Sequence): + raise ValueError(f"{name} must be a sequence") # noqa: TRY004 + return cast(Sequence[object], value) + + +def require_manifest_items( + value: object, + name: str, + item_type: type[_T], +) -> tuple[_T, ...]: + items = tuple(_require_sequence(value, name)) + if not all(isinstance(item, item_type) for item in items): + raise ValueError(f"{name} must contain {item_type.__name__}") + return tuple(cast(_T, item) for item in items) + + +def _normalize_aliases(value: object) -> tuple[TensorId, ...]: + aliases = tuple(_require_sequence(value, "aliases")) + if any(type(alias) is not str or not alias for alias in aliases): + raise ValueError("aliases must contain non-empty strings") + normalized_aliases = tuple(cast(str, alias) for alias in aliases) + if len(normalized_aliases) != len(set(normalized_aliases)): + raise ValueError("aliases must not contain duplicates") + return tuple(TensorId(alias) for alias in sorted(normalized_aliases)) + + +def _canonical_placement_fragment_id( + *, + tensor_id: TensorId, + global_offset: tuple[int, ...], + local_shape: tuple[int, ...], + nbytes: int, + rank: ParallelRank, + aliases: tuple[TensorId, ...], +) -> PlacementFragmentId: + content = { + "schema": "weight-placement-fragment", + "tensor_id": tensor_id, + "global_offset": global_offset, + "local_shape": local_shape, + "nbytes": nbytes, + "rank": asdict(rank), + "aliases": aliases, + } + encoded = json.dumps(content, sort_keys=True, separators=(",", ":")).encode() + return PlacementFragmentId(f"sha256:{hashlib.sha256(encoded).hexdigest()}") + + +def _normalize_singleton_strides_bytes( + shape: tuple[int, ...], + strides_bytes: tuple[int, ...], + itemsize: int, +) -> tuple[int, ...]: + canonical = canonical_strides_bytes(shape, itemsize) + return tuple( + expected if extent == 1 else observed + for extent, observed, expected in zip(shape, strides_bytes, canonical) + ) + + +def _validate_parallel_axis_kind(kind: ParallelAxisKind) -> None: + _require_nonempty_string(kind, "parallel axis kind") + if kind not in _PARALLEL_AXIS_ORDER: + raise ValueError(f"unsupported parallel axis kind: {kind}") + + +# Retain internal imports used by this first-stage contract package. These +# aliases do not widen the public wire or runtime contract. +_canonical_strides_bytes = canonical_strides_bytes +_require_manifest_items = require_manifest_items +_require_sha256_digest = require_sha256_digest diff --git a/mooncake-reshard/python/mooncake/reshard/weight/validation.py b/mooncake-reshard/python/mooncake/reshard/weight/validation.py new file mode 100644 index 0000000000..d1f952fff5 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/weight/validation.py @@ -0,0 +1,606 @@ +"""Cross-fragment logical and physical validation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from itertools import product as cartesian_product +from math import prod + +from ..contracts import RuntimeInstanceId, TensorId +from .topology import ParallelTopology +from .types import ( + OwnershipAxis, + ParallelRank, + PlacementFragment, + ReplicatedAxis, + RuntimeBindingFragment, + SplitAxis, + TensorDescriptor, +) + + +def _validate_fragments( + tensors: Sequence[TensorDescriptor], + fragments: Sequence[PlacementFragment], + *, + require_complete_alias_groups: bool = False, +) -> None: + tensor_by_id: dict[TensorId, TensorDescriptor] = {} + for tensor in tensors: + if tensor.tensor_id in tensor_by_id: + raise ValueError(f"duplicate tensor_id: {tensor.tensor_id}") + tensor_by_id[tensor.tensor_id] = tensor + + fragment_ids: set[str] = set() + logical_fragments: set[ + tuple[TensorId, ParallelRank, tuple[int, ...], tuple[int, ...]] + ] = set() + for fragment in fragments: + if fragment.fragment_id in fragment_ids: + raise ValueError(f"duplicate fragment_id: {fragment.fragment_id}") + fragment_ids.add(fragment.fragment_id) + logical_fragment = ( + fragment.tensor_id, + fragment.rank, + fragment.global_offset, + fragment.local_shape, + ) + if logical_fragment in logical_fragments: + raise ValueError( + "duplicate logical fragment for tensor and parallel rank: " + f"{fragment.fragment_id}" + ) + logical_fragments.add(logical_fragment) + tensor = tensor_by_id.get(fragment.tensor_id) + if tensor is None: + raise ValueError(f"unknown tensor_id: {fragment.tensor_id}") + _validate_fragment_geometry(tensor, fragment) + if require_complete_alias_groups: + _validate_complete_alias_groups(tensor_by_id, fragments) + _validate_logical_fragment_overlaps(fragments) + + +def _validate_complete_alias_groups( + tensors_by_id: Mapping[TensorId, TensorDescriptor], + fragments: Sequence[PlacementFragment], +) -> None: + """Validate alias authorization only after all placement parts are assembled. + + A placement part can legitimately name an alias member owned by another + participant. The complete placement is the first layer with the full tensor + set, so it alone can prove that every declared member exists and declares + the same group. + """ + + fragments_by_tensor: dict[TensorId, list[PlacementFragment]] = {} + alias_groups: set[tuple[TensorId, ...]] = set() + for fragment in fragments: + fragments_by_tensor.setdefault(fragment.tensor_id, []).append(fragment) + if fragment.aliases: + alias_groups.add(fragment.aliases) + + known_tensor_ids = set(tensors_by_id) + for aliases in alias_groups: + unknown_tensor_ids = sorted(set(aliases) - known_tensor_ids) + if unknown_tensor_ids: + raise ValueError( + f"alias group references unknown tensor: {unknown_tensor_ids[0]}" + ) + for tensor_id in aliases: + member_groups = { + fragment.aliases for fragment in fragments_by_tensor[tensor_id] + } + if member_groups != {aliases}: + raise ValueError( + f"alias group is not declared consistently for tensor: {tensor_id}" + ) + + +def _validate_complete_weight_placement( + tensors: Sequence[TensorDescriptor], + fragments: Sequence[PlacementFragment], + *, + topology: ParallelTopology, +) -> None: + """Validate explicit split, replication, and ownership semantics.""" + + if not tensors: + raise ValueError("global weight placement must contain tensors") + if not fragments: + raise ValueError("global weight placement must contain fragments") + + by_tensor: dict[TensorId, list[PlacementFragment]] = {} + for fragment in fragments: + by_tensor.setdefault(fragment.tensor_id, []).append(fragment) + + for tensor in tensors: + tensor_fragments = by_tensor.get(tensor.tensor_id, []) + if not tensor_fragments: + raise ValueError( + f"global placement tensor is not fully covered: {tensor.tensor_id}" + ) + _validate_tensor_axis_covers(tensor, tensor_fragments, topology=topology) + + +def _split_axis_kinds(tensor: TensorDescriptor) -> tuple[str, ...]: + return tuple( + axis.kind for axis in tensor.parallel_axes if isinstance(axis, SplitAxis) + ) + + +def _replicated_axis_kinds(tensor: TensorDescriptor) -> tuple[str, ...]: + return tuple( + axis.kind for axis in tensor.parallel_axes if isinstance(axis, ReplicatedAxis) + ) + + +def _ownership_axis_kinds(tensor: TensorDescriptor) -> tuple[str, ...]: + return tuple( + axis.kind for axis in tensor.parallel_axes if isinstance(axis, OwnershipAxis) + ) + + +def _axis_rank(rank: ParallelRank, kinds: Sequence[str]) -> tuple[tuple[str, int], ...]: + return tuple((kind, getattr(rank, kind)) for kind in kinds) + + +def _rank_matches(rank: ParallelRank, coordinates: tuple[tuple[str, int], ...]) -> bool: + return all(getattr(rank, kind) == value for kind, value in coordinates) + + +def _validate_tensor_axis_covers( + tensor: TensorDescriptor, + fragments: Sequence[PlacementFragment], + *, + topology: ParallelTopology, +) -> None: + _validate_undeclared_axis_dependencies(tensor, fragments, topology=topology) + + ownership_kinds = _ownership_axis_kinds(tensor) + replicated_kinds = _replicated_axis_kinds(tensor) + declared_owners = { + _axis_rank(fragment.rank, ownership_kinds) for fragment in fragments + } + + for owner in sorted(declared_owners): + owner_fragments = [ + fragment + for fragment in fragments + if _axis_rank(fragment.rank, ownership_kinds) == owner + ] + expected_replicas = { + _axis_rank(participant.rank, replicated_kinds) + for participant in topology.participants + if _rank_matches(participant.rank, owner) + } + actual_replicas = { + _axis_rank(fragment.rank, replicated_kinds) for fragment in owner_fragments + } + if actual_replicas != expected_replicas: + raise ValueError( + "global placement tensor is not fully covered: missing or " + "unexpected replicated-axis participant for " + f"{tensor.tensor_id}" + ) + + for replica in sorted(expected_replicas): + cover = [ + fragment + for fragment in owner_fragments + if _axis_rank(fragment.rank, replicated_kinds) == replica + ] + if not _fragments_exactly_cover_tensor(tensor, cover): + raise ValueError( + f"global placement tensor is not fully covered: {tensor.tensor_id}" + ) + _validate_split_axis_participation( + tensor, + cover, + topology=topology, + fixed_coordinates=owner + replica, + ) + _validate_split_axis_geometry(tensor, cover) + + +def _validate_undeclared_axis_dependencies( + tensor: TensorDescriptor, + fragments: Sequence[PlacementFragment], + *, + topology: ParallelTopology, +) -> None: + """Reject implicit ownership while allowing explicitly coupled TP/EP ranks.""" + + declared_axis_kinds = tuple(axis.kind for axis in tensor.parallel_axes) + varying_declared_axis_kinds = tuple( + kind + for kind in declared_axis_kinds + if len({getattr(fragment.rank, kind) for fragment in fragments}) > 1 + ) + for kind in ("dp", "pp", "ep", "tp"): + if kind in declared_axis_kinds or getattr(topology, f"{kind}_size") == 1: + continue + coordinates = {getattr(fragment.rank, kind) for fragment in fragments} + if kind in {"dp", "pp"} or len(coordinates) <= 1: + raise ValueError( + "global placement tensor requires explicit parallel semantics " + f"for active undeclared {kind} axis: {tensor.tensor_id}" + ) + if not varying_declared_axis_kinds: + raise ValueError( + "global placement tensor varies across undeclared " + f"{kind} axis: {tensor.tensor_id}" + ) + + coordinate_by_declared_rank: dict[tuple[tuple[str, int], ...], int] = {} + for fragment in fragments: + declared_rank = _axis_rank(fragment.rank, varying_declared_axis_kinds) + coordinate = getattr(fragment.rank, kind) + previous = coordinate_by_declared_rank.setdefault( + declared_rank, + coordinate, + ) + if previous != coordinate: + raise ValueError( + "global placement tensor varies independently across undeclared " + f"{kind} axis: {tensor.tensor_id}" + ) + + +def _validate_split_axis_geometry( + tensor: TensorDescriptor, + fragments: Sequence[PlacementFragment], +) -> None: + """Validate that every declared split rank owns its logical dimension.""" + + split_axes = tuple( + axis for axis in tensor.parallel_axes if isinstance(axis, SplitAxis) + ) + if not split_axes: + return + coordinate_values = tuple( + tuple(sorted({getattr(fragment.rank, axis.kind) for fragment in fragments})) + for axis in split_axes + ) + observed_coordinates = { + tuple(getattr(fragment.rank, axis.kind) for axis in split_axes) + for fragment in fragments + } + expected_coordinates = set(cartesian_product(*coordinate_values)) + if observed_coordinates != expected_coordinates: + raise ValueError( + "non-Cartesian split-axis participant mapping cannot prove " + f"rank-to-dimension ownership: {tensor.tensor_id}" + ) + + for axis in split_axes: + intervals_by_rank: dict[int, list[tuple[int, int]]] = {} + for fragment in fragments: + begin = fragment.global_offset[axis.dim] + end = begin + fragment.local_shape[axis.dim] + intervals_by_rank.setdefault(getattr(fragment.rank, axis.kind), []).append( + (begin, end) + ) + + owned_intervals = sorted( + (begin, end, rank) + for rank, intervals in intervals_by_rank.items() + for begin, end in _merge_intervals(intervals) + ) + cursor = 0 + for begin, end, _ in owned_intervals: + if begin != cursor: + raise ValueError( + "split-axis rank geometry conflicts with declared dimension: " + f"{tensor.tensor_id}: {axis.kind} -> dim {axis.dim}" + ) + cursor = end + if cursor != tensor.global_shape[axis.dim]: + raise ValueError( + "split-axis rank geometry conflicts with declared dimension: " + f"{tensor.tensor_id}: {axis.kind} -> dim {axis.dim}" + ) + + +def _merge_intervals( + intervals: Sequence[tuple[int, int]], +) -> tuple[tuple[int, int], ...]: + merged: list[list[int]] = [] + for begin, end in sorted(set(intervals)): + if not merged or begin > merged[-1][1]: + merged.append([begin, end]) + else: + merged[-1][1] = max(merged[-1][1], end) + return tuple((begin, end) for begin, end in merged) + + +def _parallel_split_rank( + tensor: TensorDescriptor, + rank: ParallelRank, +) -> tuple[tuple[str, int], ...]: + return _axis_rank(rank, _split_axis_kinds(tensor)) + + +def _validate_split_axis_participation( + tensor: TensorDescriptor, + fragments: Sequence[PlacementFragment], + *, + topology: ParallelTopology, + fixed_coordinates: tuple[tuple[str, int], ...], +) -> None: + split_axis_kinds = _split_axis_kinds(tensor) + if not split_axis_kinds: + return + + expected = { + _parallel_split_rank(tensor, participant.rank) + for participant in topology.participants + if _rank_matches(participant.rank, fixed_coordinates) + } + actual = {_parallel_split_rank(tensor, fragment.rank) for fragment in fragments} + complete_axis_ranges = all( + {dict(split_rank)[kind] for split_rank in expected} + == set(range(getattr(topology, f"{kind}_size"))) + for kind in split_axis_kinds + ) + if not complete_axis_ranges or actual != expected: + raise ValueError( + "global placement tensor is not fully covered: missing or unexpected " + f"split-axis participant for {tensor.tensor_id}" + ) + + +def _fragments_exactly_cover_tensor( + tensor: TensorDescriptor, + fragments: Sequence[PlacementFragment], +) -> bool: + boxes = tuple( + (fragment.global_offset, fragment.local_shape) for fragment in fragments + ) + if not boxes: + return False + if sum(prod(shape) for _, shape in boxes) != prod(tensor.global_shape): + return False + if any( + any( + offset < 0 or offset + extent > total + for offset, extent, total in zip( + box_offset, + box_shape, + tensor.global_shape, + ) + ) + for box_offset, box_shape in boxes + ): + return False + return not _boxes_overlap(boxes) + + +def _boxes_overlap( + boxes: Sequence[tuple[tuple[int, ...], tuple[int, ...]]], +) -> bool: + if len(boxes) < 2: + return False + ndim = len(boxes[0][0]) + sweep_dim = max( + range(ndim), + key=lambda dim: len( + {(offset[dim], offset[dim] + shape[dim]) for offset, shape in boxes} + ), + ) + ordered = sorted(boxes, key=lambda item: item[0][sweep_dim]) + active: list[tuple[tuple[int, ...], tuple[int, ...]]] = [] + for offset, shape in ordered: + begin = offset[sweep_dim] + active = [ + candidate + for candidate in active + if candidate[0][sweep_dim] + candidate[1][sweep_dim] > begin + ] + if any( + all( + left_begin < right_begin + right_extent + and right_begin < left_begin + left_extent + for left_begin, left_extent, right_begin, right_extent in zip( + candidate_offset, + candidate_shape, + offset, + shape, + ) + ) + for candidate_offset, candidate_shape in active + ): + return True + active.append((offset, shape)) + return False + + +def _validate_logical_fragment_overlaps( + fragments: Sequence[PlacementFragment], +) -> None: + by_tensor_and_rank: dict[tuple[str, ParallelRank], list[PlacementFragment]] = {} + for fragment in fragments: + by_tensor_and_rank.setdefault((fragment.tensor_id, fragment.rank), []).append( + fragment + ) + + for owner_fragments in by_tensor_and_rank.values(): + if len(owner_fragments) < 2: + continue + ndim = len(owner_fragments[0].global_offset) + sweep_dim = max( + range(ndim), + key=lambda dim: len( + { + ( + fragment.global_offset[dim], + fragment.global_offset[dim] + fragment.local_shape[dim], + ) + for fragment in owner_fragments + } + ), + ) + ordered = sorted( + owner_fragments, + key=lambda fragment: fragment.global_offset[sweep_dim], + ) + active: list[PlacementFragment] = [] + for current in ordered: + current_begin = current.global_offset[sweep_dim] + active = [ + previous + for previous in active + if previous.global_offset[sweep_dim] + previous.local_shape[sweep_dim] + > current_begin + ] + for previous in active: + if all( + previous_offset < current_offset + current_extent + and current_offset < previous_offset + previous_extent + for ( + previous_offset, + previous_extent, + current_offset, + current_extent, + ) in zip( + previous.global_offset, + previous.local_shape, + current.global_offset, + current.local_shape, + ) + ): + raise ValueError( + "logical fragment boxes overlap for tensor and " + "parallel rank: " + f"{previous.fragment_id} and {current.fragment_id}" + ) + active.append(current) + + +def _validate_fragment_geometry( + tensor: TensorDescriptor, + fragment: PlacementFragment, +) -> None: + ndim = len(tensor.global_shape) + if len(fragment.global_offset) != ndim or len(fragment.local_shape) != ndim: + raise ValueError(f"fragment rank mismatch: {fragment.fragment_id}") + for offset, extent, total in zip( + fragment.global_offset, + fragment.local_shape, + tensor.global_shape, + ): + if offset + extent > total: + raise ValueError(f"fragment is out of bounds: {fragment.fragment_id}") + + expected_nbytes = prod(fragment.local_shape) * tensor.itemsize + if fragment.nbytes != expected_nbytes: + raise ValueError( + f"fragment byte size mismatch: {fragment.fragment_id}: " + f"expected {expected_nbytes}, got {fragment.nbytes}" + ) + + +def _runtime_alias_descriptor_key(tensor: TensorDescriptor) -> tuple[object, ...]: + return ( + tensor.global_shape, + tensor.dtype, + tensor.itemsize, + tensor.shard_dims, + tensor.parallel_axes, + tensor.layer_id, + tensor.expert_id, + tensor.layout_fingerprint, + ) + + +def _is_exact_declared_runtime_alias( + left_placement: PlacementFragment, + left_binding: RuntimeBindingFragment, + right_placement: PlacementFragment, + right_binding: RuntimeBindingFragment, + tensors: Mapping[TensorId, TensorDescriptor], +) -> bool: + return ( + left_binding.address == right_binding.address + and left_binding.nbytes == right_binding.nbytes + and len(left_placement.aliases) >= 2 + and left_placement.aliases == right_placement.aliases + and left_placement.tensor_id in left_placement.aliases + and right_placement.tensor_id in left_placement.aliases + and left_placement.global_offset == right_placement.global_offset + and left_placement.local_shape == right_placement.local_shape + and _runtime_alias_descriptor_key(tensors[left_placement.tensor_id]) + == _runtime_alias_descriptor_key(tensors[right_placement.tensor_id]) + ) + + +def _validate_runtime_binding_address_ranges( + *, + instance_id: RuntimeInstanceId, + tensors: Sequence[TensorDescriptor], + placements: Sequence[PlacementFragment], + bindings: Sequence[RuntimeBindingFragment], +) -> None: + tensor_by_id = {tensor.tensor_id: tensor for tensor in tensors} + placement_by_id = { + fragment.placement_fragment_id: fragment for fragment in placements + } + by_address_space: dict[ + tuple[str, str, str], + list[tuple[PlacementFragment, RuntimeBindingFragment]], + ] = {} + for binding in bindings: + placement = placement_by_id[binding.placement_fragment_id] + address_space = (instance_id, binding.worker_id, binding.device) + by_address_space.setdefault(address_space, []).append((placement, binding)) + + for address_space, fragments in by_address_space.items(): + ordered = sorted(fragments, key=lambda item: item[1].address) + active: list[tuple[PlacementFragment, RuntimeBindingFragment]] = [] + for current_placement, current_binding in ordered: + active = [ + (previous_placement, previous_binding) + for previous_placement, previous_binding in active + if previous_binding.address + previous_binding.nbytes + > current_binding.address + ] + for previous_placement, previous_binding in active: + if _is_exact_declared_runtime_alias( + previous_placement, + previous_binding, + current_placement, + current_binding, + tensor_by_id, + ): + continue + raise ValueError( + "runtime binding address ranges overlap: " + f"{previous_binding.fragment_id} and " + f"{current_binding.fragment_id} " + f"in {address_space}" + ) + active.append((current_placement, current_binding)) + + allocation_order = sorted( + (binding for _, binding in fragments), + key=lambda binding: binding.storage_address, + ) + active_allocations: list[RuntimeBindingFragment] = [] + for current in allocation_order: + active_allocations = [ + previous + for previous in active_allocations + if previous.storage_address + previous.storage_nbytes + > current.storage_address + ] + for previous in active_allocations: + if ( + previous.storage_address == current.storage_address + and previous.storage_nbytes == current.storage_nbytes + ): + continue + raise ValueError( + "runtime binding storage allocation ranges overlap: " + f"{previous.fragment_id} and {current.fragment_id} " + f"in {address_space}" + ) + active_allocations.append(current) diff --git a/mooncake-reshard/tests/test_reshard_manifest_contracts.py b/mooncake-reshard/tests/test_reshard_manifest_contracts.py new file mode 100644 index 0000000000..95cf4e44e2 --- /dev/null +++ b/mooncake-reshard/tests/test_reshard_manifest_contracts.py @@ -0,0 +1,113 @@ +import json +from dataclasses import fields +from importlib.util import find_spec +from typing import Protocol + +import mooncake.reshard.weight.placement as weight_placement +import pytest +from mooncake import reshard +from mooncake.reshard import weight +from mooncake.reshard.contracts import ( + PlacementManifest, + ResourceKind, + ResourceManifest, + RuntimeBindingManifest, +) +from mooncake.reshard.weight import ( + WeightPlacementManifest, + WeightRuntimeBindingManifest, + weight_placement_from_json, + weight_placement_to_json, +) +from weight_manifest.helpers import ( + binding_manifest, + placement_manifest, +) + + +def test_reshard_public_api_is_resource_neutral(): + assert reshard.__all__ == [ + "ResourceKind", + "ResourceManifest", + "PlacementManifest", + "RuntimeBindingManifest", + ] + + +def test_common_manifest_contracts_are_structural_protocols(): + assert issubclass(ResourceManifest, Protocol) + assert issubclass(PlacementManifest, Protocol) + assert issubclass(RuntimeBindingManifest, Protocol) + assert PlacementManifest not in WeightPlacementManifest.__mro__ + assert RuntimeBindingManifest not in WeightRuntimeBindingManifest.__mro__ + + +def test_weight_api_has_no_combined_runtime_manifest(): + assert not hasattr(weight, "WeightRuntimeManifest") + + +def test_weight_placement_has_one_canonical_contract_without_side_aliases(): + assert not hasattr(weight_placement, "SourcePlacementManifest") + assert not hasattr(weight_placement, "TargetPlacementManifest") + + +def test_concrete_weight_manifests_own_and_validate_common_fields(): + assert [field.name for field in fields(WeightPlacementManifest)][:2] == [ + "resource_id", + "placement_id", + ] + assert [field.name for field in fields(WeightRuntimeBindingManifest)][:6] == [ + "resource_id", + "placement_id", + "placement_digest", + "instance_id", + "generation", + "lease_id", + ] + + +def test_weight_manifests_expose_common_resource_fields(): + placement = placement_manifest() + manifests = ( + placement, + binding_manifest(placement=placement), + ) + + for manifest in manifests: + assert not hasattr(manifest, "model_id") + assert manifest.resource_id == "model" + assert manifest.resource_kind is ResourceKind.MODEL_WEIGHT + + +def test_legacy_model_weight_namespace_is_not_installed(): + assert find_spec("mooncake.model_weight") is None + + +def test_kv_cache_is_reserved_without_a_manifest_implementation(): + assert ResourceKind.KV_CACHE.value == "kv_cache" + assert find_spec("mooncake.reshard.kv_cache") is None + + +def test_weight_placement_json_has_a_strict_resource_kind(): + placement = placement_manifest() + payload = json.loads(weight_placement_to_json(placement)) + + assert payload["resource_kind"] == "model_weight" + payload["resource_kind"] = "kv_cache" + + with pytest.raises(ValueError, match="resource_kind"): + weight_placement_from_json(json.dumps(payload)) + + +@pytest.mark.parametrize("resource_kind", [None, "unknown"]) +def test_weight_placement_json_rejects_missing_or_unknown_resource_kind( + resource_kind, +): + payload = json.loads(weight_placement_to_json(placement_manifest())) + if resource_kind is None: + del payload["resource_kind"] + else: + payload["resource_kind"] = resource_kind + + with pytest.raises(ValueError): + weight_placement_from_json(json.dumps(payload)) diff --git a/mooncake-reshard/tests/test_reshard_weight_module_layout.py b/mooncake-reshard/tests/test_reshard_weight_module_layout.py new file mode 100644 index 0000000000..883939b3eb --- /dev/null +++ b/mooncake-reshard/tests/test_reshard_weight_module_layout.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import mooncake.reshard.weight as model_weight +from mooncake.reshard.weight.binding import ( + validate_runtime_binding, + validate_runtime_bindings, +) +from mooncake.reshard.weight.manifest import ( + ParallelRank, + ParallelTopology, + PlacementFragment, + RuntimeBindingFragment, + TensorDescriptor, + TopologyParticipant, + WeightPlacementManifest, + WeightPlacementPart, + WeightRuntimeBindingManifest, +) +from mooncake.reshard.weight.part import WeightPlacementPart as PlacementPartContract +from mooncake.reshard.weight.placement import ( + WeightPlacementManifest as PlacementContract, +) +from mooncake.reshard.weight.runtime import ( + WeightRuntimeBindingManifest as RuntimeBindingContract, +) +from mooncake.reshard.weight.types import ( + ParallelRank as ParallelRankContract, +) +from mooncake.reshard.weight.types import ( + PlacementFragment as PlacementFragmentContract, +) +from mooncake.reshard.weight.types import ( + RuntimeBindingFragment as RuntimeBindingFragmentContract, +) +from mooncake.reshard.weight.types import TensorDescriptor as TensorContract +from mooncake.reshard.weight.topology import ( + ParallelTopology as ParallelTopologyContract, +) +from mooncake.reshard.weight.topology import ( + TopologyParticipant as TopologyParticipantContract, +) + + +def test_responsibility_modules_preserve_public_contract_identity() -> None: + assert model_weight.ParallelRank is ParallelRank is ParallelRankContract + assert model_weight.ParallelTopology is ParallelTopology is ParallelTopologyContract + assert ( + model_weight.TopologyParticipant + is TopologyParticipant + is TopologyParticipantContract + ) + assert model_weight.TensorDescriptor is TensorDescriptor is TensorContract + assert ( + model_weight.PlacementFragment is PlacementFragment is PlacementFragmentContract + ) + assert ( + model_weight.RuntimeBindingFragment + is RuntimeBindingFragment + is RuntimeBindingFragmentContract + ) + assert ( + model_weight.WeightPlacementManifest + is WeightPlacementManifest + is PlacementContract + ) + assert ( + model_weight.WeightPlacementPart is WeightPlacementPart is PlacementPartContract + ) + assert ( + model_weight.WeightRuntimeBindingManifest + is WeightRuntimeBindingManifest + is RuntimeBindingContract + ) + assert model_weight.validate_runtime_binding is validate_runtime_binding + assert model_weight.validate_runtime_bindings is validate_runtime_bindings diff --git a/mooncake-reshard/tests/test_source_package_import.py b/mooncake-reshard/tests/test_source_package_import.py new file mode 100644 index 0000000000..22ca7327a6 --- /dev/null +++ b/mooncake-reshard/tests/test_source_package_import.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import os +from pathlib import Path +import subprocess +import sys + + +def test_source_tree_package_coexists_with_installed_mooncake() -> None: + repo_root = Path(__file__).resolve().parents[2] + source_root = repo_root / "mooncake-reshard" / "python" + wheel_root = repo_root / "mooncake-wheel" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join((str(source_root), str(wheel_root))) + + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import mooncake; import mooncake.reshard; " + "assert len(tuple(mooncake.__path__)) >= 2" + ), + ], + cwd=repo_root, + env=env, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr diff --git a/mooncake-reshard/tests/weight_manifest/__init__.py b/mooncake-reshard/tests/weight_manifest/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/mooncake-reshard/tests/weight_manifest/helpers.py b/mooncake-reshard/tests/weight_manifest/helpers.py new file mode 100644 index 0000000000..aa6ca28c89 --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/helpers.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +from mooncake.reshard.weight import ( + ParallelRank, + ParallelTopology, + PlacementFragment, + RuntimeBindingFragment, + SplitAxis, + TensorDescriptor, + TopologyParticipant, + WeightPlacementManifest, + WeightPlacementPart, + WeightRuntimeBindingManifest, +) + + +MODEL_ID = "model" +REVISION = "revision" +WEIGHT_GENERATION = 7 +PLACEMENT_SET_ID = "placement-set-7" +PARTICIPANT_ID = "worker-0" + + +def parallel_topology( + *, + participants: tuple[TopologyParticipant, ...] | None = None, + **overrides, +) -> ParallelTopology: + participants = participants or ( + TopologyParticipant(PARTICIPANT_ID, ParallelRank()), + ) + values = { + "tp_size": max(item.rank.tp for item in participants) + 1, + "pp_size": max(item.rank.pp for item in participants) + 1, + "ep_size": max(item.rank.ep for item in participants) + 1, + "dp_size": max(item.rank.dp for item in participants) + 1, + "participants": participants, + } + values.update(overrides) + return ParallelTopology(**values) + + +def descriptor(**overrides) -> TensorDescriptor: + values = { + "tensor_id": "layers.2.experts.3.w1", + "global_shape": (4, 4), + "dtype": "bfloat16", + "itemsize": 2, + "shard_dims": (0,), + "layer_id": 2, + "expert_id": 3, + "layout_fingerprint": "test:qwen:bf16:v1", + "parallel_axes": (SplitAxis(kind="tp", dim=0),), + } + values.update(overrides) + return TensorDescriptor(**values) + + +def placement_fragment(**overrides) -> PlacementFragment: + values = { + "placement_fragment_id": "placement-0", + "tensor_id": "layers.2.experts.3.w1", + "global_offset": (0, 0), + "local_shape": (4, 4), + "nbytes": 32, + "rank": ParallelRank(), + } + values.update(overrides) + return PlacementFragment(**values) + + +def placement_part( + *, + topology: ParallelTopology | None = None, + participant_id: str = PARTICIPANT_ID, + rank: ParallelRank | None = None, + tensors: tuple[TensorDescriptor, ...] = (), + fragments: tuple[PlacementFragment, ...] = (), + **overrides, +) -> WeightPlacementPart: + topology = topology or parallel_topology() + rank = rank or topology.participant(participant_id).rank + values = { + "resource_id": MODEL_ID, + "revision": REVISION, + "weight_generation": WEIGHT_GENERATION, + "placement_set_id": PLACEMENT_SET_ID, + "topology_id": topology.topology_id, + "participant_id": participant_id, + "rank": rank, + "tensors": tensors, + "fragments": fragments, + } + values.update(overrides) + return WeightPlacementPart(**values) + + +def placement_manifest(**overrides) -> WeightPlacementManifest: + values = { + "resource_id": MODEL_ID, + "revision": REVISION, + "weight_generation": WEIGHT_GENERATION, + "placement_set_id": PLACEMENT_SET_ID, + "placement_id": None, + } + for field in tuple(values): + if field in overrides: + values[field] = overrides.pop(field) + + tensors = overrides.pop("tensors", (descriptor(),)) + fragments = overrides.pop("fragments", (placement_fragment(),)) + topology = overrides.pop("topology", None) + parts = overrides.pop("parts", None) + if overrides: + values.update(overrides) + + fragments_are_valid = isinstance(fragments, tuple) and all( + isinstance(item, PlacementFragment) for item in fragments + ) + tensors_are_valid = isinstance(tensors, tuple) and all( + isinstance(item, TensorDescriptor) for item in tensors + ) + if topology is None: + ranks = ( + tuple(sorted({item.rank for item in fragments}, key=_rank_key)) + if fragments_are_valid and fragments + else (ParallelRank(),) + ) + topology = parallel_topology( + participants=tuple( + TopologyParticipant(f"worker-{index}", rank) + for index, rank in enumerate(ranks) + ) + ) + + if parts is None and fragments_are_valid and tensors_are_valid: + tensor_by_id = {item.tensor_id: item for item in tensors} + generated_parts = [] + for index, participant in enumerate(topology.participants): + local_fragments = tuple( + item for item in fragments if item.rank == participant.rank + ) + local_tensor_ids = {item.tensor_id for item in local_fragments} + local_tensors = tuple( + tensor_by_id[tensor_id] + for tensor_id in sorted(local_tensor_ids) + if tensor_id in tensor_by_id + ) + if not fragments and index == 0: + local_tensors = tensors + generated_parts.append( + placement_part( + topology=topology, + participant_id=participant.participant_id, + rank=participant.rank, + tensors=local_tensors, + fragments=local_fragments, + resource_id=values["resource_id"], + revision=values["revision"], + weight_generation=values["weight_generation"], + placement_set_id=values["placement_set_id"], + ) + ) + parts = tuple(generated_parts) + elif parts is None: + participant = topology.participants[0] + parts = ( + placement_part( + topology=topology, + participant_id=participant.participant_id, + rank=participant.rank, + tensors=tensors, + fragments=fragments, + resource_id=values["resource_id"], + revision=values["revision"], + weight_generation=values["weight_generation"], + placement_set_id=values["placement_set_id"], + ), + ) + + return WeightPlacementManifest(topology=topology, parts=parts, **values) + + +def _rank_key(rank: ParallelRank) -> tuple[int, int, int, int]: + return (rank.dp, rank.tp, rank.pp, rank.ep) + + +def binding_fragment(**overrides) -> RuntimeBindingFragment: + values = { + "placement_fragment_id": "placement-0", + "fragment_id": "runtime-0", + "address": 0x1000, + "nbytes": 32, + "worker_id": "worker-0", + "endpoint": "worker-0:12345", + "device": "cuda:0", + "itemsize": 2, + "local_shape": (4, 4), + "strides_bytes": (8, 2), + "storage_address": 0x1000, + "storage_nbytes": 32, + "storage_offset_bytes": 0, + } + if "address" in overrides and "storage_address" not in overrides: + values["storage_address"] = overrides["address"] + if "nbytes" in overrides and "storage_nbytes" not in overrides: + values["storage_nbytes"] = overrides["nbytes"] + values.update(overrides) + return RuntimeBindingFragment(**values) + + +def binding_manifest( + *, + placement: WeightPlacementManifest | None = None, + **overrides, +) -> WeightRuntimeBindingManifest: + logical = placement or placement_manifest() + participant_id = overrides.pop( + "participant_id", + logical.parts[0].participant_id, + ) + placement_part = next( + (item for item in logical.parts if item.participant_id == participant_id), + None, + ) + tensor_by_id = {tensor.tensor_id: tensor for tensor in logical.tensors} + default_fragments = ( + tuple( + binding_fragment( + placement_fragment_id=fragment.placement_fragment_id, + fragment_id=f"runtime-{index}", + address=0x1000 + index * 0x100, + nbytes=fragment.nbytes, + itemsize=tensor_by_id[fragment.tensor_id].itemsize, + local_shape=fragment.local_shape, + strides_bytes=_contiguous_strides_bytes( + fragment.local_shape, + tensor_by_id[fragment.tensor_id].itemsize, + ), + ) + for index, fragment in enumerate(placement_part.fragments) + ) + if placement_part is not None + else () + ) + values = { + "resource_id": logical.resource_id, + "revision": logical.revision, + "placement_id": logical.placement_id, + "placement_digest": logical.digest, + "participant_id": participant_id, + "instance_id": "instance", + "generation": 7, + "lease_id": "lease-7", + "fragments": default_fragments, + } + values.update(overrides) + return WeightRuntimeBindingManifest(**values) + + +def _contiguous_strides_bytes(shape: tuple[int, ...], itemsize: int) -> tuple[int, ...]: + result = [] + running = itemsize + for extent in reversed(shape): + result.append(running) + running *= extent + return tuple(reversed(result)) diff --git a/mooncake-reshard/tests/weight_manifest/test_api_and_types.py b/mooncake-reshard/tests/weight_manifest/test_api_and_types.py new file mode 100644 index 0000000000..636c8e0ff9 --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/test_api_and_types.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import inspect +import typing + +import pytest + +import mooncake.reshard.weight.manifest as model_weight +from mooncake.reshard.weight import ( + OwnershipAxis, + ParallelRank, + PlacementFragment, + ReplicatedAxis, + SplitAxis, + TensorDescriptor, +) + +from .helpers import descriptor + + +def test_public_api_is_minimal_and_explicit() -> None: + assert model_weight.__all__ == [ + "ParallelRank", + "ParallelTopology", + "PlacementFragment", + "TopologyParticipant", + "WeightPlacementManifest", + "WeightPlacementPart", + "RuntimeBindingFragment", + "WeightRuntimeBindingManifest", + "SplitAxis", + "ReplicatedAxis", + "OwnershipAxis", + "TensorDescriptor", + "validate_runtime_binding", + "validate_runtime_bindings", + ] + + +def test_public_type_hints_resolve() -> None: + for name in model_weight.__all__: + value = getattr(model_weight, name) + targets = (value, value.__init__) if inspect.isclass(value) else (value,) + for target in targets: + typing.get_type_hints(target) + + +def test_tensor_descriptor_uses_canonical_shard_dims() -> None: + single_axis = descriptor() + multidim = descriptor( + shard_dims=(0, 1), + global_shape=(8, 16, 32), + expert_id=None, + parallel_axes=( + SplitAxis(kind="ep", dim=0), + SplitAxis(kind="tp", dim=1), + ), + ) + + assert single_axis.shard_dims == (0,) + assert multidim.shard_dims == (0, 1) + + +def test_tensor_descriptor_carries_explicit_parallel_axis_semantics() -> None: + value = descriptor( + global_shape=(8, 16, 32), + shard_dims=(0, 2), + expert_id=None, + parallel_axes=( + SplitAxis(kind="tp", dim=2), + SplitAxis(kind="ep", dim=0), + OwnershipAxis(kind="pp"), + ReplicatedAxis(kind="dp"), + ), + ) + + assert tuple(axis.kind for axis in value.parallel_axes) == ( + "dp", + "pp", + "ep", + "tp", + ) + assert tuple( + axis.dim for axis in value.parallel_axes if isinstance(axis, SplitAxis) + ) == (0, 2) + + +@pytest.mark.parametrize( + ("parallel_axes", "shard_dims", "message"), + [ + ((SplitAxis(kind="tp", dim=1),), (0,), "shard_dims"), + ( + (SplitAxis(kind="tp", dim=0), ReplicatedAxis(kind="tp")), + (0,), + "duplicate", + ), + ( + (SplitAxis(kind="tp", dim=0), SplitAxis(kind="ep", dim=0)), + (0,), + "share a dimension", + ), + ], +) +def test_tensor_descriptor_rejects_ambiguous_parallel_axis_semantics( + parallel_axes: tuple[object, ...], + shard_dims: tuple[int, ...], + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + descriptor( + global_shape=(8, 8), + shard_dims=shard_dims, + parallel_axes=parallel_axes, + ) + + +def test_individually_allocated_expert_rejects_ep_split_semantics() -> None: + with pytest.raises(ValueError, match="individually allocated expert"): + descriptor( + expert_id=3, + parallel_axes=(SplitAxis(kind="ep", dim=0),), + ) + + assert descriptor( + expert_id=3, + shard_dims=(), + parallel_axes=(OwnershipAxis(kind="ep"),), + ) + + +@pytest.mark.parametrize( + ("factory", "message"), + [ + (lambda: SplitAxis(kind="unknown", dim=0), "kind"), + (lambda: SplitAxis(kind="tp", dim=-1), "at least"), + (lambda: SplitAxis(kind="tp", dim=True), "integer"), + (lambda: SplitAxis(kind="pp", dim=0), "split semantics"), + (lambda: SplitAxis(kind="dp", dim=0), "split semantics"), + (lambda: SplitAxis(kind="ep", dim=1), "leading"), + (lambda: ReplicatedAxis(kind="unknown"), "kind"), + (lambda: OwnershipAxis(kind="unknown"), "kind"), + ], +) +def test_explicit_parallel_axis_rejects_invalid_schema(factory, message: str) -> None: + with pytest.raises(ValueError, match=message): + factory() + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"global_shape": ()}, "global_shape"), + ({"global_shape": (8.0, 4)}, "integer"), + ({"itemsize": True}, "integer"), + ({"shard_dims": (0, 0)}, "duplicates"), + ({"shard_dims": (1, 0)}, "sorted"), + ({"shard_dims": (2,)}, "out-of-range"), + ({"shard_dims": (True,)}, "integer"), + ({"layout_fingerprint": ""}, "layout_fingerprint"), + ({"parallel_axes": (object(),)}, "explicit axis"), + ], +) +def test_tensor_descriptor_rejects_invalid_schema( + overrides: dict, message: str +) -> None: + with pytest.raises(ValueError, match=message): + descriptor(**overrides) + + +@pytest.mark.parametrize( + "shape", + [ + {8: None, 4: None}, + {8, 4}, + frozenset((8, 4)), + (extent for extent in (8, 4)), + ], +) +def test_tensor_descriptor_rejects_unordered_or_one_shot_shape(shape) -> None: + with pytest.raises(ValueError, match="global_shape must contain integers"): + descriptor(global_shape=shape) + + +def test_tensor_descriptor_accepts_ordered_shape_sequence() -> None: + assert descriptor(global_shape=[8, 4]).global_shape == (8, 4) + + +def test_tensor_descriptor_requires_explicit_layout_fingerprint() -> None: + values = { + "tensor_id": "weight", + "global_shape": (4, 4), + "dtype": "bfloat16", + "itemsize": 2, + "shard_dims": (0,), + "parallel_axes": (SplitAxis(kind="tp", dim=0),), + } + + with pytest.raises(TypeError, match="layout_fingerprint"): + TensorDescriptor(**values) + + +def test_placement_fragment_derives_stable_id_from_logical_content() -> None: + values = { + "tensor_id": "weight", + "global_offset": (0, 0), + "local_shape": (4, 4), + "nbytes": 32, + "rank": ParallelRank(tp=1), + } + + first = PlacementFragment(**values) + second = PlacementFragment(**values) + + assert first.placement_fragment_id == second.placement_fragment_id + assert first.placement_fragment_id.startswith("sha256:") diff --git a/mooncake-reshard/tests/weight_manifest/test_binding.py b/mooncake-reshard/tests/weight_manifest/test_binding.py new file mode 100644 index 0000000000..048ffa8b62 --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/test_binding.py @@ -0,0 +1,598 @@ +from __future__ import annotations + +from dataclasses import replace +from math import prod +from types import SimpleNamespace + +import pytest + +from mooncake.reshard.weight import ( + SplitAxis, + validate_runtime_binding, +) + +from .helpers import ( + binding_fragment, + binding_manifest, + descriptor, + placement_fragment, + placement_manifest, +) + + +def test_runtime_binding_fragment_retains_owner() -> None: + owner = object() + + fragment = binding_fragment(owner=owner) + + assert fragment.owner is owner + assert fragment.device == "cuda:0" + + +def test_runtime_binding_rejects_duck_typed_manifests() -> None: + placement = placement_manifest() + binding = binding_manifest(placement=placement) + placement_fields = dict(vars(placement)) + placement_fields.update( + resource_kind=placement.resource_kind, + digest=placement.digest, + ) + duck_placement = SimpleNamespace(**placement_fields) + duck_binding = SimpleNamespace( + **vars(binding), + resource_kind=binding.resource_kind, + ) + + with pytest.raises(ValueError, match="WeightPlacementManifest"): + validate_runtime_binding(duck_placement, binding) + with pytest.raises(ValueError, match="WeightRuntimeBindingManifest"): + validate_runtime_binding(placement, duck_binding) + + +def test_runtime_binding_rejects_noncanonical_stride() -> None: + placement = placement_manifest() + binding = binding_manifest( + placement=placement, + fragments=(binding_fragment(strides_bytes=(2, 8)),), + ) + + with pytest.raises(ValueError, match="stride"): + validate_runtime_binding(placement, binding) + + +@pytest.mark.parametrize( + ("shape", "strides_bytes", "canonical_strides_bytes"), + [ + ((1, 4), (200, 2), (8, 2)), + ((4, 1), (2, 200), (2, 2)), + ((2, 1, 4), (8, 0, 2), (8, 8, 2)), + ((1,), (0,), (2,)), + ], +) +def test_binding_normalizes_stride_on_singleton_dimensions( + shape: tuple[int, ...], + strides_bytes: tuple[int, ...], + canonical_strides_bytes: tuple[int, ...], +) -> None: + nbytes = prod(shape) * 2 + tensor = descriptor( + global_shape=shape, + shard_dims=(), + parallel_axes=(), + expert_id=None, + ) + placement = placement_manifest( + tensors=(tensor,), + fragments=( + placement_fragment( + global_offset=(0,) * len(shape), + local_shape=shape, + nbytes=nbytes, + ), + ), + ) + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment( + nbytes=nbytes, + local_shape=shape, + strides_bytes=strides_bytes, + ), + ), + ) + + assert binding.fragments[0].strides_bytes == canonical_strides_bytes + assert validate_runtime_binding(placement, binding) is None + + +def test_runtime_binding_does_not_trust_runtime_itemsize() -> None: + placement = placement_manifest() + binding = binding_manifest( + placement=placement, + fragments=(binding_fragment(itemsize=1),), + ) + + with pytest.raises(ValueError, match="itemsize"): + validate_runtime_binding(placement, binding) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"storage_offset_bytes": -1}, "storage_offset_bytes"), + ({"device": ""}, "device"), + ( + { + "address": 0x1000, + "storage_address": 0x1000, + "storage_nbytes": 64, + "storage_offset_bytes": 8, + }, + "storage_address plus", + ), + ( + { + "address": 0x1010, + "storage_address": 0x1000, + "storage_nbytes": 40, + "storage_offset_bytes": 16, + }, + "storage allocation", + ), + ], +) +def test_runtime_binding_fragment_rejects_unsafe_views( + overrides: dict, message: str +) -> None: + with pytest.raises(ValueError, match=message): + binding_fragment(**overrides) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"placement_digest": ""}, "placement_digest"), + ({"placement_digest": "g" * 64}, "SHA-256"), + ({"placement_digest": "a" * 63}, "SHA-256"), + ], +) +def test_runtime_binding_requires_content_attestation( + overrides: dict, message: str +) -> None: + with pytest.raises(ValueError, match=message): + binding_manifest(**overrides) + + +@pytest.mark.parametrize( + "overrides, message", + [ + ({"resource_id": "other"}, "resource_id"), + ({"revision": "other"}, "revision"), + ({"placement_id": "other"}, "placement_id"), + ], +) +def test_binding_rejects_identity_mismatch(overrides: dict, message: str) -> None: + with pytest.raises(ValueError, match=message): + validate_runtime_binding(placement_manifest(), binding_manifest(**overrides)) + + +def test_binding_requires_exact_fragment_set_and_size() -> None: + placement = placement_manifest() + + with pytest.raises(ValueError, match="missing placement fragment"): + validate_runtime_binding( + placement, + binding_manifest(placement=placement, fragments=()), + ) + with pytest.raises(ValueError, match="unknown placement fragment"): + validate_runtime_binding( + placement, + binding_manifest( + placement=placement, + fragments=(binding_fragment(placement_fragment_id="unknown"),), + ), + ) + with pytest.raises(ValueError, match="byte size"): + validate_runtime_binding( + placement, + binding_manifest( + placement=placement, + fragments=(binding_fragment(nbytes=64),), + ), + ) + + with pytest.raises(ValueError, match="local_shape"): + validate_runtime_binding( + placement, + binding_manifest( + placement=placement, + fragments=( + binding_fragment( + local_shape=(2, 8), + strides_bytes=(16, 2), + ), + ), + ), + ) + + +def test_binding_rejects_duplicate_fragment_ids() -> None: + fragment = binding_fragment() + + with pytest.raises(ValueError, match="duplicate placement fragment"): + binding_manifest(fragments=(fragment, replace(fragment, fragment_id="other"))) + with pytest.raises(ValueError, match="duplicate runtime fragment_id"): + binding_manifest( + fragments=( + fragment, + replace( + fragment, + placement_fragment_id="placement-other", + ), + ) + ) + + +def test_binding_allows_one_rank_to_span_runtime_locations() -> None: + placement = placement_manifest( + tensors=(descriptor(global_shape=(8, 4)),), + fragments=( + placement_fragment(placement_fragment_id="left"), + placement_fragment( + placement_fragment_id="right", + global_offset=(4, 0), + ), + ), + ) + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment(placement_fragment_id="left"), + binding_fragment( + placement_fragment_id="right", + fragment_id="runtime-right", + address=0x2000, + worker_id="worker-1", + endpoint="worker-1:12345", + ), + ), + ) + + assert validate_runtime_binding(placement, binding) is None + + +@pytest.mark.parametrize("right_address", [0x1000, 0x1010]) +def test_binding_rejects_overlapping_runtime_ranges(right_address: int) -> None: + placement = placement_manifest( + tensors=( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ), + fragments=( + placement_fragment( + placement_fragment_id="a", + tensor_id="a.weight", + ), + placement_fragment( + placement_fragment_id="b", + tensor_id="b.weight", + ), + ), + ) + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment(placement_fragment_id="a", fragment_id="runtime-a"), + binding_fragment( + placement_fragment_id="b", + fragment_id="runtime-b", + address=right_address, + endpoint="worker-0:54321", + ), + ), + ) + + with pytest.raises(ValueError, match="address ranges overlap"): + validate_runtime_binding(placement, binding) + + +def test_binding_rejects_partially_overlapping_storage_allocations() -> None: + placement = placement_manifest( + tensors=( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ), + fragments=( + placement_fragment(placement_fragment_id="a", tensor_id="a.weight"), + placement_fragment(placement_fragment_id="b", tensor_id="b.weight"), + ), + ) + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment( + placement_fragment_id="a", + fragment_id="runtime-a", + address=0x1100, + storage_address=0x1000, + storage_nbytes=0x1000, + storage_offset_bytes=0x100, + ), + binding_fragment( + placement_fragment_id="b", + fragment_id="runtime-b", + address=0x2100, + storage_address=0x1800, + storage_nbytes=0x1000, + storage_offset_bytes=0x900, + ), + ), + ) + + with pytest.raises(ValueError, match="storage allocation ranges overlap"): + validate_runtime_binding(placement, binding) + + +def test_binding_allows_disjoint_views_of_same_storage_allocation() -> None: + placement = placement_manifest( + tensors=( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ), + fragments=( + placement_fragment(placement_fragment_id="a", tensor_id="a.weight"), + placement_fragment(placement_fragment_id="b", tensor_id="b.weight"), + ), + ) + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment( + placement_fragment_id="a", + fragment_id="runtime-a", + address=0x1000, + storage_address=0x1000, + storage_nbytes=0x100, + storage_offset_bytes=0, + ), + binding_fragment( + placement_fragment_id="b", + fragment_id="runtime-b", + address=0x1080, + storage_address=0x1000, + storage_nbytes=0x100, + storage_offset_bytes=0x80, + ), + ), + ) + + assert validate_runtime_binding(placement, binding) is None + + +def test_binding_treats_endpoint_as_routing_not_address_space() -> None: + placement = placement_manifest( + tensors=( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ), + fragments=( + placement_fragment(placement_fragment_id="a", tensor_id="a.weight"), + placement_fragment(placement_fragment_id="b", tensor_id="b.weight"), + ), + ) + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment(placement_fragment_id="a", fragment_id="runtime-a"), + binding_fragment( + placement_fragment_id="b", + fragment_id="runtime-b", + endpoint="worker-0:54321", + ), + ), + ) + + with pytest.raises(ValueError, match="address ranges overlap"): + validate_runtime_binding(placement, binding) + + +@pytest.mark.parametrize( + ("field", "value"), + [("worker_id", "worker-1"), ("device", "cuda:1")], +) +def test_binding_separates_worker_and_device_address_spaces( + field: str, + value: str, +) -> None: + placement = placement_manifest( + tensors=( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ), + fragments=( + placement_fragment(placement_fragment_id="a", tensor_id="a.weight"), + placement_fragment(placement_fragment_id="b", tensor_id="b.weight"), + ), + ) + second_overrides = { + "placement_fragment_id": "b", + "fragment_id": "runtime-b", + field: value, + } + if field == "worker_id": + second_overrides["endpoint"] = "worker-1:12345" + binding = binding_manifest( + placement=placement, + fragments=( + binding_fragment(placement_fragment_id="a", fragment_id="runtime-a"), + binding_fragment(**second_overrides), + ), + ) + + assert validate_runtime_binding(placement, binding) is None + + +def test_binding_allows_only_exact_compatible_declared_aliases() -> None: + aliases = ("embed.weight", "head.weight") + tensors = ( + descriptor(tensor_id="embed.weight", expert_id=None), + descriptor(tensor_id="head.weight", expert_id=None), + ) + fragments = ( + placement_fragment( + placement_fragment_id="embed", + tensor_id="embed.weight", + aliases=aliases, + ), + placement_fragment( + placement_fragment_id="head", + tensor_id="head.weight", + aliases=aliases, + ), + ) + placement = placement_manifest(tensors=tensors, fragments=fragments) + binding_fragments = ( + binding_fragment( + placement_fragment_id="embed", + fragment_id="runtime-embed", + ), + binding_fragment( + placement_fragment_id="head", + fragment_id="runtime-head", + ), + ) + + assert ( + validate_runtime_binding( + placement, + binding_manifest(placement=placement, fragments=binding_fragments), + ) + is None + ) + + incompatible = placement_manifest( + tensors=( + tensors[0], + replace(tensors[1], layout_fingerprint="different"), + ), + fragments=fragments, + ) + with pytest.raises(ValueError, match="address ranges overlap"): + validate_runtime_binding( + incompatible, + binding_manifest( + placement=incompatible, + fragments=binding_fragments, + ), + ) + + incompatible_axes = placement_manifest( + tensors=( + tensors[0], + replace( + tensors[1], + parallel_axes=(SplitAxis(kind="ep", dim=0),), + ), + ), + fragments=fragments, + ) + with pytest.raises(ValueError, match="address ranges overlap"): + validate_runtime_binding( + incompatible_axes, + binding_manifest( + placement=incompatible_axes, + fragments=binding_fragments, + ), + ) + + +def test_binding_rejects_alias_group_without_tensor_members() -> None: + aliases = ("unrelated-x", "unrelated-y") + + with pytest.raises(ValueError, match="alias"): + placement_manifest( + tensors=( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ), + fragments=( + placement_fragment( + placement_fragment_id="a", + tensor_id="a.weight", + aliases=aliases, + ), + placement_fragment( + placement_fragment_id="b", + tensor_id="b.weight", + aliases=aliases, + ), + ), + ) + + +def test_binding_validation_preserves_logical_and_physical_halves() -> None: + placement = placement_manifest() + binding = binding_manifest(placement=placement) + + assert validate_runtime_binding(placement, binding) is None + assert placement.fragments[0].global_offset == (0, 0) + assert binding.fragments[0].address == 0x1000 + + +def test_binding_validation_is_order_independent() -> None: + tensors = ( + descriptor(tensor_id="a.weight"), + descriptor(tensor_id="b.weight"), + ) + fragments = ( + placement_fragment( + placement_fragment_id="a", + tensor_id="a.weight", + ), + placement_fragment( + placement_fragment_id="b", + tensor_id="b.weight", + ), + ) + placement = placement_manifest(tensors=tensors, fragments=fragments) + bindings = ( + binding_fragment( + placement_fragment_id="a", + fragment_id="runtime-a", + address=0x1000, + ), + binding_fragment( + placement_fragment_id="b", + fragment_id="runtime-b", + address=0x2000, + ), + ) + + first = validate_runtime_binding( + placement, + binding_manifest(placement=placement, fragments=bindings), + ) + second = validate_runtime_binding( + placement, + binding_manifest( + placement=placement, + fragments=tuple(reversed(bindings)), + ), + ) + + assert first == second + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"tensors": (), "fragments": ()}, "must contain tensors"), + ], +) +def test_global_placement_must_not_be_empty( + overrides: dict, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + placement_manifest(**overrides) diff --git a/mooncake-reshard/tests/weight_manifest/test_canonical_boundary.py b/mooncake-reshard/tests/weight_manifest/test_canonical_boundary.py new file mode 100644 index 0000000000..6c2553fa11 --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/test_canonical_boundary.py @@ -0,0 +1,233 @@ +"""Contract tests for the canonical framework/core boundary.""" + +import json +from dataclasses import fields +from typing import Any, Protocol, get_type_hints + +import pytest + +from mooncake.reshard.contracts import ( + PlacementManifest, + ResourceManifest, + RuntimeBindingManifest, +) +from mooncake.reshard.weight import ( + OwnershipAxis, + ParallelRank, + ParallelTopology, + ReplicatedAxis, + RuntimeBindingFragment, + SplitAxis, + TensorDescriptor, + TopologyParticipant, + WeightPlacementPart, + WeightRuntimeBindingManifest, + weight_placement_from_json, + weight_placement_to_json, +) + +from .helpers import descriptor, placement_fragment, placement_manifest + + +def test_resource_contracts_are_structural_protocols() -> None: + assert issubclass(ResourceManifest, Protocol) + assert issubclass(PlacementManifest, Protocol) + assert issubclass(RuntimeBindingManifest, Protocol) + + +def test_core_exposes_only_typed_manifest_constructors() -> None: + assert not hasattr(WeightPlacementPart, "from_runtime_inventory") + assert not hasattr(WeightRuntimeBindingManifest, "from_runtime_inventory") + assert get_type_hints(RuntimeBindingFragment)["owner"] != Any + + +def test_tensor_descriptor_has_one_canonical_shard_representation() -> None: + assert "partition_dim" not in {field.name for field in fields(TensorDescriptor)} + + descriptor = TensorDescriptor( + tensor_id="decoder.layer.0.mlp.weight", + global_shape=(8, 16), + dtype="float16", + itemsize=2, + shard_dims=(0,), + layout_fingerprint="row-major", + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ) + + assert descriptor.shard_dims == (0,) + + +def test_parallel_axis_semantics_are_explicit() -> None: + assert SplitAxis(kind="tp", dim=1).dim == 1 + assert ReplicatedAxis(kind="dp").kind == "dp" + assert OwnershipAxis(kind="pp").kind == "pp" + + with pytest.raises(ValueError, match="EP must split"): + SplitAxis(kind="ep", dim=1) + + +def test_weight_placement_serde_accepts_only_canonical_wire_fields() -> None: + placement = placement_manifest() + encoded = weight_placement_to_json(placement) + + assert weight_placement_from_json(encoded) == placement + assert "partition_dim" not in encoded + + +@pytest.mark.parametrize("alias", ["model_id", "partition_dim", "split_dim"]) +def test_weight_placement_serde_rejects_framework_and_legacy_aliases( + alias: str, +) -> None: + payload = json.loads(weight_placement_to_json(placement_manifest())) + if alias == "model_id": + payload[alias] = payload.pop("resource_id") + elif alias == "partition_dim": + payload["tensors"][0][alias] = payload["tensors"][0]["shard_dims"][0] + else: + axis = payload["tensors"][0]["parallel_axes"][0] + axis[alias] = axis.pop("dim") + + with pytest.raises(ValueError, match="schema fields"): + weight_placement_from_json(json.dumps(payload)) + + +def test_split_axis_collectively_covers_once_with_all_split_ranks() -> None: + topology = _two_rank_topology("tp") + tensor = descriptor( + global_shape=(8, 4), + shard_dims=(0,), + expert_id=None, + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=( + placement_fragment( + placement_fragment_id="tp-0", + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(tp=0), + ), + placement_fragment( + placement_fragment_id="tp-1", + global_offset=(4, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(tp=1), + ), + ), + ) + + assert len(placement.fragments) == 2 + + +def test_replicated_axis_requires_an_independent_cover_for_every_rank() -> None: + topology = _two_rank_topology("tp") + tensor = descriptor( + shard_dims=(), + expert_id=None, + parallel_axes=(ReplicatedAxis(kind="tp"),), + ) + + with pytest.raises(ValueError, match="replicated-axis participant"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=(placement_fragment(rank=ParallelRank(tp=0)),), + ) + + +def test_replicated_axis_accepts_independent_complete_covers() -> None: + topology = _two_rank_topology("tp") + tensor = descriptor( + shard_dims=(), + expert_id=None, + parallel_axes=(ReplicatedAxis(kind="tp"),), + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=( + placement_fragment( + placement_fragment_id="replica-0", + rank=ParallelRank(tp=0), + ), + placement_fragment( + placement_fragment_id="replica-1", + rank=ParallelRank(tp=1), + ), + ), + ) + + assert len(placement.fragments) == 2 + + +def test_ownership_axis_requires_only_declared_owners() -> None: + topology = _two_rank_topology("pp") + tensor = descriptor( + shard_dims=(), + expert_id=None, + parallel_axes=(OwnershipAxis(kind="pp"),), + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=(placement_fragment(rank=ParallelRank(pp=0)),), + ) + + assert {fragment.rank.pp for fragment in placement.fragments} == {0} + + +def test_ownership_axis_rejects_cover_split_across_declared_owners() -> None: + topology = _two_rank_topology("pp") + tensor = descriptor( + global_shape=(8, 4), + shard_dims=(0,), + expert_id=None, + parallel_axes=( + OwnershipAxis(kind="pp"), + SplitAxis(kind="tp", dim=0), + ), + ) + + with pytest.raises(ValueError, match="not fully covered"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=( + placement_fragment( + placement_fragment_id="owner-0", + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(pp=0), + ), + placement_fragment( + placement_fragment_id="owner-1", + global_offset=(4, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(pp=1), + ), + ), + ) + + +def _two_rank_topology(kind: str) -> ParallelTopology: + first = ParallelRank() + second_values = {"dp": 0, "tp": 0, "pp": 0, "ep": 0} + second_values[kind] = 1 + second = ParallelRank(**second_values) + return ParallelTopology( + tp_size=2 if kind == "tp" else 1, + pp_size=2 if kind == "pp" else 1, + ep_size=2 if kind == "ep" else 1, + dp_size=2 if kind == "dp" else 1, + participants=( + TopologyParticipant("worker-0", first), + TopologyParticipant("worker-1", second), + ), + ) diff --git a/mooncake-reshard/tests/weight_manifest/test_global_placement.py b/mooncake-reshard/tests/weight_manifest/test_global_placement.py new file mode 100644 index 0000000000..a9fed095ac --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/test_global_placement.py @@ -0,0 +1,965 @@ +from __future__ import annotations + +from dataclasses import replace + +import pytest + +import mooncake.reshard.weight.serde as weight_serde +from mooncake.reshard.weight import ( + OwnershipAxis, + ParallelRank, + ParallelTopology, + PlacementFragment, + ReplicatedAxis, + RuntimeBindingFragment, + SplitAxis, + TopologyParticipant, + WeightPlacementManifest, + WeightPlacementPart, + WeightRuntimeBindingManifest, + validate_runtime_binding, + validate_runtime_bindings, + weight_placement_from_json, + weight_placement_to_json, +) + +from .helpers import ( + MODEL_ID, + PLACEMENT_SET_ID, + REVISION, + WEIGHT_GENERATION, + _contiguous_strides_bytes, + descriptor, + parallel_topology, + placement_manifest, + placement_part, +) + + +def _topology() -> ParallelTopology: + return parallel_topology( + participants=( + TopologyParticipant("worker-0", ParallelRank(tp=0)), + TopologyParticipant("worker-1", ParallelRank(tp=1)), + ), + ) + + +def _part( + participant_id: str, + rank: ParallelRank, + offset: int, + *, + placement_set_id: str = "placement-set-7", + weight_generation: int = 7, + topology_id: str | None = None, +) -> WeightPlacementPart: + topology = _topology() + tensor = descriptor( + global_shape=(8, 4), + parallel_axes=( + ReplicatedAxis(kind="dp"), + OwnershipAxis(kind="pp"), + SplitAxis(kind="tp", dim=0), + ), + ) + return placement_part( + topology=topology, + resource_id=MODEL_ID, + revision=REVISION, + weight_generation=weight_generation, + placement_set_id=placement_set_id, + topology_id=topology_id or topology.topology_id, + participant_id=participant_id, + rank=rank, + tensors=(tensor,), + fragments=( + PlacementFragment( + placement_fragment_id=f"fragment-{participant_id}", + tensor_id=tensor.tensor_id, + global_offset=(offset, 0), + local_shape=(4, 4), + nbytes=32, + rank=rank, + ), + ), + ) + + +def _manifest( + *, + topology: ParallelTopology | None = None, + parts: tuple[WeightPlacementPart, ...] | None = None, +) -> WeightPlacementManifest: + topology = topology or _topology() + selected_parts = ( + parts + if parts is not None + else ( + _part("worker-0", ParallelRank(tp=0), 0), + _part("worker-1", ParallelRank(tp=1), 4), + ) + ) + return placement_manifest( + resource_id=MODEL_ID, + revision=REVISION, + weight_generation=WEIGHT_GENERATION, + placement_set_id=PLACEMENT_SET_ID, + topology=topology, + parts=selected_parts, + ) + + +def test_weight_placement_manifest_is_one_complete_global_placement() -> None: + placement = _manifest() + + assert placement.topology.world_size == 2 + assert len(placement.parts) == 2 + assert len(placement.tensors) == 1 + assert len(placement.fragments) == 2 + assert weight_placement_from_json(weight_placement_to_json(placement)) == placement + + +def test_placement_part_rejects_unreferenced_tensor_descriptors() -> None: + with pytest.raises(ValueError, match="unreferenced tensor"): + placement_part(tensors=(descriptor(),), fragments=()) + + +def test_flat_placement_rejects_unreferenced_tensor_descriptors() -> None: + topology = parallel_topology() + referenced = descriptor(tensor_id="referenced.weight") + unreferenced = descriptor(tensor_id="unreferenced.weight") + fragment = PlacementFragment( + placement_fragment_id="referenced-fragment", + tensor_id=referenced.tensor_id, + global_offset=(0, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(), + ) + + with pytest.raises(ValueError, match="unreferenced tensor"): + WeightPlacementManifest.from_fragments( + resource_id=MODEL_ID, + revision=REVISION, + weight_generation=WEIGHT_GENERATION, + placement_set_id=PLACEMENT_SET_ID, + topology=topology, + tensors=(referenced, unreferenced), + fragments=(fragment,), + ) + + +def test_global_placement_rejects_alias_member_missing_from_tensor_inventory() -> None: + tensor = descriptor(tensor_id="tied.embedding") + fragment = PlacementFragment( + placement_fragment_id="tied-embedding", + tensor_id=tensor.tensor_id, + global_offset=(0, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(), + aliases=("tied.embedding", "tied.output"), + ) + + with pytest.raises(ValueError, match="alias group references unknown tensor"): + placement_manifest(tensors=(tensor,), fragments=(fragment,)) + + +def test_global_placement_requires_every_alias_member_to_declare_the_group() -> None: + embedding = descriptor(tensor_id="tied.embedding") + output = descriptor(tensor_id="tied.output") + aliases = (embedding.tensor_id, output.tensor_id) + fragments = ( + PlacementFragment( + placement_fragment_id="tied-embedding", + tensor_id=embedding.tensor_id, + global_offset=(0, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(), + aliases=aliases, + ), + PlacementFragment( + placement_fragment_id="tied-output", + tensor_id=output.tensor_id, + global_offset=(0, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(), + ), + ) + + with pytest.raises(ValueError, match="alias group is not declared consistently"): + placement_manifest( + tensors=(embedding, output), + fragments=fragments, + ) + + +def test_global_placement_caches_digest_for_all_participant_bindings( + monkeypatch, +) -> None: + placement = _manifest() + original = weight_serde.weight_placement_to_json + calls = 0 + + def counted_to_json(value): + nonlocal calls + calls += 1 + return original(value) + + monkeypatch.setattr(weight_serde, "weight_placement_to_json", counted_to_json) + + assert placement.digest == placement.digest + assert calls == 1 + + +def test_topology_uses_explicit_participants_instead_of_axis_product() -> None: + topology = ParallelTopology( + tp_size=2, + pp_size=1, + ep_size=2, + dp_size=1, + participants=( + TopologyParticipant("worker-0", ParallelRank(tp=0, ep=0)), + TopologyParticipant("worker-1", ParallelRank(tp=1, ep=1)), + ), + ) + + assert topology.world_size == 2 + assert topology.world_size != topology.tp_size * topology.ep_size + + +def test_global_placement_may_select_one_complete_dp_replica() -> None: + topology = ParallelTopology( + tp_size=2, + pp_size=4, + ep_size=1, + dp_size=2, + participants=( + TopologyParticipant("worker-0", ParallelRank(dp=0, tp=0, pp=1)), + TopologyParticipant("worker-1", ParallelRank(dp=0, tp=1, pp=1)), + ), + ) + placement = _manifest( + topology=topology, + parts=( + _part( + "worker-0", + ParallelRank(dp=0, tp=0, pp=1), + 0, + topology_id=topology.topology_id, + ), + _part( + "worker-1", + ParallelRank(dp=0, tp=1, pp=1), + 4, + topology_id=topology.topology_id, + ), + ), + ) + + assert placement.topology.dp_size == 2 + assert {part.rank.dp for part in placement.parts} == {0} + + +def test_global_placement_rejects_an_empty_selected_dp_replica() -> None: + topology = ParallelTopology( + tp_size=2, + pp_size=1, + ep_size=1, + dp_size=2, + participants=( + TopologyParticipant("dp0-tp0", ParallelRank(dp=0, tp=0)), + TopologyParticipant("dp0-tp1", ParallelRank(dp=0, tp=1)), + TopologyParticipant("dp1-tp0", ParallelRank(dp=1, tp=0)), + TopologyParticipant("dp1-tp1", ParallelRank(dp=1, tp=1)), + ), + ) + complete_parts = ( + _part( + "dp0-tp0", + ParallelRank(dp=0, tp=0), + 0, + topology_id=topology.topology_id, + ), + _part( + "dp0-tp1", + ParallelRank(dp=0, tp=1), + 4, + topology_id=topology.topology_id, + ), + ) + empty_parts = tuple( + WeightPlacementPart( + resource_id=MODEL_ID, + revision=REVISION, + weight_generation=WEIGHT_GENERATION, + placement_set_id=PLACEMENT_SET_ID, + topology_id=topology.topology_id, + participant_id=f"dp1-tp{tp}", + rank=ParallelRank(dp=1, tp=tp), + tensors=(), + fragments=(), + ) + for tp in range(2) + ) + + with pytest.raises(ValueError, match="replicated-axis participant"): + _manifest(topology=topology, parts=complete_parts + empty_parts) + + +def test_global_placement_rejects_missing_or_duplicate_participants() -> None: + left = _part("worker-0", ParallelRank(tp=0), 0) + right = _part("worker-1", ParallelRank(tp=1), 4) + + with pytest.raises(ValueError, match="missing topology participant"): + _manifest(parts=(left,)) + with pytest.raises(ValueError, match="duplicate placement participant"): + _manifest(parts=(left, replace(right, participant_id="worker-0"))) + + +@pytest.mark.parametrize( + ("replacement", "message"), + [ + ({"resource_id": "other-model"}, "resource_id"), + ({"revision": "other-revision"}, "revision"), + ({"placement_set_id": "other-set"}, "placement_set_id"), + ({"weight_generation": 8}, "weight_generation"), + ({"topology_id": "sha256:" + "0" * 64}, "topology_id"), + ], +) +def test_global_placement_rejects_parts_from_another_collection( + replacement: dict[str, object], message: str +) -> None: + left = _part("worker-0", ParallelRank(tp=0), 0) + right = replace(_part("worker-1", ParallelRank(tp=1), 4), **replacement) + + with pytest.raises(ValueError, match=message): + _manifest(parts=(left, right)) + + +def test_global_placement_rejects_incomplete_tensor_coverage() -> None: + left = _part("worker-0", ParallelRank(tp=0), 0) + right = _part("worker-1", ParallelRank(tp=1), 4) + incomplete = replace( + right, + fragments=( + replace( + right.fragments[0], + local_shape=(2, 4), + nbytes=16, + ), + ), + ) + + with pytest.raises(ValueError, match="not fully covered"): + _manifest(parts=(left, incomplete)) + + +def test_global_placement_requires_every_split_axis_participant() -> None: + topology = _topology() + present = _part("worker-0", ParallelRank(tp=0), 0) + present = replace( + present, + fragments=( + replace( + present.fragments[0], + local_shape=(8, 4), + nbytes=64, + ), + ), + ) + missing = WeightPlacementPart( + resource_id=MODEL_ID, + revision=REVISION, + weight_generation=WEIGHT_GENERATION, + placement_set_id=PLACEMENT_SET_ID, + topology_id=topology.topology_id, + participant_id="worker-1", + rank=ParallelRank(tp=1), + tensors=(), + fragments=(), + ) + + with pytest.raises(ValueError, match="split-axis participant"): + _manifest(topology=topology, parts=(present, missing)) + + +def test_global_placement_rejects_full_shard_on_each_split_rank() -> None: + tensor = descriptor( + tensor_id="model.layers.0.mlp.weight", + global_shape=(8, 4), + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ) + topology = _topology() + fragments = tuple( + PlacementFragment( + placement_fragment_id=f"tp{tp_rank}-full", + tensor_id=tensor.tensor_id, + global_offset=(0, 0), + local_shape=tensor.global_shape, + nbytes=64, + rank=ParallelRank(tp=tp_rank), + ) + for tp_rank in range(2) + ) + + with pytest.raises(ValueError, match="not fully covered"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_global_placement_rejects_missing_declared_split_rank() -> None: + topology = parallel_topology(tp_size=2) + + with pytest.raises(ValueError, match="split-axis participant"): + placement_manifest(topology=topology) + + +def test_global_placement_accepts_independent_tp_replicas() -> None: + tensor = descriptor( + tensor_id="model.norm.weight", + global_shape=(4,), + shard_dims=(), + parallel_axes=(ReplicatedAxis(kind="tp"),), + layer_id=None, + expert_id=None, + ) + topology = _topology() + fragments = tuple( + PlacementFragment( + placement_fragment_id=f"replica-{tp_rank}", + tensor_id=tensor.tensor_id, + global_offset=(0,), + local_shape=(4,), + nbytes=8, + rank=ParallelRank(tp=tp_rank), + ) + for tp_rank in range(2) + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + assert len(placement.fragments) == 2 + + +def test_global_placement_accepts_complete_replicas_on_multiple_pp_owners() -> None: + tensor = descriptor( + tensor_id="lm_head.weight", + global_shape=(8, 4), + layer_id=None, + expert_id=None, + parallel_axes=( + OwnershipAxis(kind="pp"), + SplitAxis(kind="tp", dim=0), + ), + ) + participants = tuple( + TopologyParticipant( + f"pp{pp_rank}-tp{tp_rank}", + ParallelRank(tp=tp_rank, pp=pp_rank), + ) + for pp_rank in range(2) + for tp_rank in range(2) + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=f"pp{pp_rank}-tp{tp_rank}", + tensor_id=tensor.tensor_id, + global_offset=(tp_rank * 4, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(tp=tp_rank, pp=pp_rank), + ) + for pp_rank in range(2) + for tp_rank in range(2) + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + assert {fragment.rank.pp for fragment in placement.fragments} == {0, 1} + + +def test_global_placement_rejects_coverage_split_across_pp_owners() -> None: + tensor = descriptor( + tensor_id="lm_head.weight", + global_shape=(8, 4), + layer_id=None, + expert_id=None, + parallel_axes=( + OwnershipAxis(kind="pp"), + SplitAxis(kind="tp", dim=0), + ), + ) + participants = ( + TopologyParticipant("pp0-tp0", ParallelRank(tp=0, pp=0)), + TopologyParticipant("pp1-tp1", ParallelRank(tp=1, pp=1)), + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + global_offset=(participant.rank.tp * 4, 0), + local_shape=(4, 4), + nbytes=32, + rank=participant.rank, + ) + for participant in participants + ) + + with pytest.raises(ValueError, match="not fully covered"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +@pytest.mark.parametrize("undeclared_kind", ["pp", "dp"]) +def test_global_placement_rejects_coverage_across_undeclared_axis( + undeclared_kind: str, +) -> None: + tensor = descriptor( + tensor_id="lm_head.weight", + global_shape=(8, 4), + layer_id=None, + expert_id=None, + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ) + rank_overrides = ( + {"tp": 0, undeclared_kind: 0}, + {"tp": 1, undeclared_kind: 1}, + ) + participants = tuple( + TopologyParticipant( + f"worker-{index}", + ParallelRank(**coordinates), + ) + for index, coordinates in enumerate(rank_overrides) + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + global_offset=(participant.rank.tp * 4, 0), + local_shape=(4, 4), + nbytes=32, + rank=participant.rank, + ) + for participant in participants + ) + + with pytest.raises(ValueError, match=f"undeclared {undeclared_kind} axis"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_global_placement_rejects_implicit_pp_owner() -> None: + tensor = descriptor( + tensor_id="lm_head.weight", + global_shape=(4, 4), + shard_dims=(), + layer_id=None, + expert_id=None, + parallel_axes=(), + ) + topology = parallel_topology( + participants=( + TopologyParticipant("pp-0", ParallelRank(pp=0)), + TopologyParticipant("pp-1", ParallelRank(pp=1)), + ) + ) + + with pytest.raises(ValueError, match="active undeclared pp axis"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=( + PlacementFragment( + placement_fragment_id="pp-1-fragment", + tensor_id=tensor.tensor_id, + global_offset=(0, 0), + local_shape=(4, 4), + nbytes=32, + rank=ParallelRank(pp=1), + ), + ), + ) + + +def test_global_placement_rejects_swapped_cartesian_ep_tp_geometry() -> None: + tensor = descriptor( + tensor_id="experts.weight", + global_shape=(4, 4), + shard_dims=(0, 1), + expert_id=None, + parallel_axes=( + SplitAxis(kind="ep", dim=0), + SplitAxis(kind="tp", dim=1), + ), + ) + participants = tuple( + TopologyParticipant( + f"ep{ep_rank}-tp{tp_rank}", + ParallelRank(ep=ep_rank, tp=tp_rank), + ) + for ep_rank in range(2) + for tp_rank in range(2) + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + # Deliberately assign TP to dim0 and EP to dim1. + global_offset=(participant.rank.tp * 2, participant.rank.ep * 2), + local_shape=(2, 2), + nbytes=8, + rank=participant.rank, + ) + for participant in participants + ) + + with pytest.raises(ValueError, match="split-axis rank geometry"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_global_placement_rejects_non_cartesian_multi_split_geometry() -> None: + tensor = descriptor( + tensor_id="experts.weight", + global_shape=(4, 4), + shard_dims=(0, 1), + expert_id=None, + parallel_axes=( + SplitAxis(kind="ep", dim=0), + SplitAxis(kind="tp", dim=1), + ), + ) + participants = ( + TopologyParticipant("ep0-tp0", ParallelRank(ep=0, tp=0)), + TopologyParticipant("ep1-tp1", ParallelRank(ep=1, tp=1)), + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + # The boxes cover the tensor only along EP/dim0. TP/dim1 is not split. + global_offset=(participant.rank.ep * 2, 0), + local_shape=(2, 4), + nbytes=16, + rank=participant.rank, + ) + for participant in participants + ) + + with pytest.raises(ValueError, match="non-Cartesian split-axis"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_global_placement_accepts_ep_rank_derived_from_tp_rank() -> None: + tensor = descriptor( + tensor_id="dense.weight", + global_shape=(8, 4), + expert_id=None, + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ) + participants = tuple( + TopologyParticipant( + f"tp-{tp_rank}", + ParallelRank(tp=tp_rank, ep=tp_rank // 2), + ) + for tp_rank in range(4) + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + global_offset=(participant.rank.tp * 2, 0), + local_shape=(2, 4), + nbytes=16, + rank=participant.rank, + ) + for participant in participants + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + assert len(placement.fragments) == 4 + + +def test_global_placement_rejects_independent_undeclared_ep_rank() -> None: + tensor = descriptor( + tensor_id="dense.weight", + global_shape=(8, 4), + expert_id=None, + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ) + participants = tuple( + TopologyParticipant( + f"ep{ep_rank}-tp{tp_rank}", + ParallelRank(ep=ep_rank, tp=tp_rank), + ) + for ep_rank in range(2) + for tp_rank in range(2) + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + global_offset=(participant.rank.tp * 4, 0), + local_shape=(4, 4), + nbytes=32, + rank=participant.rank, + ) + for participant in participants + ) + + with pytest.raises(ValueError, match="independently.*undeclared ep axis"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_global_placement_rejects_coverage_split_across_ep_owners() -> None: + tensor = descriptor( + tensor_id="experts.weight", + global_shape=(8, 4), + expert_id=None, + parallel_axes=( + OwnershipAxis(kind="ep"), + SplitAxis(kind="tp", dim=0), + ), + ) + participants = ( + TopologyParticipant("ep0-tp0", ParallelRank(ep=0, tp=0)), + TopologyParticipant("ep1-tp1", ParallelRank(ep=1, tp=1)), + ) + topology = parallel_topology(participants=participants) + fragments = tuple( + PlacementFragment( + placement_fragment_id=participant.participant_id, + tensor_id=tensor.tensor_id, + global_offset=(participant.rank.tp * 4, 0), + local_shape=(4, 4), + nbytes=32, + rank=participant.rank, + ) + for participant in participants + ) + + with pytest.raises(ValueError, match="not fully covered"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_topology_rejects_out_of_range_rank_coordinates() -> None: + with pytest.raises(ValueError, match="tp rank"): + ParallelTopology( + tp_size=2, + pp_size=1, + ep_size=1, + dp_size=1, + participants=(TopologyParticipant("worker-0", ParallelRank(tp=2)),), + ) + + +def _binding( + placement: WeightPlacementManifest, + participant_id: str, + *, + instance_id: str, +) -> WeightRuntimeBindingManifest: + part = next( + item for item in placement.parts if item.participant_id == participant_id + ) + tensor_by_id = {tensor.tensor_id: tensor for tensor in placement.tensors} + return WeightRuntimeBindingManifest( + resource_id=placement.resource_id, + revision=placement.revision, + placement_id=placement.placement_id, + placement_digest=placement.digest, + participant_id=participant_id, + instance_id=instance_id, + generation=3, + lease_id=f"lease-{participant_id}", + fragments=tuple( + RuntimeBindingFragment( + placement_fragment_id=fragment.placement_fragment_id, + fragment_id=f"runtime-{fragment.placement_fragment_id}", + address=0x1000 + index * 0x100, + nbytes=fragment.nbytes, + worker_id=participant_id, + endpoint=f"{participant_id}:12345", + device="cuda:0", + itemsize=tensor_by_id[fragment.tensor_id].itemsize, + local_shape=fragment.local_shape, + strides_bytes=_contiguous_strides_bytes( + fragment.local_shape, + tensor_by_id[fragment.tensor_id].itemsize, + ), + storage_address=0x1000 + index * 0x100, + storage_nbytes=fragment.nbytes, + storage_offset_bytes=0, + ) + for index, fragment in enumerate(part.fragments) + ), + ) + + +def test_runtime_bindings_bind_each_part_of_one_global_placement() -> None: + placement = _manifest() + bindings = ( + _binding(placement, "worker-0", instance_id="source-0"), + _binding(placement, "worker-1", instance_id="source-1"), + ) + + assert validate_runtime_binding(placement, bindings[0]) is None + assert validate_runtime_bindings(placement, bindings) is None + + +def test_runtime_binding_set_rejects_cross_participant_address_overlap() -> None: + placement = _manifest() + left = _binding(placement, "worker-0", instance_id="shared-instance") + right = _binding(placement, "worker-1", instance_id="shared-instance") + left_fragment = left.fragments[0] + right_fragment = replace( + right.fragments[0], + address=left_fragment.address, + worker_id=left_fragment.worker_id, + endpoint=left_fragment.endpoint, + device=left_fragment.device, + storage_address=left_fragment.storage_address, + storage_nbytes=left_fragment.storage_nbytes, + storage_offset_bytes=left_fragment.storage_offset_bytes, + ) + + with pytest.raises(ValueError, match="address ranges overlap"): + validate_runtime_bindings( + placement, + (left, replace(right, fragments=(right_fragment,))), + ) + + +def test_runtime_binding_set_requires_every_global_part_exactly_once() -> None: + placement = _manifest() + left = _binding(placement, "worker-0", instance_id="source-0") + right = _binding(placement, "worker-1", instance_id="source-1") + + with pytest.raises(ValueError, match="missing runtime binding participant"): + validate_runtime_bindings(placement, (left,)) + with pytest.raises(ValueError, match="duplicate runtime binding participant"): + validate_runtime_bindings( + placement, + (left, replace(right, participant_id="worker-0")), + ) + with pytest.raises(ValueError, match="unknown runtime binding participant"): + validate_runtime_bindings( + placement, + (left, right, replace(right, participant_id="worker-unknown")), + ) + + +def test_runtime_binding_set_does_not_require_empty_participants() -> None: + topology = ParallelTopology( + tp_size=2, + pp_size=2, + ep_size=1, + dp_size=1, + participants=( + TopologyParticipant("worker-0", ParallelRank(tp=0, pp=0)), + TopologyParticipant("worker-1", ParallelRank(tp=1, pp=0)), + TopologyParticipant("worker-empty", ParallelRank(tp=0, pp=1)), + ), + ) + placement = _manifest( + topology=topology, + parts=( + _part( + "worker-0", + ParallelRank(tp=0, pp=0), + 0, + topology_id=topology.topology_id, + ), + _part( + "worker-1", + ParallelRank(tp=1, pp=0), + 4, + topology_id=topology.topology_id, + ), + WeightPlacementPart( + resource_id=MODEL_ID, + revision=REVISION, + weight_generation=WEIGHT_GENERATION, + placement_set_id=PLACEMENT_SET_ID, + topology_id=topology.topology_id, + participant_id="worker-empty", + rank=ParallelRank(tp=0, pp=1), + tensors=(), + fragments=(), + ), + ), + ) + bindings = ( + _binding(placement, "worker-0", instance_id="source-0"), + _binding(placement, "worker-1", instance_id="source-1"), + ) + + assert weight_placement_from_json(weight_placement_to_json(placement)) == placement + assert validate_runtime_bindings(placement, bindings) is None + + +def test_runtime_binding_can_only_bind_its_declared_part() -> None: + placement = _manifest() + left = _binding(placement, "worker-0", instance_id="source-0") + right_part = next( + item for item in placement.parts if item.participant_id == "worker-1" + ) + wrong_fragment = replace( + left.fragments[0], + placement_fragment_id=right_part.fragments[0].placement_fragment_id, + ) + + with pytest.raises(ValueError, match="unknown placement fragment"): + validate_runtime_binding( + placement, + replace(left, fragments=(wrong_fragment,)), + ) diff --git a/mooncake-reshard/tests/weight_manifest/test_placement.py b/mooncake-reshard/tests/weight_manifest/test_placement.py new file mode 100644 index 0000000000..9fccc502c5 --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/test_placement.py @@ -0,0 +1,297 @@ +from __future__ import annotations + +import json +import pytest + +from mooncake.reshard.weight import ( + OwnershipAxis, + ParallelRank, + ReplicatedAxis, + SplitAxis, + weight_placement_from_json, + weight_placement_to_json, +) + +from .helpers import ( + descriptor, + parallel_topology, + placement_fragment, + placement_manifest, +) + + +def test_placement_round_trip_is_stable_and_address_free() -> None: + placement = placement_manifest() + + encoded = weight_placement_to_json(placement) + decoded = weight_placement_from_json(encoded) + + assert decoded == placement + assert decoded.digest == placement.digest + assert placement.placement_id == ( + "sha256:618f22d9994d8d327a7afd97f3cfcf1f6680a337d71ea3a8162a7827ed0346db" + ) + assert placement.digest == ( + "a2ade0db5b17691e6a7c16d589d0498a64359eef91589addc870590899512d98" + ) + assert encoded == weight_placement_to_json(placement) + for forbidden in ( + "address", + "endpoint", + "worker_id", + "instance_id", + '"generation":', + "lease_id", + "owner", + ): + assert forbidden not in encoded + + +def test_placement_digest_is_independent_of_inventory_order() -> None: + tensors = ( + descriptor( + tensor_id="b.weight", + shard_dims=(), + parallel_axes=(OwnershipAxis(kind="tp"),), + expert_id=None, + ), + descriptor( + tensor_id="a.weight", + shard_dims=(), + parallel_axes=(OwnershipAxis(kind="tp"),), + expert_id=None, + ), + ) + fragments = ( + placement_fragment( + placement_fragment_id="b", + tensor_id="b.weight", + rank=ParallelRank(tp=1), + ), + placement_fragment( + placement_fragment_id="a", + tensor_id="a.weight", + rank=ParallelRank(tp=0), + ), + ) + + first = placement_manifest(tensors=tensors, fragments=fragments) + second = placement_manifest( + tensors=tuple(reversed(tensors)), + fragments=tuple(reversed(fragments)), + ) + + assert first == second + assert weight_placement_to_json(first) == weight_placement_to_json(second) + assert first.digest == second.digest + + +def test_weight_generation_changes_canonical_placement_identity() -> None: + generation_1 = placement_manifest(weight_generation=1) + generation_2 = placement_manifest(weight_generation=2) + + assert generation_1.revision == generation_2.revision + assert generation_1.placement_id != generation_2.placement_id + assert generation_1.digest != generation_2.digest + assert ( + weight_placement_from_json( + weight_placement_to_json(generation_2) + ).weight_generation + == 2 + ) + + +def test_placement_uses_one_canonical_shard_representation() -> None: + placement = placement_manifest(tensors=(descriptor(shard_dims=(0,)),)) + + assert placement.tensors[0].shard_dims == (0,) + assert "partition_dim" not in weight_placement_to_json(placement) + + +@pytest.mark.parametrize("mutation", ["missing", "unknown", "nan"]) +def test_placement_json_requires_strict_schema(mutation: str) -> None: + raw = json.loads(weight_placement_to_json(placement_manifest())) + if mutation == "missing": + del raw["revision"] + elif mutation == "unknown": + raw["future_semantics"] = "required" + else: + raw["resource_id"] = float("nan") + + with pytest.raises(ValueError): + weight_placement_from_json(json.dumps(raw)) + + +@pytest.mark.parametrize( + ("path", "mutation"), + [ + (("tensors", 0), ("pop", "dtype")), + (("tensors", 0), ("set", "future_semantics")), + (("topology",), ("pop", "tp_size")), + (("topology",), ("set", "future_semantics")), + (("topology", "participants", 0), ("pop", "participant_id")), + (("topology", "participants", 0), ("set", "future_semantics")), + (("topology", "participants", 0, "rank"), ("pop", "tp")), + (("tensors", 0, "parallel_axes", 0), ("pop", "kind")), + ( + ("tensors", 0, "parallel_axes", 0), + ("set", "future_semantics"), + ), + (("parts", 0), ("pop", "participant_id")), + (("parts", 0), ("set", "future_semantics")), + (("parts", 0, "rank"), ("set", "future_semantics")), + (("parts", 0, "fragments", 0), ("pop", "nbytes")), + (("parts", 0, "fragments", 0), ("set", "future_semantics")), + (("parts", 0, "fragments", 0, "rank"), ("pop", "tp")), + ( + ("parts", 0, "fragments", 0, "rank"), + ("set", "future_semantics"), + ), + ], +) +def test_placement_json_requires_strict_nested_schema( + path: tuple, mutation: tuple[str, str] +) -> None: + raw = json.loads(weight_placement_to_json(placement_manifest())) + target = raw + for component in path: + target = target[component] + operation, field = mutation + if operation == "pop": + target.pop(field) + else: + target[field] = "unsupported" + + with pytest.raises(ValueError, match="schema"): + weight_placement_from_json(json.dumps(raw)) + + +@pytest.mark.parametrize("value", ["not-json", "[]", '"placement"']) +def test_placement_json_rejects_invalid_document(value: str) -> None: + with pytest.raises(ValueError): + weight_placement_from_json(value) + + +def test_placement_json_rejects_duplicate_object_keys() -> None: + encoded = weight_placement_to_json(placement_manifest()) + duplicated = encoded.replace( + '"resource_id":"model"', + '"resource_id":"model","resource_id":"other"', + 1, + ) + + with pytest.raises(ValueError, match="duplicate JSON field"): + weight_placement_from_json(duplicated) + + +def test_placement_json_rejects_unreferenced_tensor_descriptors() -> None: + raw = json.loads(weight_placement_to_json(placement_manifest())) + orphan = dict(raw["tensors"][0]) + orphan["tensor_id"] = "orphan.weight" + raw["tensors"].append(orphan) + + with pytest.raises(ValueError, match="unreferenced tensor"): + weight_placement_from_json(json.dumps(raw)) + + +@pytest.mark.parametrize("aliases", ["alias", {"alias": 1}, ["alias", "alias"]]) +def test_placement_json_rejects_invalid_aliases(aliases) -> None: + raw = json.loads(weight_placement_to_json(placement_manifest())) + raw["parts"][0]["fragments"][0]["aliases"] = aliases + + with pytest.raises(ValueError, match="aliases"): + weight_placement_from_json(json.dumps(raw)) + + +@pytest.mark.parametrize( + ("path", "value"), + [ + (("tensors",), {}), + (("topology", "participants"), {}), + (("parts",), 1), + (("parts", 0, "fragments"), 1), + (("tensors", 0, "global_shape"), 8), + (("tensors", 0, "shard_dims"), "0"), + (("tensors", 0, "parallel_axes"), {}), + (("parts", 0, "fragments", 0, "global_offset"), 0), + (("parts", 0, "fragments", 0, "local_shape"), None), + (("parts", 0, "fragments", 0, "rank"), []), + ], +) +def test_placement_json_rejects_wrong_container_types( + path: tuple, value: object +) -> None: + raw = json.loads(weight_placement_to_json(placement_manifest())) + target = raw + for component in path[:-1]: + target = target[component] + target[path[-1]] = value + + with pytest.raises(ValueError): + weight_placement_from_json(json.dumps(raw)) + + +def test_placement_id_must_match_canonical_logical_content() -> None: + with pytest.raises(ValueError, match="canonical logical content"): + placement_manifest(placement_id="opaque-placement-id") + + +def test_parallel_axis_semantics_participate_in_placement_identity() -> None: + tp = placement_manifest( + tensors=( + descriptor( + expert_id=None, + parallel_axes=(SplitAxis(kind="tp", dim=0),), + ), + ), + ) + ep = placement_manifest( + tensors=( + descriptor( + expert_id=None, + parallel_axes=(SplitAxis(kind="ep", dim=0),), + ), + ), + ) + + assert tp.placement_id != ep.placement_id + assert tp.digest != ep.digest + + +def test_declared_parallel_axis_sizes_participate_in_placement_identity() -> None: + tensor = descriptor( + shard_dims=(), + parallel_axes=(ReplicatedAxis(kind="tp"),), + expert_id=None, + ) + base = placement_manifest(tensors=(tensor,)) + expanded_topology = parallel_topology(tp_size=2) + expanded = placement_manifest(topology=expanded_topology, tensors=(tensor,)) + + assert base.topology.participants == expanded.topology.participants + assert base.placement_id != expanded.placement_id + assert base.digest != expanded.digest + + +def test_axis_semantic_kind_participates_in_placement_identity() -> None: + replicated = placement_manifest( + tensors=( + descriptor( + shard_dims=(), + expert_id=None, + parallel_axes=(ReplicatedAxis(kind="tp"),), + ), + ), + ) + ownership = placement_manifest( + tensors=( + descriptor( + shard_dims=(), + expert_id=None, + parallel_axes=(OwnershipAxis(kind="tp"),), + ), + ), + ) + + assert replicated.placement_id != ownership.placement_id + assert replicated.digest != ownership.digest diff --git a/mooncake-reshard/tests/weight_manifest/test_validation.py b/mooncake-reshard/tests/weight_manifest/test_validation.py new file mode 100644 index 0000000000..47a577537c --- /dev/null +++ b/mooncake-reshard/tests/weight_manifest/test_validation.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import pytest + +from mooncake.reshard.weight import ( + ParallelRank, + ReplicatedAxis, + SplitAxis, + TopologyParticipant, +) + +from .helpers import ( + binding_fragment, + binding_manifest, + descriptor, + parallel_topology, + placement_fragment, + placement_manifest, + placement_part, +) + + +@pytest.mark.parametrize( + "factory", + [ + lambda: ParallelRank(dp=True), + lambda: binding_fragment(address=4096.0), + lambda: binding_fragment(nbytes=32.0), + lambda: binding_manifest(generation=True), + ], +) +def test_contract_rejects_bool_and_float_integer_fields(factory) -> None: + with pytest.raises(ValueError, match="integer"): + factory() + + +@pytest.mark.parametrize( + ("factory", "message"), + [ + (lambda: placement_manifest(tensors=(object(),)), "tensors"), + (lambda: placement_manifest(fragments=(object(),)), "fragments"), + (lambda: binding_manifest(fragments=(object(),)), "fragments"), + ], +) +def test_manifest_collections_reject_wrong_element_types(factory, message: str) -> None: + with pytest.raises(ValueError, match=message): + factory() + + +@pytest.mark.parametrize( + ("factory", "message"), + [ + (lambda: placement_manifest(tensors=None), "tensors"), + (lambda: placement_manifest(fragments=None), "fragments"), + (lambda: binding_manifest(fragments=None), "fragments"), + ], +) +def test_manifest_collections_reject_wrong_container_types( + factory, message: str +) -> None: + with pytest.raises(ValueError, match=message): + factory() + + +@pytest.mark.parametrize( + "factory", + [ + lambda: binding_fragment(address=2**64), + lambda: binding_fragment(address=2**64 - 16, nbytes=32), + lambda: binding_fragment(nbytes=2**64), + lambda: binding_manifest(generation=2**64), + ], +) +def test_physical_contract_rejects_values_outside_u64_abi(factory) -> None: + with pytest.raises(ValueError, match="64-bit"): + factory() + + +def test_physical_contract_rejects_unrepresentable_exclusive_end() -> None: + with pytest.raises(ValueError, match="64-bit"): + binding_fragment(address=2**64 - 4, nbytes=4) + + +def _n_dim_descriptor(**overrides): + values = { + "global_shape": (8, 8), + "shard_dims": (0, 1), + "expert_id": None, + "parallel_axes": ( + SplitAxis(kind="ep", dim=0), + SplitAxis(kind="tp", dim=1), + ), + } + values.update(overrides) + return descriptor(**values) + + +def test_logical_validation_rejects_same_owner_n_dim_overlap() -> None: + tensor = _n_dim_descriptor() + + with pytest.raises(ValueError, match="logical fragment boxes overlap"): + placement_part( + tensors=(tensor,), + fragments=( + placement_fragment( + placement_fragment_id="left", + local_shape=(6, 8), + nbytes=96, + ), + placement_fragment( + placement_fragment_id="right", + global_offset=(4, 0), + local_shape=(4, 8), + nbytes=64, + ), + ), + ) + + +def test_logical_validation_accepts_adjacent_n_dim_boxes() -> None: + tensor = _n_dim_descriptor() + fragments = tuple( + placement_fragment( + placement_fragment_id=f"box-{row}-{column}", + global_offset=(row, column), + local_shape=(4, 4), + nbytes=32, + ) + for row in (0, 4) + for column in (0, 4) + ) + + placement = placement_manifest(tensors=(tensor,), fragments=fragments) + + assert len(placement.fragments) == 4 + + +def test_complete_placement_validation_is_python_39_compatible() -> None: + placement = placement_manifest() + + assert placement.fragments + + +@pytest.mark.parametrize( + ("tensor", "fragment", "message"), + [ + ( + _n_dim_descriptor(), + placement_fragment( + global_offset=(7, 0), + local_shape=(2, 8), + nbytes=32, + ), + "out of bounds", + ), + ( + descriptor(), + placement_fragment(nbytes=31), + "byte size mismatch", + ), + ], +) +def test_logical_validation_rejects_invalid_fragment_geometry( + tensor, fragment, message: str +) -> None: + with pytest.raises(ValueError, match=message): + placement_part(tensors=(tensor,), fragments=(fragment,)) + + +def test_replicated_tensor_accepts_multiple_physical_fragments() -> None: + tensor = descriptor( + global_shape=(2, 4), + shard_dims=(), + expert_id=None, + parallel_axes=(ReplicatedAxis(kind="tp"),), + ) + fragments = ( + placement_fragment( + placement_fragment_id="row-0", + global_offset=(0, 0), + local_shape=(1, 4), + nbytes=8, + ), + placement_fragment( + placement_fragment_id="row-1", + global_offset=(1, 0), + local_shape=(1, 4), + nbytes=8, + ), + ) + + placement = placement_manifest(tensors=(tensor,), fragments=fragments) + + assert placement.fragments == fragments + + +def test_split_tensor_accepts_storage_segmentation_on_non_split_axis() -> None: + tensor = descriptor(global_shape=(4, 4)) + topology = parallel_topology( + participants=( + TopologyParticipant("worker-0", ParallelRank(tp=0)), + TopologyParticipant("worker-1", ParallelRank(tp=1)), + ), + ) + fragments = tuple( + placement_fragment( + placement_fragment_id=f"tp-{tp_rank}-column-{column}", + global_offset=(tp_rank * 2, column), + local_shape=(2, 2), + nbytes=8, + rank=ParallelRank(tp=tp_rank), + ) + for tp_rank in range(2) + for column in (0, 2) + ) + + placement = placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + assert placement.fragments == fragments + + +def test_global_placement_rejects_split_axis_cross_rank_overlap() -> None: + tensor = descriptor(global_shape=(8, 4)) + topology = parallel_topology( + participants=( + TopologyParticipant("worker-0", ParallelRank(tp=0)), + TopologyParticipant("worker-1", ParallelRank(tp=1)), + ), + ) + fragments = ( + placement_fragment( + placement_fragment_id="rank-0", + local_shape=(8, 4), + nbytes=64, + rank=ParallelRank(tp=0), + ), + placement_fragment( + placement_fragment_id="rank-1", + local_shape=(8, 4), + nbytes=64, + rank=ParallelRank(tp=1), + ), + ) + + with pytest.raises(ValueError, match="not fully covered"): + placement_manifest( + topology=topology, + tensors=(tensor,), + fragments=fragments, + ) + + +def test_partial_part_is_allowed_but_global_placement_is_rejected() -> None: + topology = parallel_topology() + tensor = descriptor(global_shape=(8, 4)) + part = placement_part( + topology=topology, + tensors=(tensor,), + fragments=( + placement_fragment( + local_shape=(4, 4), + nbytes=32, + ), + ), + ) + + assert part.fragments[0].local_shape == (4, 4) + with pytest.raises(ValueError, match="not fully covered"): + placement_manifest(topology=topology, parts=(part,)) diff --git a/mooncake-reshard/typecheck/negative/invalid_contract_usage.py b/mooncake-reshard/typecheck/negative/invalid_contract_usage.py new file mode 100644 index 0000000000..a9640175ca --- /dev/null +++ b/mooncake-reshard/typecheck/negative/invalid_contract_usage.py @@ -0,0 +1,13 @@ +"""Negative static checks for canonical reshard contract categories. + +This file is intentionally invalid. The type-check script requires pyright to +reject it after the production contract has passed strict checking. +""" + +from mooncake.reshard.contracts import ParticipantId, PlacementId +from mooncake.reshard.weight.types import SplitAxis + + +participant_id = ParticipantId("participant-0") +placement_id: PlacementId = participant_id +SplitAxis(kind="pp", dim=0) diff --git a/mooncake-reshard/typecheck/negative/pyrightconfig.json b/mooncake-reshard/typecheck/negative/pyrightconfig.json new file mode 100644 index 0000000000..f3c7a4df7d --- /dev/null +++ b/mooncake-reshard/typecheck/negative/pyrightconfig.json @@ -0,0 +1,6 @@ +{ + "include": ["invalid_contract_usage.py"], + "extraPaths": ["../../python"], + "pythonVersion": "3.10", + "typeCheckingMode": "strict" +} diff --git a/mooncake-wheel/mooncake/__init__.py b/mooncake-wheel/mooncake/__init__.py index f1bb787f50..4532fdbf4c 100644 --- a/mooncake-wheel/mooncake/__init__.py +++ b/mooncake-wheel/mooncake/__init__.py @@ -1,6 +1,9 @@ """Mooncake public Python package.""" +from pkgutil import extend_path + from mooncake.buffer_pool import BufferPool, RegisteredBufferPool -__all__ = ["BufferPool", "RegisteredBufferPool"] +__path__ = extend_path(__path__, __name__) +__all__ = ["BufferPool", "RegisteredBufferPool"] diff --git a/requirements.txt b/requirements.txt index 11f5e714c9..fe41512416 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ # Then enable hooks: pre-commit install pre-commit==3.7.1 ruff==0.6.9 +pyright==1.1.411 codespell==2.2.6 cmake-format==0.6.13 # clang-format is provided via system package (e.g., sudo apt-get install -y clang-format-20) or toolchain diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index bfbcc3c3b1..9050f07569 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -163,6 +163,17 @@ if [ "$NPU_BUILD" = "1" ]; then fi echo "Building wheel package..." +# Stage the reshard Python package for the combined Mooncake wheel. The tracked +# source of truth remains in the top-level module. +RESHARD_SOURCE_DIR="mooncake-reshard/python/mooncake/reshard" +RESHARD_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake/reshard" +cleanup_reshard_staging() { + rm -rf "${RESHARD_STAGING_DIR}" +} +trap cleanup_reshard_staging EXIT +rm -rf "${RESHARD_STAGING_DIR}" +cp -R "${RESHARD_SOURCE_DIR}" "${RESHARD_STAGING_DIR}" + # Build the wheel package cd mooncake-wheel @@ -176,6 +187,7 @@ WHEEL_DIR="$(pwd)" cleanup_wheel_metadata_state() { [[ -f "${WHEEL_DIR}/pyproject.toml.backup" ]] && mv "${WHEEL_DIR}/pyproject.toml.backup" "${WHEEL_DIR}/pyproject.toml" rm -f "${WHEEL_DIR}/README.md" + cleanup_reshard_staging } trap cleanup_wheel_metadata_state EXIT diff --git a/scripts/check_reshard_types.sh b/scripts/check_reshard_types.sh new file mode 100755 index 0000000000..903fa1c18e --- /dev/null +++ b/scripts/check_reshard_types.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# Validate the canonical reshard contracts and their required negative examples. +set -euo pipefail + +pyright --project mooncake-reshard/pyrightconfig.json + +if negative_output=$(pyright --project mooncake-reshard/typecheck/negative/pyrightconfig.json 2>&1); then + echo "Expected invalid reshard contract examples to fail static checking." + exit 1 +fi + +printf '%s\n' "${negative_output}" +printf '%s\n' "${negative_output}" | grep -F "ParticipantId" >/dev/null +printf '%s\n' "${negative_output}" | grep -F "SplitAxisKind" >/dev/null diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index 68cfd181d6..8ea02eaae4 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -32,10 +32,12 @@ sudo apt-get install -y $SYSTEM_PACKAGES echo "Verifying that import succeeds after installation..." python -c "import mooncake.engine" && echo "Success: Import succeeded after installation" || { echo "ERROR: Import failed after installation!"; exit 1; } +python -c "import mooncake.reshard.weight" && echo "Success: Reshard import succeeded after installation" || { echo "ERROR: Reshard import failed after installation!"; exit 1; } echo "Running import structure test..." # Run the import structure test cp -r mooncake-wheel/tests test_env/ +cp -r mooncake-reshard/tests test_env/reshard_tests cd test_env pip install torch==2.11.0 numpy python -c "import mooncake._fast_copy" @@ -45,6 +47,10 @@ python tests/test_import_structure.py echo "Running mooncake config test..." python tests/test_mooncake_config.py +echo "Running reshard contract tests..." +python -m pip install pytest +python -m pytest reshard_tests -q + echo "Verifying mooncake_master entry point..." # Check if the mooncake_master entry point is installed and executable which mooncake_master || { echo "ERROR: mooncake_master entry point not found!"; exit 1; }