dabsn 0.1.0__py3-none-any.whl

dabsn/__init__.py

@@ -0,0 +1,86 @@
+"""Canonical DABSN framework."""
+
+from .config import (
+    DABSN_ARCH,
+    DABSNConfig,
+    DABSNLayerSpec,
+    DABSNPretrainConfig,
+    parse_dabsn_layer_specs,
+    resolve_dabsn_layers,
+    resolve_layer_geometries,
+)
+from .core import DABSNCore
+from .pretrain import pretrain_next_token
+from .read import DABSNRead
+from .checkpoint import dabsn_config_dict, inspect_dabsn, load_dabsn, save_dabsn
+from .model import (
+    DABSNBackbone,
+    DABSNBlock,
+    DABSNModel,
+    DABSNSequenceLM,
+    DABSNTaskModel,
+    build_dabsn,
+    build_dabsn_from_config,
+    dabsn_adamw_param_groups,
+)
+from .runtime import (
+    DABSNSequenceModule,
+    DistributedState,
+    cleanup_distributed,
+    load_distributed_optimizer,
+    load_sharded_model_checkpoint,
+    load_sharded_training_checkpoint,
+    prepare_distributed_model,
+    save_distributed_dabsn,
+    save_sharded_training_checkpoint,
+    setup_distributed,
+    train,
+    train_step,
+    evaluate,
+    export_dabsn,
+    infer,
+    verify_gradients,
+)
+
+__all__ = [
+    "DABSN_ARCH",
+    "DABSNConfig",
+    "DABSNCore",
+    "DABSNBackbone",
+    "DABSNBlock",
+    "DABSNLayerSpec",
+    "DABSNPretrainConfig",
+    "DABSNRead",
+    "DABSNModel",
+    "DABSNSequenceLM",
+    "DABSNTaskModel",
+    "DABSNSequenceModule",
+    "DistributedState",
+    "build_dabsn",
+    "build_dabsn_from_config",
+    "dabsn_adamw_param_groups",
+    "dabsn_config_dict",
+    "cleanup_distributed",
+    "evaluate",
+    "export_dabsn",
+    "infer",
+    "load_dabsn",
+    "load_distributed_optimizer",
+    "load_sharded_model_checkpoint",
+    "load_sharded_training_checkpoint",
+    "inspect_dabsn",
+    "parse_dabsn_layer_specs",
+    "resolve_dabsn_layers",
+    "resolve_layer_geometries",
+    "prepare_distributed_model",
+    "pretrain_next_token",
+    "save_dabsn",
+    "save_distributed_dabsn",
+    "save_sharded_training_checkpoint",
+    "setup_distributed",
+    "train",
+    "train_step",
+    "verify_gradients",
+]
+
+__version__ = "0.1.0"

dabsn/adapters.py

