288 lines
8.9 KiB
Python
288 lines
8.9 KiB
Python
"""Dense-Llama GGUF ownership helpers.
|
|
|
|
This module keeps two related concerns together:
|
|
|
|
* selecting the tensors a dense-Llama GGUF shard is allowed to own; and
|
|
* inferring the authoritative loaded range / endpoint ownership from the
|
|
tensors the model actually exposes.
|
|
|
|
The first is used by the range-aware loader seam. The second is used by the
|
|
doctor/admission/reporting path so the tracker sees what the model loaded, not
|
|
what a CLI flag claimed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import Any, Iterable, Mapping
|
|
|
|
_BLOCK_RE = re.compile(r"^blk\.(\d+)\.")
|
|
|
|
_HEAD_TENSOR_NAMES = {
|
|
"token_embd.weight",
|
|
"token_embd.bias",
|
|
"tok_embeddings.weight",
|
|
"tok_embeddings.bias",
|
|
"embed_tokens.weight",
|
|
"embed_tokens.bias",
|
|
}
|
|
|
|
_TAIL_TENSOR_NAMES = {
|
|
"output_norm.weight",
|
|
"output_norm.bias",
|
|
"output.weight",
|
|
"output.bias",
|
|
"lm_head.weight",
|
|
"lm_head.bias",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DenseLlamaShardOwnership:
|
|
"""Authoritative ownership for one dense-Llama shard."""
|
|
|
|
start_layer: int
|
|
end_layer: int
|
|
owns_embedding: bool
|
|
owns_final_head: bool
|
|
tensor_names: tuple[str, ...] = ()
|
|
source_artifact_hash: str | None = None
|
|
slice_artifact_hash: str | None = None
|
|
derivative_slice: bool = False
|
|
final_artifact_semantics: bool = True
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.start_layer < 0:
|
|
raise ValueError("start_layer must be non-negative")
|
|
if self.end_layer < self.start_layer:
|
|
raise ValueError("end_layer must be >= start_layer")
|
|
if self.derivative_slice:
|
|
if not self.source_artifact_hash or not self.slice_artifact_hash:
|
|
raise ValueError(
|
|
"temporary derivative sub-GGUFs must carry source and slice hashes"
|
|
)
|
|
if self.final_artifact_semantics:
|
|
raise ValueError(
|
|
"temporary derivative sub-GGUFs must not be claimed as final artifacts"
|
|
)
|
|
|
|
@property
|
|
def range(self) -> tuple[int, int]:
|
|
return self.start_layer, self.end_layer
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"start_layer": self.start_layer,
|
|
"end_layer": self.end_layer,
|
|
"owns_embedding": self.owns_embedding,
|
|
"owns_final_head": self.owns_final_head,
|
|
"tensor_names": list(self.tensor_names),
|
|
"source_artifact_hash": self.source_artifact_hash,
|
|
"slice_artifact_hash": self.slice_artifact_hash,
|
|
"derivative_slice": self.derivative_slice,
|
|
"final_artifact_semantics": self.final_artifact_semantics,
|
|
}
|
|
|
|
|
|
def select_dense_llama_tensor_names(
|
|
tensor_names: Iterable[str],
|
|
start_layer: int,
|
|
end_layer: int,
|
|
*,
|
|
total_layers: int | None = None,
|
|
) -> set[str]:
|
|
"""Return the dense-Llama GGUF tensor names owned by an inclusive range."""
|
|
if start_layer < 0:
|
|
raise ValueError("start_layer must be non-negative")
|
|
if end_layer < start_layer:
|
|
raise ValueError("end_layer must be greater than or equal to start_layer")
|
|
|
|
selected: set[str] = set()
|
|
for tensor_name in tensor_names:
|
|
if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
|
|
selected.add(tensor_name)
|
|
return selected
|
|
|
|
|
|
def infer_dense_llama_ownership(
|
|
tensor_names: Iterable[str],
|
|
*,
|
|
total_layers: int | None = None,
|
|
source_artifact_hash: str | None = None,
|
|
slice_artifact_hash: str | None = None,
|
|
derivative_slice: bool = False,
|
|
final_artifact_semantics: bool = True,
|
|
) -> DenseLlamaShardOwnership:
|
|
"""Infer authoritative loaded range and endpoint ownership from tensors."""
|
|
names = tuple(str(name) for name in tensor_names if isinstance(name, str))
|
|
if not names:
|
|
raise ValueError("tensor inventory is empty")
|
|
|
|
block_layers = sorted(
|
|
{
|
|
layer
|
|
for name in names
|
|
if (layer := _layer_index(name)) is not None
|
|
}
|
|
)
|
|
if not block_layers:
|
|
raise ValueError("tensor inventory does not contain any blk.N.* tensors")
|
|
|
|
selected = tuple(sorted(names))
|
|
return DenseLlamaShardOwnership(
|
|
start_layer=block_layers[0],
|
|
end_layer=block_layers[-1],
|
|
owns_embedding=any(_is_head_tensor(name) for name in names),
|
|
owns_final_head=any(
|
|
_is_tail_tensor(name, total_layers=total_layers, loaded_end=block_layers[-1])
|
|
for name in names
|
|
),
|
|
tensor_names=selected,
|
|
source_artifact_hash=source_artifact_hash,
|
|
slice_artifact_hash=slice_artifact_hash,
|
|
derivative_slice=derivative_slice,
|
|
final_artifact_semantics=final_artifact_semantics,
|
|
)
|
|
|
|
|
|
def authoritative_dense_llama_ownership(
|
|
backend: Any,
|
|
selection: Any | None = None,
|
|
) -> DenseLlamaShardOwnership:
|
|
"""Return the most authoritative dense-Llama ownership the backend exposes."""
|
|
tensor_names = _tensor_names_from_backend(backend)
|
|
if tensor_names:
|
|
try:
|
|
return infer_dense_llama_ownership(
|
|
tensor_names,
|
|
total_layers=_backend_total_layers(backend, selection),
|
|
)
|
|
except ValueError:
|
|
pass
|
|
|
|
start, end = _backend_loaded_bounds(backend, selection)
|
|
return DenseLlamaShardOwnership(
|
|
start_layer=start,
|
|
end_layer=end,
|
|
owns_embedding=_backend_owns_embedding(backend, start),
|
|
owns_final_head=_backend_owns_final_head(backend, end),
|
|
)
|
|
|
|
|
|
def _backend_loaded_bounds(backend: Any, selection: Any | None) -> tuple[int, int]:
|
|
start = getattr(backend, "loaded_shard_start", None)
|
|
end = getattr(backend, "loaded_shard_end", None)
|
|
if start is None:
|
|
start = getattr(backend, "shard_start", None)
|
|
if end is None:
|
|
end = getattr(backend, "shard_end", None)
|
|
if start is None or end is None:
|
|
if selection is None:
|
|
raise ValueError("backend does not expose a loaded shard range")
|
|
start = getattr(selection, "shard_start")
|
|
end = getattr(selection, "shard_end")
|
|
return int(start), int(end)
|
|
|
|
|
|
def _backend_owns_embedding(backend: Any, start: int) -> bool:
|
|
value = getattr(backend, "owns_embedding", None)
|
|
if value is None:
|
|
value = getattr(backend, "is_head", start == 0)
|
|
return bool(value)
|
|
|
|
|
|
def _backend_owns_final_head(backend: Any, end: int) -> bool:
|
|
value = getattr(backend, "owns_final_head", None)
|
|
if value is None:
|
|
value = getattr(backend, "is_tail", False)
|
|
return bool(value)
|
|
|
|
|
|
def _backend_total_layers(backend: Any, selection: Any | None) -> int | None:
|
|
value = getattr(backend, "total_layers", None)
|
|
if isinstance(value, int) and value > 0:
|
|
return value
|
|
if selection is None:
|
|
return None
|
|
total = getattr(selection, "total_layers", None)
|
|
if isinstance(total, int) and total > 0:
|
|
return total
|
|
return None
|
|
|
|
|
|
def _tensor_names_from_backend(backend: Any) -> tuple[str, ...]:
|
|
for attr in ("loaded_tensor_names", "tensor_names", "tensor_inventory"):
|
|
value = getattr(backend, attr, None)
|
|
names = _normalise_tensor_names(value)
|
|
if names:
|
|
return names
|
|
return ()
|
|
|
|
|
|
def _normalise_tensor_names(value: Any) -> tuple[str, ...]:
|
|
if value is None:
|
|
return ()
|
|
if isinstance(value, Mapping):
|
|
items = value.keys()
|
|
else:
|
|
try:
|
|
items = list(value)
|
|
except TypeError:
|
|
return ()
|
|
names = [str(item) for item in items if isinstance(item, str) and item.strip()]
|
|
return tuple(names)
|
|
|
|
|
|
def _tensor_belongs_to_range(
|
|
tensor_name: str,
|
|
start_layer: int,
|
|
end_layer: int,
|
|
total_layers: int | None,
|
|
) -> bool:
|
|
layer = _layer_index(tensor_name)
|
|
if layer is not None:
|
|
return start_layer <= layer <= end_layer
|
|
|
|
if start_layer == 0 and _is_head_tensor(tensor_name):
|
|
return True
|
|
|
|
if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(
|
|
tensor_name, total_layers=total_layers, loaded_end=end_layer
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def _layer_index(tensor_name: str) -> int | None:
|
|
match = _BLOCK_RE.match(tensor_name)
|
|
if match is None:
|
|
return None
|
|
return int(match.group(1))
|
|
|
|
|
|
def _is_head_tensor(tensor_name: str) -> bool:
|
|
lowered = tensor_name.lower()
|
|
return lowered in _HEAD_TENSOR_NAMES or any(
|
|
lowered.startswith(prefix)
|
|
for prefix in ("token_embd.", "tok_embeddings.", "embed_tokens.")
|
|
)
|
|
|
|
|
|
def _is_tail_tensor(
|
|
tensor_name: str,
|
|
*,
|
|
total_layers: int | None,
|
|
loaded_end: int,
|
|
) -> bool:
|
|
lowered = tensor_name.lower()
|
|
if lowered in _TAIL_TENSOR_NAMES:
|
|
return True
|
|
if total_layers is not None and loaded_end >= total_layers - 1:
|
|
return any(
|
|
lowered.startswith(prefix)
|
|
for prefix in ("output_norm.", "final_norm.", "norm.")
|
|
)
|
|
return False
|