ExtraltodeusandClaude Opus 4.8
d9773394e3
The fit corpus was limited to three hardcoded choices (wikitext, Semantic-Harmless, mixed), and any other id was rejected. Now any HuggingFace dataset id works, and any number of them can be ticked to fit on an equal-parts mix, shuffled. n_prompts now counts training SEQUENCES (what the fit iterates over) instead of source rows: each dataset is packed up to its quota, so the number entered is exactly what runs, regardless of the dataset. The fixed dropdown becomes a checkable dataset library persisted in localStorage. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
437 lines
17 KiB
Python
437 lines
17 KiB
Python
import hashlib
|
|
import json
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timezone
|
|
|
|
from jlens.lens import JacobianLens
|
|
|
|
import config
|
|
from core.gpus import gpu_stats
|
|
|
|
FITS_DIR = config.DATA_DIR / "fits"
|
|
WORKER = config.ROOT / "scripts" / "fit_worker.py"
|
|
|
|
# Fit corpus. Any HuggingFace dataset id works as-is: wikitext is the default
|
|
# and keeps a dedicated streamed path, every other id goes through the generic
|
|
# loader below. Several ids = an equal-parts mix.
|
|
DATASET_WIKITEXT = "Salesforce/wikitext-103-raw-v1"
|
|
|
|
# Seed shared by the row sampling and the mix shuffle: a continued fit that asks
|
|
# for skip+n rows deterministically extends the sequence it drew the first n
|
|
# from (sample(skip+n) then drop the head).
|
|
_SAMPLE_SEED = 1729
|
|
|
|
|
|
def _slug(dataset):
|
|
"""Short, filename-safe tag from a dataset id's last path segment, e.g.
|
|
``heretic-org/Semantic-Harmless`` -> ``semantic-harmless``."""
|
|
tail = dataset.rstrip("/").split("/")[-1].lower()
|
|
return re.sub(r"[^a-z0-9]+", "-", tail).strip("-")[:24] or "dataset"
|
|
|
|
|
|
def _text_column(features):
|
|
"""Column to fit on: prefer ``text``, else the first string-valued column."""
|
|
from datasets import Value
|
|
|
|
if "text" in features:
|
|
return "text"
|
|
for name, feat in features.items():
|
|
if isinstance(feat, Value) and feat.dtype == "string":
|
|
return name
|
|
raise ValueError(
|
|
"dataset exposes no text column to fit on "
|
|
f"(columns: {', '.join(features) or 'none'})"
|
|
)
|
|
|
|
|
|
def _pack(texts, count, target=350):
|
|
"""Pack ``texts`` into ~``target``-char sequences, stopping as soon as
|
|
``count`` sequences are ready. jlens skips the first 16 positions of every
|
|
sequence as attention sinks, so short unpacked prompts would almost all be
|
|
dropped as too short. Returns fewer than ``count`` only if ``texts`` runs
|
|
out (the caller decides whether that is an error)."""
|
|
packs, cur = [], ""
|
|
for text in texts:
|
|
text = (text or "").strip()
|
|
if not text:
|
|
continue
|
|
cur = f"{cur}\n\n{text}" if cur else text
|
|
if len(cur) >= target:
|
|
packs.append(cur)
|
|
cur = ""
|
|
if len(packs) >= count:
|
|
return packs
|
|
if cur and len(packs) < count:
|
|
# trailing remainder: keep it so a just-large-enough dataset still fills
|
|
# its quota (jlens tolerates a slightly-short final sequence)
|
|
packs.append(cur)
|
|
return packs
|
|
|
|
|
|
def _load_split(dataset):
|
|
"""``dataset``'s ``train`` split, or its first split if it has no ``train``."""
|
|
from datasets import load_dataset
|
|
|
|
try:
|
|
return load_dataset(dataset, split="train")
|
|
except ValueError:
|
|
dd = load_dataset(dataset)
|
|
return dd[next(iter(dd))]
|
|
|
|
|
|
def _load_one(dataset, n, skip):
|
|
"""``n`` training SEQUENCES from a single ``dataset`` id, skipping the first
|
|
``skip`` (continue-from: the new sequences must not overlap the base lens's).
|
|
|
|
wikitext keeps its historical path (first records >=600 chars, streamed —
|
|
one record already is one sequence). Any other HF dataset is loaded whole,
|
|
its text column shuffled with a fixed seed, then PACKED into ~350-char
|
|
sequences until skip+n are ready (median instruct prompt ~10 tokens, so
|
|
several rows per sequence). ``n``/``skip`` count OUTPUT sequences, so the
|
|
number the user asks for is exactly what the fit iterates over — not source
|
|
rows, whose count varies per dataset."""
|
|
if n <= 0:
|
|
return []
|
|
if dataset == DATASET_WIKITEXT:
|
|
from jlens.examples import load_wikitext_prompts
|
|
|
|
# load skip + n then keep the tail: the new sequences don't overlap the
|
|
# base lens's
|
|
return load_wikitext_prompts(skip + n)[skip:]
|
|
import random
|
|
|
|
ds = _load_split(dataset)
|
|
col = _text_column(ds.features)
|
|
texts = [r[col] for r in ds]
|
|
random.Random(_SAMPLE_SEED).shuffle(texts)
|
|
packs = _pack(texts, skip + n)
|
|
if len(packs) < skip + n:
|
|
raise ValueError(
|
|
f"{dataset}: {len(texts)} rows pack into only {len(packs)} sequences, "
|
|
f"{skip + n} requested — lower n_prompts"
|
|
)
|
|
return packs[skip:skip + n]
|
|
|
|
|
|
def _load_corpus(datasets, n, skip=0):
|
|
"""``n`` training SEQUENCES drawn from ``datasets`` (a list of HF dataset
|
|
ids). A single id loads that dataset; several are mixed in EQUAL parts — n
|
|
and skip are each split across them (the first datasets take the rounding
|
|
remainder) and the union is shuffled so multi-GPU slices stay mixed. Because
|
|
the count is in sequences, ``n`` is exactly what the fit iterates over."""
|
|
datasets = list(datasets)
|
|
if len(datasets) == 1:
|
|
return _load_one(datasets[0], n, skip)
|
|
import random
|
|
|
|
k = len(datasets)
|
|
prompts = []
|
|
for i, ds in enumerate(datasets):
|
|
prompts += _load_one(ds, n // k + int(i < n % k), skip // k + int(i < skip % k))
|
|
random.Random(_SAMPLE_SEED).shuffle(prompts)
|
|
return prompts
|
|
|
|
def _default_dim_batch(device):
|
|
"""Default dim_batch scaled to the device's VRAM.
|
|
Measured on a 4B bf16 fit: 8 fits in 16 GB, 4 in 12 GB."""
|
|
try:
|
|
total = gpu_stats()[int(device.split(":")[1])]["vram_total"]
|
|
return 8 if total >= 15 * 2**30 else 4
|
|
except Exception:
|
|
return 4
|
|
|
|
|
|
def _now():
|
|
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
|
|
|
|
class FitManager:
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self._procs = []
|
|
self.state = {"state": "idle"}
|
|
self.on_progress = None
|
|
|
|
def _emit(self):
|
|
if self.on_progress:
|
|
self.on_progress(dict(self.state))
|
|
|
|
def start(self, *, model_id, source, n_prompts=100, dtype="bf16", quant=None,
|
|
devices=("cuda:0",), name=None, dim_batch=None,
|
|
max_seq_len=128, source_layers=None, model_revision=None,
|
|
continue_from=None, datasets=(DATASET_WIKITEXT,)):
|
|
with self._lock:
|
|
if self.state.get("state") == "running":
|
|
raise ValueError("a fitting is already in progress")
|
|
if not devices:
|
|
raise ValueError("at least one device required")
|
|
datasets = [d.strip() for d in datasets if d and d.strip()]
|
|
if not datasets:
|
|
raise ValueError("at least one dataset required")
|
|
skip_prompts = 0
|
|
base_lens = None
|
|
if continue_from:
|
|
base_lens = JacobianLens.load(continue_from)
|
|
# new prompts: skip those already seen by the base lens
|
|
skip_prompts = base_lens.n_prompts
|
|
if source_layers is None:
|
|
source_layers = list(base_lens.source_layers)
|
|
if name is None:
|
|
base = model_id.split("/")[-1]
|
|
if len(datasets) > 1:
|
|
base += "_mixed"
|
|
elif datasets[0] != DATASET_WIKITEXT:
|
|
base += "_" + _slug(datasets[0])
|
|
total = n_prompts + skip_prompts
|
|
name = f"{base}_n{total}" if continue_from else f"{base}_n{n_prompts}"
|
|
params = {
|
|
"model_id": model_id,
|
|
"source": source,
|
|
"model_revision": model_revision,
|
|
"dtype": dtype,
|
|
"quant": quant,
|
|
"n_prompts": n_prompts,
|
|
"datasets": datasets,
|
|
"devices": list(devices),
|
|
"dim_batch": dim_batch,
|
|
"max_seq_len": max_seq_len,
|
|
"source_layers": source_layers,
|
|
"continue_from": continue_from,
|
|
"skip_prompts": skip_prompts,
|
|
}
|
|
self.state = {
|
|
"state": "running",
|
|
"name": name,
|
|
"phase": "corpus",
|
|
"total": n_prompts,
|
|
"done": 0,
|
|
"workers": [],
|
|
"eta_seconds": None,
|
|
"started_at": _now(),
|
|
"params": params,
|
|
"error": None,
|
|
}
|
|
self._procs = []
|
|
threading.Thread(target=self._run, args=(name, params), daemon=True).start()
|
|
return dict(self.state)
|
|
|
|
def stop(self):
|
|
with self._lock:
|
|
for proc in self._procs:
|
|
if proc.poll() is None:
|
|
proc.terminate()
|
|
if self.state.get("state") == "running":
|
|
self.state["state"] = "stopping"
|
|
self._emit()
|
|
return dict(self.state)
|
|
|
|
def _run(self, name, params):
|
|
try:
|
|
job_dir = FITS_DIR / name
|
|
job_dir.mkdir(parents=True, exist_ok=True)
|
|
corpus_path = job_dir / "corpus.json"
|
|
if corpus_path.exists():
|
|
prompts = json.loads(corpus_path.read_text(encoding="utf-8"))
|
|
else:
|
|
prompts = _load_corpus(
|
|
params.get("datasets", [DATASET_WIKITEXT]),
|
|
params["n_prompts"],
|
|
params.get("skip_prompts", 0),
|
|
)
|
|
corpus_path.write_text(
|
|
json.dumps(prompts, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
devices = params["devices"]
|
|
n = len(prompts)
|
|
if len(devices) == 2:
|
|
cut = int(n * 0.65)
|
|
slices = [prompts[:cut], prompts[cut:]]
|
|
else:
|
|
slices = [prompts]
|
|
|
|
self.state.update(phase="fitting", total=n)
|
|
workers = []
|
|
started = time.perf_counter()
|
|
for i, (device, chunk) in enumerate(zip(devices, slices)):
|
|
slice_path = job_dir / f"slice{i}.json"
|
|
if not slice_path.exists():
|
|
slice_path.write_text(json.dumps(chunk, ensure_ascii=False), encoding="utf-8")
|
|
dim_batch = params["dim_batch"] or _default_dim_batch(device)
|
|
cmd = [
|
|
sys.executable, "-X", "utf8", str(WORKER),
|
|
"--model", params["source"],
|
|
"--device", device,
|
|
"--dtype", params["dtype"],
|
|
"--prompts", str(slice_path),
|
|
"--checkpoint", str(job_dir / f"ckpt{i}.pt"),
|
|
"--out", str(job_dir / f"lens{i}.pt"),
|
|
"--dim-batch", str(dim_batch),
|
|
"--max-seq-len", str(params["max_seq_len"]),
|
|
]
|
|
if params["quant"]:
|
|
cmd += ["--quant", params["quant"]]
|
|
if params["source_layers"]:
|
|
cmd += ["--source-layers", json.dumps(params["source_layers"])]
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
encoding="utf-8",
|
|
cwd=str(config.ROOT),
|
|
)
|
|
self._procs.append(proc)
|
|
worker_state = {
|
|
"device": device,
|
|
"done": 0,
|
|
"total": len(chunk),
|
|
"dim_batch": dim_batch,
|
|
"state": "loading",
|
|
"elapsed": 0.0,
|
|
# [done, elapsed] of the last 10 updates: the ETA follows the
|
|
# RECENT pace (throughput can degrade mid-fit, e.g. VRAM
|
|
# saturated — a global average would then freeze the ETA)
|
|
"hist": [],
|
|
}
|
|
workers.append(worker_state)
|
|
threading.Thread(
|
|
target=self._read_worker, args=(proc, worker_state, started), daemon=True
|
|
).start()
|
|
self.state["workers"] = workers
|
|
self._emit()
|
|
|
|
stderr_tails = [""] * len(self._procs)
|
|
|
|
def drain_err(index, proc):
|
|
data = proc.stderr.read()
|
|
stderr_tails[index] = (data or "")[-2000:]
|
|
|
|
drainers = [
|
|
threading.Thread(target=drain_err, args=(i, p), daemon=True)
|
|
for i, p in enumerate(self._procs)
|
|
]
|
|
for t in drainers:
|
|
t.start()
|
|
for proc in self._procs:
|
|
proc.wait()
|
|
for t in drainers:
|
|
t.join()
|
|
failed = [i for i, p in enumerate(self._procs) if p.returncode != 0]
|
|
if self.state.get("state") == "stopping":
|
|
self.state.update(state="stopped")
|
|
self._emit()
|
|
return
|
|
if failed:
|
|
detail = " | ".join(stderr_tails[i].strip().splitlines()[-1] if stderr_tails[i].strip() else "?" for i in failed)
|
|
raise RuntimeError(f"worker(s) {failed} failed: {detail}")
|
|
|
|
self.state.update(phase="merge")
|
|
self._emit()
|
|
partials = [
|
|
JacobianLens.load(str(job_dir / f"lens{i}.pt"))
|
|
for i in range(len(slices))
|
|
]
|
|
merged = JacobianLens.merge(partials) if len(partials) > 1 else partials[0]
|
|
if params.get("continue_from"):
|
|
base_lens = JacobianLens.load(params["continue_from"])
|
|
if base_lens.source_layers != merged.source_layers:
|
|
raise RuntimeError(
|
|
"cannot continue: the source layers differ from the base lens "
|
|
f"({base_lens.source_layers[0]}..{base_lens.source_layers[-1]} vs "
|
|
f"{merged.source_layers[0]}..{merged.source_layers[-1]})"
|
|
)
|
|
# weighted average by n_prompts = equivalent to a fit over the union
|
|
merged = JacobianLens.merge([base_lens, merged])
|
|
out_dir = config.LENSES_DIR / name
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
lens_path = out_dir / "lens.pt"
|
|
merged.save(str(lens_path))
|
|
meta = {
|
|
"name": name,
|
|
"model_id": params["model_id"],
|
|
"model_revision": params["model_revision"],
|
|
"model_source": params["source"],
|
|
"d_model": merged.d_model,
|
|
"source_layers": [merged.source_layers[0], merged.source_layers[-1]],
|
|
"dtype": params["dtype"],
|
|
"quant": params["quant"],
|
|
"n_prompts": merged.n_prompts,
|
|
"corpus": (
|
|
"mixed: " + " + ".join(params["datasets"]) + " (equal parts)"
|
|
if len(params["datasets"]) > 1
|
|
else params["datasets"][0]
|
|
),
|
|
"max_seq_len": params["max_seq_len"],
|
|
"devices": params["devices"],
|
|
"continued_from": params.get("continue_from"),
|
|
"config_hash": hashlib.sha1(
|
|
json.dumps(params, sort_keys=True).encode()
|
|
).hexdigest()[:16],
|
|
"created_at": _now(),
|
|
"fit_seconds": round(time.perf_counter() - started, 1),
|
|
}
|
|
(out_dir / "meta.json").write_text(
|
|
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
|
)
|
|
self.state.update(
|
|
state="done",
|
|
phase="done",
|
|
lens_path=str(lens_path),
|
|
meta=meta,
|
|
eta_seconds=0,
|
|
)
|
|
self._emit()
|
|
except Exception as exc:
|
|
self.state.update(state="error", error=str(exc))
|
|
self._emit()
|
|
|
|
def _read_worker(self, proc, worker_state, started):
|
|
for line in proc.stdout:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
event = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if event["event"] == "loading":
|
|
worker_state["state"] = "loading"
|
|
elif event["event"] in ("progress", "resume"):
|
|
worker_state["state"] = "fitting"
|
|
worker_state["done"] = event["done"]
|
|
worker_state["total"] = event["total"]
|
|
worker_state["elapsed"] = round(time.perf_counter() - started, 1)
|
|
hist = worker_state.setdefault("hist", [])
|
|
hist.append([worker_state["done"], worker_state["elapsed"]])
|
|
del hist[:-10]
|
|
elif event["event"] == "done":
|
|
worker_state["state"] = "done"
|
|
worker_state["done"] = worker_state["total"]
|
|
self._refresh_totals(started)
|
|
self._emit()
|
|
|
|
def _refresh_totals(self, started):
|
|
workers = self.state.get("workers", [])
|
|
self.state["done"] = sum(w["done"] for w in workers)
|
|
etas = []
|
|
for w in workers:
|
|
if not (w["done"] > 0 and w["elapsed"] > 0 and w["done"] < w["total"]):
|
|
continue
|
|
hist = w.get("hist") or []
|
|
if len(hist) >= 2 and hist[-1][1] > hist[0][1] and hist[-1][0] > hist[0][0]:
|
|
# pace over the last 10 updates (sliding window)
|
|
rate = (hist[-1][0] - hist[0][0]) / (hist[-1][1] - hist[0][1])
|
|
else:
|
|
rate = w["done"] / w["elapsed"]
|
|
etas.append((w["total"] - w["done"]) / rate)
|
|
# multi-GPU: the fit ETA = the slowest worker
|
|
self.state["eta_seconds"] = round(max(etas), 0) if etas else None
|
|
self.state["vram"] = [
|
|
{"index": g["index"], "used_gb": round(g["vram_used"] / 2**30, 1)}
|
|
for g in gpu_stats()
|
|
]
|