@@ -0,0 +1,118 @@
+"""Built-in DABSN adapters and registration APIs for task-specific modules."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class IdentityInputAdapter(nn.Module):
+    """Already-vectorized sequence or field input."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.output_dim = dim
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return inputs.float()
+
+
+class LinearInputAdapter(nn.Module):
+    """Continuous features projected to the DABSN input width."""
+
+    def __init__(self, input_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self.output_dim = output_dim
+        self.proj = nn.Linear(input_dim, output_dim)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.proj(inputs.float())
+
+
+class ByteInputAdapter(nn.Module):
+    """Byte ids ``[B,T]`` to vectors ``[B,T,D]``."""
+
+    def __init__(self, output_dim: int, vocab_size: int = 256) -> None:
+        super().__init__()
+        self.output_dim = output_dim
+        self.emb = nn.Embedding(vocab_size, output_dim)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.emb(inputs.long())
+
+
+class LinearOutputHead(nn.Module):
+    """Per-position linear output head."""
+
+    def __init__(self, input_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(input_dim, out_dim)
+
+    def forward(self, hidden: Tensor) -> Tensor:
+        return self.proj(hidden)
+
+
+class ByteOutputHead(LinearOutputHead):
+    def __init__(self, input_dim: int, out_dim: int = 256) -> None:
+        del out_dim
+        super().__init__(input_dim, 256)
+
+
+InputBuilder = Callable[[int, int | None], nn.Module]
+OutputBuilder = Callable[[int, int], nn.Module]
+
+_INPUT_ADAPTERS: dict[str, InputBuilder] = {
+    "identity": lambda input_dim, output_dim: IdentityInputAdapter(
+        input_dim if output_dim is None else output_dim
+    ),
+    "tensor": lambda input_dim, output_dim: IdentityInputAdapter(
+        input_dim if output_dim is None else output_dim
+    ),
+    "linear": lambda input_dim, output_dim: LinearInputAdapter(
+        input_dim,
+        input_dim if output_dim is None else output_dim,
+    ),
+    "byte": lambda input_dim, output_dim: ByteInputAdapter(
+        input_dim if output_dim is None else output_dim
+    ),
+}
+
+_OUTPUT_HEADS: dict[str, OutputBuilder] = {
+    "linear": LinearOutputHead,
+    "field": LinearOutputHead,
+    "token": LinearOutputHead,
+    "byte": ByteOutputHead,
+}
+
+
+def register_input_adapter(kind: str, builder: InputBuilder) -> None:
+    """Register an input-adapter builder under a case-insensitive name."""
+    _INPUT_ADAPTERS[kind.lower()] = builder
+
+
+def register_output_head(kind: str, builder: OutputBuilder) -> None:
+    """Register an output-head builder under a case-insensitive name."""
+    _OUTPUT_HEADS[kind.lower()] = builder
+
+
+def build_input_adapter(
+    kind: str,
+    input_dim: int,
+    output_dim: int | None = None,
+) -> nn.Module:
+    """Construct a registered input adapter."""
+    normalized = kind.lower()
+    if normalized not in _INPUT_ADAPTERS:
+        raise ValueError(f"unknown DABSN input adapter: {normalized}")
+    return _INPUT_ADAPTERS[normalized](input_dim, output_dim)
+
+
+def build_output_head(kind: str, input_dim: int, out_dim: int) -> nn.Module:
+    """Construct a registered output head."""
+    normalized = kind.lower()
+    if normalized not in _OUTPUT_HEADS:
+        raise ValueError(f"unknown DABSN output head: {normalized}")
+    return _OUTPUT_HEADS[normalized](input_dim, out_dim)

dabsn/checkpoint.py

@@ -0,0 +1,231 @@
+"""Self-describing, non-pickle DABSN model checkpoints."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Mapping
+
+import torch
+from torch import Tensor
+
+from .config import DABSNConfig, DABSNLayerSpec
+from .model import (
+    DABSNModel,
+    DABSNSequenceLM,
+    DABSNTaskModel,
+    build_dabsn_from_config,
+)
+
+Model = DABSNModel | DABSNTaskModel | DABSNSequenceLM
+
+CHECKPOINT_FORMAT = "dabsn-model"
+CHECKPOINT_VERSION = 1
+
+
+def _introspect_config(model: DABSNModel | DABSNTaskModel) -> DABSNConfig:
+    if isinstance(model, DABSNTaskModel):
+        body = model.body
+        backbone = body.backbone
+        return DABSNConfig(
+            input_dim=int(model.raw_input_dim),
+            out_dim=int(model.out_dim),
+            hidden_dim=int(model.model_input_dim),
+            input_adapter=model.input_adapter_kind,
+            output_adapter=body.output_adapter_kind,
+            layers=[spec.to_metadata() for spec in backbone.layer_specs],
+        )
+    return DABSNConfig(
+        input_dim=int(model.input_dim),
+        out_dim=int(model.out_dim),
+        hidden_dim=int(model.input_dim),
+        input_adapter="identity",
+        output_adapter=model.output_adapter_kind,
+        layers=[spec.to_metadata() for spec in model.backbone.layer_specs],
+    )
+
+
+def dabsn_config_dict(model: Model) -> dict[str, object]:
+    """Return the complete architecture needed to reconstruct ``model``."""
+
+    if isinstance(model, DABSNSequenceLM):
+        return {
+            "model_kind": "sequence_lm",
+            "vocab": int(model.vocab),
+            "hidden_dim": int(model.hidden_dim),
+            "depth": int(model.depth),
+            "state_dim": model.state_dim,
+            "layers": [spec.to_metadata() for spec in model.layers],
+            "tie_embeddings": bool(model.tie_embeddings),
+            "grad_checkpoint": bool(model.backbone.grad_checkpoint),
+        }
+    config = getattr(model, "_dabsn_config", None) or _introspect_config(model)
+    data = {name: getattr(config, name) for name in DABSNConfig.__dataclass_fields__}
+    data["layers"] = [
+        spec.to_metadata() if isinstance(spec, DABSNLayerSpec) else dict(spec)
+        for spec in config.layer_specs()
+    ]
+    data["model_kind"] = "task" if isinstance(model, DABSNTaskModel) else "model"
+    data["grad_checkpoint"] = bool(model.backbone.grad_checkpoint)
+    return data
+
+
+def build_dabsn_from_checkpoint_config(config: Mapping[str, object]) -> Model:
+    """Construct a model from clean checkpoint metadata without loading weights."""
+
+    if config.get("model_kind") == "sequence_lm":
+        return DABSNSequenceLM(
+            vocab=int(config["vocab"]),
+            hidden_dim=int(config["hidden_dim"]),
+            depth=int(config["depth"]),
+            layers=config["layers"],
+            state_dim=config.get("state_dim"),
+            tie_embeddings=bool(config.get("tie_embeddings", False)),
+            grad_checkpoint=bool(config.get("grad_checkpoint", False)),
+        )
+    config_fields = {
+        name: value
+        for name, value in config.items()
+        if name in DABSNConfig.__dataclass_fields__
+    }
+    return build_dabsn_from_config(
+        DABSNConfig(**config_fields),
+        grad_checkpoint=bool(config.get("grad_checkpoint", False)),
+    )
+
+
+def _storage_identity(tensor: Tensor) -> tuple[int, int, tuple[int, ...], tuple[int, ...]]:
+    try:
+        pointer = tensor.untyped_storage().data_ptr()
+    except AttributeError:
+        pointer = tensor.storage().data_ptr()
+    return pointer, tensor.storage_offset(), tuple(tensor.shape), tuple(tensor.stride())
+
+
+def _deduplicate_state_dict(
+    state_dict: Mapping[str, Tensor],
+) -> tuple[dict[str, Tensor], dict[str, str]]:
+    saved: dict[str, Tensor] = {}
+    shared: dict[str, str] = {}
+    seen: dict[tuple[int, int, tuple[int, ...], tuple[int, ...]], str] = {}
+    for name, tensor in state_dict.items():
+        identity = _storage_identity(tensor)
+        if identity in seen:
+            shared[name] = seen[identity]
+        else:
+            seen[identity] = name
+            saved[name] = tensor.detach().cpu().contiguous()
+    return saved, shared
+
+
+def save_dabsn_state(
+    state_dict: Mapping[str, Tensor],
+    config: Mapping[str, object],
+    path: str | Path,
+    *,
+    extra: Mapping[str, object] | None = None,
+) -> None:
+    """Atomically write a clean DABSN SafeTensors checkpoint."""
+
+    from safetensors.torch import save_file
+
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    tensors, shared = _deduplicate_state_dict(state_dict)
+    metadata = {
+        "format": CHECKPOINT_FORMAT,
+        "version": str(CHECKPOINT_VERSION),
+        "config": json.dumps(dict(config), sort_keys=True, separators=(",", ":")),
+        "extra": json.dumps(dict(extra or {}), sort_keys=True, separators=(",", ":")),
+        "shared": json.dumps(shared, sort_keys=True, separators=(",", ":")),
+    }
+    temporary = destination.with_name(destination.name + ".tmp")
+    try:
+        save_file(tensors, str(temporary), metadata=metadata)
+        os.replace(temporary, destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def save_dabsn(
+    model: Model,
+    path: str | Path,
+    *,
+    extra: Mapping[str, object] | None = None,
+) -> None:
+    """Save ``model`` as a self-describing SafeTensors checkpoint."""
+
+    save_dabsn_state(model.state_dict(), dabsn_config_dict(model), path, extra=extra)
+
+
+def inspect_dabsn(path: str | Path) -> dict[str, object]:
+    """Read clean checkpoint metadata without allocating model tensors."""
+
+    from safetensors import safe_open
+
+    source = Path(path)
+    try:
+        with safe_open(str(source), framework="pt") as checkpoint:
+            metadata = checkpoint.metadata() or {}
+    except Exception as exc:
+        raise ValueError(f"{source} is not a readable SafeTensors checkpoint") from exc
+    if metadata.get("format") != CHECKPOINT_FORMAT:
+        raise ValueError(f"{source} is not a {CHECKPOINT_FORMAT} checkpoint")
+    version = int(metadata.get("version", "0"))
+    if version != CHECKPOINT_VERSION:
+        raise ValueError(
+            f"unsupported {CHECKPOINT_FORMAT} version {version}; expected {CHECKPOINT_VERSION}"
+        )
+    try:
+        config = json.loads(metadata["config"])
+        extra = json.loads(metadata.get("extra", "{}"))
+        shared = json.loads(metadata.get("shared", "{}"))
+    except (KeyError, TypeError, json.JSONDecodeError) as exc:
+        raise ValueError(f"{source} has invalid DABSN metadata") from exc
+    if not isinstance(config, dict) or not isinstance(extra, dict) or not isinstance(shared, dict):
+        raise ValueError(f"{source} has invalid DABSN metadata objects")
+    return {
+        "format": CHECKPOINT_FORMAT,
+        "version": version,
+        "config": config,
+        "extra": extra,
+        "shared": shared,
+    }
+
+
+def load_dabsn(
+    path: str | Path,
+    *,
+    map_location: str | torch.device | None = None,
+    strict: bool = True,
+) -> Model:
+    """Load only the clean DABSN SafeTensors format; no legacy pickle fallback."""
+
+    from safetensors.torch import load_file
+
+    source = Path(path)
+    metadata = inspect_dabsn(source)
+    device = torch.device("cpu" if map_location is None else map_location)
+    model = build_dabsn_from_checkpoint_config(metadata["config"]).to(device)
+    state_dict = load_file(str(source), device=str(device))
+    for duplicate, original in metadata["shared"].items():
+        if original not in state_dict:
+            raise ValueError(
+                f"checkpoint shared tensor {duplicate!r} refers to missing {original!r}"
+            )
+        state_dict[duplicate] = state_dict[original]
+    model.load_state_dict(state_dict, strict=strict)
+    return model
+
+
+__all__ = [
+    "CHECKPOINT_FORMAT",
+    "CHECKPOINT_VERSION",
+    "build_dabsn_from_checkpoint_config",
+    "dabsn_config_dict",
+    "inspect_dabsn",
+    "load_dabsn",
+    "save_dabsn",
+    "save_dabsn_state",
+]