feat: checkpoint distributed gguf runtime stories
This commit is contained in:
Dobromir Popov
287
packages/node/meshnet_node/gguf_ownership.py
Normal file
287
packages/node/meshnet_node/gguf_ownership.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user