File size: 8,514 Bytes
70aae54 | 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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | import hashlib
import json
import os
import random
import tempfile
from pathlib import Path
import pyarrow.parquet as pq
import torch
import torch.distributed as dist
from filelock import FileLock
def parquet_files(data_dir, split):
if not data_dir:
raise ValueError("A local dataset directory is required.")
root = Path(data_dir).expanduser()
if not root.is_dir():
raise FileNotFoundError(f"Local dataset directory does not exist: {root}")
files = sorted((root / "data").glob(f"{split}-*.parquet"))
if not files:
raise FileNotFoundError(f"No {split} parquet shards under {root / 'data'}")
return files
def file_sha256(path):
path = Path(path)
before = path.stat()
digest = hashlib.sha256()
with path.open("rb") as handle:
for block in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(block)
after = path.stat()
if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns):
raise RuntimeError(f"File changed while hashing: {path}")
return digest.hexdigest()
def atomic_save(value, path):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory(prefix=".pending-", dir=path.parent) as tmp:
pending = Path(tmp) / path.name
torch.save(value, pending)
os.replace(pending, path)
def _metadata(data_dir, split, n, seed, seq_len):
return {
"version": 1, "split": split, "seed": int(seed),
"seq_len": int(seq_len), "segments": int(n),
"encoding": "utf8-ignore+LF;byte-offset=2;overlap=1",
"files": [{"name": p.name, "sha256": file_sha256(p)} for p in parquet_files(data_dir, split)],
}
def _tensor_sha256(t):
return hashlib.sha256(t.numpy().astype("<i2", copy=False).tobytes()).hexdigest()
def _load_cache(path, metadata):
value = torch.load(path, map_location="cpu", weights_only=True)
if not isinstance(value, dict) or value.get("metadata") != metadata:
raise ValueError(f"Cache provenance is missing or mismatched: {path}. Use a new output directory or move the old cache aside.")
t = value.get("t")
if (
not isinstance(t, torch.Tensor) or t.dtype != torch.int16
or tuple(t.shape) != (metadata["segments"], metadata["seq_len"] + 1)
or int(t.min()) < 2 or int(t.max()) > 257
or value.get("tensor_sha256") != _tensor_sha256(t)
):
raise ValueError(f"Cache tensor failed integrity validation: {path}")
return t
def _segments(data_dir, split, seed, seq_len):
files = parquet_files(data_dir, split)
rng = random.Random(seed)
buf = bytearray()
while True:
order = list(files)
rng.shuffle(order)
nonempty = False
for path in order:
for rb in pq.ParquetFile(path).iter_batches(batch_size=256, columns=["text"]):
for text in rb.column(0).to_pylist():
if not text:
continue
nonempty = True
buf.extend(text.encode("utf-8", errors="ignore"))
buf.extend(b"\n")
while len(buf) >= seq_len + 1:
yield torch.frombuffer(bytearray(buf[:seq_len + 1]), dtype=torch.uint8).long().add_(2)
del buf[:seq_len]
if not nonempty:
raise ValueError(f"No nonempty text in the {split} split.")
@torch.no_grad()
def build_eval_cache(data_dir, cache_dir, split, n, seed, seq_len, rank=0, world_size=1, distributed=False):
if n <= 0 or seq_len <= 0:
raise ValueError("Cache size and sequence length must be positive.")
root = Path(cache_dir)
path = root / f"slimpj_{split}_L{seq_len}_N{n}.pt"
status, tensor = [None], None
if not distributed or rank == 0:
try:
root.mkdir(parents=True, exist_ok=True)
with FileLock(str(path) + ".lock", timeout=300):
metadata = _metadata(data_dir, split, n, seed, seq_len)
if path.exists():
tensor = _load_cache(path, metadata)
else:
stream = iter(_segments(data_dir, split, seed, seq_len))
tensor = torch.stack([next(stream).to(torch.int16) for _ in range(n)])
if _metadata(data_dir, split, n, seed, seq_len) != metadata:
raise RuntimeError("Dataset changed during cache generation.")
atomic_save({"t": tensor, "metadata": metadata, "tensor_sha256": _tensor_sha256(tensor)}, path)
status[0] = {"metadata": metadata}
except Exception as exc:
if not distributed:
raise
status[0] = {"error": f"{type(exc).__name__}: {exc}"}
if distributed:
dist.broadcast_object_list(status, src=0)
if "error" in status[0]:
raise RuntimeError(status[0]["error"])
error = None
if rank != 0:
try:
tensor = _load_cache(path, status[0]["metadata"])
except Exception as exc:
error = f"Rank {rank}: {type(exc).__name__}: {exc}"
errors = [None] * world_size
dist.all_gather_object(errors, error)
if any(errors):
raise RuntimeError("; ".join(x for x in errors if x))
return tensor
def hankel_kernel(length, eps=1e-6):
s = torch.arange(2, 2 * length + 1, dtype=torch.float32, device="cpu")
a = 2 / (s ** 3 - s + eps)
return torch.cat([a[length - 1:], torch.zeros(1), a[:length - 1]])
def hmv(g_f, x):
length, p = x.shape
padded = torch.zeros(p, 2 * length, dtype=x.dtype, device=x.device)
padded[:, :length] = x.flip(0).T
return torch.fft.irfft(torch.fft.rfft(padded, dim=-1) * g_f[None], n=2 * length, dim=-1)[:, :length].T
def _basis_digest(sigma, phi):
digest = hashlib.sha256()
for tensor in (sigma, phi):
digest.update(tensor.numpy().astype("<f4", copy=False).tobytes())
return digest.hexdigest()
@torch.no_grad()
def hankel_basis(cache_dir, length, channels, iters=60, oversample=16, tol=1e-6):
if not 1 <= channels <= length or iters <= 0 or oversample < 0 or tol <= 0:
raise ValueError("Invalid Hankel basis configuration.")
metadata = dict(version=2, length=length, channels=channels, iters=iters, oversample=oversample, tol=tol, seed=0)
signature = hashlib.sha256(json.dumps(metadata, sort_keys=True).encode()).hexdigest()[:12]
root = Path(cache_dir)
root.mkdir(parents=True, exist_ok=True)
path = root / f"hankel_v2_L{length}_K{channels}_{signature}.pt"
with FileLock(str(path) + ".lock", timeout=300):
if path.exists():
saved = torch.load(path, map_location="cpu", weights_only=True)
sigma, phi = saved["sigma"], saved["phi"]
if (
saved.get("metadata") != metadata or sigma.shape != (channels,)
or phi.shape != (length, channels)
or not torch.isfinite(sigma).all() or not torch.isfinite(phi).all()
or saved.get("sha256") != _basis_digest(sigma, phi)
):
raise ValueError(f"Invalid Hankel cache: {path}")
return sigma, phi
generator = torch.Generator(device="cpu").manual_seed(0)
q = torch.randn(length, channels + oversample, generator=generator, device="cpu", dtype=torch.float32)
q, _ = torch.linalg.qr(q, mode="reduced")
gf = torch.fft.rfft(hankel_kernel(length))
for step in range(1, iters + 1):
q, _ = torch.linalg.qr(hmv(gf, q), mode="reduced")
if step % max(1, iters // 4) == 0 or step == iters:
values, vectors = torch.linalg.eigh(q.T @ hmv(gf, q))
order = values.argsort(descending=True)
values, vectors = values[order], vectors[:, order]
q = q @ vectors
residual = (hmv(gf, q[:, :channels]) - q[:, :channels] * values[:channels][None]).norm(dim=0)
if float((residual / (values[:channels].abs() + 1e-12)).max()) < tol:
break
values, vectors = torch.linalg.eigh(q.T @ hmv(gf, q))
order = values.argsort(descending=True)
sigma, phi = values[order][:channels], (q @ vectors[:, order])[:, :channels]
atomic_save({"sigma": sigma, "phi": phi, "metadata": metadata, "sha256": _basis_digest(sigma, phi)}, path)
return sigma, phi
|