Spaces:
Sleeping
Sleeping
| """Hub format of a DSSD head set: `heads.safetensors` + `dssd_config.json`. | |
| heads.safetensors, for every head k (bf16 except the token ids): | |
| h{k}.norm.weight [d], h{k}.proj.weight [B_k + 1, d] (last row = reject class ⊥), | |
| h{k}.reject_bias [1], h{k}.vocab_ids [B_k] (int32). | |
| dssd_config.json: format_version, base_model, base_revision, head_layers, budgets, | |
| reject_class, dtype, default_domain, domains {name: {policy, stop}}, provenance. | |
| A policy uses the bench format (scripts/compose_final.py): type "threshold", thresholds, | |
| cum_gamma, acc_maps {k: [knots, values]}, max_draft, tree. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from pathlib import Path | |
| import torch | |
| from .heads import ExitHead | |
| FORMAT_VERSION = 1 | |
| HEADS_FILE = "heads.safetensors" | |
| CONFIG_FILE = "dssd_config.json" | |
| BASE_DIMS = ("hidden_size", "vocab_size", "num_hidden_layers") | |
| def save_heads(out_dir: str | Path, heads: list[ExitHead], layers: list[int], config: dict, | |
| model_config=None) -> None: | |
| """Write heads.safetensors and dssd_config.json (format, layers, budgets and reject class | |
| are filled in here; every other key comes from `config`). With the base model's config, its | |
| dimensions are recorded and checked at load time.""" | |
| from safetensors.torch import save_file | |
| out = Path(out_dir) | |
| out.mkdir(parents=True, exist_ok=True) | |
| tensors = {} | |
| for k, h in enumerate(heads): | |
| if h.adapter is not None or h.vocab_ids is None or not h.reject: | |
| raise ValueError("only pruned heads with a reject class and no adapter are supported") | |
| tensors[f"h{k}.norm.weight"] = h.norm.weight.detach().to(torch.bfloat16).contiguous() | |
| tensors[f"h{k}.proj.weight"] = h.proj.weight.detach().to(torch.bfloat16).contiguous() | |
| tensors[f"h{k}.reject_bias"] = h.reject_bias.detach().to(torch.bfloat16).contiguous() | |
| tensors[f"h{k}.vocab_ids"] = h.vocab_ids.to(torch.int32).contiguous() | |
| save_file({n: t.cpu() for n, t in tensors.items()}, str(out / HEADS_FILE)) | |
| conf = {"format_version": FORMAT_VERSION, **config, "head_layers": [int(x) for x in layers], | |
| "budgets": [h.n_vocab for h in heads], "reject_class": True, "dtype": "bfloat16"} | |
| if model_config is not None: | |
| conf["base_dims"] = {k: int(getattr(model_config, k)) for k in BASE_DIMS} | |
| (out / CONFIG_FILE).write_text(json.dumps(conf, indent=1, ensure_ascii=False)) | |
| def _resolve(path_or_repo: str, revision: str | None, token: str | None) -> Path: | |
| if os.path.isdir(path_or_repo): | |
| return Path(path_or_repo) | |
| from huggingface_hub import snapshot_download | |
| return Path(snapshot_download(path_or_repo, revision=revision, token=token, | |
| allow_patterns=[HEADS_FILE, CONFIG_FILE])) | |
| def read_config(path_or_repo: str, revision: str | None = None, token: str | None = None) -> dict: | |
| return json.loads((_resolve(path_or_repo, revision, token) / CONFIG_FILE).read_text()) | |
| def load_heads(path_or_repo: str, model_config, device="cpu", dtype=torch.bfloat16, | |
| revision: str | None = None, token: str | None = None) -> tuple[list[ExitHead], list[int], dict]: | |
| """Heads, their layers and the config from a local directory or a Hub repo. Checks that | |
| they fit the base model (hidden size, vocabulary, number of layers).""" | |
| from safetensors.torch import load_file | |
| root = _resolve(path_or_repo, revision, token) | |
| conf = json.loads((root / CONFIG_FILE).read_text()) | |
| if conf.get("format_version") != FORMAT_VERSION: | |
| raise ValueError(f"unsupported format_version {conf.get('format_version')!r} (expected {FORMAT_VERSION})") | |
| for k, v in conf.get("base_dims", {}).items(): # heads trained on another base model | |
| if int(getattr(model_config, k)) != v: | |
| raise ValueError(f"base model {k} {getattr(model_config, k)} != {v} recorded with the heads") | |
| layers = [int(x) for x in conf["head_layers"]] | |
| n_layers = model_config.num_hidden_layers | |
| if any(not 0 < x < n_layers for x in layers): | |
| raise ValueError(f"head layer out of range for a {n_layers}-layer base model: {layers}") | |
| t = load_file(str(root / HEADS_FILE)) | |
| heads = [] | |
| for k in range(len(layers)): | |
| w, ids = t[f"h{k}.proj.weight"], t[f"h{k}.vocab_ids"].long() | |
| if w.shape[1] != model_config.hidden_size: | |
| raise ValueError(f"head {k}: hidden_size {w.shape[1]} != base model {model_config.hidden_size}") | |
| if int(ids.max()) >= model_config.vocab_size: | |
| raise ValueError(f"head {k}: token id {int(ids.max())} >= base model vocab_size {model_config.vocab_size}") | |
| h = ExitHead(model_config.hidden_size, model_config.vocab_size, model_config.rms_norm_eps, | |
| vocab_ids=ids, reject=True) | |
| h.load_state_dict({"norm.weight": t[f"h{k}.norm.weight"], "proj.weight": w, | |
| "reject_bias": t[f"h{k}.reject_bias"], "vocab_ids": ids}) | |
| heads.append(h.to(device=device, dtype=dtype).eval()) | |
| return heads, layers, conf | |
| def heads_from_pt(path: str, model_config) -> tuple[list[ExitHead], list[int]]: | |
| """Pruned heads from a training checkpoint {'heads': [state_dict], 'meta': {'layers'}} | |
| (scripts/train_pruned.py, merge_pruned_heads.py), dtype preserved.""" | |
| ck = torch.load(path, map_location="cpu", weights_only=False) | |
| layers = [int(x) for x in ck["meta"]["layers"]] | |
| heads = [] | |
| for k in range(len(layers)): | |
| sd = ck["heads"][k] | |
| if any(n.startswith("adapter.") for n in sd) or sd.get("vocab_ids") is None or "reject_bias" not in sd: | |
| raise ValueError(f"head {k}: only pruned heads with a reject class and no adapter are supported") | |
| n_out = sd["proj.weight"].shape[0] - 1 | |
| h = ExitHead(model_config.hidden_size, model_config.vocab_size, model_config.rms_norm_eps, | |
| vocab_ids=sd["vocab_ids"][:n_out], reject=True) | |
| h.load_state_dict(sd) | |
| heads.append(h.to(sd["proj.weight"].dtype).eval()) | |
| return heads, layers | |