| import argparse |
| import gc |
| import json |
| import math |
| import os |
| import random |
| import warnings |
| from contextlib import contextmanager |
| from pathlib import Path |
| from runtime_utils import atomic_save, build_eval_cache, parquet_files as _parquet_files, hankel_basis |
|
|
| import torch |
| import torch.utils.checkpoint |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torch.distributed as dist |
| from torch.utils.data import DataLoader, Dataset, IterableDataset, get_worker_info |
|
|
| try: |
| import pyarrow.parquet as pq |
| except ImportError: |
| pq = None |
|
|
|
|
| CFG = { |
| "seed": 86, |
| "vocab_size": 258, |
| "seq_len": 2048, |
| "d_model": 1262, |
| "n_layers": 20, |
| "K_max": 32, |
| "batch_size": 16, |
| "grad_accum_steps": 2, |
| "grad_ckpt": False, |
| "gate_bias_init": -2.0, |
| "dropout": 0.1, |
| "stoch_depth": 0.1, |
| "lr": 3e-4, |
| "warmup_ratio": 0.01, |
| "min_lr_ratio": 0.1, |
| "weight_decay": 0.1, |
| "grad_clip": 1.0, |
| "target_tokens": 30_000_000_000, |
| "tokens_between_val": 1_000_000_000, |
| "budget_warmup_tokens": 1_000_000_000, |
| "budget_k_min": 2, |
| "budget_full_every": 8, |
| "budget_ks": [2, 3, 4, 5, 6, 8, 12, 16, 24, 32], |
| "eval_k": [2, 3, 4, 5, 6, 8, 12, 16, 24, 32], |
| "K_chunk": 8, |
| "eval_batches": 200, |
| "test_batches": 250, |
| "filter_iters": 60, |
| "filter_oversample": 16, |
| "filter_tol": 1e-6, |
| "num_workers": 4, |
| "val_segments": 1000, |
| "test_segments": 1000, |
| } |
|
|
|
|
| WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1)) |
| RANK = int(os.environ.get("RANK", 0)) |
| LOCAL_RANK = int(os.environ.get("LOCAL_RANK", 0)) |
| IS_DIST = WORLD_SIZE > 1 |
| IS_MASTER = RANK == 0 |
|
|
|
|
| if torch.cuda.is_available(): |
| DEVICE = torch.device(f"cuda:{LOCAL_RANK}") |
| elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available() and not IS_DIST: |
| DEVICE = torch.device("mps") |
| else: |
| DEVICE = torch.device("cpu") |
| AMP_DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
| def _amp_dtype(): |
| if not torch.cuda.is_available(): |
| return torch.float32 |
| try: |
| return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| except Exception: |
| return torch.float16 |
|
|
|
|
| AMP_DTYPE = _amp_dtype() |
|
|
|
|
| @contextmanager |
| def autocast(): |
| if torch.cuda.is_available(): |
| with torch.amp.autocast(device_type=AMP_DEVICE, dtype=AMP_DTYPE): |
| yield |
| else: |
| yield |
|
|
|
|
| def set_seed(s): |
| random.seed(s) |
| torch.manual_seed(s) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(s) |
|
|
|
|
| def count_params(m): |
| return sum(p.numel() for p in m.parameters() if p.requires_grad) |
|
|
|
|
| def cleanup(): |
| gc.collect() |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
|
|
| def emit(msg): |
| if IS_MASTER: |
| print(msg, flush=True) |
|
|
|
|
| def _b2t(b): |
| return torch.frombuffer(bytearray(b), dtype=torch.uint8).to(torch.int64).add_(2) |
|
|
|
|
| class StreamingByteDataset(IterableDataset): |
| def __init__(self, data_dir, split, seed, distributed=True): |
| if pq is None: |
| raise ImportError("Local parquet loading requires pyarrow.") |
| self.data_dir = data_dir |
| self.split = split |
| self.seed = seed |
| self.seg = CFG["seq_len"] + 1 |
| self.files = _parquet_files(data_dir, split) |
| self.distributed = distributed |
|
|
| def __iter__(self): |
| wi = get_worker_info() |
| wid = 0 if wi is None else wi.id |
| nw = 1 if wi is None else wi.num_workers |
| cid = RANK * nw + wid if self.distributed else wid |
| ctot = WORLD_SIZE * nw if self.distributed else nw |
| if len(self.files) < ctot: |
| raise ValueError(f"Need at least {ctot} parquet shards for the selected ranks/workers; found {len(self.files)}. Reduce num_workers or ranks.") |
| shard = self.files[cid::ctot] if ctot > 1 else self.files |
| if not shard: |
| return |
| rng = random.Random(self.seed + 13 * cid) |
| buf = bytearray() |
| while True: |
| order = list(shard) |
| rng.shuffle(order) |
| nonempty = False |
| for fp in order: |
| pf = pq.ParquetFile(fp) |
| for rb in pf.iter_batches(batch_size=256, columns=["text"]): |
| for t in rb.column(0).to_pylist(): |
| if not t: |
| continue |
| nonempty = True |
| buf.extend(t.encode("utf-8", errors="ignore")) |
| buf.extend(b"\n") |
| while len(buf) >= self.seg: |
| yield _b2t(bytes(buf[: self.seg])) |
| del buf[: self.seg - 1] |
| if not nonempty: |
| raise ValueError(f"No nonempty text in {self.split} shards.") |
|
|
|
|
| class CachedSegments(Dataset): |
| def __init__(self, t): |
| self.t = t |
|
|
| def __len__(self): |
| return self.t.shape[0] |
|
|
| def __getitem__(self, i): |
| return self.t[i].to(torch.long) |
|
|
|
|
| @torch.no_grad() |
| def build_cache(data_dir, output_dir, split, n, seed): |
| return build_eval_cache(data_dir, Path(output_dir) / "eval_cache", split, n, seed, CFG["seq_len"], RANK, WORLD_SIZE, IS_DIST) |
|
|
|
|
| @torch.no_grad() |
| def hankel_top_k(cache_dir, L, K): |
| return hankel_basis(cache_dir, L, K, CFG["filter_iters"], CFG["filter_oversample"], CFG["filter_tol"]) |
|
|
|
|
| class DropPath(nn.Module): |
| def __init__(self, p): |
| super().__init__() |
| self.p = float(p) |
|
|
| def forward(self, x): |
| if not self.training or self.p <= 0: |
| return x |
| shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| mask = torch.rand(shape, device=x.device, dtype=x.dtype) < (1 - self.p) |
| return x * mask / (1 - self.p) |
|
|
|
|
| class ESSMLayer(nn.Module): |
| def __init__(self, cache_dir, d, L, K, dp=0.0): |
| super().__init__() |
| self.K, self.L = K, L |
| sigma, phi = hankel_top_k(cache_dir, L, K) |
| self.register_buffer("sigma", sigma) |
| with torch.no_grad(): |
| phi_f = torch.fft.rfft(phi.T.float(), n=2 * L, dim=-1) |
| self.register_buffer("phi_f_ri", torch.view_as_real(phi_f).contiguous()) |
| self.M = nn.Parameter(torch.randn(K, d, d) * 1e-3) |
| self.gate = nn.Sequential(nn.Linear(d, d // 2), nn.GELU(), nn.Linear(d // 2, K)) |
| nn.init.constant_(self.gate[-1].bias, float(CFG["gate_bias_init"])) |
| self.dp = DropPath(dp) |
|
|
| def forward(self, u, K_rt=None): |
| B, L, D = u.shape |
| n = 2 * L |
| K = self.K if K_rt is None else max(CFG["budget_k_min"], min(self.K, int(K_rt))) |
| w = (torch.sigmoid(self.gate(u)[:, :, :K].float()) / math.sqrt(K)).to(u.dtype) |
| chunk = int(CFG.get("K_chunk", 0)) |
| chunked = (not self.training) and (0 < chunk < K) |
| if not chunked: |
| with torch.amp.autocast(device_type=AMP_DEVICE, enabled=False): |
| xf = torch.fft.rfft(u.float().permute(0, 2, 1), n=n, dim=-1) |
| pf = torch.view_as_complex(self.phi_f_ri[:K].contiguous()) |
| Uo = torch.fft.irfft(xf.unsqueeze(2) * pf.unsqueeze(0).unsqueeze(0), n=n, dim=-1)[..., :L] |
| Uo = Uo.permute(0, 3, 2, 1) |
| if Uo.dtype != u.dtype: |
| Uo = Uo.to(dtype=u.dtype, memory_format=torch.contiguous_format) |
| y = torch.einsum("blkd,kdo->blo", Uo * w.unsqueeze(-1), self.M[:K].to(u.dtype)) |
| else: |
| with torch.amp.autocast(device_type=AMP_DEVICE, enabled=False): |
| xf = torch.fft.rfft(u.float().permute(0, 2, 1), n=n, dim=-1) |
| pf_full = torch.view_as_complex(self.phi_f_ri[:K].contiguous()) |
| y = torch.zeros(B, L, D, device=u.device, dtype=u.dtype) |
| for s in range(0, K, chunk): |
| e = min(s + chunk, K) |
| with torch.amp.autocast(device_type=AMP_DEVICE, enabled=False): |
| pf_c = pf_full[s:e] |
| Uo_c = torch.fft.irfft(xf.unsqueeze(2) * pf_c.unsqueeze(0).unsqueeze(0), n=n, dim=-1)[..., :L] |
| Uo_c = Uo_c.permute(0, 3, 2, 1) |
| if Uo_c.dtype != u.dtype: |
| Uo_c = Uo_c.to(dtype=u.dtype, memory_format=torch.contiguous_format) |
| y = y + torch.einsum("blkd,kdo->blo", Uo_c * w[:, :, s:e].unsqueeze(-1), self.M[s:e].to(u.dtype)) |
| del Uo_c |
| return self.dp(y), K |
|
|
|
|
| class ESSSMLanguageModel(nn.Module): |
| def __init__(self, cache_dir): |
| super().__init__() |
| d = CFG["d_model"] |
| L = CFG["seq_len"] |
| K = CFG["K_max"] |
| nl = CFG["n_layers"] |
| self.K_max = K |
| self.d_model = d |
| self.use_ckpt = CFG["grad_ckpt"] |
| self.emb = nn.Embedding(CFG["vocab_size"], d, padding_idx=0) |
| self.drop = nn.Dropout(CFG["dropout"]) |
| sd = torch.linspace(0, CFG["stoch_depth"], nl).tolist() |
| self.layers = nn.ModuleList([ESSMLayer(cache_dir, d, L, K, sd[i]) for i in range(nl)]) |
| self.n1 = nn.ModuleList([nn.LayerNorm(d) for _ in range(nl)]) |
| self.n2 = nn.ModuleList([nn.LayerNorm(d) for _ in range(nl)]) |
| self.ffn = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(d, d * 4), |
| nn.GELU(), |
| nn.Dropout(CFG["dropout"]), |
| nn.Linear(d * 4, d), |
| ) for _ in range(nl) |
| ]) |
| self.ffn_dp = nn.ModuleList([DropPath(sd[i]) for i in range(nl)]) |
| self.ln_f = nn.LayerNorm(d) |
| self.head = nn.Linear(d, CFG["vocab_size"], bias=False) |
| nn.init.normal_(self.emb.weight, 0, 0.02) |
| self.head.weight = self.emb.weight |
|
|
| def _block(self, x, i, K_rt): |
| y, _ = self.layers[i](self.n1[i](x), K_rt=K_rt) |
| x = x + y |
| x = x + self.ffn_dp[i](self.ffn[i](self.n2[i](x))) |
| return x |
|
|
| def forward(self, tok, K_rt=None): |
| x = self.drop(self.emb(tok)) |
| for i in range(len(self.layers)): |
| if self.use_ckpt and self.training: |
| x = torch.utils.checkpoint.checkpoint(self._block, x, i, K_rt, use_reentrant=False) |
| else: |
| x = self._block(x, i, K_rt) |
| return self.head(self.ln_f(x)) |
|
|
|
|
| def build_compact(cache_dir, source, K): |
| old = CFG["K_max"] |
| try: |
| CFG["K_max"] = int(K) |
| compact = ESSSMLanguageModel(cache_dir) |
| finally: |
| CFG["K_max"] = old |
| with torch.no_grad(): |
| compact.emb.weight.copy_(source.emb.weight.detach().to(compact.emb.weight.device)) |
| compact.ln_f.load_state_dict(source.ln_f.state_dict()) |
| for d, s in zip(compact.n1, source.n1): |
| d.load_state_dict(s.state_dict()) |
| for d, s in zip(compact.n2, source.n2): |
| d.load_state_dict(s.state_dict()) |
| for d, s in zip(compact.ffn, source.ffn): |
| d.load_state_dict(s.state_dict()) |
| for dl, sl in zip(compact.layers, source.layers): |
| dl.sigma.copy_(sl.sigma[:K].detach().to(dl.sigma.device)) |
| dl.phi_f_ri.copy_(sl.phi_f_ri[:K].detach().to(dl.phi_f_ri.device)) |
| dl.M.copy_(sl.M[:K].detach().to(dl.M.device)) |
| dl.gate[0].load_state_dict(sl.gate[0].state_dict()) |
| dl.gate[2].weight.copy_(sl.gate[2].weight[:K].detach().to(dl.gate[2].weight.device)) |
| dl.gate[2].bias.copy_(sl.gate[2].bias[:K].detach().to(dl.gate[2].bias.device)) |
| compact.head.weight = compact.emb.weight |
| compact.use_ckpt = False |
| return compact |
|
|
|
|
| class BudgetSampler: |
| def __init__(self): |
| self.rng = random.Random(CFG["seed"] + 777) |
| self.ks = sorted(CFG["budget_ks"]) |
| self._eligible = [k for k in self.ks if k >= CFG["budget_k_min"]] |
| self._k_max = CFG["K_max"] |
| self._lo = max(CFG["budget_k_min"], min(self.ks)) |
| self._hi = max(self.ks) |
| self._log_ratio = math.log(self._hi / self._lo) |
|
|
| def state_dict(self): |
| return {"rng": self.rng.getstate()} |
|
|
| def load_state_dict(self, state): |
| if state and "rng" in state: |
| self.rng.setstate(state["rng"]) |
|
|
| def _draw(self, step, tokens_seen): |
| if tokens_seen < CFG["budget_warmup_tokens"]: |
| return self._k_max |
| fe = int(CFG["budget_full_every"]) |
| if fe > 0 and step % fe == 0: |
| return self._k_max |
| u = self.rng.random() |
| target = self._lo * math.exp(u * self._log_ratio) |
| return min(self._eligible, key=lambda k: abs(k - target)) |
|
|
| def sample(self, step, tokens_seen): |
| if IS_DIST: |
| if IS_MASTER: |
| k = self._draw(step, tokens_seen) |
| buf = torch.tensor([k], dtype=torch.long, device=DEVICE) |
| else: |
| buf = torch.empty(1, dtype=torch.long, device=DEVICE) |
| dist.broadcast(buf, src=0) |
| return int(buf.item()) |
| return self._draw(step, tokens_seen) |
|
|
|
|
| def build_optimizer(model): |
| decay, no_decay, seen = [], [], set() |
| for n, p in model.named_parameters(): |
| if not p.requires_grad or id(p) in seen: |
| continue |
| seen.add(id(p)) |
| if p.ndim == 1 or any(k in n.lower() for k in ("bias", "norm", "emb", "ln")): |
| no_decay.append(p) |
| else: |
| decay.append(p) |
| return torch.optim.AdamW( |
| [{"params": decay, "weight_decay": CFG["weight_decay"]}, |
| {"params": no_decay, "weight_decay": 0.0}], |
| lr=CFG["lr"], betas=(0.9, 0.95), eps=1e-8, |
| ) |
|
|
|
|
| def make_scaler(enabled): |
| try: |
| return torch.amp.GradScaler(AMP_DEVICE, enabled=enabled) |
| except TypeError: |
| return torch.cuda.amp.GradScaler(enabled=enabled) |
|
|
|
|
| class CompiledInference(nn.Module): |
| def __init__(self, model, enabled=False): |
| super().__init__() |
| self.model = model.eval() |
| self.compiled = None |
| self.checked_inputs = set() |
| self.execution = "eager" |
| if enabled: |
| try: |
| device = next(model.parameters()).device |
| mode = "reduce-overhead" if device.type == "cuda" else "default" |
| self.compiled = torch.compile(model, mode=mode, dynamic=False, fullgraph=False) |
| self.execution = "pending" |
| except Exception as exc: |
| self._fallback(exc) |
| self.eval() |
|
|
| def _fallback(self, exc): |
| self.compiled = None |
| self.execution = "fallback" |
| detail = str(exc).splitlines()[0] if str(exc) else type(exc).__name__ |
| warnings.warn(f"Compiled inference fell back to eager execution: {detail}", RuntimeWarning) |
|
|
| def forward(self, tokens): |
| if self.training: |
| raise RuntimeError("CompiledInference only supports evaluation.") |
| if self.compiled is None: |
| return self.model(tokens) |
| try: |
| compiler = getattr(torch, "compiler", None) |
| if tokens.is_cuda and hasattr(compiler, "cudagraph_mark_step_begin"): |
| compiler.cudagraph_mark_step_begin() |
| output = self.compiled(tokens) |
| key = (tuple(tokens.shape), tuple(tokens.stride()), tokens.device, tokens.dtype) |
| if key not in self.checked_inputs: |
| reference = self.model(tokens) |
| if output.shape != reference.shape or output.dtype != reference.dtype: |
| raise RuntimeError("Compiled output shape or dtype differs from eager output.") |
| if not torch.isfinite(output).all() or not torch.isfinite(reference).all(): |
| raise RuntimeError("Non-finite output during compiled-inference validation.") |
| error = (output.float() - reference.float()).norm() |
| error = float(error / reference.float().norm().clamp_min(1e-12)) |
| tolerance = 1e-4 if reference.dtype == torch.float32 else 0.005 |
| if not math.isfinite(error) or error > tolerance: |
| raise RuntimeError(f"Compiled output relative L2 error {error:.6g} exceeds {tolerance}.") |
| self.checked_inputs.add(key) |
| self.execution = "compiled" |
| return output |
| except Exception as exc: |
| self._fallback(exc) |
| return self.model(tokens) |
|
|
|
|
| def checkpoint_selection(mode="full", budgets=None, weights=None, full_bpb_ceiling=None): |
| if mode not in ("full", "multi"): |
| raise ValueError("Checkpoint selection must be 'full' or 'multi'.") |
| if mode == "full": |
| if budgets is not None or weights is not None or full_bpb_ceiling is not None: |
| raise ValueError("Budget weights and the full-budget ceiling require --checkpoint-selection multi.") |
| return {"mode": "full", "budgets": [CFG["K_max"]], "weights": [1.0], "full_bpb_ceiling": None} |
| ks = [int(k) for k in budgets.split(",")] if budgets is not None else list(CFG["eval_k"]) |
| ws = [float(w) for w in weights.split(",")] if weights is not None else [1.0] * len(ks) |
| if not ks or len(ks) != len(set(ks)) or any(k < CFG["budget_k_min"] or k > CFG["K_max"] for k in ks): |
| raise ValueError("Selection budgets must be unique supported integer budgets.") |
| if len(ws) != len(ks) or any(not math.isfinite(w) or w <= 0 for w in ws) or not math.isfinite(sum(ws)): |
| raise ValueError("Selection weights must be finite, positive, and match the budget count.") |
| if full_bpb_ceiling is not None and (not math.isfinite(full_bpb_ceiling) or full_bpb_ceiling <= 0): |
| raise ValueError("The full-budget BPB ceiling must be finite and positive.") |
| return {"mode": "multi", "budgets": ks, "weights": ws, "full_bpb_ceiling": full_bpb_ceiling} |
|
|
|
|
| def checkpoint_score(metrics, selection): |
| if any(not math.isfinite(value) for value in metrics.values()): |
| raise ValueError("Non-finite validation BPB cannot select a checkpoint.") |
| total = sum(selection["weights"]) |
| return sum(metrics[k] * (w / total) for k, w in zip(selection["budgets"], selection["weights"])) |
|
|
|
|
| @torch.no_grad() |
| def eval_bpb(model, loader, K=None, elastic=True, max_batches=None, distributed=True): |
| model.eval() |
| nll, tok = 0.0, 0 |
| max_batches = max_batches if max_batches else CFG["eval_batches"] |
| for i, b in enumerate(loader): |
| if i >= max_batches: |
| break |
| b = b.to(DEVICE, non_blocking=True) |
| x, y = b[:, :-1], b[:, 1:] |
| with autocast(): |
| lo = model(x, K_rt=K) if elastic else model(x) |
| loss = F.cross_entropy(lo.reshape(-1, lo.size(-1)), y.reshape(-1)) |
| nll += loss.item() * y.numel() |
| tok += y.numel() |
| if IS_DIST and distributed: |
| stats = torch.tensor([nll, float(tok)], device=DEVICE) |
| dist.all_reduce(stats, op=dist.ReduceOp.SUM) |
| nll, tok = stats[0].item(), stats[1].item() |
| if tok <= 0: |
| raise ValueError("Evaluation produced no target tokens.") |
| value = (nll / tok) / math.log(2) |
| if not math.isfinite(value): |
| raise FloatingPointError("Evaluation BPB is non-finite.") |
| return value |
|
|
|
|
| def train_loop(model, train_dl, val_dl, ckpt_dir, selection=None): |
| selection = checkpoint_selection() if selection is None else selection |
| model.to(DEVICE) |
| if hasattr(model, "use_ckpt"): |
| model.use_ckpt = CFG["grad_ckpt"] |
| core = model.module if IS_DIST and hasattr(model, "module") else model |
|
|
| target_tokens = CFG["target_tokens"] |
| tokens_per_step = CFG["batch_size"] * CFG["grad_accum_steps"] * CFG["seq_len"] * WORLD_SIZE |
| total_steps = max(1, target_tokens // tokens_per_step) |
| wu = max(1, int(total_steps * CFG["warmup_ratio"])) |
|
|
| opt = build_optimizer(model) |
|
|
| def lrf(s): |
| if s < wu: |
| return (s + 1) / wu |
| p = (s - wu) / max(1, total_steps - wu) |
| return CFG["min_lr_ratio"] + (1 - CFG["min_lr_ratio"]) * 0.5 * (1 + math.cos(math.pi * p)) |
|
|
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lrf) |
| use_scaler = AMP_DTYPE == torch.float16 |
| scaler = make_scaler(use_scaler) |
| bud = BudgetSampler() |
|
|
| it = iter(train_dl) |
| g = 0 |
| tokens_seen = 0 |
| next_val = CFG["tokens_between_val"] |
| best_bpb = float("inf") |
| best_score = float("inf") |
| selected_path = None |
|
|
| while tokens_seen < target_tokens: |
| model.train() |
| K_cur = bud.sample(g, tokens_seen) |
| opt.zero_grad(set_to_none=True) |
| step_tokens = 0 |
| for _ in range(CFG["grad_accum_steps"]): |
| try: |
| batch = next(it) |
| except StopIteration: |
| it = iter(train_dl) |
| batch = next(it) |
| batch = batch.to(DEVICE, non_blocking=True) |
| x, y = batch[:, :-1], batch[:, 1:] |
| with autocast(): |
| logits = model(x, K_rt=K_cur) |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) |
| scaled = loss / CFG["grad_accum_steps"] |
| if use_scaler: |
| scaler.scale(scaled).backward() |
| else: |
| scaled.backward() |
| step_tokens += y.numel() |
| if use_scaler: |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), CFG["grad_clip"]) |
| if use_scaler: |
| scaler.step(opt) |
| scaler.update() |
| else: |
| opt.step() |
| sched.step() |
| g += 1 |
| tokens_seen += step_tokens * WORLD_SIZE |
|
|
| if tokens_seen >= next_val or tokens_seen >= target_tokens: |
| next_val += CFG["tokens_between_val"] |
| ks = list(dict.fromkeys([CFG["K_max"]] + selection["budgets"])) |
| metrics = { |
| k: eval_bpb(model, val_dl, K=k, elastic=True, max_batches=CFG["eval_batches"]) |
| for k in ks |
| } |
| vbpb = metrics[CFG["K_max"]] |
| score = checkpoint_score(metrics, selection) |
| ceiling = selection["full_bpb_ceiling"] |
| eligible = ceiling is None or vbpb <= ceiling |
| emit(json.dumps({"step": g, "tokens_seen": tokens_seen, "validation_bpb": metrics, |
| "selection_score": score, "eligible": eligible}, sort_keys=True)) |
| metadata = { |
| "cfg": dict(CFG), |
| "step": g, |
| "tokens_seen": tokens_seen, |
| "validation_bpb": metrics, |
| "selection_protocol": selection, |
| } |
| if vbpb < best_bpb - 1e-4: |
| best_bpb = vbpb |
| if selection["mode"] == "full" and ckpt_dir is not None: |
| selected_path = ckpt_dir / "best.pt" |
| if IS_MASTER and ckpt_dir is not None: |
| atomic_save({"model": core.state_dict(), **metadata, "selected_by": "full", |
| "selection_score": vbpb}, str(ckpt_dir / "best.pt")) |
| if selection["mode"] == "multi" and eligible and score < best_score - 1e-4: |
| best_score = score |
| if ckpt_dir is not None: |
| selected_path = ckpt_dir / "best_multibudget.pt" |
| if IS_MASTER: |
| atomic_save({"model": core.state_dict(), **metadata, "selected_by": "multi", |
| "selection_score": score}, str(selected_path)) |
| if IS_DIST: |
| dist.barrier() |
| if selection["mode"] == "multi" and not math.isfinite(best_score): |
| raise RuntimeError("No checkpoint satisfied the prescribed full-budget BPB ceiling. The full-budget best is saved as best.pt.") |
| return selected_path |
|
|
|
|
| def maybe_init_dist(): |
| global IS_DIST |
| if DEVICE.type == "cuda": |
| torch.cuda.set_device(DEVICE) |
| if WORLD_SIZE > 1 and not dist.is_initialized(): |
| backend = "nccl" if torch.cuda.is_available() else "gloo" |
| dist.init_process_group(backend=backend, rank=RANK, world_size=WORLD_SIZE) |
| IS_DIST = True |
|
|
|
|
| def main(): |
| pa = argparse.ArgumentParser() |
| pa.add_argument("--data-dir", type=str, required=True) |
| pa.add_argument("--output-dir", type=str, required=True) |
| pa.add_argument("--filter-cache-dir", type=str, default="./cache/filter") |
| pa.add_argument("--inference-only", action="store_true") |
| pa.add_argument("--checkpoint", type=str, default=None) |
| pa.add_argument("--compile-inference", action="store_true") |
| pa.add_argument("--checkpoint-selection", choices=("full", "multi"), default="full") |
| pa.add_argument("--selection-budgets", type=str) |
| pa.add_argument("--selection-weights", type=str) |
| pa.add_argument("--full-bpb-ceiling", type=float) |
| pa.add_argument("--num-workers", type=int, default=None) |
| pa.add_argument("--local-rank", "--local_rank", type=int, default=None) |
| args = pa.parse_args() |
| if args.checkpoint and not args.inference_only: |
| pa.error("--checkpoint is for --inference-only; training resume is not implemented by this entrypoint.") |
| if args.num_workers is not None: |
| if args.num_workers < 0: |
| pa.error("--num-workers must be nonnegative") |
| CFG["num_workers"] = args.num_workers |
| try: |
| selection = checkpoint_selection(args.checkpoint_selection, args.selection_budgets, |
| args.selection_weights, args.full_bpb_ceiling) |
| except ValueError as exc: |
| pa.error(str(exc)) |
|
|
| maybe_init_dist() |
| set_seed(CFG["seed"] + RANK) |
| if torch.cuda.is_available(): |
| torch.backends.cudnn.benchmark = True |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
|
|
| out = Path(args.output_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| cache_dir = Path(args.filter_cache_dir) |
| cache_dir.mkdir(parents=True, exist_ok=True) |
|
|
| val_t = build_cache(args.data_dir, args.output_dir, "validation", CFG["val_segments"], CFG["seed"] + 111) |
| val_dl = DataLoader(CachedSegments(val_t), batch_size=CFG["batch_size"], shuffle=False, num_workers=0) |
|
|
| model = ESSSMLanguageModel(args.filter_cache_dir) |
| if IS_DIST: |
| model = model.to(DEVICE) |
| model = torch.nn.parallel.DistributedDataParallel( |
| model, device_ids=[LOCAL_RANK] if DEVICE.type == "cuda" else None, |
| output_device=LOCAL_RANK if DEVICE.type == "cuda" else None, |
| find_unused_parameters=False, |
| ) |
|
|
| ckpt_dir = out / "checkpoints" |
| ckpt_dir.mkdir(parents=True, exist_ok=True) |
|
|
| selected_path = ckpt_dir / ("best_multibudget.pt" if selection["mode"] == "multi" else "best.pt") |
| if not args.inference_only: |
| train_ds = StreamingByteDataset(args.data_dir, "train", CFG["seed"]) |
| train_dl = DataLoader( |
| train_ds, |
| batch_size=CFG["batch_size"], |
| num_workers=CFG["num_workers"], |
| pin_memory=True, |
| collate_fn=torch.stack, |
| drop_last=True, |
| persistent_workers=CFG["num_workers"] > 0, |
| ) |
| selected_path = train_loop(model, train_dl, val_dl, ckpt_dir, selection=selection) |
| checkpoint_path = Path(args.checkpoint) if args.checkpoint else selected_path |
| if checkpoint_path is None or not checkpoint_path.is_file(): |
| raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") |
| ck = torch.load(str(checkpoint_path), map_location="cpu") |
| (model.module if IS_DIST else model).load_state_dict(ck["model"]) |
| del ck |
| test_t = build_cache(args.data_dir, args.output_dir, "test", CFG["test_segments"], CFG["seed"] + 222) |
| test_dl = DataLoader(CachedSegments(test_t), batch_size=CFG["batch_size"], shuffle=False, num_workers=0) |
|
|
| core = model.module if IS_DIST else model |
| core.use_ckpt = False |
|
|
| if IS_MASTER: |
| core = core.cpu() |
| cleanup() |
| for K in CFG["eval_k"]: |
| compact = build_compact(args.filter_cache_dir, core, K) |
| compact = compact.to(DEVICE) |
| compiler = getattr(torch, "compiler", None) |
| if args.compile_inference and hasattr(compiler, "reset"): |
| compiler.reset() |
| inference = CompiledInference(compact, enabled=args.compile_inference) |
| v = eval_bpb(inference, val_dl, elastic=False, max_batches=CFG["eval_batches"], distributed=False) |
| t = eval_bpb(inference, test_dl, elastic=False, max_batches=CFG["test_batches"], distributed=False) |
| emit(f"K={K} val={v:.4f} test={t:.4f} execution={inference.execution}") |
| del inference, compact |
| cleanup() |
|
|
| if IS_DIST: |
| dist.barrier() |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
