485 lines
19 KiB
Python
485 lines
19 KiB
Python
"""Architecture-defined boundary input/output for distributed Shards (DGR-006).
|
|
|
|
A public-network Shard is a contiguous range of transformer layers (RALPH runtime
|
|
decision #1). For disjoint processes to reproduce whole-model execution, every
|
|
Shard must agree on *exactly* what boundary state it consumes and emits:
|
|
|
|
* The **head** owns token embedding: it accepts token IDs and turns them into the
|
|
residual stream. No other Shard may embed tokens.
|
|
* **Middle and tail** Shards bypass token embedding entirely; they accept the named
|
|
boundary bundle (the residual stream handed over by the previous range).
|
|
* A **non-tail** Shard emits the *unnormalized* architecture-defined residual /
|
|
boundary — before the final norm, before the LM head, and before any tail-only
|
|
row pruning — so the next range sees precisely the state the whole model would
|
|
have at that layer index.
|
|
* The **tail** owns the final norm + LM head and turns the residual into logits or
|
|
a sampled token through an explicit sampling contract.
|
|
|
|
This module is deliberately backend-agnostic. It enforces the boundary *contract*
|
|
and defers the arithmetic to a ``ShardComputation`` (a duck-typed object exposing
|
|
``embed_tokens`` / ``run_layers`` / ``final_norm`` / ``lm_head``). The pinned
|
|
llama.cpp worker (DGR-008) and the reference PyTorch backend both satisfy that
|
|
protocol, and the numpy reference model in the tests proves whole-model versus
|
|
two-range parity without any download, GPU, or API credit.
|
|
|
|
The adapter **fails closed** for uncertified architectures: only architectures
|
|
that have passed real certification (dense Llama-family first, per RALPH runtime
|
|
decision #13) are accepted. Everything else raises rather than silently guessing a
|
|
tensor layout — Qwen3/Qwen3-MoE stays registered-but-dark until DGR-015 certifies
|
|
its own adapter.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
# The boundary bundle wire schema version. This is the ``boundary_schema_version``
|
|
# carried by ``runtime_recipe.RuntimeRecipeIdentity``; a receiver refuses a bundle
|
|
# whose schema it does not implement (forward/backward compatibility is a routing
|
|
# concern, not a silent reinterpretation).
|
|
BOUNDARY_SCHEMA_VERSION = 1
|
|
|
|
|
|
class BoundaryAdapterError(RuntimeError):
|
|
"""Base class for boundary-contract violations."""
|
|
|
|
|
|
class UncertifiedArchitectureError(BoundaryAdapterError):
|
|
"""Raised when a boundary adapter is requested for an uncertified architecture.
|
|
|
|
Failing closed here is a safety property: an unknown architecture has an
|
|
unknown tensor layout, so guessing where the residual boundary lives would
|
|
silently corrupt distributed output. The architecture must pass real
|
|
certification first.
|
|
"""
|
|
|
|
|
|
class BoundaryContractError(BoundaryAdapterError):
|
|
"""Raised when a Shard is fed the wrong boundary input for its role.
|
|
|
|
Examples: a head handed a residual bundle instead of token IDs, a middle
|
|
Shard handed token IDs it must not embed, or a boundary bundle whose
|
|
architecture / schema / seam layer does not match the receiving range.
|
|
"""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArchitectureBoundary:
|
|
"""The architecture-defined boundary description for one certified adapter.
|
|
|
|
These fields are what makes the boundary *architecture-defined* rather than a
|
|
hardcoded assumption: the residual tensor name, whether the tail normalizes
|
|
before the LM head, and whether row pruning is a tail-only concern all come
|
|
from here.
|
|
"""
|
|
|
|
adapter: str
|
|
boundary_tensor_name: str
|
|
boundary_schema_version: int
|
|
normalizes_before_head: bool
|
|
prunes_rows_at_tail: bool
|
|
|
|
|
|
# Certified architectures only. Dense Llama-family is first (RALPH runtime decision
|
|
# #13 + native discipline). Aliases map the many spellings a runtime recipe /
|
|
# GGUF / HF config may use onto the single canonical adapter id. Anything not in
|
|
# this table fails closed.
|
|
_DENSE_LLAMA = ArchitectureBoundary(
|
|
adapter="dense-llama",
|
|
boundary_tensor_name="residual_stream",
|
|
boundary_schema_version=BOUNDARY_SCHEMA_VERSION,
|
|
normalizes_before_head=True,
|
|
prunes_rows_at_tail=True,
|
|
)
|
|
|
|
_CERTIFIED_ARCHITECTURES: dict[str, ArchitectureBoundary] = {
|
|
"dense-llama": _DENSE_LLAMA,
|
|
"dense_llama": _DENSE_LLAMA,
|
|
"llama": _DENSE_LLAMA,
|
|
"llamaforcausallm": _DENSE_LLAMA,
|
|
"llamamodel": _DENSE_LLAMA,
|
|
}
|
|
|
|
|
|
def certified_architecture(name: Any) -> ArchitectureBoundary:
|
|
"""Return the certified boundary description for ``name`` or fail closed.
|
|
|
|
``name`` may be the canonical adapter id (``dense-llama``), an HF architecture
|
|
class (``LlamaForCausalLM``), or a GGUF/config ``model_type`` (``llama``).
|
|
Uncertified architectures raise ``UncertifiedArchitectureError``.
|
|
"""
|
|
if not isinstance(name, str) or not name.strip():
|
|
raise UncertifiedArchitectureError(
|
|
"architecture adapter must be a non-empty string; "
|
|
"the boundary adapter refuses to guess a tensor layout"
|
|
)
|
|
key = name.strip().lower()
|
|
boundary = _CERTIFIED_ARCHITECTURES.get(key)
|
|
if boundary is None:
|
|
raise UncertifiedArchitectureError(
|
|
f"architecture {name!r} is not certified for the boundary adapter; "
|
|
f"certified adapters: {sorted(set(v.adapter for v in _CERTIFIED_ARCHITECTURES.values()))}. "
|
|
"Uncertified architectures stay registered-but-dark until real "
|
|
"certification passes."
|
|
)
|
|
return boundary
|
|
|
|
|
|
def is_certified_architecture(name: Any) -> bool:
|
|
"""Return True when ``name`` maps to a certified boundary adapter."""
|
|
try:
|
|
certified_architecture(name)
|
|
except UncertifiedArchitectureError:
|
|
return False
|
|
return True
|
|
|
|
|
|
class ShardRole(str, Enum):
|
|
"""Where a contiguous layer range sits in the whole model."""
|
|
|
|
HEAD = "head"
|
|
MIDDLE = "middle"
|
|
TAIL = "tail"
|
|
FULL = "full"
|
|
|
|
@property
|
|
def owns_embedding(self) -> bool:
|
|
return self in (ShardRole.HEAD, ShardRole.FULL)
|
|
|
|
@property
|
|
def owns_final_head(self) -> bool:
|
|
return self in (ShardRole.TAIL, ShardRole.FULL)
|
|
|
|
|
|
def role_for_range(start_layer: int, end_layer: int, total_layers: int) -> ShardRole:
|
|
"""Classify a contiguous inclusive layer range within a model of ``total_layers``."""
|
|
if total_layers <= 0:
|
|
raise ValueError("total_layers must be positive")
|
|
if start_layer < 0 or end_layer < start_layer:
|
|
raise ValueError("require 0 <= start_layer <= end_layer")
|
|
if end_layer > total_layers - 1:
|
|
raise ValueError(
|
|
f"end_layer {end_layer} exceeds last layer index {total_layers - 1}"
|
|
)
|
|
is_head = start_layer == 0
|
|
is_tail = end_layer >= total_layers - 1
|
|
if is_head and is_tail:
|
|
return ShardRole.FULL
|
|
if is_head:
|
|
return ShardRole.HEAD
|
|
if is_tail:
|
|
return ShardRole.TAIL
|
|
return ShardRole.MIDDLE
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BoundaryBundle:
|
|
"""The versioned named-tensor bundle handed between adjacent Shard ranges.
|
|
|
|
``residual`` is the *unnormalized* architecture-defined residual stream with
|
|
every position row intact (no tail-only pruning). ``next_layer`` is the layer
|
|
index the receiving range must start at — it is the overlap-safe effective
|
|
start of the seam, so a receiver can reject a bundle meant for a different cut.
|
|
"""
|
|
|
|
architecture_adapter: str
|
|
schema_version: int
|
|
tensor_name: str
|
|
residual: np.ndarray
|
|
positions: np.ndarray
|
|
next_layer: int
|
|
normalized: bool = False
|
|
|
|
def named_tensor_fields(self) -> dict[str, Any]:
|
|
"""Return the wire-shaped description of the residual tensor.
|
|
|
|
These are exactly the fields the DGR-002 ``NamedTensor`` carries (name,
|
|
shape, dtype, byte order, raw bytes), so a worker can serialize this
|
|
bundle into the gRPC protobuf without re-deriving them.
|
|
"""
|
|
residual = np.ascontiguousarray(self.residual)
|
|
return {
|
|
"name": self.tensor_name,
|
|
"shape": list(residual.shape),
|
|
"dtype": residual.dtype.name,
|
|
"byte_order": _byte_order(residual.dtype),
|
|
"data": residual.tobytes(),
|
|
}
|
|
|
|
def pack(self) -> dict[str, Any]:
|
|
"""Serialize the bundle to a transport-agnostic dict (proves the seam).
|
|
|
|
The residual and positions are carried as raw little/big-endian bytes plus
|
|
shape/dtype so that a truly disjoint process can reconstruct the exact
|
|
array — this is what lets two OS processes reproduce whole-model math.
|
|
"""
|
|
residual = np.ascontiguousarray(self.residual)
|
|
positions = np.ascontiguousarray(self.positions)
|
|
return {
|
|
"architecture_adapter": self.architecture_adapter,
|
|
"schema_version": self.schema_version,
|
|
"tensor_name": self.tensor_name,
|
|
"next_layer": self.next_layer,
|
|
"normalized": self.normalized,
|
|
"residual": {
|
|
"shape": list(residual.shape),
|
|
"dtype": residual.dtype.str,
|
|
"data": residual.tobytes(),
|
|
},
|
|
"positions": {
|
|
"shape": list(positions.shape),
|
|
"dtype": positions.dtype.str,
|
|
"data": positions.tobytes(),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def unpack(cls, payload: dict[str, Any]) -> "BoundaryBundle":
|
|
"""Reconstruct a bundle produced by :meth:`pack`."""
|
|
residual = _array_from_wire(payload["residual"])
|
|
positions = _array_from_wire(payload["positions"])
|
|
return cls(
|
|
architecture_adapter=payload["architecture_adapter"],
|
|
schema_version=int(payload["schema_version"]),
|
|
tensor_name=payload["tensor_name"],
|
|
residual=residual,
|
|
positions=positions,
|
|
next_layer=int(payload["next_layer"]),
|
|
normalized=bool(payload.get("normalized", False)),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SamplingContract:
|
|
"""Explicit contract for turning tail logits into a token.
|
|
|
|
The tail never hides the sampling decision inside the adapter: the contract is
|
|
a first-class value so the head/route can reproduce it and so greedy decoding
|
|
is deterministic by construction. Only greedy is certified here; temperature /
|
|
top-p are declared but must be requested explicitly and are out of scope for
|
|
the deterministic parity gate.
|
|
"""
|
|
|
|
mode: str = "greedy"
|
|
temperature: float = 1.0
|
|
top_p: float = 1.0
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.mode not in ("greedy",):
|
|
raise BoundaryContractError(
|
|
f"sampling mode {self.mode!r} is not certified; only 'greedy' is "
|
|
"deterministic and supported by the boundary adapter today"
|
|
)
|
|
|
|
@classmethod
|
|
def greedy(cls) -> "SamplingContract":
|
|
return cls(mode="greedy")
|
|
|
|
def sample(self, last_logits: np.ndarray) -> int:
|
|
"""Return the next token id from the final-position logits row."""
|
|
logits = np.asarray(last_logits)
|
|
if logits.ndim == 2:
|
|
# (batch, vocab) — parity harness uses batch size 1.
|
|
logits = logits[0]
|
|
if logits.ndim != 1:
|
|
raise BoundaryContractError(
|
|
"sampling expects the pruned final-position logits row"
|
|
)
|
|
return int(np.argmax(logits))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TailOutput:
|
|
"""What a tail Shard emits: the sampled token plus the pruned logits row."""
|
|
|
|
token_id: int
|
|
logits: np.ndarray
|
|
sampling: SamplingContract
|
|
|
|
|
|
@dataclass
|
|
class BoundaryAdapter:
|
|
"""Enforces the architecture-defined boundary contract for one Shard range.
|
|
|
|
Construction fails closed for uncertified architectures. The adapter derives
|
|
the Shard's role from its range and drives a duck-typed ``ShardComputation``.
|
|
"""
|
|
|
|
computation: Any
|
|
sampling: SamplingContract = field(default_factory=SamplingContract.greedy)
|
|
architecture: ArchitectureBoundary = field(init=False)
|
|
role: ShardRole = field(init=False)
|
|
start_layer: int = field(init=False)
|
|
end_layer: int = field(init=False)
|
|
total_layers: int = field(init=False)
|
|
|
|
def __post_init__(self) -> None:
|
|
arch_name = getattr(self.computation, "architecture_adapter", None)
|
|
self.architecture = certified_architecture(arch_name)
|
|
self.start_layer = int(getattr(self.computation, "start_layer"))
|
|
self.end_layer = int(getattr(self.computation, "end_layer"))
|
|
self.total_layers = int(getattr(self.computation, "total_layers"))
|
|
self.role = role_for_range(
|
|
self.start_layer, self.end_layer, self.total_layers
|
|
)
|
|
|
|
@property
|
|
def is_head(self) -> bool:
|
|
return self.role.owns_embedding
|
|
|
|
@property
|
|
def is_tail(self) -> bool:
|
|
return self.role.owns_final_head
|
|
|
|
def forward(
|
|
self,
|
|
*,
|
|
token_ids: Any | None = None,
|
|
boundary: BoundaryBundle | None = None,
|
|
) -> BoundaryBundle | TailOutput:
|
|
"""Run one prefill/decode pass for this range and emit its boundary output.
|
|
|
|
Head/full ranges require ``token_ids``; middle/tail ranges require the
|
|
``boundary`` bundle. Non-tail ranges return a :class:`BoundaryBundle`;
|
|
tail/full ranges return a :class:`TailOutput` through the sampling
|
|
contract.
|
|
"""
|
|
hidden, positions = self._ingest(token_ids, boundary)
|
|
hidden = self.computation.run_layers(hidden, positions=positions)
|
|
if self.is_tail:
|
|
return self._emit_tail(hidden)
|
|
return self._emit_boundary(hidden, positions)
|
|
|
|
# -- input side -----------------------------------------------------------
|
|
|
|
def _ingest(
|
|
self, token_ids: Any | None, boundary: BoundaryBundle | None
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if self.role.owns_embedding:
|
|
return self._ingest_tokens(token_ids, boundary)
|
|
return self._ingest_boundary(token_ids, boundary)
|
|
|
|
def _ingest_tokens(
|
|
self, token_ids: Any | None, boundary: BoundaryBundle | None
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if token_ids is None:
|
|
raise BoundaryContractError(
|
|
"the head owns token embedding and must receive token IDs"
|
|
)
|
|
if boundary is not None:
|
|
raise BoundaryContractError(
|
|
"the head owns token embedding; it must not receive a boundary "
|
|
"bundle from an upstream range"
|
|
)
|
|
ids = np.asarray(token_ids)
|
|
if ids.ndim == 1:
|
|
ids = ids[None, :]
|
|
if ids.ndim != 2:
|
|
raise BoundaryContractError("token IDs must be (seq,) or (batch, seq)")
|
|
hidden = np.asarray(self.computation.embed_tokens(ids))
|
|
positions = np.broadcast_to(
|
|
np.arange(ids.shape[1], dtype=np.int64), ids.shape
|
|
).copy()
|
|
return hidden, positions
|
|
|
|
def _ingest_boundary(
|
|
self, token_ids: Any | None, boundary: BoundaryBundle | None
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if token_ids is not None:
|
|
raise BoundaryContractError(
|
|
"middle/tail Shards bypass token embedding; they must not receive "
|
|
"token IDs"
|
|
)
|
|
if boundary is None:
|
|
raise BoundaryContractError(
|
|
"middle/tail Shards must receive the named boundary bundle"
|
|
)
|
|
self._check_boundary(boundary)
|
|
return np.asarray(boundary.residual), np.asarray(boundary.positions)
|
|
|
|
def _check_boundary(self, boundary: BoundaryBundle) -> None:
|
|
if certified_architecture(boundary.architecture_adapter) is not self.architecture:
|
|
raise BoundaryContractError(
|
|
f"boundary bundle architecture {boundary.architecture_adapter!r} "
|
|
f"does not match this Shard's adapter {self.architecture.adapter!r}"
|
|
)
|
|
if boundary.schema_version != self.architecture.boundary_schema_version:
|
|
raise BoundaryContractError(
|
|
f"boundary schema v{boundary.schema_version} is not supported by "
|
|
f"this Shard (expects v{self.architecture.boundary_schema_version})"
|
|
)
|
|
if boundary.tensor_name != self.architecture.boundary_tensor_name:
|
|
raise BoundaryContractError(
|
|
f"boundary tensor {boundary.tensor_name!r} is not the "
|
|
f"architecture-defined {self.architecture.boundary_tensor_name!r}"
|
|
)
|
|
if boundary.normalized:
|
|
raise BoundaryContractError(
|
|
"boundary bundle is normalized; a Shard range must receive the "
|
|
"UNNORMALIZED architecture-defined residual"
|
|
)
|
|
if boundary.next_layer != self.start_layer:
|
|
raise BoundaryContractError(
|
|
f"boundary hands over at layer {boundary.next_layer} but this "
|
|
f"Shard starts at layer {self.start_layer}"
|
|
)
|
|
|
|
# -- output side ----------------------------------------------------------
|
|
|
|
def _emit_boundary(
|
|
self, hidden: np.ndarray, positions: np.ndarray
|
|
) -> BoundaryBundle:
|
|
# A non-tail Shard emits the unnormalized residual with every position row
|
|
# intact: no final norm, no LM head, no tail-only row pruning. next_layer
|
|
# is the receiver's overlap-safe effective start.
|
|
return BoundaryBundle(
|
|
architecture_adapter=self.architecture.adapter,
|
|
schema_version=self.architecture.boundary_schema_version,
|
|
tensor_name=self.architecture.boundary_tensor_name,
|
|
residual=np.asarray(hidden),
|
|
positions=np.asarray(positions),
|
|
next_layer=self.end_layer + 1,
|
|
normalized=False,
|
|
)
|
|
|
|
def _emit_tail(self, hidden: np.ndarray) -> TailOutput:
|
|
hidden = np.asarray(hidden)
|
|
# Tail-only row pruning: only the final position is needed to sample the
|
|
# next token, so the LM head runs on the pruned row. A non-tail Shard is
|
|
# forbidden from doing this (it must forward every row).
|
|
if self.architecture.prunes_rows_at_tail:
|
|
last_hidden = hidden[:, -1:, :]
|
|
else: # pragma: no cover - no certified architecture takes this path yet
|
|
last_hidden = hidden
|
|
if self.architecture.normalizes_before_head:
|
|
last_hidden = np.asarray(self.computation.final_norm(last_hidden))
|
|
logits = np.asarray(self.computation.lm_head(last_hidden))
|
|
last_logits = logits[:, -1, :]
|
|
token_id = self.sampling.sample(last_logits)
|
|
return TailOutput(
|
|
token_id=token_id, logits=last_logits, sampling=self.sampling
|
|
)
|
|
|
|
|
|
def _byte_order(dtype: np.dtype) -> str:
|
|
order = dtype.byteorder
|
|
if order == "<":
|
|
return "little"
|
|
if order == ">":
|
|
return "big"
|
|
# '=' native, '|' not applicable (single byte)
|
|
import sys
|
|
|
|
return sys.byteorder if order in ("=", "|") else "little"
|
|
|
|
|
|
def _array_from_wire(field_payload: dict[str, Any]) -> np.ndarray:
|
|
array = np.frombuffer(
|
|
field_payload["data"], dtype=np.dtype(field_payload["dtype"])
|
|
)
|
|
return array.reshape(field_payload["shape"]).copy()
|