File size: 5,127 Bytes
3e363e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """Reloadable inference artifact export."""
from __future__ import annotations
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any
import torch
from graph.config import GraphConfig, PATConfig
from graph.model import GraphModel
def load_checkpoint_file(
path: str | Path, *, map_location: str | torch.device = "cpu"
) -> dict[str, Any]:
"""Load a GRM checkpoint with ``weights_only=True``.
Checkpoints written by :class:`~graph.training.trainer.GraphTrainer` contain
only tensors and plain containers, so unpickling arbitrary Python objects is
never required. Loading model artifacts with ``weights_only=False`` would
permit code execution from a checkpoint file, which matters because the
registry promotion flow (PLAN section 40) moves artifacts between machines.
"""
# A checkpoint is a heterogeneous payload from outside the process; this is
# the external boundary where dict[str, Any] is the honest type.
payload = torch.load(path, map_location=map_location, weights_only=True)
if not isinstance(payload, dict):
raise TypeError(f"{path}: checkpoint is not a dict payload")
return payload
def save_inference_artifact(
model: GraphModel, directory: str | Path, *, use_safetensors: bool = False
) -> Path:
root = Path(directory)
root.mkdir(parents=True, exist_ok=True)
if use_safetensors:
try:
from safetensors.torch import save_file
except (
ImportError
) as exc: # pragma: no cover - depends on optional export dependency
raise RuntimeError(
"safetensors export requested but safetensors is not installed"
) from exc
weights = root / "model.safetensors"
save_file(
{key: value.detach().cpu() for key, value in model.state_dict().items()},
str(weights),
)
weight_format = "safetensors"
else:
weights = root / "model.pt"
torch.save(model.state_dict(), weights)
weight_format = "torch"
config = model.config.to_dict()
config["weight_format"] = weight_format
(root / "config.json").write_text(
json.dumps(config, indent=2, sort_keys=True), encoding="utf-8"
)
return weights
def load_inference_artifact(
directory: str | Path, *, map_location: str | torch.device = "cpu"
) -> GraphModel:
root = Path(directory)
raw = json.loads((root / "config.json").read_text(encoding="utf-8"))
weight_format = raw.pop("weight_format", "torch")
pat = PATConfig(**raw.pop("pat"))
raw.pop("training", None)
raw.pop("loss", None)
config = GraphConfig(pat=pat, **raw)
model = GraphModel(config)
if weight_format == "safetensors":
try:
from safetensors.torch import load_file
except (
ImportError
) as exc: # pragma: no cover - depends on optional export dependency
raise RuntimeError(
"safetensors artifact requires the safetensors package"
) from exc
state = load_file(str(root / "model.safetensors"), device=str(map_location))
else:
state = torch.load(
root / "model.pt", map_location=map_location, weights_only=True
)
load_state_dict_compat(model, state)
model.eval()
return model
def load_state_dict_compat(
model: GraphModel, state: Mapping[str, torch.Tensor]
) -> None:
"""Load checkpoints across explicitly supported PAT shape migrations.
The structured PAT token expansion added type/attribute tables and a
larger position table. Older registered smoke artifacts remain
loadable by zero-initializing only those new inputs and padding the old
position rows; every other missing or incompatible tensor fails loudly.
"""
expected = model.state_dict()
migrated: dict[str, torch.Tensor] = {}
missing_allowed = {
"pat.type_embedding.weight",
"pat.attribute_embedding.weight",
}
for key, value in state.items():
if key not in expected:
raise RuntimeError(f"checkpoint contains unexpected tensor: {key}")
if value.shape == expected[key].shape:
migrated[key] = value
continue
if (
key == "pat.path_position_embedding.weight"
and value.ndim == 2
and value.shape[1] == expected[key].shape[1]
and value.shape[0] < expected[key].shape[0]
):
padded = torch.zeros_like(expected[key])
padded[: value.shape[0]] = value
migrated[key] = padded
continue
raise RuntimeError(
f"checkpoint tensor shape mismatch for {key}: {tuple(value.shape)} != {tuple(expected[key].shape)}"
)
for key, value in expected.items():
if key not in migrated:
if key in missing_allowed:
migrated[key] = torch.zeros_like(value)
else:
raise RuntimeError(f"checkpoint is missing tensor: {key}")
model.load_state_dict(migrated, strict=True)
|