489 lines
20 KiB
Python
489 lines
20 KiB
Python
"""Architecture-defined boundary input/output and dense-Llama parity (DGR-006).
|
|
|
|
These tests prove the boundary contract with a *pure-numpy* dense-Llama reference
|
|
model: no download, no GPU, no torch, no API credit. The reference implements the
|
|
same ``ShardComputation`` duck type the real llama.cpp/PyTorch backends expose, so
|
|
whole-model execution and a two-range (or three-range) split are the exact same
|
|
arithmetic applied to the exact same float32 residual stream. Splitting the layer
|
|
stack at a seam and shipping the *unnormalized* residual bundle across a simulated
|
|
process boundary must reproduce the whole-model tokens bit-for-bit.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from meshnet_node.boundary_adapter import (
|
|
BOUNDARY_SCHEMA_VERSION,
|
|
BoundaryAdapter,
|
|
BoundaryBundle,
|
|
BoundaryContractError,
|
|
SamplingContract,
|
|
ShardRole,
|
|
TailOutput,
|
|
UncertifiedArchitectureError,
|
|
certified_architecture,
|
|
is_certified_architecture,
|
|
role_for_range,
|
|
)
|
|
|
|
# Documented parity tolerance. The split path applies the identical layer
|
|
# functions in the identical order to the identical float32 arrays, so the
|
|
# residual seam is bit-exact in practice; the tolerance is a conservative guard.
|
|
PARITY_ATOL = 1e-6
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Pure-numpy dense-Llama reference model (test fixture, not production).
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class _ReferenceDenseLlama:
|
|
"""A tiny deterministic dense-Llama: RMSNorm, RoPE attention, SwiGLU MLP."""
|
|
|
|
architecture_adapter = "dense-llama"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
vocab: int = 48,
|
|
hidden: int = 32,
|
|
n_layers: int = 6,
|
|
n_heads: int = 4,
|
|
intermediate: int = 64,
|
|
rms_eps: float = 1e-6,
|
|
rope_theta: float = 10000.0,
|
|
seed: int = 20260715,
|
|
) -> None:
|
|
assert hidden % n_heads == 0
|
|
self.vocab = vocab
|
|
self.hidden = hidden
|
|
self.n_layers = n_layers
|
|
self.n_heads = n_heads
|
|
self.head_dim = hidden // n_heads
|
|
assert self.head_dim % 2 == 0
|
|
self.rms_eps = rms_eps
|
|
self.rope_theta = rope_theta
|
|
|
|
rng = np.random.default_rng(seed)
|
|
|
|
def w(*shape: int) -> np.ndarray:
|
|
return (rng.standard_normal(shape) * 0.08).astype(np.float32)
|
|
|
|
self.embed = w(vocab, hidden)
|
|
self.layers = []
|
|
for _ in range(n_layers):
|
|
self.layers.append(
|
|
{
|
|
"in_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
|
"q": w(hidden, hidden),
|
|
"k": w(hidden, hidden),
|
|
"v": w(hidden, hidden),
|
|
"o": w(hidden, hidden),
|
|
"post_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
|
"gate": w(intermediate, hidden),
|
|
"up": w(intermediate, hidden),
|
|
"down": w(hidden, intermediate),
|
|
}
|
|
)
|
|
self.final_ln = (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32)
|
|
self.lm_head_w = w(vocab, hidden)
|
|
|
|
inv_freq = 1.0 / (
|
|
rope_theta ** (np.arange(0, self.head_dim, 2, dtype=np.float32) / self.head_dim)
|
|
)
|
|
self.inv_freq = inv_freq.astype(np.float32)
|
|
|
|
# -- primitive ops -----------------------------------------------------
|
|
def _rmsnorm(self, x: np.ndarray, weight: np.ndarray) -> np.ndarray:
|
|
variance = np.mean(x.astype(np.float32) ** 2, axis=-1, keepdims=True)
|
|
normed = x / np.sqrt(variance + self.rms_eps)
|
|
return (normed * weight).astype(np.float32)
|
|
|
|
def _rope(self, positions: np.ndarray):
|
|
# positions: (batch, seq) -> cos/sin: (batch, seq, head_dim)
|
|
angles = positions[..., None].astype(np.float32) * self.inv_freq[None, None, :]
|
|
emb = np.concatenate([angles, angles], axis=-1)
|
|
return np.cos(emb).astype(np.float32), np.sin(emb).astype(np.float32)
|
|
|
|
@staticmethod
|
|
def _rotate_half(x: np.ndarray) -> np.ndarray:
|
|
half = x.shape[-1] // 2
|
|
return np.concatenate([-x[..., half:], x[..., :half]], axis=-1)
|
|
|
|
def _apply_rope(self, t: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
|
|
# t: (batch, n_heads, seq, head_dim); cos/sin: (batch, seq, head_dim)
|
|
cos = cos[:, None, :, :]
|
|
sin = sin[:, None, :, :]
|
|
return t * cos + self._rotate_half(t) * sin
|
|
|
|
def _attention(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
|
|
batch, seq, _ = x.shape
|
|
q = (x @ layer["q"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
|
k = (x @ layer["k"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
|
v = (x @ layer["v"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
|
q = q.transpose(0, 2, 1, 3)
|
|
k = k.transpose(0, 2, 1, 3)
|
|
v = v.transpose(0, 2, 1, 3)
|
|
cos, sin = self._rope(positions)
|
|
q = self._apply_rope(q, cos, sin)
|
|
k = self._apply_rope(k, cos, sin)
|
|
scores = (q @ k.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
|
|
causal = np.triu(np.full((seq, seq), -1e30, dtype=np.float32), k=1)
|
|
scores = scores + causal[None, None, :, :]
|
|
scores = scores - scores.max(axis=-1, keepdims=True)
|
|
weights = np.exp(scores)
|
|
weights = weights / weights.sum(axis=-1, keepdims=True)
|
|
out = weights @ v
|
|
out = out.transpose(0, 2, 1, 3).reshape(batch, seq, self.hidden)
|
|
return (out @ layer["o"].T).astype(np.float32)
|
|
|
|
def _mlp(self, x: np.ndarray, layer: dict) -> np.ndarray:
|
|
gate = x @ layer["gate"].T
|
|
up = x @ layer["up"].T
|
|
silu = gate * (1.0 / (1.0 + np.exp(-gate)))
|
|
return ((silu * up) @ layer["down"].T).astype(np.float32)
|
|
|
|
def _run_layer(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
|
|
h = x + self._attention(self._rmsnorm(x, layer["in_ln"]), layer, positions)
|
|
h = h + self._mlp(self._rmsnorm(h, layer["post_ln"]), layer)
|
|
return h.astype(np.float32)
|
|
|
|
|
|
class _ReferenceShard:
|
|
"""A contiguous inclusive layer range of the reference model.
|
|
|
|
Satisfies the ``ShardComputation`` duck type used by ``BoundaryAdapter``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: _ReferenceDenseLlama,
|
|
start_layer: int,
|
|
end_layer: int,
|
|
*,
|
|
architecture_adapter: str | None = None,
|
|
) -> None:
|
|
self._model = model
|
|
self.start_layer = start_layer
|
|
self.end_layer = end_layer
|
|
self.total_layers = model.n_layers
|
|
self.architecture_adapter = architecture_adapter or model.architecture_adapter
|
|
|
|
def embed_tokens(self, token_ids: np.ndarray) -> np.ndarray:
|
|
return self._model.embed[np.asarray(token_ids)]
|
|
|
|
def run_layers(self, hidden: np.ndarray, *, positions: np.ndarray) -> np.ndarray:
|
|
h = np.asarray(hidden, dtype=np.float32)
|
|
for idx in range(self.start_layer, self.end_layer + 1):
|
|
h = self._model._run_layer(h, self._model.layers[idx], positions)
|
|
return h
|
|
|
|
def final_norm(self, hidden: np.ndarray) -> np.ndarray:
|
|
return self._model._rmsnorm(np.asarray(hidden, dtype=np.float32), self._model.final_ln)
|
|
|
|
def lm_head(self, hidden: np.ndarray) -> np.ndarray:
|
|
return np.asarray(hidden, dtype=np.float32) @ self._model.lm_head_w.T
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Whole-model and split reference drivers.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _whole_model_next_token(model: _ReferenceDenseLlama, token_ids: list[int]) -> TailOutput:
|
|
shard = _ReferenceShard(model, 0, model.n_layers - 1)
|
|
adapter = BoundaryAdapter(shard)
|
|
result = adapter.forward(token_ids=np.asarray(token_ids)[None, :])
|
|
assert isinstance(result, TailOutput)
|
|
return result
|
|
|
|
|
|
def _split_next_token(
|
|
model: _ReferenceDenseLlama,
|
|
token_ids: list[int],
|
|
cut_points: list[int],
|
|
*,
|
|
through_wire: bool = True,
|
|
) -> TailOutput:
|
|
"""Run the model as N contiguous ranges, shipping the bundle across each seam.
|
|
|
|
``cut_points`` are the last (inclusive) layer of each non-final range.
|
|
"""
|
|
bounds = _ranges_from_cuts(cut_points, model.n_layers)
|
|
boundary: BoundaryBundle | None = None
|
|
result: BoundaryBundle | TailOutput | None = None
|
|
for i, (start, end) in enumerate(bounds):
|
|
shard = _ReferenceShard(model, start, end)
|
|
adapter = BoundaryAdapter(shard)
|
|
if i == 0:
|
|
result = adapter.forward(token_ids=np.asarray(token_ids)[None, :])
|
|
else:
|
|
assert isinstance(boundary, BoundaryBundle)
|
|
incoming = BoundaryBundle.unpack(boundary.pack()) if through_wire else boundary
|
|
result = adapter.forward(boundary=incoming)
|
|
if isinstance(result, BoundaryBundle):
|
|
boundary = result
|
|
assert isinstance(result, TailOutput)
|
|
return result
|
|
|
|
|
|
def _ranges_from_cuts(cut_points: list[int], n_layers: int) -> list[tuple[int, int]]:
|
|
bounds: list[tuple[int, int]] = []
|
|
start = 0
|
|
for cut in cut_points:
|
|
bounds.append((start, cut))
|
|
start = cut + 1
|
|
bounds.append((start, n_layers - 1))
|
|
return bounds
|
|
|
|
|
|
def _greedy_generate(next_token_fn, prompt: list[int], n_new: int) -> list[int]:
|
|
tokens = list(prompt)
|
|
generated: list[int] = []
|
|
for _ in range(n_new):
|
|
out = next_token_fn(tokens)
|
|
tokens.append(out.token_id)
|
|
generated.append(out.token_id)
|
|
return generated
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Certification / fail-closed.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_dense_llama_and_aliases_are_certified():
|
|
"Dense Llama-family identifiers all resolve to the one certified adapter.\n\nTags: node, boundary"
|
|
for name in ("dense-llama", "llama", "LlamaForCausalLM", "LlamaModel"):
|
|
boundary = certified_architecture(name)
|
|
assert boundary.adapter == "dense-llama"
|
|
assert boundary.boundary_tensor_name == "residual_stream"
|
|
assert is_certified_architecture(name)
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["qwen3", "qwen3-moe", "mixtral", "gpt2", "", None, 123])
|
|
def test_uncertified_architectures_fail_closed(name):
|
|
"Uncertified architectures raise instead of guessing a tensor layout.\n\nTags: node, boundary"
|
|
assert not is_certified_architecture(name)
|
|
with pytest.raises(UncertifiedArchitectureError):
|
|
certified_architecture(name)
|
|
|
|
|
|
def test_adapter_construction_fails_closed_for_uncertified_backend():
|
|
"Building the adapter over an uncertified computation fails closed.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
shard = _ReferenceShard(model, 0, 2, architecture_adapter="qwen3-moe")
|
|
with pytest.raises(UncertifiedArchitectureError):
|
|
BoundaryAdapter(shard)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Roles.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_role_classification():
|
|
"Range endpoints map to head/middle/tail/full roles.\n\nTags: node, boundary"
|
|
assert role_for_range(0, 2, 6) is ShardRole.HEAD
|
|
assert role_for_range(2, 3, 6) is ShardRole.MIDDLE
|
|
assert role_for_range(4, 5, 6) is ShardRole.TAIL
|
|
assert role_for_range(0, 5, 6) is ShardRole.FULL
|
|
assert ShardRole.HEAD.owns_embedding and not ShardRole.HEAD.owns_final_head
|
|
assert ShardRole.TAIL.owns_final_head and not ShardRole.TAIL.owns_embedding
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Input-side contract.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_head_accepts_token_ids_and_owns_embedding():
|
|
"The head embeds token IDs and refuses an upstream boundary bundle.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
|
out = head.forward(token_ids=[1, 2, 3])
|
|
assert isinstance(out, BoundaryBundle)
|
|
|
|
# Head owns embedding: a residual bundle from upstream is a contract error.
|
|
bundle = out
|
|
with pytest.raises(BoundaryContractError, match="head owns token embedding"):
|
|
head.forward(boundary=bundle)
|
|
|
|
|
|
def test_middle_and_tail_bypass_embedding_and_require_the_bundle():
|
|
"Middle/tail Shards reject token IDs and demand the named boundary bundle.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
tail = BoundaryAdapter(_ReferenceShard(model, 3, 5))
|
|
with pytest.raises(BoundaryContractError, match="bypass token embedding"):
|
|
tail.forward(token_ids=[1, 2, 3])
|
|
with pytest.raises(BoundaryContractError, match="must receive the named boundary bundle"):
|
|
tail.forward()
|
|
|
|
|
|
def test_boundary_seam_layer_mismatch_is_rejected():
|
|
"A bundle handed to the wrong range (seam layer mismatch) is rejected.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
|
bundle = head.forward(token_ids=[1, 2, 3])
|
|
assert isinstance(bundle, BoundaryBundle)
|
|
assert bundle.next_layer == 3
|
|
|
|
# A range that starts at layer 4 must not accept a bundle cut at layer 3.
|
|
wrong = BoundaryAdapter(_ReferenceShard(model, 4, 5))
|
|
with pytest.raises(BoundaryContractError, match="starts at layer 4"):
|
|
wrong.forward(boundary=bundle)
|
|
|
|
|
|
def test_normalized_bundle_is_rejected():
|
|
"A normalized residual is not the architecture-defined boundary.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
|
bundle = head.forward(token_ids=[1, 2, 3])
|
|
assert isinstance(bundle, BoundaryBundle)
|
|
normalized = BoundaryBundle(
|
|
architecture_adapter=bundle.architecture_adapter,
|
|
schema_version=bundle.schema_version,
|
|
tensor_name=bundle.tensor_name,
|
|
residual=bundle.residual,
|
|
positions=bundle.positions,
|
|
next_layer=bundle.next_layer,
|
|
normalized=True,
|
|
)
|
|
tail = BoundaryAdapter(_ReferenceShard(model, 3, 5))
|
|
with pytest.raises(BoundaryContractError, match="UNNORMALIZED"):
|
|
tail.forward(boundary=normalized)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Output-side contract.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_non_tail_emits_unnormalized_full_row_boundary():
|
|
"A non-tail Shard emits the unnormalized residual with every position row.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
tokens = [3, 7, 1, 9, 2]
|
|
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
|
bundle = head.forward(token_ids=tokens)
|
|
assert isinstance(bundle, BoundaryBundle)
|
|
assert bundle.normalized is False
|
|
assert bundle.tensor_name == "residual_stream"
|
|
assert bundle.schema_version == BOUNDARY_SCHEMA_VERSION
|
|
assert bundle.next_layer == 3
|
|
# No tail-only row pruning: all sequence positions are forwarded.
|
|
assert bundle.residual.shape == (1, len(tokens), model.hidden)
|
|
assert bundle.positions.shape == (1, len(tokens))
|
|
|
|
# The emitted residual must be exactly the whole model's residual after layer 2
|
|
# (i.e. before any final norm) — prove it is NOT normalized.
|
|
positions = np.arange(len(tokens))[None, :]
|
|
hidden = model.embed[np.asarray(tokens)][None, :]
|
|
for idx in range(0, 3):
|
|
hidden = model._run_layer(hidden, model.layers[idx], positions)
|
|
assert np.allclose(bundle.residual, hidden, atol=0)
|
|
assert not np.allclose(bundle.residual, model._rmsnorm(hidden, model.final_ln))
|
|
|
|
|
|
def test_tail_emits_pruned_logits_through_the_sampling_contract():
|
|
"The tail prunes to the final row and samples through an explicit contract.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
out = _whole_model_next_token(model, [4, 8, 15, 16, 23])
|
|
assert isinstance(out, TailOutput)
|
|
assert out.logits.shape == (1, model.vocab) # tail-only row pruning to last row
|
|
assert out.sampling.mode == "greedy"
|
|
assert 0 <= out.token_id < model.vocab
|
|
assert out.token_id == int(np.argmax(out.logits[0]))
|
|
|
|
|
|
def test_sampling_contract_rejects_uncertified_modes():
|
|
"Only the certified greedy sampling mode is accepted.\n\nTags: node, boundary"
|
|
with pytest.raises(BoundaryContractError):
|
|
SamplingContract(mode="top_p")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# The core parity gate.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_two_range_prefill_parity_matches_whole_model():
|
|
"Whole-model vs two-range prefill produce the same next-token logits and token.\n\nTags: node, boundary, parity"
|
|
model = _ReferenceDenseLlama()
|
|
prompt = [5, 12, 3, 41, 7, 19, 2, 33]
|
|
|
|
whole = _whole_model_next_token(model, prompt)
|
|
split = _split_next_token(model, prompt, cut_points=[2])
|
|
|
|
assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
|
|
assert whole.token_id == split.token_id
|
|
|
|
|
|
def test_three_range_prefill_parity_exercises_the_middle_role():
|
|
"A head/middle/tail split reproduces whole-model prefill through two seams.\n\nTags: node, boundary, parity"
|
|
model = _ReferenceDenseLlama()
|
|
prompt = [9, 1, 44, 6, 30, 11]
|
|
|
|
whole = _whole_model_next_token(model, prompt)
|
|
split = _split_next_token(model, prompt, cut_points=[1, 3])
|
|
|
|
assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
|
|
assert whole.token_id == split.token_id
|
|
|
|
|
|
def test_two_range_greedy_decode_parity_matches_whole_model():
|
|
"Whole-model vs two-range greedy decode produce identical token sequences.\n\nTags: node, boundary, parity"
|
|
model = _ReferenceDenseLlama()
|
|
prompt = [2, 17, 8, 25]
|
|
n_new = 12
|
|
|
|
whole_tokens = _greedy_generate(
|
|
lambda toks: _whole_model_next_token(model, toks), prompt, n_new
|
|
)
|
|
split_tokens = _greedy_generate(
|
|
lambda toks: _split_next_token(model, toks, cut_points=[2]), prompt, n_new
|
|
)
|
|
|
|
assert whole_tokens == split_tokens
|
|
assert len(whole_tokens) == n_new
|
|
|
|
|
|
def test_boundary_bundle_wire_round_trip_is_exact():
|
|
"Packing and unpacking the boundary bundle reconstructs the exact arrays.\n\nTags: node, boundary"
|
|
model = _ReferenceDenseLlama()
|
|
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
|
bundle = head.forward(token_ids=[1, 2, 3, 4])
|
|
assert isinstance(bundle, BoundaryBundle)
|
|
|
|
restored = BoundaryBundle.unpack(bundle.pack())
|
|
assert np.array_equal(restored.residual, bundle.residual)
|
|
assert np.array_equal(restored.positions, bundle.positions)
|
|
assert restored.next_layer == bundle.next_layer
|
|
assert restored.architecture_adapter == bundle.architecture_adapter
|
|
|
|
fields = bundle.named_tensor_fields()
|
|
assert fields["name"] == "residual_stream"
|
|
assert fields["shape"] == [1, 4, model.hidden]
|
|
assert fields["byte_order"] in ("little", "big")
|
|
|
|
|
|
def test_alias_architecture_still_parity_matches():
|
|
"A Shard advertised as 'llama' interoperates with the canonical adapter.\n\nTags: node, boundary, parity"
|
|
model = _ReferenceDenseLlama()
|
|
prompt = [7, 3, 22, 5]
|
|
|
|
whole = _whole_model_next_token(model, prompt)
|
|
|
|
# Head advertises 'LlamaForCausalLM', tail advertises 'llama'; both certify to
|
|
# the same canonical adapter, so the seam contract still matches.
|
|
head = BoundaryAdapter(_ReferenceShard(model, 0, 2, architecture_adapter="LlamaForCausalLM"))
|
|
bundle = head.forward(token_ids=np.asarray(prompt)[None, :])
|
|
assert isinstance(bundle, BoundaryBundle)
|
|
tail = BoundaryAdapter(_ReferenceShard(model, 3, 5, architecture_adapter="llama"))
|
|
split = tail.forward(boundary=BoundaryBundle.unpack(bundle.pack()))
|
|
assert isinstance(split, TailOutput)
|
|
|
|
assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
|
|
assert whole.token_id == split.token_id
|