| """Evaluation-only native base plus FP32 low-rank delta. Never merges weights. |
| |
| TP convention: column/output sharding slices B rows; row/input sharding slices |
| A columns, with partials reduced by the existing worker. Install once BEFORE |
| capture. Every trained norm/Markov/confidence tensor is consumed or we refuse. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import asdict |
| import hashlib |
| import json |
| from pathlib import Path |
| from typing import Callable, Mapping, Protocol, cast |
|
|
| import torch |
| from torch import Tensor |
| from torch.nn import functional as F |
|
|
| from ...native.arms import ArmConfig |
| from .cache_guard import DependencyGuard |
| from .model import Config |
|
|
|
|
| class NativeLinear(Protocol): |
| in_features: int |
| out_features: int |
| def __call__(self, x: Tensor, out_f32: bool = False, out: Tensor | None = None) -> Tensor: ... |
| def nbytes(self) -> int: ... |
|
|
|
|
| class DeltaLinear: |
| def __init__(self, base: NativeLinear, a: Tensor, b: Tensor): |
| if a.ndim != 2 or b.shape != (base.out_features, a.shape[0]) or a.shape[1] != base.in_features: |
| raise ValueError("adapter shard shape differs from native base") |
| self.base, self.a, self.b = base, a.detach().float().contiguous(), b.detach().float().contiguous() |
| self.in_features, self.out_features = base.in_features, base.out_features |
|
|
| def __call__(self, x: Tensor, out_f32: bool = False, out: Tensor | None = None) -> Tensor: |
| |
| |
| value = self.base(x, out_f32=True).float()+F.linear(F.linear(x.float(), self.a), self.b) |
| value = value if out_f32 else value.to(x.dtype) |
| if out is not None: |
| out.copy_(value) |
| return out |
| return value |
|
|
| def nbytes(self) -> int: |
| return self.base.nbytes()+4*(self.a.numel()+self.b.numel()) |
|
|
|
|
| class GroupedDelta: |
| def __init__(self, weight: Tensor, factors: list[tuple[Tensor, Tensor]]): |
| if weight.ndim != 3 or len(factors) != weight.shape[0]: |
| raise ValueError("grouped adapter ownership mismatch") |
| self.weight, self.factors = weight, factors |
| if any(a.ndim != 2 or a.shape[1] != weight.shape[2] or b.shape != (weight.shape[1], a.shape[0]) for a, b in factors): |
| raise ValueError("grouped adapter factor shape mismatch") |
|
|
| def __call__(self, x: Tensor) -> Tensor: |
| grouped = x.reshape(x.shape[0], len(self.factors), -1).float() |
| base = torch.einsum("tgd,god->tgo", grouped, self.weight.float()) |
| delta = torch.stack([F.linear(F.linear(grouped[:, g], a), b) for g, (a, b) in enumerate(self.factors)], 1) |
| return (base+delta).reshape(x.shape[0], -1).to(x.dtype) |
|
|
|
|
| class Stage(Protocol): |
| wq_a: NativeLinear |
| wq_b: NativeLinear |
| wkv: NativeLinear |
| wo_b: NativeLinear |
| sh_w1: NativeLinear |
| sh_w2: NativeLinear |
| sh_w3: NativeLinear |
| wo_a_n: object |
| adapter_wo_a: Callable[[Tensor], Tensor] | None |
| shared_pair: object |
| attn_norm: Tensor |
| ffn_norm: Tensor |
| q_norm: Tensor |
| kv_norm: Tensor |
|
|
|
|
| class Markov(Protocol): |
| w1: Tensor |
| w2: Tensor |
| cache: Tensor |
| slot_of_token: Tensor |
|
|
|
|
| class Rank(Protocol): |
| g0: int |
| g1: int |
| h0: int |
| h1: int |
| i0: int |
| i1: int |
| v0: int |
| v1: int |
| dev: torch.device |
| stages: list[Stage] |
| main_norm: Tensor |
| norm: Tensor |
| conf_w: Tensor |
| markov: Markov |
| mrows: object |
| ring_kv: Tensor |
| ring_pos: Tensor |
|
|
|
|
| class Worker(Protocol): |
| cfg: ArmConfig |
| ALL: tuple[int, ...] |
| ranks: list[Rank] |
| graphs: dict[object, object] |
| margin_graphs: dict[object, object] |
| insert_graphs: dict[object, object] |
| adapter_guard: DependencyGuard |
|
|
|
|
| CacheFactory = Callable[[Tensor, Tensor, list[int], int], object] |
|
|
|
|
| def install(model: Worker, values: Mapping[str, Tensor], config: Config, rank: int, |
| cache_factory: CacheFactory, hot_tokens: list[int]) -> None: |
| if (model.cfg.dense != "bf16" or model.cfg.head != "bf16" or model.cfg.combine != "official" |
| or not model.cfg.wo_a_qdq or model.cfg.kv_qdq != "official" or model.cfg.dense_kernel != "n1"): |
| raise ValueError("adapter evaluation requires native BF16 dense/head and official QDQ/combine") |
| if model.graphs or model.margin_graphs or model.insert_graphs or hasattr(model, "adapter_guard"): |
| raise ValueError("install a single immutable adapter before graph capture; recreate worker") |
| expected: dict[str, tuple[int, ...]] = {} |
| projections = {"attn.wq_a": (config.qr, config.dim), "attn.wq_b": (config.heads*config.hd, config.qr), |
| "attn.wkv": (config.hd, config.dim), "attn.wo_b": (config.dim, config.groups*config.outr), |
| "ffn.shared_experts.w1": (config.inter, config.dim), "ffn.shared_experts.w3": (config.inter, config.dim), |
| "ffn.shared_experts.w2": (config.dim, config.inter)} |
| def tensor_key(name: str) -> str: |
| return name.replace(".", "__") |
| expected["t.mtp__0__main_norm__weight"] = (config.dim,) |
| for name, shape in {"norm.weight": (config.dim,), "markov_head.embed.weight": (config.vocab, config.markov), |
| "markov_head.head.weight": (config.vocab, config.markov), |
| "confidence_head.proj.weight": (1, config.dim+config.markov)}.items(): |
| expected["t."+tensor_key(f"mtp.{config.stages-1}."+name)] = shape |
| for i in range(config.stages): |
| for name, (out, width) in projections.items(): |
| stem = f"stages.{i}.dense.{tensor_key(name)}." |
| expected[stem+"a"], expected[stem+"b"] = (rank, width), (out, rank) |
| for group in range(config.groups): |
| stem = f"stages.{i}.wo_a.{group}." |
| expected[stem+"a"] = (rank, config.heads//config.groups*config.hd) |
| expected[stem+"b"] = (config.outr, rank) |
| for name, width in (("attn_norm", config.dim), ("ffn_norm", config.dim), ("attn.q_norm", config.qr), ("attn.kv_norm", config.hd)): |
| expected[f"stages.{i}.t."+tensor_key(f"mtp.{i}.{name}.weight")] = (width,) |
| if set(expected) != set(values) or rank not in (32, 64): |
| raise ValueError("adapter trainable coverage/rank mismatch; no ignored tensors") |
| for name, shape in expected.items(): |
| if values[name].shape != shape or not bool(torch.isfinite(values[name]).all()): |
| raise ValueError(f"adapter shape/nonfinite mismatch: {name}") |
| |
| changes: list[tuple[object, str, object]] = [] |
| dependencies: dict[str, Tensor] = {} |
| for r in model.ALL: |
| rk = model.ranks[r] |
| def on_device(name: str) -> Tensor: |
| tensor = values[name].detach().to(device=rk.dev, dtype=torch.float32).clone().contiguous() |
| dependencies[f"{r}/{name}"] = tensor |
| return tensor |
| def replace(owner: object, attribute: str, name: str, selection: slice | None = None, flatten: bool = False) -> None: |
| t = on_device(name) |
| if selection is not None: |
| t = t[selection].contiguous() |
| if flatten: |
| t = t.flatten() |
| changes.append((owner, attribute, t)) |
| replace(rk, "main_norm", "t.mtp__0__main_norm__weight") |
| last = "t."+tensor_key(f"mtp.{config.stages-1}.") |
| replace(rk, "norm", last+"norm__weight") |
| replace(rk, "conf_w", last+"confidence_head__proj__weight", flatten=True) |
| w1 = on_device(last+"markov_head__embed__weight") |
| w2 = on_device(last+"markov_head__head__weight")[rk.v0:rk.v1].contiguous() |
| |
| changes.extend(((rk.markov, "w1", w1), (rk.markov, "w2", w2), |
| (rk, "mrows", cache_factory(w1, w2, hot_tokens, model.cfg.markov_rows)))) |
| for i, stage in enumerate(rk.stages): |
| for op, attribute in (("attn.wq_a", "wq_a"), ("attn.wq_b", "wq_b"), ("attn.wkv", "wkv"), |
| ("attn.wo_b", "wo_b"), ("ffn.shared_experts.w1", "sh_w1"), |
| ("ffn.shared_experts.w2", "sh_w2"), ("ffn.shared_experts.w3", "sh_w3")): |
| stem = f"stages.{i}.dense.{tensor_key(op)}." |
| a, b = on_device(stem+"a"), on_device(stem+"b") |
| if op == "attn.wq_b": |
| b = b[rk.h0*config.hd:rk.h1*config.hd] |
| elif op == "attn.wo_b": |
| a = a[:, rk.g0*config.outr:rk.g1*config.outr] |
| elif op in ("ffn.shared_experts.w1", "ffn.shared_experts.w3"): |
| b = b[rk.i0:rk.i1] |
| elif op == "ffn.shared_experts.w2": |
| a = a[:, rk.i0:rk.i1] |
| changes.append((stage, attribute, DeltaLinear(cast(NativeLinear, getattr(stage, attribute)), a, b))) |
| |
| weight = cast(Tensor, getattr(stage.wo_a_n, "w")) |
| factors = [(on_device(f"stages.{i}.wo_a.{g}.a"), on_device(f"stages.{i}.wo_a.{g}.b")) for g in range(rk.g0, rk.g1)] |
| changes.append((stage, "adapter_wo_a", GroupedDelta(weight.reshape(rk.g1-rk.g0, config.outr, -1), factors))) |
| changes.append((stage, "shared_pair", None)) |
| for op, attribute in (("attn_norm", "attn_norm"), ("ffn_norm", "ffn_norm"), ("attn.q_norm", "q_norm"), ("attn.kv_norm", "kv_norm")): |
| replace(stage, attribute, f"stages.{i}.t."+tensor_key(f"mtp.{i}.{op}.weight")) |
| with torch.no_grad(): |
| for owner, attribute, value in changes: |
| setattr(owner, attribute, value) |
| |
| |
| def current_dependencies() -> dict[str, Tensor]: |
| current = dict(dependencies) |
| for i, (owner, attribute, _) in enumerate(changes): |
| value = getattr(owner, attribute) |
| if isinstance(value, Tensor): |
| current[f"installed/{i}"] = value |
| elif isinstance(value, DeltaLinear): |
| current[f"delta/{i}/a"], current[f"delta/{i}/b"] = value.a, value.b |
| elif isinstance(value, GroupedDelta): |
| for g, (a, b) in enumerate(value.factors): |
| current[f"group/{i}/{g}/a"], current[f"group/{i}/{g}/b"] = a, b |
| return current |
| def invalidate() -> None: |
| model.graphs.clear() |
| model.margin_graphs.clear() |
| model.insert_graphs.clear() |
| for r in model.ALL: |
| rk = model.ranks[r] |
| rk.ring_kv.zero_() |
| rk.ring_pos.fill_(-1) |
| rk.markov.cache.zero_() |
| rk.markov.slot_of_token.fill_(-1) |
| rk.mrows = None |
| |
| for r in model.ALL: |
| model.ranks[r].ring_kv.zero_() |
| model.ranks[r].ring_pos.fill_(-1) |
| model.ranks[r].markov.cache.zero_() |
| model.ranks[r].markov.slot_of_token.fill_(-1) |
| model.adapter_guard = DependencyGuard(current_dependencies(), invalidate, current_dependencies) |
|
|
|
|
| def adapted_identity(base: Mapping[str, object], checkpoint_sha256: str, config: Config) -> dict[str, object]: |
| from .pilot import sources |
| from .data import sha256 |
| runtime_sources = sources() |
| proxy_root = Path(__file__).resolve().parents[2]/"spark/remote_dspark" |
| runtime_sources.update({"process_proxy/"+p.name: sha256(p) for p in sorted(proxy_root.iterdir()) |
| if p.is_file() and p.suffix in (".py", ".c")}) |
| def identity(role: str) -> str: |
| return hashlib.sha256(json.dumps({"schema": "dsv41.d2.native-adapter.v1", "role": role, "base": dict(base), |
| "adapter_sha256": checkpoint_sha256, "config": asdict(config), "runtime_sources": runtime_sources, |
| "semantics": "native-base-plus-fp32-delta-before-rounding;immutable;official-QDQ"}, sort_keys=True).encode()).hexdigest() |
| return {"schema": "dsv41.d2.adapted-native-profile.v1", "source_receipt": dict(base), |
| "native_base_profile": "native-hybrid-head-hot", "arm": "d2-adapted-native", "profile_label": "adapted; not unchanged source-faithful", |
| "adapter_sha256": checkpoint_sha256, "draft_head": "official BF16 head with FP32 logits", "drafter_identity": identity("drafter"), "markov_identity": identity("markov")} |
|
|
