561 lines
24 KiB
Python
561 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
REFERENCE ONLY - build a small test "slice" of a HF safetensors model without downloading the full weights.
|
|
|
|
The heuristics below (layer/expert detection, config patching) cover common layouts only. New models often use tensor names, shapes or config keys that this script does not know about. Always run --inspect first, check the plan, and adapt the script to the model.
|
|
|
|
How it works: each safetensors shard starts with an 8-byte header length + JSON header, which gives the exact byte span of every tensor. We fetch only the headers, then use HTTP range requests to fetch the tensors we keep:
|
|
1. keep only some layers (per numeric "family" in tensor names whose count matches a layer count in config.json), renumbered to 0..n-1
|
|
2. keep only the first E experts (per-expert tensors are dropped, stacked expert tensors and router weights are row-sliced along dim 0)
|
|
3. optionally keep only the first V rows of vocab-sized tensors (off by default, see --vocab-size)
|
|
config.json is patched to match. Other small repo files are copied as-is.
|
|
|
|
Usage:
|
|
python slice_model.py Qwen/Qwen3-0.6B --inspect
|
|
python slice_model.py Qwen/Qwen3-30B-A3B --num-layers 2 --num-experts 8 --dry-run
|
|
python slice_model.py Qwen/Qwen3-30B-A3B --keep-layers 0,3 --output ../models/qwen3moe-slice
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import concurrent.futures as cf
|
|
import fnmatch
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import requests
|
|
from huggingface_hub import HfApi, get_token, hf_hub_url
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
logger = logging.getLogger("slice_model")
|
|
|
|
INDEX_FILE = "model.safetensors.index.json"
|
|
CONFIG_FILE = "config.json"
|
|
|
|
# files not worth copying into a slice
|
|
AUX_BLACKLIST = [
|
|
# weights in other formats (the slice replaces them)
|
|
"*.bin", "*.pth", "*.pt", "*.ckpt", "*.h5", "*.msgpack",
|
|
"*.onnx", "*.tflite", "*.mlmodel", "*.mlpackage/*", "*.gguf", "*.nemo",
|
|
"onnx/*", "openvino/*", "coreml/*", "original/*", "metal/*",
|
|
# media / demo assets
|
|
"*.png", "*.jpg", "*.jpeg", "*.gif", "*.webp", "*.avif", "*.svg",
|
|
"*.mp4", "*.webm", "*.mov", "*.wav", "*.mp3", "*.flac", "*.ogg", "*.pdf",
|
|
".gitattributes",
|
|
]
|
|
AUX_MAX_SIZE = 50 * 1024 * 1024
|
|
|
|
LAYER_CONTAINERS = {"layers", "layer", "blocks", "block", "h"}
|
|
EXPERT_ROW_KEYWORDS = ("expert", "router", "gate", "score")
|
|
|
|
LAYER_KEYS = {
|
|
"num_hidden_layers", "num_layers", "n_layer", "n_layers", "num_layer",
|
|
"num_decoder_layers", "encoder_layers", "decoder_layers", "depth",
|
|
"num_encoder_layers", "num_attention_layers", "layer_count",
|
|
}
|
|
VOCAB_KEYS = {"vocab_size", "padded_vocab_size"}
|
|
NEXTN_KEYS = {"num_nextn_predict_layers", "num_mtp_layers"}
|
|
TOPK_SUBSTR = ("per_tok", "topk", "top_k", "moe_k")
|
|
EXPERT_KEY_EXCLUDE = ("per_tok", "topk", "top_k", "shared", "group", "intermediate", "size", "dim", "dtype")
|
|
|
|
|
|
def parse_size(text: str) -> int:
|
|
m = re.fullmatch(r"(\d+(?:\.\d+)?)\s*([kKmMgGtT]?)[bB]?", text.strip())
|
|
if not m:
|
|
raise argparse.ArgumentTypeError(f"invalid size: {text!r}")
|
|
mult = {"": 1, "k": 1024, "m": 1024**2, "g": 1024**3, "t": 1024**4}[m.group(2).lower()]
|
|
return int(float(m.group(1)) * mult)
|
|
|
|
|
|
def human(n: float) -> str:
|
|
for unit in ("B", "KB", "MB", "GB"):
|
|
if n < 1024 or unit == "GB":
|
|
return f"{n:.1f}{unit}" if unit != "B" else f"{int(n)}B"
|
|
n /= 1024
|
|
return f"{n}GB"
|
|
|
|
|
|
# --------------------------------------------------------------------------------------
|
|
# HTTP
|
|
# --------------------------------------------------------------------------------------
|
|
|
|
_tls = threading.local()
|
|
|
|
|
|
def _session() -> requests.Session:
|
|
s = getattr(_tls, "s", None)
|
|
if s is None:
|
|
s = requests.Session()
|
|
_tls.s = s
|
|
return s
|
|
|
|
|
|
class HubClient:
|
|
def __init__(self, repo_id: str, revision: str, token: str | None):
|
|
self.repo_id = repo_id
|
|
self.revision = revision
|
|
self.headers = {"Authorization": f"Bearer {token}"} if token else {}
|
|
|
|
def url(self, filename: str) -> str:
|
|
return hf_hub_url(self.repo_id, filename, revision=self.revision)
|
|
|
|
def _request(self, filename: str, headers: dict, stream: bool = False, retries: int = 5):
|
|
last: Exception | None = None
|
|
for attempt in range(1, retries + 1):
|
|
try:
|
|
r = _session().get(self.url(filename), headers=headers, stream=stream, timeout=90, allow_redirects=True)
|
|
if r.status_code in (429, 500, 502, 503, 504):
|
|
raise requests.HTTPError(f"HTTP {r.status_code}", response=r)
|
|
r.raise_for_status()
|
|
return r
|
|
except Exception as e: # noqa: BLE001
|
|
last = e
|
|
if attempt == retries or "404" in str(e) or "401" in str(e) or "403" in str(e):
|
|
break
|
|
time.sleep(2.0 * attempt)
|
|
raise RuntimeError(f"{self.repo_id}/{filename}: {last}") from last
|
|
|
|
def get_range(self, filename: str, start: int, end_inclusive: int) -> bytes:
|
|
h = dict(self.headers)
|
|
h["Range"] = f"bytes={start}-{end_inclusive}"
|
|
return self._request(filename, h).content
|
|
|
|
def get_json(self, filename: str) -> Any:
|
|
return json.loads(self._request(filename, dict(self.headers)).content)
|
|
|
|
def download_file(self, filename: str, dest: Path):
|
|
r = self._request(filename, dict(self.headers), stream=True)
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(dest, "wb") as f:
|
|
for chunk in r.iter_content(chunk_size=4 * 1024 * 1024):
|
|
if chunk:
|
|
f.write(chunk)
|
|
|
|
def stream_range_into_fd(self, filename: str, start: int, end_inclusive: int, fd: int, write_offset: int):
|
|
h = dict(self.headers)
|
|
h["Range"] = f"bytes={start}-{end_inclusive}"
|
|
r = self._request(filename, h, stream=True)
|
|
pos = write_offset
|
|
for chunk in r.iter_content(chunk_size=8 * 1024 * 1024):
|
|
if chunk:
|
|
os.pwrite(fd, chunk, pos)
|
|
pos += len(chunk)
|
|
if pos - write_offset != end_inclusive - start + 1:
|
|
raise RuntimeError(f"short read: {filename} bytes={start}-{end_inclusive}")
|
|
|
|
|
|
def fetch_shard_header(client: HubClient, filename: str) -> tuple[dict, int]:
|
|
(header_len,) = struct.unpack("<Q", client.get_range(filename, 0, 7))
|
|
header = json.loads(client.get_range(filename, 8, 8 + header_len - 1))
|
|
header.pop("__metadata__", None)
|
|
return header, 8 + header_len
|
|
|
|
|
|
# --------------------------------------------------------------------------------------
|
|
# config.json analysis / patching
|
|
# --------------------------------------------------------------------------------------
|
|
|
|
def walk_config(cfg: Any):
|
|
"""Yield (container_dict, key, value) for every leaf in nested dicts."""
|
|
if isinstance(cfg, dict):
|
|
for k, v in cfg.items():
|
|
if isinstance(v, dict):
|
|
yield from walk_config(v)
|
|
else:
|
|
yield cfg, k, v
|
|
|
|
|
|
@dataclass
|
|
class ConfigInfo:
|
|
cfg: dict | None
|
|
layer_counts: set[int] = field(default_factory=set)
|
|
nextn_counts: set[int] = field(default_factory=set)
|
|
expert_counts: set[int] = field(default_factory=set)
|
|
vocab_sizes: set[int] = field(default_factory=set)
|
|
first_k_dense: int = 0
|
|
|
|
|
|
def analyze_config(cfg: dict | None) -> ConfigInfo:
|
|
info = ConfigInfo(cfg=cfg)
|
|
if cfg is None:
|
|
return info
|
|
for _d, key, val in walk_config(cfg):
|
|
if not isinstance(val, int) or isinstance(val, bool) or val <= 0:
|
|
continue
|
|
if key in LAYER_KEYS:
|
|
info.layer_counts.add(val)
|
|
elif key in NEXTN_KEYS:
|
|
info.nextn_counts.add(val)
|
|
elif key in VOCAB_KEYS:
|
|
info.vocab_sizes.add(val)
|
|
elif key == "first_k_dense_replace":
|
|
info.first_k_dense = val
|
|
elif "expert" in key and val > 1 and not any(x in key for x in EXPERT_KEY_EXCLUDE):
|
|
info.expert_counts.add(val)
|
|
return info
|
|
|
|
|
|
def patch_config(cfg: dict, families: dict[str, Family], new_vocab: int | None, orig_vocabs: set[int]) -> list[str]:
|
|
"""Patch counts in-place to match the slice. Returns log lines."""
|
|
log = []
|
|
by_card: dict[int, Family] = {}
|
|
for fam in families.values():
|
|
if fam.kept is not None:
|
|
by_card.setdefault(fam.card, fam)
|
|
by_card.setdefault(fam.regular_total, fam)
|
|
max_new_experts = max((f.new_count for f in families.values() if f.kind == "expert" and f.kept is not None), default=None)
|
|
|
|
for d, key, val in list(walk_config(cfg)):
|
|
if isinstance(val, int) and not isinstance(val, bool):
|
|
fam = by_card.get(val)
|
|
if key in LAYER_KEYS and fam is not None and fam.kind == "layer":
|
|
d[key] = fam.new_regular
|
|
elif key == "first_k_dense_replace":
|
|
lfam = next((f for f in families.values() if f.kind == "layer"), None)
|
|
if lfam is None:
|
|
continue
|
|
assert lfam.kept is not None
|
|
d[key] = sum(1 for i in lfam.kept if i < val)
|
|
elif "expert" in key and fam is not None and fam.kind == "expert" and not any(x in key for x in EXPERT_KEY_EXCLUDE):
|
|
d[key] = fam.new_count
|
|
elif "expert" in key and any(x in key for x in TOPK_SUBSTR) and max_new_experts is not None and val > max_new_experts:
|
|
d[key] = max_new_experts
|
|
elif key in VOCAB_KEYS and new_vocab is not None and val in orig_vocabs:
|
|
d[key] = new_vocab
|
|
if d[key] != val:
|
|
log.append(f"{key}: {val} -> {d[key]}")
|
|
elif isinstance(val, list) and val and all(not isinstance(x, (dict, list)) for x in val):
|
|
# per-layer lists, e.g. layer_types
|
|
fam = by_card.get(len(val))
|
|
if fam is not None and fam.kind == "layer" and len(val) in (fam.card, fam.regular_total):
|
|
assert fam.kept is not None
|
|
d[key] = [val[i] for i in fam.kept if i < len(val)]
|
|
log.append(f"{key}: list[{len(val)}] -> list[{len(d[key])}]")
|
|
return log
|
|
|
|
|
|
# --------------------------------------------------------------------------------------
|
|
# tensor name analysis
|
|
# --------------------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class Family:
|
|
"""One numeric slot in tensor names, e.g. 'model.layers.#'."""
|
|
prefix: str
|
|
kind: str # "layer" | "expert" | "other"
|
|
indices: set[int] = field(default_factory=set)
|
|
kept: list[int] | None = None # None = keep all, no renumbering
|
|
renumber: dict[int, int] = field(default_factory=dict)
|
|
regular_total: int = 0 # layer count excluding MTP tail
|
|
new_regular: int = 0
|
|
|
|
@property
|
|
def card(self) -> int:
|
|
return len(self.indices)
|
|
|
|
@property
|
|
def new_count(self) -> int:
|
|
return len(self.kept) if self.kept is not None else self.card
|
|
|
|
|
|
def name_slots(name: str):
|
|
"""Yield (slot_pos, family_prefix, index) for each numeric path component."""
|
|
tokens = name.split(".")
|
|
norm: list[str] = []
|
|
for i, tok in enumerate(tokens):
|
|
if tok.isdigit():
|
|
yield i, ".".join(norm), int(tok)
|
|
norm.append("#")
|
|
else:
|
|
norm.append(tok)
|
|
|
|
|
|
def build_families(names: list[str], cfg_info: ConfigInfo) -> dict[str, Family]:
|
|
families: dict[str, Family] = {}
|
|
for name in names:
|
|
tokens = name.split(".")
|
|
for i, prefix, idx in name_slots(name):
|
|
container = tokens[i - 1] if i > 0 else ""
|
|
fam = families.get(prefix)
|
|
if fam is None:
|
|
fam = Family(prefix=prefix, kind="expert" if "expert" in container else "other")
|
|
families[prefix] = fam
|
|
fam.indices.add(idx)
|
|
|
|
acceptable_layer_cards: dict[int, int] = {} # card -> regular layer count
|
|
for c in cfg_info.layer_counts:
|
|
acceptable_layer_cards[c] = c
|
|
for m in cfg_info.nextn_counts:
|
|
acceptable_layer_cards[c + m] = c
|
|
|
|
for fam in families.values():
|
|
if fam.kind == "expert":
|
|
continue
|
|
container = fam.prefix.rsplit(".", 1)[-1]
|
|
if cfg_info.cfg is not None:
|
|
if fam.card in acceptable_layer_cards:
|
|
fam.kind = "layer"
|
|
fam.regular_total = acceptable_layer_cards[fam.card]
|
|
elif fam.card >= 4 and container in LAYER_CONTAINERS:
|
|
fam.kind = "layer"
|
|
fam.regular_total = fam.card
|
|
return families
|
|
|
|
|
|
def plan_families(families: dict[str, Family], num_layers: int, keep_layers: list[int] | None, num_experts: int, first_k_dense: int):
|
|
for fam in families.values():
|
|
if fam.kind == "layer":
|
|
if keep_layers is not None:
|
|
kept = [i for i in keep_layers if i < fam.regular_total]
|
|
else:
|
|
base = min(num_layers, fam.regular_total)
|
|
kept = list(range(base))
|
|
# keep one MoE layer for models whose first k layers are dense
|
|
if 0 < first_k_dense < fam.regular_total and first_k_dense >= base >= 1:
|
|
kept = list(range(base - 1)) + [first_k_dense]
|
|
fam.new_regular = len(kept)
|
|
# keep the MTP tail (layer indices beyond the regular count)
|
|
kept += [i for i in sorted(fam.indices) if i >= fam.regular_total]
|
|
fam.kept = sorted(set(kept))
|
|
elif fam.kind == "expert":
|
|
fam.kept = sorted(fam.indices)[:num_experts]
|
|
fam.new_regular = len(fam.kept)
|
|
else:
|
|
fam.kept = None
|
|
continue
|
|
fam.renumber = {old: new for new, old in enumerate(fam.kept)}
|
|
|
|
|
|
def slice_tensor_name(name: str, families: dict[str, Family]) -> str | None:
|
|
"""Return the renumbered output name, or None if the tensor is dropped."""
|
|
tokens = name.split(".")
|
|
for i, prefix, idx in name_slots(name):
|
|
fam = families[prefix]
|
|
if fam.kept is None:
|
|
continue
|
|
if idx not in fam.renumber:
|
|
return None
|
|
tokens[i] = str(fam.renumber[idx])
|
|
return ".".join(tokens)
|
|
|
|
|
|
# --------------------------------------------------------------------------------------
|
|
# slicing
|
|
# --------------------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class TensorPlan:
|
|
src_name: str
|
|
out_name: str
|
|
dtype: str
|
|
shape: list[int]
|
|
orig_shape: list[int]
|
|
size: int
|
|
src_file: str
|
|
seg_start: int # absolute byte range in src_file
|
|
seg_end: int # inclusive
|
|
|
|
|
|
def pick_weight_files(filenames: list[str]) -> tuple[list[str], bool]:
|
|
"""Return (weight files, has_index). Preference: index -> model.safetensors -> any root safetensors."""
|
|
if INDEX_FILE in filenames:
|
|
return [], True
|
|
if "model.safetensors" in filenames:
|
|
return ["model.safetensors"], False
|
|
root_st = [f for f in filenames if f.endswith(".safetensors") and "/" not in f]
|
|
if "consolidated.safetensors" in root_st:
|
|
return ["consolidated.safetensors"], False
|
|
return root_st[:1], False
|
|
|
|
|
|
def row_slice(shape: list[int], size: int, keep_rows: int) -> tuple[list[int], int] | None:
|
|
# only a prefix of dim 0 is a single contiguous byte range
|
|
if len(shape) < 1 or shape[0] <= keep_rows or size % shape[0] != 0:
|
|
return None
|
|
row_bytes = size // shape[0]
|
|
return [keep_rows] + shape[1:], keep_rows * row_bytes
|
|
|
|
|
|
def build_plan(args, families, headers, weight_files, src_file_of, cfg_info) -> list[TensorPlan]:
|
|
plan_families(families, args.num_layers, args.keep_layers, args.num_experts, cfg_info.first_k_dense)
|
|
plans = []
|
|
for fname in weight_files:
|
|
header, data_start = headers[fname]
|
|
for name, meta in header.items():
|
|
out_name = slice_tensor_name(name, families)
|
|
if out_name is None:
|
|
continue
|
|
start, end = meta["data_offsets"]
|
|
shape, size = list(meta["shape"]), end - start
|
|
new_rows = None
|
|
if args.vocab_size and shape and shape[0] > args.vocab_size \
|
|
and any(shape[0] >= vs and shape[0] - vs <= 1024 for vs in cfg_info.vocab_sizes):
|
|
new_rows = args.vocab_size
|
|
elif shape and shape[0] in cfg_info.expert_counts and shape[0] > args.num_experts \
|
|
and any(k in name for k in EXPERT_ROW_KEYWORDS):
|
|
new_rows = args.num_experts
|
|
if new_rows is not None:
|
|
sliced = row_slice(shape, size, new_rows)
|
|
if sliced is not None:
|
|
shape, size = sliced
|
|
plans.append(TensorPlan(name, out_name, meta["dtype"], shape, list(meta["shape"]), size, src_file_of[name],
|
|
data_start + start, data_start + start + size - 1))
|
|
return plans
|
|
|
|
|
|
def print_inspect(cfg_info: ConfigInfo, families: dict[str, Family], plans: list[TensorPlan], all_names: list[str]):
|
|
logger.info("=== config analysis ===")
|
|
logger.info(f"layer counts: {sorted(cfg_info.layer_counts)} nextn: {sorted(cfg_info.nextn_counts)} "
|
|
f"experts: {sorted(cfg_info.expert_counts)} vocab: {sorted(cfg_info.vocab_sizes)} first_k_dense: {cfg_info.first_k_dense}")
|
|
logger.info("\n=== numeric families in tensor names ===")
|
|
for fam in sorted(families.values(), key=lambda f: f.prefix):
|
|
kept = "all" if fam.kept is None else f"{fam.kept[:16]}{'...' if len(fam.kept) > 16 else ''}"
|
|
logger.info(f"{fam.prefix:60s} kind={fam.kind:6s} count={fam.card:4d} kept={kept}")
|
|
logger.info("\n=== kept tensors (first index of each family only) ===")
|
|
seen = set()
|
|
for t in plans:
|
|
key = re.sub(r"\.\d+\.", ".#.", t.out_name)
|
|
if key in seen and t.shape == t.orig_shape:
|
|
continue
|
|
seen.add(key)
|
|
mark = f" (sliced from {t.orig_shape})" if t.shape != t.orig_shape else ""
|
|
logger.info(f"{t.out_name:70s} {t.dtype:5s} {t.shape}{mark}")
|
|
kept_src = {t.src_name for t in plans}
|
|
dropped = [n for n in all_names if n not in kept_src]
|
|
logger.info(f"\n{len(plans)}/{len(all_names)} tensors kept, {len(dropped)} dropped")
|
|
|
|
|
|
def write_slice(out_dir: Path, client: HubClient, plans: list[TensorPlan], jobs: int):
|
|
header_out: dict[str, Any] = {}
|
|
offset = 0
|
|
for t in plans:
|
|
header_out[t.out_name] = {"dtype": t.dtype, "shape": t.shape, "data_offsets": [offset, offset + t.size]}
|
|
offset += t.size
|
|
header_out["__metadata__"] = {"format": "pt"}
|
|
hbytes = json.dumps(header_out, separators=(",", ":")).encode()
|
|
hbytes += b" " * ((-(len(hbytes) + 8)) % 8)
|
|
data_start = 8 + len(hbytes)
|
|
|
|
shard_path = out_dir / "model.safetensors"
|
|
with open(shard_path, "wb") as f:
|
|
f.write(struct.pack("<Q", len(hbytes)))
|
|
f.write(hbytes)
|
|
f.truncate(data_start + offset)
|
|
fd = os.open(shard_path, os.O_WRONLY)
|
|
try:
|
|
def fetch(t: TensorPlan):
|
|
dst = data_start + header_out[t.out_name]["data_offsets"][0]
|
|
client.stream_range_into_fd(t.src_file, t.seg_start, t.seg_end, fd, dst)
|
|
|
|
with cf.ThreadPoolExecutor(max_workers=jobs) as ex:
|
|
list(ex.map(fetch, plans))
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("repo", help="HF repo id, e.g. Qwen/Qwen3-0.6B")
|
|
ap.add_argument("--revision", default=None)
|
|
ap.add_argument("--output", type=Path, default=None, help="output dir (default: ./tmp/hf_slices/{user}--{model})")
|
|
ap.add_argument("--inspect", action="store_true", help="print config analysis, tensor families and the plan, write nothing")
|
|
ap.add_argument("--dry-run", action="store_true", help="plan and estimate size only, write nothing")
|
|
ap.add_argument("--num-layers", type=int, default=2, help="keep the first N layers (default: 2)")
|
|
ap.add_argument("--keep-layers", default=None, help="explicit comma-separated layer indices to keep, overrides --num-layers")
|
|
ap.add_argument("--num-experts", type=int, default=8)
|
|
ap.add_argument("--vocab-size", type=int, default=0,
|
|
help="row-slice vocab-sized tensors to this size (default: 0 = off; tokenizer is NOT sliced, so this breaks conversion for most tokenizers)")
|
|
ap.add_argument("--max-size", type=parse_size, default="1G", help="abort if the sliced weights exceed this (default: 1G)")
|
|
ap.add_argument("--hf-token", default=None, help="HF token (default: HF_TOKEN env var or 'hf auth login' token)")
|
|
ap.add_argument("--jobs", type=int, default=8, help="concurrent HTTP requests")
|
|
args = ap.parse_args()
|
|
args.keep_layers = sorted({int(x) for x in args.keep_layers.split(",")}) if args.keep_layers else None
|
|
|
|
token = args.hf_token or get_token()
|
|
api = HfApi(token=token)
|
|
info = api.model_info(args.repo, revision=args.revision, files_metadata=True)
|
|
client = HubClient(args.repo, info.sha or "main", token)
|
|
siblings = {s.rfilename: (s.size or 0) for s in info.siblings or []}
|
|
filenames = sorted(siblings)
|
|
|
|
weight_files, has_index = pick_weight_files(filenames)
|
|
weight_map: dict[str, str] = {}
|
|
if has_index:
|
|
weight_map = client.get_json(INDEX_FILE)["weight_map"]
|
|
weight_files = sorted(set(weight_map.values()))
|
|
elif not weight_files:
|
|
logger.error("no .safetensors found")
|
|
return 1
|
|
|
|
cfg = client.get_json(CONFIG_FILE) if CONFIG_FILE in filenames else None
|
|
cfg_info = analyze_config(cfg)
|
|
|
|
with cf.ThreadPoolExecutor(max_workers=args.jobs) as ex:
|
|
headers = dict(zip(weight_files, ex.map(lambda f: fetch_shard_header(client, f), weight_files)))
|
|
all_names = [n for fname in weight_files for n in headers[fname][0]]
|
|
src_file_of = weight_map if has_index else {n: weight_files[0] for n in all_names}
|
|
|
|
families = build_families(all_names, cfg_info)
|
|
plans = build_plan(args, families, headers, weight_files, src_file_of, cfg_info)
|
|
est = sum(t.size for t in plans)
|
|
|
|
if args.inspect:
|
|
print_inspect(cfg_info, families, plans, all_names)
|
|
logger.info(f"{args.repo}@{info.sha}: {len(plans)}/{len(all_names)} tensors, {human(est)}")
|
|
if not any(f.kind == "layer" for f in families.values()):
|
|
logger.warning("no layer family detected - check --inspect output and adapt the script")
|
|
if est > args.max_size:
|
|
logger.error(f"slice is {human(est)} > --max-size {human(args.max_size)}")
|
|
return 1
|
|
if args.inspect or args.dry_run:
|
|
return 0
|
|
|
|
out_dir = args.output or Path("tmp") / "hf_slices" / args.repo.replace("/", "--")
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
write_slice(out_dir, client, plans, args.jobs)
|
|
|
|
skipped_aux = []
|
|
for fname in filenames:
|
|
if fname in weight_files or fname in (INDEX_FILE, CONFIG_FILE) or fname.endswith(".safetensors"):
|
|
continue
|
|
if any(fnmatch.fnmatch(fname, p) or fnmatch.fnmatch(Path(fname).name, p) for p in AUX_BLACKLIST):
|
|
continue
|
|
if siblings[fname] > AUX_MAX_SIZE:
|
|
skipped_aux.append(f"{fname} ({human(siblings[fname])})")
|
|
continue
|
|
client.download_file(fname, out_dir / fname)
|
|
for f in skipped_aux:
|
|
logger.warning(f"skipped large aux file {f}")
|
|
|
|
patch_log = []
|
|
if cfg is not None:
|
|
new_vocab = args.vocab_size if args.vocab_size and any(vs > args.vocab_size for vs in cfg_info.vocab_sizes) else None
|
|
patch_log = patch_config(cfg, families, new_vocab, cfg_info.vocab_sizes)
|
|
with open(out_dir / CONFIG_FILE, "w") as f:
|
|
json.dump(cfg, f, indent=2)
|
|
f.write("\n")
|
|
for line in patch_log:
|
|
logger.info(f"config: {line}")
|
|
|
|
with open(out_dir / ".slice_meta.json", "w") as f:
|
|
json.dump({
|
|
"repo": args.repo, "revision": info.sha,
|
|
"tensors": len(plans), "total_tensors": len(all_names), "size": est,
|
|
"config_patches": patch_log, "skipped_aux": skipped_aux,
|
|
}, f, indent=2)
|
|
logger.info(f"wrote {out_dir}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|