| |
| """hf_do_ckpt.py (Hart, 09.11) — wagi HF Slayer1B 1.07 -> ckpt trenera r24.2 dla --init-from (bieg Slayer1B 1.08, 8xB200). |
| |
| Eksport (eksport_hf_1b.py): stan modelu fp32 z kluczami trenera, BEZ head.weight (wiązanie z tok.weight), |
| config.json = ckpt["config"] + metadane. Tu odwrotność: head.weight = tok.weight, zapis {"model", "config", "step"}. |
| Bez sieci, bez optymalizatora (--init-from i tak bierze tylko wagi: r24.2 :1128-1140). |
| |
| Kontrole (STOP przy rozjeździe): |
| 1. sha256 model.safetensors / config.json / trenera = piny; |
| 2. wiązanie w config.json = {"tok.weight": "head.weight"} i brak head.weight w safetensors; |
| 3. strict load do klasy GPT trenera (jak nauczyciel KD, r24.2 :1160-1161) + te same klucze configu, które sprawdza --init-from; |
| 4. po zapisie: ponowny torch.load(weights_only=True) i każdy tensor == safetensors bit w bit. |
| Użycie (węzeł): python hf_do_ckpt.py --pkg <katalog HF 1.07> --trainer train_gpt_ref_r24_2.py --out init_107.pt |
| """ |
| import argparse |
| import hashlib |
| import importlib.util |
| import json |
| import os |
| import sys |
| import types |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
| SHA_MODEL = "d826b4a11bf4eaccee10fe43ec19a4f6d515dfc5f3e5c859adb631f1a0b76ad5" |
| SHA_CONFIG = "9e2062501a87d0b87d7ef75ba8bc5484f3b1afdf42081e727a74dc02f5a4afce" |
| SHA_TRENER = "46f83b51fcec18ffc3273d80fb9eab29e182c52d8667800dab27b80f7ff18afc" |
| KLUCZE_INIT = ("vocab", "n_layer", "n_embd", "n_head") |
| KLUCZE_DFLT = (("loop", 1), ("xsa", 0), ("rope_dims", 0), ("ffn", "gelu"), ("imha_t", 0), ("kv_heads", 0)) |
|
|
|
|
| def stop(msg): |
| print(f"STOP: {msg}", file=sys.stderr) |
| sys.exit(2) |
|
|
|
|
| def sha(path): |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 24), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("--pkg", required=True, help="katalog z model.safetensors i config.json") |
| ap.add_argument("--trainer", required=True, help="train_gpt_ref_r24_2.py") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--sha-model", default=SHA_MODEL) |
| ap.add_argument("--sha-config", default=SHA_CONFIG) |
| ap.add_argument("--sha-trener", default=SHA_TRENER) |
| a = ap.parse_args() |
|
|
| if os.path.exists(a.out): |
| stop(f"{a.out} istnieje") |
| model_p, config_p = os.path.join(a.pkg, "model.safetensors"), os.path.join(a.pkg, "config.json") |
| for path, want in ((model_p, a.sha_model), (config_p, a.sha_config), (a.trainer, a.sha_trener)): |
| got = sha(path) |
| if got != want: |
| stop(f"{path}: sha {got[:12]} != pin {want[:12]}") |
|
|
| sd = load_file(model_p, device="cpu") |
| with open(config_p, encoding="utf-8") as f: |
| cfg = json.load(f) |
| if cfg.get("tied_weights") != {"tok.weight": "head.weight"} or "head.weight" in sd or "tok.weight" not in sd: |
| stop("nieoczekiwane wiązanie wag (oczekiwane tied_weights tok.weight -> head.weight, bez head.weight w pliku)") |
| for key in KLUCZE_INIT: |
| if cfg.get(key) is None: |
| stop(f"config.json bez '{key}' (wymagane przez --init-from)") |
| sd["head.weight"] = sd["tok.weight"] |
|
|
| spec = importlib.util.spec_from_file_location("trener_r24_2", a.trainer) |
| trener = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(trener) |
| with torch.random.fork_rng(devices=[]): |
| model = trener.GPT(cfg["vocab"], cfg["n_layer"], cfg["n_embd"], cfg["n_head"], cfg["block"], |
| types.SimpleNamespace(**cfg)) |
| model.load_state_dict(sd, strict=True) |
|
|
| state = {"model": model.state_dict(), "config": cfg, "step": int(cfg.get("step", 0)), |
| "hf_init": {"model_sha256": a.sha_model, "config_sha256": a.sha_config}} |
| tmp = a.out + ".tmp" |
| torch.save(state, tmp) |
| try: |
| ck = torch.load(tmp, map_location="cpu", weights_only=True) |
| c0 = ck["config"] |
| for key in KLUCZE_INIT: |
| if c0[key] != cfg[key]: |
| stop(f"config po zapisie: {key} {c0[key]} != {cfg[key]}") |
| dflt = {key: c0.get(key, d) for key, d in KLUCZE_DFLT} |
| if set(ck["model"]) != set(sd): |
| stop(f"klucze po zapisie != safetensors+head ({len(ck['model'])} vs {len(sd)})") |
| rozne = [k for k, v in sd.items() if not torch.equal(ck["model"][k], v)] |
| if rozne: |
| stop(f"tensory po zapisie różne od safetensors: {rozne[:5]}") |
| except SystemExit: |
| os.remove(tmp) |
| raise |
| os.replace(tmp, a.out) |
| n = sum(v.numel() for k, v in sd.items() if k != "head.weight") |
| print(f"HF->CKPT OK {a.out} sha {sha(a.out)} | krok {state['step']} | P {n:,} | klucze {len(sd)} | " |
| f"config --init-from {dict((k, cfg[k]) for k in KLUCZE_INIT)} {dflt}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
