|
|
|
|
| |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| |
|
|
|
|
|
|
| |
|
|
| |
|
|
|
|
|
|
| from __future__ import annotations
|
| import argparse, json, math, pathlib, random, time, os, sys, threading, hashlib, re, subprocess, shutil
|
| from pathlib import Path
|
| from contextlib import nullcontext
|
| from typing import Dict, Any, List, Optional, Tuple
|
| from datetime import datetime, timezone
|
| STATUS_SCRIPT_PATH = Path(__file__).resolve()
|
| STATUS_DEFAULT_LOG = STATUS_SCRIPT_PATH.parent / "train.log"
|
| STATUS_DEFAULT_SAVE_DIR = STATUS_SCRIPT_PATH.parent / "ckpts_expansion"
|
| _STATUS_PROGRESS_RE = re.compile(
|
| r"^\[(?P<percent>\d+(?:\.\d+)?)%\]\s+"
|
| r"(?P<seen>[\d,]+)/(?P<target>[\d,]+)\s+tok\s+\|\s+"
|
| r"(?P<tok_s>[\d.]+)\s+tok/s\s+\|\s+"
|
| r"loss=(?P<loss>-?[\d.]+)\s+B=(?P<batch>\d+)\s+L=(?P<block>\d+)\s*$"
|
| )
|
| _STATUS_DELTA_RE = re.compile(r"\[delta\]\s+saved\s+(?P<name>\S+?\.pt)\s+\((?P<sha>[0-9a-f]+)\.\.\.\)")
|
| _STATUS_STEP_RE = re.compile(r"step(?P<step>\d+)")
|
|
|
|
|
| def _status_iso(ts: Optional[float]) -> Optional[str]:
|
| if ts is None:
|
| return None
|
| return datetime.fromtimestamp(ts, tz=timezone.utc).astimezone().isoformat(timespec="seconds")
|
|
|
|
|
| def _status_human_duration(seconds: Optional[float]) -> Optional[str]:
|
| if seconds is None:
|
| return None
|
| total = max(0, int(seconds))
|
| days, rem = divmod(total, 86400)
|
| hours, rem = divmod(rem, 3600)
|
| minutes, secs = divmod(rem, 60)
|
| parts = []
|
| if days:
|
| parts.append(f"{days}d")
|
| if hours or parts:
|
| parts.append(f"{hours}h")
|
| if minutes or parts:
|
| parts.append(f"{minutes}m")
|
| parts.append(f"{secs}s")
|
| return " ".join(parts)
|
|
|
|
|
| def _status_format_int(value: Optional[int]) -> str:
|
| return "?" if value is None else f"{value:,}"
|
|
|
|
|
| def _status_parse_step(text: str) -> Optional[int]:
|
| match = _STATUS_STEP_RE.search(text)
|
| return int(match.group("step")) if match else None
|
|
|
|
|
| def _status_resolve_ckpt_path(raw_path: str, base_dir: Path) -> Path:
|
| ckpt_path = Path(raw_path)
|
| return ckpt_path if ckpt_path.is_absolute() else (base_dir / ckpt_path).resolve()
|
|
|
|
|
| def _status_read_cmdline(proc_dir: Path) -> Optional[List[str]]:
|
| try:
|
| data = (proc_dir / "cmdline").read_bytes().split(b"\0")
|
| return [item.decode("utf-8", errors="ignore") for item in data if item]
|
| except Exception:
|
| return None
|
|
|
|
|
| def _status_resolve_proc_arg(proc_dir: Path, raw_arg: str) -> Optional[Path]:
|
| try:
|
| arg_path = Path(raw_arg)
|
| if arg_path.is_absolute():
|
| return arg_path.resolve()
|
| cwd = Path(os.readlink(proc_dir / "cwd"))
|
| return (cwd / arg_path).resolve()
|
| except Exception:
|
| return None
|
|
|
|
|
| def _status_proc_uptime(proc_dir: Path) -> Optional[float]:
|
| try:
|
| proc_uptime = float((Path("/proc") / "uptime").read_text().split()[0])
|
| stat_text = (proc_dir / "stat").read_text()
|
| after = stat_text[stat_text.rfind(")") + 2:].split()
|
| start_ticks = float(after[19])
|
| clock_ticks = os.sysconf(os.sysconf_names["SC_CLK_TCK"])
|
| return max(0.0, proc_uptime - (start_ticks / clock_ticks))
|
| except Exception:
|
| return None
|
|
|
|
|
| def _status_find_trainers(script_path: Path) -> List[Dict[str, Any]]:
|
| matches: List[Dict[str, Any]] = []
|
| for proc_dir in Path("/proc").iterdir():
|
| if not proc_dir.name.isdigit():
|
| continue
|
| args = _status_read_cmdline(proc_dir)
|
| if not args or "train" not in args:
|
| continue
|
| resolved_script = None
|
| for arg in args:
|
| if Path(arg).name != script_path.name:
|
| continue
|
| candidate = _status_resolve_proc_arg(proc_dir, arg)
|
| if candidate == script_path:
|
| resolved_script = candidate
|
| break
|
| if resolved_script is None:
|
| continue
|
| uptime_seconds = _status_proc_uptime(proc_dir)
|
| try:
|
| cwd = str(Path(os.readlink(proc_dir / "cwd")))
|
| except Exception:
|
| cwd = None
|
| matches.append({
|
| "pid": int(proc_dir.name),
|
| "cmdline": " ".join(args),
|
| "args": args,
|
| "cwd": cwd,
|
| "uptime_seconds": round(uptime_seconds, 3) if uptime_seconds is not None else None,
|
| "uptime_human": _status_human_duration(uptime_seconds),
|
| })
|
| return sorted(matches, key=lambda item: item["pid"])
|
|
|
|
|
| def _status_parse_progress_line(line: str) -> Optional[Dict[str, Any]]:
|
| match = _STATUS_PROGRESS_RE.match(line.strip())
|
| if not match:
|
| return None
|
| tok_per_sec = float(match.group("tok_s"))
|
| loss = float(match.group("loss"))
|
| return {
|
| "raw_line": line.strip(),
|
| "percent": float(match.group("percent")),
|
| "seen_tokens": int(match.group("seen").replace(",", "")),
|
| "target_tokens": int(match.group("target").replace(",", "")),
|
| "tok_per_sec": int(tok_per_sec) if tok_per_sec.is_integer() else tok_per_sec,
|
| "loss": loss,
|
| "batch": int(match.group("batch")),
|
| "block": int(match.group("block")),
|
| }
|
|
|
|
|
| def _status_parse_delta_line(line: str) -> Optional[Dict[str, Any]]:
|
| match = _STATUS_DELTA_RE.search(line)
|
| if not match:
|
| return None
|
| name = match.group("name")
|
| return {
|
| "raw_line": line.strip(),
|
| "name": name,
|
| "step": _status_parse_step(name),
|
| "sha_prefix": match.group("sha"),
|
| "source": "log",
|
| }
|
|
|
|
|
| def _status_scan_log(log_path: Path) -> tuple[Dict[str, Any], Optional[Dict[str, Any]], Optional[Dict[str, Any]], List[str]]:
|
| now = time.time()
|
| info: Dict[str, Any] = {
|
| "path": str(log_path),
|
| "exists": log_path.exists(),
|
| "mtime": None,
|
| "mtime_iso": None,
|
| "age_seconds": None,
|
| "age_human": None,
|
| "size_bytes": None,
|
| }
|
| warnings: List[str] = []
|
| if not log_path.exists():
|
| warnings.append(f"train log missing: {log_path}")
|
| return info, None, None, warnings
|
| try:
|
| st = log_path.stat()
|
| info["mtime"] = st.st_mtime
|
| info["mtime_iso"] = _status_iso(st.st_mtime)
|
| info["age_seconds"] = round(max(0.0, now - st.st_mtime), 3)
|
| info["age_human"] = _status_human_duration(info["age_seconds"])
|
| info["size_bytes"] = st.st_size
|
| except Exception as exc:
|
| warnings.append(f"failed to stat train log: {exc}")
|
| last_progress = None
|
| last_delta = None
|
| try:
|
| with log_path.open("r", encoding="utf-8", errors="ignore") as handle:
|
| for raw_line in handle:
|
| line = raw_line.rstrip("\n")
|
| progress = _status_parse_progress_line(line)
|
| if progress is not None:
|
| last_progress = progress
|
| delta = _status_parse_delta_line(line)
|
| if delta is not None:
|
| last_delta = delta
|
| except Exception as exc:
|
| warnings.append(f"failed to read train log: {exc}")
|
| return info, last_progress, last_delta, warnings
|
|
|
|
|
| def _status_latest_full_checkpoint(save_dir: Path, base_dir: Path) -> tuple[Dict[str, Any], List[str]]:
|
| latest_path = save_dir / "latest.json"
|
| info: Dict[str, Any] = {
|
| "metadata_path": str(latest_path),
|
| "exists": latest_path.exists(),
|
| "raw_path": None,
|
| "checkpoint_path": None,
|
| "checkpoint_name": None,
|
| "checkpoint_exists": None,
|
| "step": None,
|
| "checkpoint_mtime": None,
|
| "checkpoint_mtime_iso": None,
|
| }
|
| warnings: List[str] = []
|
| if not latest_path.exists():
|
| warnings.append(f"latest.json missing: {latest_path}")
|
| return info, warnings
|
| try:
|
| payload = json.loads(latest_path.read_text(encoding="utf-8"))
|
| except Exception as exc:
|
| warnings.append(f"failed to parse latest.json: {exc}")
|
| return info, warnings
|
| raw_path = payload.get("path")
|
| info["raw_path"] = raw_path
|
| info["step"] = payload.get("step")
|
| if raw_path:
|
| ckpt_path = _status_resolve_ckpt_path(raw_path, base_dir)
|
| info["checkpoint_path"] = str(ckpt_path)
|
| info["checkpoint_name"] = ckpt_path.name
|
| info["checkpoint_exists"] = ckpt_path.exists()
|
| if ckpt_path.exists():
|
| try:
|
| st = ckpt_path.stat()
|
| info["checkpoint_mtime"] = st.st_mtime
|
| info["checkpoint_mtime_iso"] = _status_iso(st.st_mtime)
|
| except Exception as exc:
|
| warnings.append(f"failed to stat full checkpoint: {exc}")
|
| else:
|
| warnings.append(f"latest.json points to missing checkpoint: {ckpt_path}")
|
| return info, warnings
|
|
|
|
|
| def _status_newest_delta(save_dir: Path) -> tuple[Optional[Dict[str, Any]], List[str]]:
|
| warnings: List[str] = []
|
| if not save_dir.exists():
|
| warnings.append(f"save dir missing: {save_dir}")
|
| return None, warnings
|
| try:
|
| candidates = [item for item in save_dir.glob("*_delta_step*.pt") if item.is_file()]
|
| except Exception as exc:
|
| warnings.append(f"failed to list delta checkpoints: {exc}")
|
| return None, warnings
|
| if not candidates:
|
| warnings.append(f"no delta checkpoints found in {save_dir}")
|
| return None, warnings
|
| newest = max(candidates, key=lambda item: item.stat().st_mtime)
|
| st = newest.stat()
|
| return {
|
| "path": str(newest),
|
| "name": newest.name,
|
| "step": _status_parse_step(newest.name),
|
| "mtime": st.st_mtime,
|
| "mtime_iso": _status_iso(st.st_mtime),
|
| "size_bytes": st.st_size,
|
| "source": "disk",
|
| }, warnings
|
|
|
|
|
| def _status_gpu_info() -> tuple[Optional[Dict[str, Any]], List[str]]:
|
| warnings: List[str] = []
|
| try:
|
| result = subprocess.run(
|
| [
|
| "nvidia-smi",
|
| "--query-gpu=name,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw",
|
| "--format=csv,noheader,nounits",
|
| ],
|
| capture_output=True,
|
| text=True,
|
| timeout=5,
|
| check=False,
|
| )
|
| except FileNotFoundError:
|
| return None, warnings
|
| except Exception as exc:
|
| warnings.append(f"failed to query GPU status: {exc}")
|
| return None, warnings
|
| if result.returncode != 0:
|
| warnings.append(result.stderr.strip() or "nvidia-smi returned non-zero exit status")
|
| return None, warnings
|
| lines = [line.strip() for line in result.stdout.splitlines() if line.strip()]
|
| if not lines:
|
| return None, warnings
|
| if len(lines) > 1:
|
| warnings.append("multiple GPUs detected; reporting the first GPU only")
|
| parts = [part.strip() for part in lines[0].split(",")]
|
| if len(parts) != 6:
|
| warnings.append(f"unexpected nvidia-smi format: {lines[0]}")
|
| return None, warnings
|
|
|
| def _parse_int(raw: str) -> Optional[int]:
|
| try:
|
| return int(float(raw))
|
| except Exception:
|
| return None
|
|
|
| def _parse_float(raw: str) -> Optional[float]:
|
| try:
|
| return float(raw)
|
| except Exception:
|
| return None
|
|
|
| return {
|
| "name": parts[0],
|
| "utilization_gpu": _parse_int(parts[1]),
|
| "memory_used_mib": _parse_int(parts[2]),
|
| "memory_total_mib": _parse_int(parts[3]),
|
| "temperature_c": _parse_int(parts[4]),
|
| "power_draw_w": _parse_float(parts[5]),
|
| }, warnings
|
|
|
|
|
| def _status_choose_delta(from_log: Optional[Dict[str, Any]], from_disk: Optional[Dict[str, Any]], warnings: List[str]) -> Optional[Dict[str, Any]]:
|
| if from_log and from_disk:
|
| log_step = from_log.get("step")
|
| disk_step = from_disk.get("step")
|
| if log_step is not None and disk_step is not None:
|
| if log_step != disk_step:
|
| warnings.append(
|
| f"log delta step {log_step} and newest on-disk delta step {disk_step} differ; using the newer step"
|
| )
|
| if disk_step >= log_step:
|
| merged = dict(from_disk)
|
| merged["source"] = "disk+log" if disk_step == log_step else "disk"
|
| if disk_step == log_step:
|
| merged["sha_prefix"] = from_log.get("sha_prefix")
|
| return merged
|
| return dict(from_log)
|
| return dict(from_disk)
|
| if from_disk:
|
| return dict(from_disk)
|
| if from_log:
|
| return dict(from_log)
|
| return None
|
|
|
|
|
| def _collect_status(log_path: Path, save_dir: Path) -> tuple[Dict[str, Any], int]:
|
| checked_at = time.time()
|
| requested_save_dir = save_dir.expanduser()
|
| log_path = log_path.expanduser()
|
| status: Dict[str, Any] = {
|
| "checked_at": checked_at,
|
| "checked_at_iso": _status_iso(checked_at),
|
| "running": False,
|
| "process": None,
|
| "progress": None,
|
| "delta_checkpoint": None,
|
| "delta_from_log": None,
|
| "delta_on_disk": None,
|
| "latest_full_checkpoint": None,
|
| "log": None,
|
| "gpu": None,
|
| "save_dir": {
|
| "requested_path": str(requested_save_dir),
|
| "path": str(requested_save_dir),
|
| "exists": requested_save_dir.exists(),
|
| "source": "requested",
|
| },
|
| "warnings": [],
|
| }
|
| warnings = status["warnings"]
|
|
|
| matches = _status_find_trainers(STATUS_SCRIPT_PATH)
|
| if len(matches) > 1:
|
| status["error"] = "multiple active n.py train processes found"
|
| status["processes"] = matches
|
| return status, 1
|
| if matches:
|
| status["running"] = True
|
| status["process"] = matches[0]
|
|
|
| save_dir = requested_save_dir
|
| if status["process"] and status["process"].get("cwd"):
|
| proc_cwd = Path(status["process"]["cwd"])
|
| alt_save_dir = (proc_cwd / requested_save_dir.name).resolve()
|
| if alt_save_dir != requested_save_dir and alt_save_dir.exists():
|
| requested_delta, _ = _status_newest_delta(requested_save_dir)
|
| requested_full, _ = _status_latest_full_checkpoint(requested_save_dir, STATUS_SCRIPT_PATH.parent)
|
| alt_delta, _ = _status_newest_delta(alt_save_dir)
|
| alt_full, _ = _status_latest_full_checkpoint(alt_save_dir, proc_cwd)
|
| requested_score = int(requested_delta is not None) + int(bool(requested_full.get("checkpoint_exists")))
|
| alt_score = int(alt_delta is not None) + int(bool(alt_full.get("checkpoint_exists")))
|
| if alt_score > requested_score:
|
| save_dir = alt_save_dir
|
| status["save_dir"] = {
|
| "requested_path": str(requested_save_dir),
|
| "path": str(save_dir),
|
| "exists": save_dir.exists(),
|
| "source": "process_cwd_fallback",
|
| }
|
| warnings.append(
|
| f"using process cwd save dir fallback: {save_dir} (requested {requested_save_dir})"
|
| )
|
|
|
| log_info, progress, delta_from_log, log_warnings = _status_scan_log(log_path)
|
| warnings.extend(log_warnings)
|
| status["log"] = log_info
|
| status["progress"] = progress
|
| status["delta_from_log"] = delta_from_log
|
|
|
| latest_base_dir = STATUS_SCRIPT_PATH.parent
|
| if status["save_dir"].get("source") == "process_cwd_fallback" and status["process"] and status["process"].get("cwd"):
|
| latest_base_dir = Path(status["process"]["cwd"])
|
| latest_full, latest_warnings = _status_latest_full_checkpoint(save_dir, latest_base_dir)
|
| warnings.extend(latest_warnings)
|
| status["latest_full_checkpoint"] = latest_full
|
|
|
| delta_on_disk, delta_warnings = _status_newest_delta(save_dir)
|
| warnings.extend(delta_warnings)
|
| status["delta_on_disk"] = delta_on_disk
|
| status["delta_checkpoint"] = _status_choose_delta(delta_from_log, delta_on_disk, warnings)
|
|
|
| gpu, gpu_warnings = _status_gpu_info()
|
| warnings.extend(gpu_warnings)
|
| status["gpu"] = gpu
|
|
|
| if status["running"] and log_info.get("age_seconds") is not None and log_info["age_seconds"] > 600:
|
| warnings.append(f"train log appears stale while trainer is running ({log_info['age_human']} old)")
|
| if log_info.get("exists") and progress is None:
|
| warnings.append("no parseable progress line found in train log")
|
| latest_step = latest_full.get("step") if latest_full else None
|
| delta_step = status["delta_checkpoint"].get("step") if status["delta_checkpoint"] else None
|
| if latest_step is not None and delta_step is not None and latest_step < delta_step:
|
| warnings.append(f"latest.json step {latest_step} lags newest delta step {delta_step}")
|
| if not status["running"] and progress is None:
|
| warnings.append("no active trainer process found")
|
|
|
| return status, 0
|
|
|
|
|
| def _format_status_text(status: Dict[str, Any]) -> str:
|
| lines = [f"AGILLM status @ {status.get('checked_at_iso')}"]
|
| if status.get("error"):
|
| lines.append(f"Error: {status['error']}")
|
| for proc in status.get("processes", []):
|
| lines.append(f"- pid {proc.get('pid')}: {proc.get('cmdline')}")
|
| return "\n".join(lines)
|
|
|
| process = status.get("process")
|
| if status.get("running") and process:
|
| lines.append(f"Process: RUNNING | pid {process.get('pid')} | uptime {process.get('uptime_human') or 'unknown'}")
|
| lines.append(f"Cmd: {process.get('cmdline')}")
|
| else:
|
| lines.append("Process: NOT RUNNING")
|
|
|
| progress = status.get("progress")
|
| if progress:
|
| lines.append(
|
| "Progress: "
|
| f"{progress['percent']:.1f}% | "
|
| f"{_status_format_int(progress['seen_tokens'])}/{_status_format_int(progress['target_tokens'])} tok | "
|
| f"{progress['tok_per_sec']} tok/s | loss {progress['loss']:.3f} | "
|
| f"B={progress['batch']} L={progress['block']}"
|
| )
|
| else:
|
| lines.append("Progress: unavailable")
|
|
|
| log_info = status.get("log") or {}
|
| if log_info.get("exists"):
|
| lines.append(
|
| f"Log: {log_info.get('path')} | updated {log_info.get('age_human') or 'unknown'} ago | "
|
| f"mtime {log_info.get('mtime_iso')}"
|
| )
|
| else:
|
| lines.append(f"Log: missing ({log_info.get('path')})")
|
|
|
| delta = status.get("delta_checkpoint")
|
| if delta:
|
| line = f"Delta: {delta.get('name')} | step {delta.get('step')} | source {delta.get('source')}"
|
| if delta.get("path"):
|
| line += f" | {delta['path']}"
|
| lines.append(line)
|
| else:
|
| lines.append("Delta: unavailable")
|
|
|
| latest_full = status.get("latest_full_checkpoint") or {}
|
| if latest_full.get("exists"):
|
| lines.append(
|
| f"Latest full: step {latest_full.get('step')} | {latest_full.get('checkpoint_path') or latest_full.get('raw_path')}"
|
| )
|
| else:
|
| lines.append(f"Latest full: unavailable ({latest_full.get('metadata_path')})")
|
|
|
| gpu = status.get("gpu")
|
| if gpu:
|
| lines.append(
|
| f"GPU: {gpu.get('name')} | {gpu.get('utilization_gpu')}% | "
|
| f"{gpu.get('memory_used_mib')}/{gpu.get('memory_total_mib')} MiB | "
|
| f"{gpu.get('temperature_c')}C | {gpu.get('power_draw_w')} W"
|
| )
|
|
|
| warnings = status.get("warnings") or []
|
| if warnings:
|
| lines.append("Warnings:")
|
| lines.extend(f"- {warning}" for warning in warnings)
|
| return "\n".join(lines)
|
|
|
|
|
| def _emit_status(log_path: Path, save_dir: Path, as_json: bool) -> int:
|
| status, exit_code = _collect_status(log_path, save_dir)
|
| if as_json:
|
| print(json.dumps(status, indent=2, sort_keys=True))
|
| else:
|
| print(_format_status_text(status))
|
| return exit_code
|
|
|
|
|
| def _run_status_command(argv: List[str]) -> int:
|
| parser = argparse.ArgumentParser(prog=f"{STATUS_SCRIPT_PATH.name} status", description="Read-only training status")
|
| parser.add_argument("--json", dest="json_output", action="store_true", help="Emit machine-readable JSON")
|
| parser.add_argument("--log", type=Path, default=STATUS_DEFAULT_LOG, help="Path to the training log")
|
| parser.add_argument("--save_dir", type=Path, default=STATUS_DEFAULT_SAVE_DIR, help="Checkpoint directory")
|
| args = parser.parse_args(argv)
|
| return _emit_status(args.log, args.save_dir, args.json_output)
|
|
|
|
|
| def _maybe_handle_status_fastpath() -> None:
|
| if len(sys.argv) > 1 and sys.argv[1] == "status":
|
| raise SystemExit(_run_status_command(sys.argv[2:]))
|
|
|
|
|
| _maybe_handle_status_fastpath()
|
|
|
| import torch
|
|
|
|
|
| class SafeProgress:
|
| def __init__(self, total, initial=0, unit="tok", print_every_tokens=1_000_000):
|
| self.total, self.n, self.unit = total, initial, unit
|
| self.initial = initial
|
| self.last_print, self.postfix = initial, {}
|
| self.print_every_tokens = print_every_tokens
|
| self.start_time = time.time()
|
| def update(self, n=1):
|
| self.n += n
|
| if self.n - self.last_print >= self.print_every_tokens:
|
| self._print(); self.last_print = self.n
|
| def set_postfix(self, **kwargs):
|
|
|
|
|
| self.postfix = kwargs
|
| def _resolve(self, v):
|
| if hasattr(v, "item") and not isinstance(v, (int, float, str, bytes)):
|
| try:
|
| return f"{float(v.item()):.3f}"
|
| except Exception:
|
| return str(v)
|
| return str(v)
|
| def _print(self):
|
| elapsed = __import__('time').time() - self.start_time
|
| rate = (self.n - self.initial) / elapsed if elapsed > 0 else 0
|
| pct = 100 * self.n / self.total if self.total > 0 else 0
|
| pf = ' '.join(f"{k}={self._resolve(v)}" for k, v in self.postfix.items())
|
| print(f"[{pct:.1f}%] {self.n:,}/{self.total:,} {self.unit} | {rate:.0f} tok/s | {pf}")
|
| def close(self): self._print(); print("Done.")
|
|
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from datasets import load_dataset, DownloadConfig
|
| from transformers import AutoTokenizer, PreTrainedTokenizerFast, logging as hf_log
|
|
|
|
|
|
|
| HOT_CONFIG_PATH = Path("/workspace/hot_config.json")
|
| _hot_config_cache = {"mtime": 0, "data": {}}
|
|
|
| def get_hot_config() -> dict:
|
| """Load hot_config.json with caching, return empty dict if missing"""
|
| try:
|
| if HOT_CONFIG_PATH.exists():
|
| mtime = HOT_CONFIG_PATH.stat().st_mtime
|
| if mtime > _hot_config_cache["mtime"]:
|
| with open(HOT_CONFIG_PATH) as f:
|
| _hot_config_cache["data"] = json.load(f)
|
| _hot_config_cache["mtime"] = mtime
|
| return _hot_config_cache["data"]
|
| except Exception as e:
|
| print(f"[hot_config] Error loading: {e}")
|
| return {}
|
|
|
| def get_hot_datasets(default_sources: str) -> str:
|
| """Get datasets from hot_config if present, else use default"""
|
| cfg = get_hot_config()
|
| if "datasets" in cfg and cfg["datasets"]:
|
| hot_ds = cfg["datasets"]
|
| if isinstance(hot_ds, list):
|
| hot_ds = ",".join(hot_ds)
|
| print(f"[hot_config] Using hot datasets: {hot_ds}")
|
| return hot_ds
|
| return default_sources
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| class Colors:
|
| RESET = "\033[0m"
|
| BOLD = "\033[1m"
|
| PROMPT = "\033[36m"
|
| GEN = "\033[0m"
|
| INFO = "\033[90m"
|
| WARN = "\033[93m"
|
|
|
|
|
| hf_log.set_verbosity_error()
|
| DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| torch.backends.cuda.matmul.allow_tf32 = True
|
| try:
|
| torch.set_float32_matmul_precision("high")
|
| except Exception:
|
| pass
|
|
|
| TOKENIZER_BACKEND = os.environ.get("TOKENIZER_BACKEND", "hf").strip().lower()
|
| TOKENIZER_ID = os.environ.get("TOKENIZER_ID", "deepseek-ai/DeepSeek-V4-Pro")
|
| TIKTOKEN_ENCODING = os.environ.get("TIKTOKEN_ENCODING", "o200k_base")
|
|
|
|
|
| class TiktokenTokenizerAdapter:
|
| """Tiny compatibility wrapper for OpenAI tiktoken encodings."""
|
| backend_name = "tiktoken"
|
|
|
| def __init__(self, encoding_name: str):
|
| import tiktoken
|
| self.encoding_name = encoding_name
|
| self.enc = tiktoken.get_encoding(encoding_name)
|
| self.eos_token_id = getattr(self.enc, "eot_token", None)
|
| self.sep_token_id = self.eos_token_id
|
| self.pad_token = "<|endoftext|>"
|
| self.pad_token_id = self.eos_token_id
|
| self.vocab_size = self.enc.n_vocab
|
|
|
| def encode(self, text: str, *args, **kwargs):
|
| return self.enc.encode(text, allowed_special="all")
|
|
|
| def decode(self, ids, *args, **kwargs):
|
| return self.enc.decode([int(i) for i in ids])
|
|
|
| def get_vocab(self):
|
| return {str(i): i for i in range(self.enc.n_vocab)}
|
|
|
| def add_special_tokens(self, *_args, **_kwargs):
|
| return 0
|
|
|
| def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False):
|
| rendered = []
|
| for msg in messages:
|
| role = msg.get("role", "user")
|
| content = msg.get("content", "")
|
| rendered.append(f"{role}: {content}")
|
| if add_generation_prompt:
|
| rendered.append("assistant:")
|
| text = "\n".join(rendered)
|
| return self.encode(text) if tokenize else text
|
|
|
|
|
| def _load_tokenizer():
|
| if TOKENIZER_BACKEND in {"tiktoken", "openai"}:
|
| print(f"[tokenizer] backend=tiktoken encoding={TIKTOKEN_ENCODING}")
|
| return TiktokenTokenizerAdapter(TIKTOKEN_ENCODING)
|
| if TOKENIZER_BACKEND not in {"hf", "huggingface", "transformers"}:
|
| print(f"[tokenizer] WARNING: unknown TOKENIZER_BACKEND={TOKENIZER_BACKEND!r}; falling back to Hugging Face")
|
| print(f"[tokenizer] backend=hf id={TOKENIZER_ID}")
|
| try:
|
| tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True, trust_remote_code=True)
|
| except Exception as e:
|
| print(f"[tokenizer] AutoTokenizer failed ({type(e).__name__}: {e})")
|
| print("[tokenizer] Falling back to direct tokenizer.json load")
|
| from huggingface_hub import hf_hub_download
|
| tokenizer_file = hf_hub_download(TOKENIZER_ID, "tokenizer.json")
|
| tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_file)
|
| if tokenizer.pad_token is None:
|
| tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
|
| tokenizer.backend_name = "hf"
|
| return tokenizer
|
|
|
|
|
| tok = _load_tokenizer()
|
|
|
|
|
|
|
|
|
|
|
| def _fix_tokenizer_space_mismatch(tokenizer):
|
| try:
|
| if not hasattr(tokenizer, "backend_tokenizer"):
|
| return
|
| import json as _json
|
| from tokenizers import Tokenizer as _Tokenizer
|
| bt = tokenizer.backend_tokenizer
|
| tj = _json.loads(bt.to_str())
|
| pre = tj.get("pre_tokenizer", {})
|
| needs_fix = (pre.get("type") == "Metaspace" and pre.get("replacement") == "\u2581")
|
| if not needs_fix:
|
| return
|
|
|
| vocab = tj.get("model", {}).get("vocab", {})
|
| has_gpt2_space = any(k.startswith("\u0120") for k in list(vocab.keys())[:500])
|
| if not has_gpt2_space:
|
| return
|
|
|
| tj["pre_tokenizer"]["replacement"] = "\u0120"
|
|
|
| for step in tj.get("decoder", {}).get("decoders", []):
|
| if step.get("type") == "Replace":
|
| pat = step.get("pattern", {})
|
| if pat.get("String") == "\u2581":
|
| pat["String"] = "\u0120"
|
|
|
| fixed = _Tokenizer.from_str(_json.dumps(tj))
|
| tokenizer.backend_tokenizer = fixed
|
|
|
| test_ids = tokenizer.encode("hello world")
|
| test_dec = tokenizer.decode(test_ids, skip_special_tokens=True)
|
| if "hello world" in test_dec:
|
| print("[tokenizer] Fixed Ä /â– space mismatch")
|
| else:
|
| print(f"[tokenizer] WARNING: fix applied but decode test failed: {repr(test_dec)}")
|
| except Exception as e:
|
| print(f"[tokenizer] Could not fix space mismatch: {e}")
|
|
|
| _fix_tokenizer_space_mismatch(tok)
|
|
|
|
|
|
|
| def _tokenizer_health_check(tokenizer):
|
| import transformers as _tf
|
| ver = _tf.__version__
|
| print(f"[tokenizer] transformers={ver}, tokenizers={__import__('tokenizers').__version__}")
|
|
|
| try:
|
| from packaging.version import Version
|
| if Version(ver) >= Version('5.0.0'):
|
| print(f'[tokenizer] WARNING: transformers {ver} may have Metaspace bug — verify carefully')
|
| except ImportError:
|
| pass
|
|
|
| tests = [
|
| 'Water boils at one hundred degrees',
|
| 'The quick brown fox jumps over the lazy dog',
|
| 'Hello world! This is a test sentence with spaces.',
|
| ]
|
| for text in tests:
|
| ids = tokenizer.encode(text)
|
| decoded = tokenizer.decode(ids, skip_special_tokens=True)
|
| if ' ' not in decoded:
|
| print(f'[tokenizer] FATAL: Roundtrip lost all spaces!')
|
| print(f' Input: {repr(text)}')
|
| print(f' Encoded: {ids[:20]}...')
|
| print(f' Decoded: {repr(decoded)}')
|
| print(f'[tokenizer] ABORTING — fix tokenizer before training!')
|
| sys.exit(1)
|
|
|
| if text.lower().split()[:3] != decoded.lower().split()[:3]:
|
| print(f'[tokenizer] WARNING: Roundtrip diverged:')
|
| print(f' Input: {repr(text[:60])}')
|
| print(f' Decoded: {repr(decoded[:60])}')
|
| print(f'[tokenizer] Health check PASSED — spaces preserved in roundtrip')
|
|
|
| _tokenizer_health_check(tok)
|
|
|
|
|
| def _tokenizer_vocab_size(tokenizer) -> int:
|
| if isinstance(tokenizer, TiktokenTokenizerAdapter):
|
| return tokenizer.vocab_size
|
| vocab_size = getattr(tokenizer, "vocab_size", None)
|
| try:
|
| actual = max(tokenizer.get_vocab().values()) + 1
|
| except Exception:
|
| actual = 0
|
| if isinstance(vocab_size, int) and vocab_size > 0:
|
| return max(vocab_size, actual)
|
| return actual
|
|
|
|
|
| def _tokenizer_eos_id(tokenizer) -> int:
|
| eos = getattr(tokenizer, "eos_token_id", None)
|
| if eos is not None:
|
| return eos
|
| sep = getattr(tokenizer, "sep_token_id", None)
|
| if sep is not None:
|
| return sep
|
| pad = getattr(tokenizer, "pad_token_id", None)
|
| if pad is not None:
|
| return pad
|
| return _tokenizer_vocab_size(tokenizer) - 1
|
|
|
|
|
| VOCAB, EOS = _tokenizer_vocab_size(tok), _tokenizer_eos_id(tok)
|
| print(f"[tokenizer] vocab={VOCAB:,} eos={EOS}")
|
|
|
|
|
| PRESETS: Dict[str, Dict[str, int]] = {
|
| "femto_1x": dict(d=16, layers=1, heads=1, rank=16),
|
| "femto_12x": dict(d=16, layers=1, heads=1, rank=192),
|
| "femto_24x": dict(d=16, layers=1, heads=1, rank=384),
|
| "pico_1x": dict(d=32, layers=1, heads=2, rank=16),
|
| "pico_3x": dict(d=32, layers=1, heads=2, rank=48),
|
| "pico_6x": dict(d=32, layers=1, heads=2, rank=96),
|
| "pico_12x": dict(d=32, layers=1, heads=2, rank=192),
|
| "pico_24x": dict(d=32, layers=1, heads=2, rank=384),
|
| "pico_48x": dict(d=32, layers=1, heads=2, rank=768),
|
| "nano_1x": dict(d=64, layers=2, heads=4, rank=16),
|
| "nano_3x": dict(d=64, layers=2, heads=4, rank=48),
|
| "nano_6x": dict(d=64, layers=2, heads=4, rank=96),
|
| "nano_12x": dict(d=64, layers=2, heads=4, rank=192),
|
| "nano_24x": dict(d=64, layers=2, heads=4, rank=384),
|
| "nano_48x": dict(d=64, layers=2, heads=4, rank=768),
|
| "nano_96x": dict(d=64, layers=2, heads=4, rank=1536),
|
| "micro_3x": dict(d=128, layers=4, heads=8, rank=48),
|
| "micro_6x": dict(d=128, layers=4, heads=8, rank=96),
|
| "micro_12x": dict(d=128, layers=4, heads=8, rank=192),
|
| "micro_24x": dict(d=128, layers=4, heads=8, rank=384),
|
| "small": dict(d=512, layers=8, heads=16, rank=64),
|
| "smallx2": dict(d=512, layers=16, heads=16, rank=64),
|
| "base": dict(d=768, layers=12, heads=24, rank=96),
|
| "base18": dict(d=768, layers=18, heads=24, rank=96),
|
| "large": dict(d=1024, layers=24, heads=16, rank=128),
|
| }
|
|
|
| DEFAULT_BLOCK = 1122
|
| DEFAULT_BATCH = 4
|
| SAT_BLOCK = 2
|
| LR_CORE, LR_HEAD = 5e-5, 2e-4
|
| EMIT_LAMBDA = 0.1
|
| DEFAULT_SAVE_SEC = 24 * 3600
|
| DEFAULT_DELTA_STEPS = 500
|
| DEFAULT_MAX_DELTAS = 5
|
| CKDIR = pathlib.Path("ckpts_expansion")
|
|
|
| DEFAULT_PRETRAIN_SOURCES = "OpenTransformer/goddess-crawl,OpenTransformer/agillm-crawl-data,OpenTransformer/web-crawl-2026,OpenTransformer/web-crawl-clean-v2,OpenTransformer/scraped-web-data,OpenTransformer/turbo-crawl,OpenTransformer/sft-data-clean,OpenTransformer/web-crawl-v1,HuggingFaceFW/fineweb,wikimedia/wikipedia:20231101.en,allenai/c4:en"
|
| DEFAULT_AFTER_SFT_SOURCES = "mlabonne/opc-sft-stage2-chat,HuggingFaceH4/ultrachat_200k"
|
| DEFAULT_AFTER_SFT_BLOCK = 1122
|
|
|
|
|
| def get_uk_time() -> str:
|
| utc_now = datetime.now(timezone.utc)
|
| year = utc_now.year
|
| march_last = datetime(year, 3, 31, 1, 0, tzinfo=timezone.utc)
|
| while march_last.weekday() != 6:
|
| march_last = march_last.replace(day=march_last.day - 1)
|
| oct_last = datetime(year, 10, 31, 1, 0, tzinfo=timezone.utc)
|
| while oct_last.weekday() != 6:
|
| oct_last = oct_last.replace(day=oct_last.day - 1)
|
| if march_last <= utc_now < oct_last:
|
| uk_offset = 1
|
| tz_name = "BST"
|
| else:
|
| uk_offset = 0
|
| tz_name = "GMT"
|
| from datetime import timedelta
|
| uk_time = utc_now + timedelta(hours=uk_offset)
|
| return uk_time.strftime(f'%Y-%m-%d %H:%M:%S {tz_name}')
|
|
|
|
|
| def rng_state():
|
| if DEV.type == "cuda":
|
| try:
|
| return torch.cuda.get_rng_state(DEV)
|
| except TypeError:
|
| return torch.cuda.get_rng_state()
|
| return torch.get_rng_state()
|
|
|
| def _is_probably_ckpt(path: pathlib.Path) -> bool:
|
| try:
|
| return path.is_file() and path.suffix == ".pt" and not path.name.endswith(".pt.tmp") and path.stat().st_size > (1<<20)
|
| except Exception:
|
| return False
|
|
|
| def _resolve_ckpt(path: pathlib.Path) -> pathlib.Path | None:
|
| try:
|
| if path.is_dir():
|
| cands = sorted([p for p in path.glob("*.pt") if _is_probably_ckpt(p)],
|
| key=lambda p: p.stat().st_mtime, reverse=True)
|
| return cands[0] if cands else None
|
| if path.suffix == ".tmp":
|
| solid = path.with_suffix("")
|
| return solid if _is_probably_ckpt(solid) else _resolve_ckpt(path.parent)
|
| return path if _is_probably_ckpt(path) else _resolve_ckpt(path.parent)
|
| except Exception:
|
| return None
|
|
|
| def _try_load(path: pathlib.Path, map_location="cpu"):
|
| try:
|
| return torch.load(path, map_location="cpu")
|
| except Exception as e:
|
| print(f"[ckpt-skip] {path} not usable: {e}")
|
| return None
|
|
|
| def _prune_checkpoints(save_dir: pathlib.Path, phase_name: str, max_ckpts: int):
|
| if max_ckpts is None or max_ckpts <= 0:
|
| return
|
| try:
|
| pattern = f"{phase_name}_step*.pt"
|
| ckpts = sorted(
|
| [p for p in save_dir.glob(pattern) if _is_probably_ckpt(p)],
|
| key=lambda p: p.stat().st_mtime
|
| )
|
| excess = len(ckpts) - max_ckpts
|
| if excess > 0:
|
| for p in ckpts[:excess]:
|
| try:
|
| p.unlink()
|
| print(f" [prune] deleted old {p.name}")
|
| except Exception:
|
| pass
|
| except Exception as e:
|
| print(f"[ckpt-prune] error: {e}")
|
|
|
| def print_expansion_info(cfg: dict, tie_weights: bool = False):
|
| d_k = cfg["d"] // cfg["heads"]
|
| rank = cfg["rank"]
|
| ratio = rank / d_k
|
| regime = "COMPRESSION" if ratio < 1 else ("IDENTITY" if ratio == 1 else "EXPANSION")
|
| tie_str = "YES" if tie_weights else "NO"
|
| print(f"┌─────────────────────────────────────────â”")
|
| print(f"│ TUNEABLE ATTENTION CONFIG │")
|
| print(f"├─────────────────────────────────────────┤")
|
| print(f"│ d_model: {cfg['d']:4d} heads: {cfg['heads']:2d} d_k: {d_k:3d} │")
|
| print(f"│ layers: {cfg['layers']:4d} tie_weights: {tie_str:3s} │")
|
| print(f"│ rank: {rank:4d} ratio: {ratio:.1f}x [{regime:11s}] │")
|
| print(f"└─────────────────────────────────────────┘")
|
|
|
|
|
| try:
|
| from torch.amp import autocast as _ac, GradScaler
|
| except ImportError:
|
| from torch.cuda.amp import autocast as _ac, GradScaler
|
|
|
| def _auto_amp_dtype():
|
| if DEV.type == "cuda":
|
| try:
|
| if torch.cuda.is_bf16_supported(): return torch.bfloat16
|
| return torch.float16
|
| except Exception: return torch.float16
|
| return torch.float32
|
|
|
| def amp(enabled: bool):
|
| return nullcontext() if not (enabled and DEV.type == "cuda") else _ac(device_type="cuda", dtype=_auto_amp_dtype())
|
|
|
|
|
| def _coerce_role(r: str) -> str:
|
| r = (r or "").lower()
|
| if r in {"user", "human", "customer"}: return "user"
|
| if r in {"assistant", "gpt", "bot"}: return "assistant"
|
| if r in {"system", "context"}: return "system"
|
| return r or "user"
|
|
|
| def _render_chat_text_from_ex(ex: dict, messages_key: str, add_generation_prompt: bool) -> Optional[str]:
|
| msgs = ex.get(messages_key)
|
| if msgs is None:
|
| for alt in ("conversations", "dialog", "turns"):
|
| if isinstance(ex.get(alt), list):
|
| msgs = ex[alt]; break
|
| if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict):
|
| try:
|
| norm = []
|
| for m in msgs:
|
| role = _coerce_role(m.get("role", "")); content = m.get("content", m.get("text", ""))
|
| if not isinstance(content, str): continue
|
| norm.append({"role": role, "content": content})
|
| if not norm: return None
|
| return tok.apply_chat_template(norm, tokenize=False, add_generation_prompt=add_generation_prompt)
|
| except Exception: return None
|
| for a, b in (("prompt", "response"), ("instruction", "output"), ("question", "answer")):
|
| if isinstance(ex.get(a), str) and isinstance(ex.get(b), str):
|
| return f"User: {ex[a]}\nAssistant: {ex[b]}"
|
| return None
|
|
|
| def _open_stream_one(ds_name: str, seed: int, streaming: bool = True):
|
| dc = DownloadConfig(max_retries=5, use_etag=True, resume_download=True)
|
| if ":" in ds_name: base, config = ds_name.split(":", 1)
|
| else: base, config = ds_name, None
|
| if not streaming:
|
| print(f"[download] Downloading {ds_name} (non-streaming)...")
|
| if base == "json":
|
| data_files = {"train": config}
|
| ds = load_dataset("json", data_files=data_files, split="train", streaming=streaming, download_config=dc)
|
| else:
|
| ds = load_dataset(base, config, split="train", streaming=streaming, download_config=dc) if config else \
|
| load_dataset(base, split="train", streaming=streaming, download_config=dc)
|
| if streaming:
|
| return iter(ds.shuffle(buffer_size=1000, seed=seed))
|
| else:
|
| print(f"[download] Got {len(ds):,} examples. Shuffling...")
|
| ds = ds.shuffle(seed=seed)
|
| return iter(ds)
|
|
|
| def token_stream(ds_names: str, target: int, seed: int = 42,
|
| chat: bool = False, chat_messages_key: str = "messages",
|
| sft_add_generation_prompt: bool = False, dataset_field_text: str = "text",
|
| streaming: bool = True):
|
| ds_names = get_hot_datasets(ds_names)
|
| sources = [s.strip() for s in ds_names.split(",") if s.strip()]
|
| if not sources: return
|
| src_idx = 0; emitted = 0; it = None; attempts = 0; backoff_base = 2.0
|
| while emitted < target:
|
| try:
|
| if it is None: it = _open_stream_one(sources[src_idx], seed, streaming=streaming)
|
| ex = next(it)
|
| text = None
|
| if isinstance(ex, dict):
|
| if chat:
|
| text = _render_chat_text_from_ex(ex, chat_messages_key, sft_add_generation_prompt)
|
| if text is None:
|
| if dataset_field_text and isinstance(ex.get(dataset_field_text), str):
|
| text = ex[dataset_field_text]
|
| elif isinstance(ex.get("text"), str):
|
| text = ex["text"]
|
| if not isinstance(text, str):
|
| attempts = 0; continue
|
| enc = tok.encode(text)
|
| if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
|
| enc = enc + [EOS]
|
| for t in enc:
|
| yield t
|
| emitted += 1
|
| if emitted >= target: return
|
| attempts = 0
|
| except StopIteration:
|
| it = None; src_idx = (src_idx + 1) % len(sources)
|
| except Exception as e:
|
| attempts += 1
|
| sleep_s = min(60.0, backoff_base ** min(attempts, 6))
|
| print(f"[stream-retry] {sources[src_idx]} error: {type(e).__name__}, sleeping {sleep_s:.1f}s")
|
| time.sleep(sleep_s); it = None
|
| if attempts % 5 == 0 and len(sources) > 1:
|
| src_idx = (src_idx + 1) % len(sources)
|
|
|
|
|
| _ALIBI_CACHE_MAX = 64
|
| _alibi_slopes_cache: Dict[int, torch.Tensor] = {}
|
| _alibi_bias_cache: Dict[Tuple[int, int], torch.Tensor] = {}
|
|
|
| def _alibi_slopes(n_heads: int):
|
| cached = _alibi_slopes_cache.get(n_heads)
|
| if cached is not None:
|
| return cached
|
| def pow2slopes(n):
|
| start = 2 ** (-2 ** -(math.log2(n) - 3))
|
| ratio = start
|
| return [start * (ratio ** i) for i in range(n)]
|
| if math.log2(n_heads).is_integer(): vals = pow2slopes(n_heads)
|
| else:
|
| closest = 2 ** math.floor(math.log2(n_heads))
|
| vals = pow2slopes(closest)
|
| extra = pow2slopes(2 * closest)
|
| vals += extra[0::2][: n_heads - closest]
|
| t = torch.tensor(vals, device=DEV).view(1, n_heads, 1, 1)
|
| _alibi_slopes_cache[n_heads] = t
|
| return t
|
|
|
| def alibi_bias(n_heads: int, n_tokens: int):
|
| key = (n_heads, n_tokens)
|
| cached = _alibi_bias_cache.get(key)
|
| if cached is not None:
|
| return cached
|
| i = torch.arange(n_tokens, device=DEV).view(1, 1, n_tokens, 1)
|
| j = torch.arange(n_tokens, device=DEV).view(1, 1, 1, n_tokens)
|
| dist = (j - i).clamp_min(0)
|
| bias = -_alibi_slopes(n_heads) * dist
|
| if len(_alibi_bias_cache) < _ALIBI_CACHE_MAX:
|
| _alibi_bias_cache[key] = bias
|
| return bias
|
|
|
|
|
| def _effective_psd_rank(M: torch.Tensor, eps_ratio: float = 1e-6) -> int:
|
| """Numerical rank of a symmetric PSD matrix via eigenvalue threshold.
|
| Used to detect post-train how many of U's r columns are actually exercised.
|
| """
|
| eigvals = torch.linalg.eigvalsh(M.detach().to(torch.float64))
|
| if eigvals.numel() == 0:
|
| return 0
|
| threshold = eps_ratio * eigvals.abs().max().clamp_min(torch.finfo(eigvals.dtype).tiny)
|
| return int((eigvals > threshold).sum().item())
|
|
|
|
|
| def compact_attention_U(mha: "TuneableAttentionMHA", eps_ratio: float = 1e-6) -> Dict[str, int]:
|
| """Constructive proof of CONSEQUENCE 2 (see file header).
|
|
|
| Replace mha.U in R^{dk x r} with U' in R^{dk x r*} satisfying U' U'^T == U U^T,
|
| where r* is the numerical rank of U U^T. Attention scores (and therefore the
|
| network output and gradients flowing back through it) are unchanged.
|
|
|
| Construction:
|
| M = U U^T # symmetric PSD, dk x dk
|
| M = Q diag(lambda) Q^T # eigendecomposition
|
| keep i where lambda_i > eps_ratio * max(lambda)
|
| U' = Q[:, keep] * sqrt(lambda[keep])
|
|
|
| Returns {"r_old": r, "r_new": r*, "params_saved": dk*(r - r*)}.
|
| """
|
| with torch.no_grad():
|
| U = mha.U.detach()
|
| dk, r_old = U.shape
|
| if r_old <= 1:
|
| return {"r_old": r_old, "r_new": r_old, "params_saved": 0}
|
|
|
| M = U.to(torch.float64) @ U.to(torch.float64).T
|
| eigvals, eigvecs = torch.linalg.eigh(M)
|
| threshold = eps_ratio * eigvals.abs().max().clamp_min(1e-30)
|
| keep_mask = eigvals > threshold
|
| r_new = int(keep_mask.sum().item())
|
| if r_new == 0:
|
| r_new = 1
|
| keep_mask = torch.zeros_like(keep_mask, dtype=torch.bool); keep_mask[-1] = True
|
| if r_new >= r_old:
|
| return {"r_old": r_old, "r_new": r_old, "params_saved": 0}
|
|
|
| kept_vals = eigvals[keep_mask].clamp_min(0.0).sqrt()
|
| kept_vecs = eigvecs[:, keep_mask]
|
| U_new = (kept_vecs * kept_vals).to(U.dtype)
|
|
|
| mha.U = nn.Parameter(U_new.contiguous())
|
| mha.r = r_new
|
| return {"r_old": r_old, "r_new": r_new, "params_saved": dk * (r_old - r_new)}
|
|
|
|
|
| def compact_checkpoint_U(src_path: str, dst_path: str, eps_ratio: float = 1e-6) -> Dict[str, Any]:
|
| """Load a checkpoint, compact every attention layer's U, save to dst.
|
|
|
| The result is functionally identical to the original (within fp32 noise) but
|
| smaller. Use for deploying a trained model where you don't need the over-
|
| parameterized U for further training.
|
| """
|
| ck = torch.load(src_path, map_location="cpu", weights_only=False)
|
| cfg = ck.get("cfg")
|
| if cfg is None:
|
| raise ValueError(f"checkpoint {src_path} has no 'cfg'; cannot reconstruct model")
|
| tie = ck.get("tie_weights", False)
|
| core = Encoder(cfg, tie_weights=tie)
|
| ar_h = ARHead(cfg["d"], tie_weights=tie, embedding_weight=core.emb.weight if tie else None)
|
| sat_h = SATHead(cfg["d"], mode="var")
|
| core.load_state_dict(ck["core"])
|
| ar_h.load_state_dict(ck["ar"])
|
| sat_h.load_state_dict(ck["sat"])
|
|
|
| total_old = 0; effective_new = 0; report = []
|
| for li, blk in enumerate(core.blocks):
|
| before = blk.mha.U.numel()
|
| info = compact_attention_U(blk.mha, eps_ratio=eps_ratio)
|
| after = blk.mha.U.numel()
|
| total_old += before; effective_new += after
|
| report.append({"layer": li, **info})
|
|
|
|
|
|
|
| target_rank = max(blk.mha.U.shape[1] for blk in core.blocks)
|
| stored_new = 0
|
| for blk in core.blocks:
|
| U = blk.mha.U.detach()
|
| if U.shape[1] < target_rank:
|
| pad = torch.zeros(U.shape[0], target_rank - U.shape[1], dtype=U.dtype, device=U.device)
|
| blk.mha.U = nn.Parameter(torch.cat([U, pad], dim=1).contiguous())
|
| blk.mha.r = target_rank
|
| stored_new += blk.mha.U.numel()
|
|
|
|
|
| out = dict(ck)
|
| out["cfg"] = dict(cfg)
|
| out["cfg"]["rank"] = target_rank
|
| out["core"] = core.state_dict()
|
| out["compacted"] = True
|
| out["compact_eps_ratio"] = eps_ratio
|
| out["compact_target_rank"] = target_rank
|
| out["compact_report"] = report
|
|
|
|
|
| out.pop("opt", None); out.pop("scaler", None)
|
|
|
| torch.save(out, dst_path)
|
| return {
|
| "src": src_path, "dst": dst_path,
|
| "U_params_old": total_old,
|
| "U_params_new": stored_new,
|
| "U_params_effective_new": effective_new,
|
| "U_params_saved": total_old - stored_new,
|
| "target_rank": target_rank,
|
| "by_layer": report,
|
| }
|
|
|
|
|
|
|
| class TuneableAttentionMHA(nn.Module):
|
| def __init__(self, d: int, h: int, r: int, use_relpos: bool = True, use_sdpa: bool = True):
|
| super().__init__()
|
| assert d % h == 0
|
| self.h, self.dk, self.r = h, d // h, r
|
| self.use_relpos = use_relpos
|
| self.use_sdpa = bool(use_sdpa and hasattr(F, "scaled_dot_product_attention"))
|
| self.q = nn.Linear(d, d, bias=False)
|
| self.k = nn.Linear(d, d, bias=False)
|
| self.v = nn.Linear(d, d, bias=False)
|
| self.U = nn.Parameter(torch.randn(self.dk, r))
|
| nn.init.orthogonal_(self.U)
|
| self.proj = nn.Linear(h * self.dk, d, bias=False)
|
| self.drop = nn.Dropout(0.1)
|
|
|
| def _proj_qk(self, x):
|
| B, N, _ = x.shape
|
| return (x.view(B, N, self.h, self.dk).transpose(1, 2) @ self.U)
|
|
|
| def _reshape_qk_base(self, x):
|
| B, N, _ = x.shape
|
| return x.view(B, N, self.h, self.dk).transpose(1, 2)
|
|
|
| def _reshape_v(self, x):
|
| B, N, _ = x.shape
|
| return x.view(B, N, self.h, self.dk).transpose(1, 2)
|
|
|
| def forward(self, x, mask=None, rel_bias_tokens=None, kv_cache=None, use_cache=False):
|
| q_len = x.size(1)
|
| v_new = self._reshape_v(self.v(x))
|
|
|
| if self.r > self.dk:
|
| q = self._reshape_qk_base(self.q(x))
|
| k_new = self._reshape_qk_base(self.k(x))
|
| if kv_cache is None:
|
| k, v = k_new, v_new
|
| else:
|
| k_cached, v_cached = kv_cache
|
| if use_cache:
|
| k = torch.cat([k_cached, k_new], dim=2)
|
| v = torch.cat([v_cached, v_new], dim=2)
|
| else:
|
| k, v = k_new, v_new
|
|
|
|
|
|
|
|
|
| metric = self.U @ self.U.T
|
| q_scores = q @ metric
|
| else:
|
| q = self._proj_qk(self.q(x))
|
| k_new = self._proj_qk(self.k(x))
|
| if kv_cache is None:
|
| k, v = k_new, v_new
|
| else:
|
| k_cached, v_cached = kv_cache
|
| if use_cache:
|
| k = torch.cat([k_cached, k_new], dim=2)
|
| v = torch.cat([v_cached, v_new], dim=2)
|
| else:
|
| k, v = k_new, v_new
|
| q_scores = q
|
|
|
| attn_bias = None
|
| if self.use_relpos and rel_bias_tokens is not None:
|
| attn_bias = alibi_bias(self.h, rel_bias_tokens)[:, :, -q_len:, :]
|
| if mask is not None:
|
| attn_bias = mask if attn_bias is None else (attn_bias + mask)
|
| if self.use_sdpa:
|
| z = F.scaled_dot_product_attention(
|
| q_scores, k, v,
|
| attn_mask=attn_bias,
|
| dropout_p=0.0,
|
| is_causal=False,
|
| scale=1.0 / math.sqrt(self.dk),
|
| )
|
| else:
|
| att = (q_scores @ k.transpose(-1, -2)) / math.sqrt(self.dk)
|
| if attn_bias is not None:
|
| att = att + attn_bias
|
| z = att.softmax(-1) @ v
|
| z = z.transpose(1, 2).reshape(x.size(0), x.size(1), -1)
|
| out = self.drop(self.proj(z))
|
| return (out, (k, v)) if use_cache else out
|
|
|
|
|
| class Block(nn.Module):
|
| def __init__(self, d: int, h: int, r: int, use_sdpa: bool = True):
|
| super().__init__()
|
| self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
|
| self.mha = TuneableAttentionMHA(d, h, r, use_sdpa=use_sdpa)
|
| self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.ReLU(), nn.Linear(4 * d, d))
|
|
|
| def forward(self, x, mask, kv=None, use_cache=False, total_seq_len=None):
|
| if use_cache:
|
| y, new_kv = self.mha(self.ln1(x), mask, rel_bias_tokens=total_seq_len, kv_cache=kv, use_cache=True)
|
| x = x + y
|
| x = x + self.ff(self.ln2(x))
|
| return x, new_kv
|
| n = x.size(1)
|
| x = x + self.mha(self.ln1(x), mask, rel_bias_tokens=n)
|
| return x + self.ff(self.ln2(x))
|
|
|
|
|
| class Encoder(nn.Module):
|
| def __init__(self, cfg, tie_weights: bool = False, use_sdpa: bool = True):
|
| super().__init__()
|
| d, l, h, r = cfg["d"], cfg["layers"], cfg["heads"], cfg["rank"]
|
| self.emb = nn.Embedding(VOCAB, d)
|
| self.blocks = nn.ModuleList([Block(d, h, r, use_sdpa=use_sdpa) for _ in range(l)])
|
| self.ln = nn.LayerNorm(d)
|
| self.tie_weights = tie_weights
|
| self.use_sdpa = bool(use_sdpa and hasattr(F, "scaled_dot_product_attention"))
|
|
|
| def forward(self, ids, mask, kv_caches=None, use_cache=False, total_seq_len=None):
|
| x = self.emb(ids)
|
| if not use_cache:
|
| for blk in self.blocks:
|
| x = blk(x, mask)
|
| return self.ln(x)
|
| new_kvs = []
|
| for i, blk in enumerate(self.blocks):
|
| kv = kv_caches[i] if kv_caches else None
|
| x, kv_out = blk(x, mask, kv, use_cache=True, total_seq_len=total_seq_len)
|
| new_kvs.append(kv_out)
|
| return self.ln(x), new_kvs
|
|
|
|
|
| class ARHead(nn.Module):
|
| def __init__(self, d, tie_weights: bool = False, embedding_weight: nn.Parameter = None):
|
| super().__init__()
|
| self.tie_weights = tie_weights
|
| if tie_weights and embedding_weight is not None:
|
| self.proj = nn.Linear(d, VOCAB, bias=False)
|
| self.proj.weight = embedding_weight
|
| else:
|
| self.proj = nn.Linear(d, VOCAB)
|
|
|
| def forward(self, h):
|
| return self.proj(h)
|
|
|
|
|
| class SATHead(nn.Module):
|
| def __init__(self, d, mode="var"):
|
| super().__init__()
|
| self.proj = nn.Linear(d, VOCAB)
|
| self.gate = nn.Linear(d, 2) if mode == "var" else None
|
| def forward(self, h_last):
|
| return self.proj(h_last), (self.gate(h_last[:, 0]) if self.gate else None)
|
|
|
|
|
|
|
| _MASK_CACHE_MAX = 64
|
| _causal_mask_cache: Dict[int, torch.Tensor] = {}
|
| _sat_mask_cache: Dict[Tuple[int, int], torch.Tensor] = {}
|
|
|
| def causal_mask(n):
|
| if n in _causal_mask_cache:
|
| return _causal_mask_cache[n]
|
| m = torch.triu(torch.full((1, 1, n, n), float("-inf"), device=DEV), 1)
|
| if len(_causal_mask_cache) < _MASK_CACHE_MAX:
|
| _causal_mask_cache[n] = m
|
| return m
|
|
|
| def sat_mask(n, block=SAT_BLOCK):
|
| key = (n, block)
|
| if key in _sat_mask_cache:
|
| return _sat_mask_cache[key]
|
| idx = torch.arange(n, device=DEV)
|
| grp = idx.unsqueeze(0) // block
|
| allow = (grp.T == grp) | (grp.T > grp)
|
| m = torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
|
| if len(_sat_mask_cache) < _MASK_CACHE_MAX:
|
| _sat_mask_cache[key] = m
|
| return m
|
|
|
| def sat_mask_cached(new_len: int, cached_len: int, block=SAT_BLOCK):
|
| total_len = cached_len + new_len
|
| mask = torch.zeros((1, 1, new_len, total_len), device=DEV)
|
| return mask
|
|
|
| def causal_mask_cached(new_len: int, cached_len: int):
|
| total_len = cached_len + new_len
|
| mask = torch.zeros((1, 1, new_len, total_len), device=DEV)
|
| if new_len > 1:
|
| mask[:, :, :, cached_len:] = torch.triu(
|
| torch.full((new_len, new_len), float("-inf"), device=DEV), 1
|
| ).view(1, 1, new_len, new_len)
|
| return mask
|
|
|
|
|
|
|
|
|
|
|
| _delta_lock = threading.Lock()
|
| _delta_thread: Optional[threading.Thread] = None
|
|
|
| def _sha256_file(path: pathlib.Path) -> str:
|
| """Compute SHA256 of a file for integrity verification."""
|
| h = hashlib.sha256()
|
| with open(path, "rb") as f:
|
| for chunk in iter(lambda: f.read(1 << 20), b""):
|
| h.update(chunk)
|
| return h.hexdigest()
|
|
|
| def _do_delta_save(tensors: dict, path: pathlib.Path, meta: dict):
|
| """Background worker: write weight-only checkpoint + checksum."""
|
| try:
|
| path.parent.mkdir(exist_ok=True, parents=True)
|
| tmp = path.with_suffix(path.suffix + ".dtmp")
|
| torch.save({"weights": tensors, **meta}, tmp, _use_new_zipfile_serialization=False)
|
| digest = _sha256_file(tmp)
|
| tmp.replace(path)
|
|
|
| path.with_suffix(".sha256").write_text(f"{digest} {path.name}\n")
|
| print(f" [delta] saved {path.name} ({digest[:12]}...)")
|
| except Exception as e:
|
| print(f" [delta] FAILED {path.name}: {e}")
|
|
|
| def save_delta(core, ar_h, sat_h, step: int, seen_tok: int, save_dir: pathlib.Path, phase_name: str):
|
| """Save weight-only delta in background thread. Non-blocking."""
|
| global _delta_thread
|
|
|
| if _delta_thread is not None and _delta_thread.is_alive():
|
| _delta_thread.join(timeout=60)
|
|
|
| with _delta_lock:
|
| tensors = {
|
| "core": {k: v.detach().cpu() for k, v in core.state_dict().items()},
|
| "ar": {k: v.detach().cpu() for k, v in ar_h.state_dict().items()},
|
| "sat": {k: v.detach().cpu() for k, v in sat_h.state_dict().items()},
|
| }
|
| meta = {"step": step, "seen_tok": seen_tok, "wall_time": time.time(), "delta": True, **_tokenizer_checkpoint_meta()}
|
| path = save_dir / f"{phase_name}_delta_step{step:08d}.pt"
|
| _delta_thread = threading.Thread(target=_do_delta_save, args=(tensors, path, meta), daemon=True)
|
| _delta_thread.start()
|
|
|
| def _prune_deltas(save_dir: pathlib.Path, phase_name: str, max_deltas: int):
|
| """Keep only the most recent max_deltas delta files."""
|
| if max_deltas is None or max_deltas <= 0:
|
| return
|
| try:
|
| pattern = f"{phase_name}_delta_step*.pt"
|
| deltas = sorted(
|
| [p for p in save_dir.glob(pattern) if p.stat().st_size > 0],
|
| key=lambda p: p.stat().st_mtime
|
| )
|
| excess = len(deltas) - max_deltas
|
| if excess > 0:
|
| for p in deltas[:excess]:
|
| try:
|
| p.unlink()
|
| sha = p.with_suffix(".sha256")
|
| if sha.exists(): sha.unlink()
|
| print(f" [delta-prune] deleted {p.name}")
|
| except Exception:
|
| pass
|
| except Exception as e:
|
| print(f" [delta-prune] error: {e}")
|
|
|
| def load_delta(path: pathlib.Path, core, ar_h, sat_h):
|
| """Load weight-only delta. Returns (step, seen_tok) or raises."""
|
|
|
| sha_path = path.with_suffix(".sha256")
|
| if sha_path.exists():
|
| expected = sha_path.read_text().split()[0]
|
| actual = _sha256_file(path)
|
| if expected != actual:
|
| raise ValueError(f"Checksum mismatch for {path.name}: expected {expected[:12]}... got {actual[:12]}...")
|
| print(f" [delta] checksum OK for {path.name}")
|
| ck = torch.load(path, map_location="cpu", weights_only=False)
|
| if not ck.get("delta"):
|
| raise ValueError(f"{path.name} is not a delta checkpoint")
|
| core.load_state_dict(ck["weights"]["core"])
|
| ar_h.load_state_dict(ck["weights"]["ar"])
|
| sat_h.load_state_dict(ck["weights"]["sat"])
|
| return ck.get("step", 0), ck.get("seen_tok", 0)
|
|
|
| def _flush_delta():
|
| """Wait for any in-flight delta save to complete."""
|
| global _delta_thread
|
| if _delta_thread is not None and _delta_thread.is_alive():
|
| print(" [delta] flushing in-flight write...")
|
| _delta_thread.join(timeout=120)
|
|
|
| def _tokenizer_checkpoint_meta() -> Dict[str, Any]:
|
| meta = {
|
| "tokenizer_backend": TOKENIZER_BACKEND,
|
| "tokenizer_id": TOKENIZER_ID,
|
| "tokenizer_vocab": VOCAB,
|
| "tokenizer_eos": EOS,
|
| }
|
| if isinstance(tok, TiktokenTokenizerAdapter):
|
| meta["tiktoken_encoding"] = tok.encoding_name
|
| meta["tiktoken_version"] = __import__("tiktoken").__version__
|
| return meta
|
| if hasattr(tok, "backend_tokenizer"):
|
| try:
|
| meta["tokenizer_json"] = tok.backend_tokenizer.to_str()
|
| except Exception as e:
|
| meta["tokenizer_json_error"] = str(e)
|
| try:
|
| meta["transformers_version"] = __import__("transformers").__version__
|
| meta["tokenizers_version"] = __import__("tokenizers").__version__
|
| except Exception:
|
| pass
|
| return meta
|
|
|
| def _warn_tokenizer_checkpoint_mismatch(ck: Dict[str, Any]):
|
| ck_backend = ck.get("tokenizer_backend")
|
| ck_id = ck.get("tokenizer_id")
|
| ck_vocab = ck.get("tokenizer_vocab")
|
| if ck_backend and ck_backend != TOKENIZER_BACKEND:
|
| print(f"[tokenizer] WARNING: checkpoint backend={ck_backend}, runtime backend={TOKENIZER_BACKEND}")
|
| if ck_id and ck_id != TOKENIZER_ID and TOKENIZER_BACKEND not in {"tiktoken", "openai"}:
|
| print(f"[tokenizer] WARNING: checkpoint tokenizer={ck_id}, runtime tokenizer={TOKENIZER_ID}")
|
| if ck_vocab and ck_vocab != VOCAB:
|
| print(f"[tokenizer] WARNING: checkpoint vocab={ck_vocab:,}, runtime vocab={VOCAB:,}; embedding/head shapes may not load")
|
|
|
| def _restore_tokenizer_from_checkpoint(ck: Dict[str, Any]):
|
| _warn_tokenizer_checkpoint_mismatch(ck)
|
| if "tokenizer_json" in ck and hasattr(tok, "backend_tokenizer"):
|
| try:
|
| from tokenizers import Tokenizer as _Tokenizer
|
| tok.backend_tokenizer = _Tokenizer.from_str(ck["tokenizer_json"])
|
| print("[tokenizer] Restored Hugging Face tokenizer JSON from checkpoint")
|
| except Exception as e:
|
| print(f"[tokenizer] WARNING: could not restore tokenizer from checkpoint: {e}")
|
| if "transformers_version" in ck:
|
| import transformers as _tf
|
| if ck["transformers_version"] != _tf.__version__:
|
| print(f"[tokenizer] WARNING: checkpoint saved with transformers={ck['transformers_version']}, now running {_tf.__version__}")
|
|
|
| def save_ckpt(path: pathlib.Path, core, ar_h, sat_h, opt, scaler, meta):
|
| path.parent.mkdir(exist_ok=True, parents=True)
|
| tmp = path.with_suffix(path.suffix + ".tmp")
|
| state = {
|
| "core": core.state_dict(), "ar": ar_h.state_dict(), "sat": sat_h.state_dict(),
|
| "opt": opt.state_dict(), "scaler": scaler.state_dict(),
|
| "cfg": meta.get("cfg"), **_tokenizer_checkpoint_meta(),
|
| "tie_weights": meta.get("tie_weights", False),
|
| **{k: v for k, v in meta.items() if k not in ("cfg", "tie_weights")}
|
| }
|
| torch.save(state, tmp, _use_new_zipfile_serialization=False)
|
| tmp.replace(path)
|
| (path.parent / "latest.json").write_text(json.dumps({"path": str(path), "step": meta["step"]}))
|
| print(f"\n✓ saved checkpoint {path.name}")
|
|
|
| def load_ckpt(path, core, ar_h, sat_h, opt, scaler):
|
| p = _resolve_ckpt(path) or path
|
| ck = _try_load(p, map_location="cpu")
|
| if ck is None: raise FileNotFoundError(f"No valid checkpoint at {p}")
|
| core.load_state_dict(ck["core"])
|
| ar_h.load_state_dict(ck["ar"])
|
| sat_h.load_state_dict(ck["sat"])
|
| opt.load_state_dict(ck["opt"])
|
| scaler.load_state_dict(ck["scaler"])
|
| _restore_tokenizer_from_checkpoint(ck)
|
| return ck.get("step", 0), ck.get("seen_tok", 0), ck.get("wall_time", time.time())
|
|
|
| def _safe_load_any(path: pathlib.Path, tgt: nn.Module, key: str | None = None):
|
| p = _resolve_ckpt(path) or path
|
| if not p.exists(): return 0
|
| ck = _try_load(p, map_location="cpu")
|
| if ck is None: return 0
|
| sd = ck.get(key, ck) if key else ck
|
| if isinstance(sd, dict) and "state_dict" in sd: sd = sd["state_dict"]
|
| tgt_sd = tgt.state_dict()
|
| filt = {k: v for k, v in sd.items() if k in tgt_sd and v.shape == tgt_sd[k].shape}
|
| if filt: tgt.load_state_dict(filt, strict=False)
|
| return len(filt)
|
|
|
| def infer_cfg_from_ckpt(path: pathlib.Path):
|
| """Return (cfg, ckpt_meta) where ckpt_meta exposes side-channel flags like
|
| tie_weights that aren't part of cfg itself. Both may be None on miss."""
|
| p = _resolve_ckpt(path) or path
|
| if not p.exists(): return None, None
|
| sd = _try_load(p, map_location="cpu")
|
| if sd is None: return None, None
|
| cfg = dict(sd["cfg"]) if "cfg" in sd else None
|
| meta = {"tie_weights": bool(sd.get("tie_weights", False))}
|
| return cfg, meta
|
|
|
|
|
|
|
| def _parse_grow_plan(s: str) -> List[int]:
|
| return sorted(set([int(x.strip()) for x in s.split(",") if x.strip() and int(x.strip()) >= 128]))
|
|
|
| def _count_enabled_params(*modules) -> int:
|
| seen_data_ptrs = set()
|
| total = 0
|
| for m in modules:
|
| if m is None:
|
| continue
|
| for p in m.parameters():
|
| if p.data_ptr() not in seen_data_ptrs:
|
| seen_data_ptrs.add(p.data_ptr())
|
| total += p.numel()
|
| return total
|
|
|
|
|
| def _params_excluding(params, *exclude_modules):
|
|
|
|
|
| excluded_ids = set()
|
| for m in exclude_modules:
|
| if m is None:
|
| continue
|
| for p in m.parameters():
|
| excluded_ids.add(id(p))
|
| return [p for p in params if id(p) not in excluded_ids]
|
|
|
|
|
| def _phase_freeze(core: nn.Module, *, freeze_core: bool, unfreeze_ln: bool, train_emb: bool):
|
| for p in core.parameters(): p.requires_grad = not freeze_core
|
| if freeze_core:
|
| if unfreeze_ln:
|
| for blk in core.blocks:
|
| for p in blk.ln1.parameters(): p.requires_grad = True
|
| for p in blk.ln2.parameters(): p.requires_grad = True
|
| for p in core.ln.parameters(): p.requires_grad = True
|
| if train_emb:
|
| for p in core.emb.parameters(): p.requires_grad = True
|
|
|
| def _sat_window_starts(n_tokens: int, window: int, stride: int, max_windows: int, device=None) -> torch.Tensor:
|
| """Start offsets for SAT future-window training.
|
|
|
| A source window [s, s+window) predicts the next window
|
| [s+window, s+2*window), matching SAT inference's "last window drafts next
|
| window" contract.
|
| """
|
| stride = max(1, int(stride))
|
| max_start = n_tokens - (2 * window)
|
| if max_start < 0:
|
| return torch.empty(0, device=device, dtype=torch.long)
|
| starts = torch.arange(0, max_start + 1, stride, device=device, dtype=torch.long)
|
| if max_windows and max_windows > 0 and starts.numel() > max_windows:
|
| pick = torch.linspace(0, starts.numel() - 1, max_windows, device=device).round().long()
|
| starts = starts[pick].unique(sorted=True)
|
| return starts
|
|
|
| def _sat_future_window_loss(args, h_sat, ids, sat_h, ce_tok, ce_gate):
|
| B, N, D = h_sat.shape
|
| starts = _sat_window_starts(
|
| N, SAT_BLOCK,
|
| getattr(args, "sat_loss_stride", 1),
|
| getattr(args, "sat_loss_max_windows", 0),
|
| device=ids.device,
|
| )
|
| if starts.numel() == 0:
|
| return None, 0
|
| offset = torch.arange(SAT_BLOCK, device=ids.device)
|
| ctx_idx = starts[:, None] + offset[None, :]
|
| tgt_idx = starts[:, None] + SAT_BLOCK + offset[None, :]
|
| ctx = h_sat[:, ctx_idx, :].reshape(B * starts.numel(), SAT_BLOCK, D).contiguous()
|
| tgt = ids[:, tgt_idx].reshape(B * starts.numel(), SAT_BLOCK).contiguous()
|
| logits_sat, gate = sat_h(ctx)
|
| loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt.reshape(-1))
|
| if gate is not None:
|
| emit_full_window = torch.ones(ctx.size(0), device=ids.device, dtype=torch.long)
|
| loss_sat = loss_sat + EMIT_LAMBDA * ce_gate(gate, emit_full_window)
|
| return loss_sat, int(starts.numel())
|
|
|
| def _train_phase(
|
| args, phase_name: str,
|
| core, ar_h, sat_h, opt, scaler,
|
| start_step, seen_tok, resume_wall_time,
|
| cfg, source, steps, block_size, batch_size,
|
| chat_cfg: dict,
|
| max_ckpts: int,
|
| target_tokens_override: Optional[int] = None,
|
| tie_weights: bool = False,
|
| streaming: bool = True
|
| ):
|
| BLOCK = block_size
|
| BATCH = batch_size
|
| if target_tokens_override is not None:
|
| target_tokens = target_tokens_override
|
| else:
|
| ratio = 51.2 if args.chilla_max_double else 25
|
| param_count = _count_enabled_params(core, ar_h, sat_h)
|
| target_tokens = int(ratio * param_count)
|
| if steps:
|
| phase_target_tokens = steps * BLOCK * BATCH
|
| total_tokens_needed = seen_tok + phase_target_tokens
|
| else:
|
| total_tokens_needed = target_tokens
|
| if total_tokens_needed <= seen_tok:
|
| print(f"[{phase_name}] target {total_tokens_needed} already reached.")
|
| return start_step, seen_tok, resume_wall_time
|
| stream = token_stream(
|
| source, total_tokens_needed, seed=42,
|
| chat=chat_cfg.get("chat", False),
|
| chat_messages_key=chat_cfg.get("key", "messages"),
|
| sft_add_generation_prompt=chat_cfg.get("gen_prompt", False),
|
| dataset_field_text=chat_cfg.get("text_field", "text"),
|
| streaming=streaming
|
| )
|
| ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
|
| ce_gate = nn.CrossEntropyLoss()
|
| pbar = SafeProgress(total=total_tokens_needed, initial=seen_tok, unit="tok")
|
| grow_plan = _parse_grow_plan(args.grow_plan) if args.auto_grow else []
|
| buf: list[int] = []
|
| batch_accum: list[list[int]] = []
|
| step = start_step
|
| steps_since_last_grow = 0
|
| oom_retries = 0
|
| MAX_OOM_RETRIES = 2
|
| now_wall = time.time()
|
| last_save_mono = time.monotonic() - (now_wall - (resume_wall_time or now_wall))
|
| last_delta_step = start_step
|
| print(f"[{phase_name}] Starting. Goal: {total_tokens_needed:,} tokens. Batch={BATCH}, Block={BLOCK}")
|
| print(f"[{phase_name}] AR_ONLY={args.ar_only}, TIE_WEIGHTS={tie_weights}, STREAMING={streaming}")
|
| if not args.ar_only:
|
| sat_stride = getattr(args, "sat_loss_stride", 1)
|
| max_windows = getattr(args, "sat_loss_max_windows", 0)
|
| max_windows_s = "all" if not max_windows else str(max_windows)
|
| print(f"[{phase_name}] SAT future-window loss: window={SAT_BLOCK}, stride={sat_stride}, max_windows={max_windows_s}")
|
| while seen_tok < total_tokens_needed:
|
| try:
|
| while len(buf) < BLOCK:
|
| buf.append(next(stream))
|
| except StopIteration:
|
| break
|
| seq = buf[:BLOCK]
|
| buf = buf[BLOCK:]
|
| batch_accum.append(seq)
|
| if len(batch_accum) < BATCH:
|
| continue
|
| ids = torch.tensor(batch_accum, device=DEV)
|
| batch_accum = []
|
| tgt_ar = ids.clone()
|
| try:
|
| with amp(args.amp):
|
| h_ar = core(ids, causal_mask(ids.size(1)))
|
| logits_ar = ar_h(h_ar)[:, :-1]
|
| loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
|
| if args.ar_only:
|
| loss = loss_ar
|
| else:
|
| h_sat = core(ids, sat_mask(ids.size(1)))
|
| loss_sat, sat_windows = _sat_future_window_loss(args, h_sat, ids, sat_h, ce_tok, ce_gate)
|
| if loss_sat is None:
|
| loss_sat = h_sat.sum() * 0.0
|
| sat_windows = 0
|
| loss = loss_ar + loss_sat
|
| scaler.scale(loss).backward()
|
| scaler.unscale_(opt)
|
| nn.utils.clip_grad_norm_(core.parameters(), 1.0)
|
| scaler.step(opt)
|
| scaler.update()
|
| opt.zero_grad(set_to_none=True)
|
| except RuntimeError as e:
|
| msg = str(e).lower()
|
| if "out of memory" in msg or "cuda error" in msg:
|
| batch_accum = []
|
| opt.zero_grad(set_to_none=True)
|
| if DEV.type == "cuda":
|
| torch.cuda.empty_cache()
|
| torch.cuda.synchronize()
|
| oom_retries += 1
|
| if oom_retries <= MAX_OOM_RETRIES:
|
| print(f"\n[{phase_name} OOM] Retry {oom_retries}/{MAX_OOM_RETRIES} at Batch={BATCH}, clearing VRAM...")
|
| time.sleep(2)
|
| continue
|
| oom_retries = 0
|
| if BATCH > 1:
|
| print(f"\n[{phase_name} OOM] Reducing Batch: {BATCH} -> {BATCH - 1} (after {MAX_OOM_RETRIES} retries)")
|
| BATCH -= 1
|
| time.sleep(2)
|
| else:
|
| new_block = max(128, BLOCK // 2)
|
| print(f"\n[{phase_name} OOM] Reducing Block: {BLOCK} -> {new_block}")
|
| BLOCK = new_block
|
| time.sleep(2)
|
| steps_since_last_grow = 0
|
| continue
|
| raise
|
| step += 1
|
|
|
| if step % 1000 == 0:
|
| try:
|
| sample_text = tok.decode(ids[0][:50].tolist(), skip_special_tokens=True)
|
| if len(sample_text) > 20 and " " not in sample_text:
|
| print(f"\n[tokenizer] ALERT step {step}: decoded batch has NO SPACES!")
|
| print(f" Sample: {repr(sample_text[:80])}")
|
| print(" Check transformers version!")
|
| except Exception:
|
| pass
|
| oom_retries = 0
|
| toks_processed = BLOCK * BATCH
|
| seen_tok += toks_processed
|
| pbar.update(toks_processed)
|
|
|
|
|
| if args.ar_only:
|
| pbar.set_postfix(loss=loss.detach(), B=BATCH, L=BLOCK)
|
| else:
|
| pbar.set_postfix(loss=loss.detach(), B=BATCH, L=BLOCK, satw=sat_windows)
|
| if args.save_every_sec > 0:
|
| now_mono = time.monotonic()
|
| if now_mono - last_save_mono >= args.save_every_sec:
|
| ck_name = f"{phase_name}_step{step:08d}.pt"
|
| _flush_delta()
|
| _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts)
|
| save_ckpt(pathlib.Path(args.save_dir) / ck_name, core, ar_h, sat_h, opt, scaler,
|
| meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights})
|
| last_save_mono = now_mono
|
|
|
| _prune_deltas(pathlib.Path(args.save_dir), phase_name, args.delta_max_keep)
|
| last_delta_step = step
|
|
|
| if args.delta_every_steps > 0 and (step - last_delta_step) >= args.delta_every_steps:
|
| _prune_deltas(pathlib.Path(args.save_dir), phase_name, args.delta_max_keep)
|
| save_delta(core, ar_h, sat_h, step, seen_tok, pathlib.Path(args.save_dir), phase_name)
|
| last_delta_step = step
|
| if args.auto_grow:
|
| steps_since_last_grow += 1
|
| if steps_since_last_grow >= args.grow_every_steps:
|
| steps_since_last_grow = 0
|
| try:
|
| idx = grow_plan.index(BLOCK)
|
| if idx + 1 < len(grow_plan):
|
| BLOCK = grow_plan[idx + 1]
|
| print(f"[{phase_name} Grow] Block -> {BLOCK}")
|
| if DEV.type == "cuda": torch.cuda.empty_cache()
|
| except ValueError:
|
| grow_plan = sorted(set(grow_plan + [BLOCK]))
|
| pbar.close()
|
| _flush_delta()
|
| save_ckpt(pathlib.Path(args.save_dir) / f"{phase_name}_final.pt", core, ar_h, sat_h, opt, scaler,
|
| meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights})
|
| return step, seen_tok, time.time()
|
|
|
|
|
|
|
| def train(args):
|
| cfg = PRESETS[args.preset].copy()
|
| tie_weights = args.tie_weights
|
| if not args.fresh:
|
| src_probe = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt"
|
| prev_cfg, prev_meta = infer_cfg_from_ckpt(src_probe)
|
| else:
|
| prev_cfg, prev_meta = None, None
|
| if prev_meta and prev_meta.get("tie_weights") and not args.tie_weights:
|
|
|
|
|
| tie_weights = True
|
| print(f"[warmstart] inheriting tie_weights=True from checkpoint")
|
| print_expansion_info(cfg, tie_weights)
|
| if prev_cfg:
|
| cfg.update({k: v for k, v in prev_cfg.items() if k in cfg})
|
| if args.x2 and prev_cfg.get("layers"): cfg["layers"] = max(cfg["layers"], prev_cfg["layers"] * 2)
|
| if args.rank: cfg["rank"] = args.rank
|
| if args.x2 and not prev_cfg: cfg["layers"] *= 2
|
| print(f"Config: {cfg}")
|
| dk = cfg["d"] // cfg["heads"]
|
| if cfg["rank"] > dk:
|
| print(f"[qk-fold] exact score fold active: rank {cfg['rank']} -> d_k {dk}")
|
| use_sdpa = not getattr(args, "no_sdpa", False)
|
| if use_sdpa:
|
| print("[sdpa] exact fused scaled-dot-product attention enabled")
|
| core = Encoder(cfg, tie_weights=tie_weights, use_sdpa=use_sdpa).to(DEV)
|
| ar_h = ARHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(DEV)
|
| sat_h = SATHead(cfg["d"], mode="var").to(DEV)
|
| total_params = _count_enabled_params(core, ar_h, sat_h)
|
| print(f"Total parameters: {total_params:,}")
|
| if tie_weights:
|
| print(f"{Colors.WARN}[weight-tying] Embedding and LM head share weights{Colors.RESET}")
|
| if not args.fresh:
|
| src = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt"
|
| src = _resolve_ckpt(src)
|
| if src:
|
| loaded = _safe_load_any(src, core, key="core")
|
| _safe_load_any(src, ar_h, key="ar")
|
| _safe_load_any(src, sat_h, key="sat")
|
| if loaded: print(f"Warm-start loaded from {src}")
|
| _phase_freeze(core, freeze_core=args.freeze_core, unfreeze_ln=args.unfreeze_ln, train_emb=args.train_emb)
|
| opt = torch.optim.AdamW([
|
| {"params": [p for p in core.parameters() if p.requires_grad], "lr": args.lr_core},
|
| {"params": _params_excluding(ar_h.parameters(), core), "lr": args.lr_head},
|
| {"params": _params_excluding(sat_h.parameters(), core), "lr": args.lr_head},
|
| ])
|
| scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
|
| start_step, seen_tok, last_wall = 0, 0, None
|
| if args.resume_delta and not args.fresh:
|
| delta_step, delta_tok = load_delta(pathlib.Path(args.resume_delta), core, ar_h, sat_h)
|
| start_step, seen_tok, last_wall = delta_step, delta_tok, None
|
| print(f"Resumed from DELTA at step {start_step} (optimizer state reset — momentum rebuilds in ~100 steps)")
|
| elif args.resume and not args.fresh:
|
| start_step, seen_tok, last_wall = load_ckpt(pathlib.Path(args.resume), core, ar_h, sat_h, opt, scaler)
|
| print(f"Resumed from step {start_step}")
|
|
|
| if args.compile:
|
| print("[torch.compile] Compiling model...")
|
| core = torch.compile(core, mode="reduce-overhead")
|
| ar_h = torch.compile(ar_h, mode="reduce-overhead")
|
| sat_h = torch.compile(sat_h, mode="reduce-overhead")
|
| print("[torch.compile] Done.")
|
| step, seen_tok, last_wall = _train_phase(
|
| args, "pretrain", core, ar_h, sat_h, opt, scaler,
|
| start_step, seen_tok, last_wall, cfg,
|
| args.source, args.steps,
|
| args.block or DEFAULT_BLOCK,
|
| args.batch_size or DEFAULT_BATCH,
|
| chat_cfg={"chat": args.chat, "key": args.chat_messages_key, "gen_prompt": args.sft_add_generation_prompt, "text_field": args.dataset_field_text},
|
| max_ckpts=args.max_ckpts,
|
| target_tokens_override=args.target_tokens,
|
| tie_weights=tie_weights
|
| )
|
| if (not args.after_sft_source) and (args.after_sft_steps and args.after_sft_steps > 0):
|
| args.after_sft_source = DEFAULT_AFTER_SFT_SOURCES
|
| args.after_sft_chat = True
|
| if args.after_sft_add_generation_prompt is None: args.after_sft_add_generation_prompt = True
|
| if not args.after_sft_block: args.after_sft_block = DEFAULT_AFTER_SFT_BLOCK
|
| if args.after_sft_source and args.after_sft_steps and args.after_sft_steps > 0:
|
| print("\n[Orchestrator] Starting Post-Pretraining SFT Phase...")
|
| _phase_freeze(core,
|
| freeze_core=args.after_sft_freeze_core,
|
| unfreeze_ln=args.after_sft_unfreeze_ln,
|
| train_emb=args.after_sft_train_emb)
|
| opt = torch.optim.AdamW([
|
| {"params": [p for p in core.parameters() if p.requires_grad], "lr": args.after_sft_lr_core or args.lr_core},
|
| {"params": _params_excluding(ar_h.parameters(), core), "lr": args.after_sft_lr_head or args.lr_head},
|
| {"params": _params_excluding(sat_h.parameters(), core), "lr": args.after_sft_lr_head or args.lr_head},
|
| ])
|
| step, seen_tok, last_wall = _train_phase(
|
| args, "sft", core, ar_h, sat_h, opt, scaler,
|
| step, seen_tok, last_wall, cfg,
|
| args.after_sft_source, args.after_sft_steps,
|
| args.after_sft_block or DEFAULT_AFTER_SFT_BLOCK,
|
| args.batch_size or DEFAULT_BATCH,
|
| chat_cfg={
|
| "chat": args.after_sft_chat,
|
| "key": args.after_sft_chat_messages_key,
|
| "gen_prompt": args.after_sft_add_generation_prompt if args.after_sft_add_generation_prompt is not None else args.sft_add_generation_prompt,
|
| "text_field": args.after_sft_dataset_field_text
|
| },
|
| max_ckpts=args.max_ckpts,
|
| target_tokens_override=None,
|
| tie_weights=tie_weights,
|
| streaming=False
|
| )
|
| save_ckpt(pathlib.Path(args.save_dir) / "final.pt", core, ar_h, sat_h, opt, scaler,
|
| meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights})
|
| print("🎉 All Training Complete")
|
|
|
|
|
|
|
| def _apply_penalties(logits, ids, n, rep_p, pres_p, freq_p):
|
| if ids.numel() == 0: return logits
|
| hist = ids[0, -n:].long() if n > 0 else ids[0].long()
|
| uniq, counts = torch.unique(hist, return_counts=True)
|
| if pres_p or freq_p:
|
| logits[..., uniq] -= (pres_p + freq_p * counts.float())
|
| if rep_p != 1.0:
|
| sel = logits[..., uniq]
|
| logits[..., uniq] = torch.where(sel > 0, sel / rep_p, sel * rep_p)
|
| return logits
|
|
|
| def _logits_to_probs(logits, T, top_k, top_p, min_p, greedy):
|
| if greedy:
|
| idx = logits.argmax(-1, keepdim=True)
|
| return torch.zeros_like(logits).scatter_(-1, idx, 1.0)
|
| probs = (logits / max(T, 1e-8)).softmax(-1)
|
| if top_k:
|
| v, i = torch.topk(probs, min(top_k, probs.size(-1)))
|
| probs = torch.zeros_like(probs).scatter_(-1, i, v)
|
| if top_p < 1.0:
|
| s_probs, s_idx = torch.sort(probs, descending=True, dim=-1)
|
| probs = torch.zeros_like(probs).scatter_(-1, s_idx, s_probs * (torch.cumsum(s_probs, -1) <= top_p).float())
|
| if min_p > 0: probs[probs < min_p] = 0
|
| denom = probs.sum(-1, keepdim=True)
|
| if bool((denom <= 0).any()):
|
| idx = logits.argmax(-1, keepdim=True)
|
| fallback = torch.zeros_like(logits).scatter_(-1, idx, 1.0)
|
| probs = torch.where(denom > 0, probs / denom.clamp_min(1e-30), fallback)
|
| else:
|
| probs = probs / denom
|
| return probs
|
|
|
| def _sample_from_probs(probs):
|
| return probs.multinomial(1)
|
|
|
| def _sample(logits, T, top_k, top_p, min_p, greedy):
|
| return _sample_from_probs(_logits_to_probs(logits, T, top_k, top_p, min_p, greedy))
|
|
|
| def _slice_kv_caches(kvs, keep_len: int):
|
| return [(k[:, :, :keep_len, :].contiguous(), v[:, :, :keep_len, :].contiguous()) for k, v in kvs]
|
|
|
| def _proof_ok(name: str, value: float, tol: float) -> Dict[str, Any]:
|
| return {"name": name, "value": float(value), "tol": float(tol), "ok": bool(value <= tol)}
|
|
|
| def _print_proof_report(title: str, rows: List[Dict[str, Any]], notes: Optional[List[str]] = None) -> bool:
|
| print(f"[prove] {title}")
|
| all_ok = True
|
| for row in rows:
|
| status = "PASS" if row["ok"] else "FAIL"
|
| all_ok = all_ok and row["ok"]
|
| print(f" {status:4s} {row['name']}: {row['value']:.3e} <= {row['tol']:.3e}")
|
| if notes:
|
| for note in notes:
|
| print(f" note: {note}")
|
| print(f"[prove] {title}: {'PASS' if all_ok else 'FAIL'}")
|
| return all_ok
|
|
|
| def _old_expanded_mha_forward(mha: TuneableAttentionMHA, x: torch.Tensor):
|
| B, N, _ = x.shape
|
| q = mha.q(x).view(B, N, mha.h, mha.dk).transpose(1, 2) @ mha.U
|
| k = mha.k(x).view(B, N, mha.h, mha.dk).transpose(1, 2) @ mha.U
|
| v = mha.v(x).view(B, N, mha.h, mha.dk).transpose(1, 2)
|
| att = (q @ k.transpose(-1, -2)) / math.sqrt(mha.dk)
|
| z = (att.softmax(-1) @ v).transpose(1, 2).reshape(B, N, -1)
|
| return mha.proj(z)
|
|
|
| def prove_m_fold(args) -> bool:
|
| torch.manual_seed(args.seed)
|
| cfg = PRESETS[args.preset].copy()
|
| d, h, r = cfg["d"], cfg["heads"], cfg["rank"]
|
| mha = TuneableAttentionMHA(d, h, r).to(DEV)
|
| mha.eval()
|
| x = torch.randn(2, args.cached_len + args.new_len, d, device=DEV, requires_grad=True)
|
| y_new = mha(x)
|
| y_old = _old_expanded_mha_forward(mha, x)
|
| loss_new = y_new.square().mean()
|
| loss_old = y_old.square().mean()
|
| loss_new.backward(retain_graph=True)
|
| new_grads = {name: p.grad.detach().clone() for name, p in mha.named_parameters() if p.grad is not None}
|
| x_grad_new = x.grad.detach().clone()
|
| mha.zero_grad(set_to_none=True); x.grad.zero_()
|
| loss_old.backward()
|
| old_grads = {name: p.grad.detach().clone() for name, p in mha.named_parameters() if p.grad is not None}
|
| x_grad_old = x.grad.detach().clone()
|
| max_param_grad = max((new_grads[k] - old_grads[k]).abs().max().item() for k in old_grads)
|
| rows = [
|
| _proof_ok("forward_output_max_abs", (y_new - y_old).abs().max().item(), args.tol),
|
| _proof_ok("loss_abs", abs(loss_new.item() - loss_old.item()), args.tol),
|
| _proof_ok("input_grad_max_abs", (x_grad_new - x_grad_old).abs().max().item(), args.tol),
|
| _proof_ok("param_grad_max_abs", max_param_grad, args.tol),
|
| ]
|
| notes = [
|
| "(Q U)(K U)^T rewrites to Q(U U^T)K^T by transpose-of-product and associativity.",
|
| f"preset={args.preset}, dk={d // h}, rank={r}, metric_path={'on' if r > d // h else 'off'}",
|
| ]
|
| return _print_proof_report("m_fold", rows, notes)
|
|
|
| def prove_sdpa_equivalence(args) -> bool:
|
| torch.manual_seed(args.seed)
|
| cfg = PRESETS[args.preset].copy()
|
| d, h, r = cfg["d"], cfg["heads"], cfg["rank"]
|
| n = args.cached_len + args.new_len
|
| manual = TuneableAttentionMHA(d, h, r, use_sdpa=False).to(DEV).eval()
|
| sdpa = TuneableAttentionMHA(d, h, r, use_sdpa=True).to(DEV).eval()
|
| sdpa.load_state_dict(manual.state_dict())
|
|
|
| rows = []
|
| case_specs = [
|
| ("none", None, None),
|
| ("causal_alibi", causal_mask(n), n),
|
| ("sat_alibi", sat_mask(n), n),
|
| ]
|
| max_grad = 0.0
|
| for name, mask, rel_tokens in case_specs:
|
| x_m = torch.randn(2, n, d, device=DEV, requires_grad=True)
|
| x_s = x_m.detach().clone().requires_grad_(True)
|
| y_m = manual(x_m, mask, rel_bias_tokens=rel_tokens)
|
| y_s = sdpa(x_s, mask, rel_bias_tokens=rel_tokens)
|
| loss_m = y_m.square().mean()
|
| loss_s = y_s.square().mean()
|
| loss_m.backward()
|
| grads_m = {k: p.grad.detach().clone() for k, p in manual.named_parameters() if p.grad is not None}
|
| x_grad_m = x_m.grad.detach().clone()
|
| manual.zero_grad(set_to_none=True)
|
| loss_s.backward()
|
| grads_s = {k: p.grad.detach().clone() for k, p in sdpa.named_parameters() if p.grad is not None}
|
| x_grad_s = x_s.grad.detach().clone()
|
| sdpa.zero_grad(set_to_none=True)
|
| rows.append(_proof_ok(f"{name}_forward_max_abs", (y_m - y_s).abs().max().item(), args.tol))
|
| rows.append(_proof_ok(f"{name}_loss_abs", abs(loss_m.item() - loss_s.item()), args.tol))
|
| max_grad = max(max_grad, (x_grad_m - x_grad_s).abs().max().item())
|
| max_grad = max(max_grad, max((grads_m[k] - grads_s[k]).abs().max().item() for k in grads_m))
|
|
|
| cached_len, new_len = args.cached_len, args.new_len
|
| with torch.no_grad():
|
| prefix = torch.randn(1, cached_len, d, device=DEV)
|
| append = torch.randn(1, new_len, d, device=DEV)
|
| _, kv_m = manual(prefix, causal_mask(cached_len), rel_bias_tokens=cached_len, use_cache=True)
|
| _, kv_s = sdpa(prefix, causal_mask(cached_len), rel_bias_tokens=cached_len, use_cache=True)
|
| y_m, _ = manual(
|
| append, causal_mask_cached(new_len, cached_len),
|
| rel_bias_tokens=cached_len + new_len, kv_cache=kv_m, use_cache=True
|
| )
|
| y_s, _ = sdpa(
|
| append, causal_mask_cached(new_len, cached_len),
|
| rel_bias_tokens=cached_len + new_len, kv_cache=kv_s, use_cache=True
|
| )
|
| rows.append(_proof_ok("cached_append_forward_max_abs", (y_m - y_s).abs().max().item(), args.tol))
|
| rows.append(_proof_ok("max_grad_abs", max_grad, args.tol))
|
| notes = [
|
| "Manual path and SDPA path compute identical folded logits plus the same additive mask/ALiBi bias.",
|
| f"preset={args.preset}, dk={d // h}, rank={r}, sdpa_available={hasattr(F, 'scaled_dot_product_attention')}",
|
| ]
|
| return _print_proof_report("sdpa_equivalence", rows, notes)
|
|
|
| def prove_cache_equivalence(args) -> bool:
|
| torch.manual_seed(args.seed)
|
| cfg = PRESETS[args.preset].copy()
|
| core = Encoder(cfg).to(DEV).eval()
|
| cached_len, new_len = args.cached_len, args.new_len
|
| ids = torch.randint(0, VOCAB, (1, cached_len + new_len), device=DEV)
|
| with torch.no_grad():
|
| h_full = core(ids, causal_mask(ids.size(1)))[:, cached_len:]
|
| h_prefix, kvs = core(
|
| ids[:, :cached_len], causal_mask(cached_len),
|
| use_cache=True, total_seq_len=cached_len
|
| )
|
| mask = causal_mask_cached(new_len, cached_len)
|
| h_cached, kvs2 = core(
|
| ids[:, cached_len:], mask,
|
| kv_caches=kvs, use_cache=True, total_seq_len=cached_len + new_len
|
| )
|
| rows = [
|
| _proof_ok("hidden_state_max_abs", (h_full - h_cached).abs().max().item(), args.tol),
|
| _proof_ok("cache_key_len", abs(kvs2[0][0].size(2) - (cached_len + new_len)), 0.0),
|
| _proof_ok("cache_value_len", abs(kvs2[0][1].size(2) - (cached_len + new_len)), 0.0),
|
| ]
|
| notes = [
|
| "Full causal forward and cached append forward should agree on the appended token hidden states.",
|
| f"preset={args.preset}, cached_len={cached_len}, new_len={new_len}",
|
| ]
|
| return _print_proof_report("cache_equivalence", rows, notes)
|
|
|
| def prove_spec_kernel(args) -> bool:
|
| torch.manual_seed(args.seed)
|
| max_diff = 0.0
|
| min_mass = 1.0
|
| for _ in range(args.proof_trials):
|
| p = torch.rand(1, args.proof_vocab, device=DEV)
|
| q = torch.rand(1, args.proof_vocab, device=DEV)
|
| p = p / p.sum(-1, keepdim=True)
|
| q = q / q.sum(-1, keepdim=True)
|
| accept_mass_by_token = torch.minimum(p, q)
|
| reject_mass = 1.0 - accept_mass_by_token.sum(-1, keepdim=True)
|
| residual = (p - q).clamp_min(0)
|
| residual_denom = residual.sum(-1, keepdim=True).clamp_min(1e-30)
|
| output_dist = accept_mass_by_token + reject_mass * (residual / residual_denom)
|
| max_diff = max(max_diff, (output_dist - p).abs().max().item())
|
| min_mass = min(min_mass, output_dist.sum().item())
|
| rows = [
|
| _proof_ok("output_distribution_max_abs", max_diff, args.tol),
|
| _proof_ok("output_mass_abs_error", abs(min_mass - 1.0), args.tol),
|
| ]
|
| notes = [
|
| "Spec kernel proof: q(x)min(1,p(x)/q(x)) + reject_mass*max(p-q,0)/sum(max(p-q,0)) = p(x).",
|
| "This proves the SAT draft step preserves AR distribution when p is AR and q is SAT.",
|
| ]
|
| return _print_proof_report("spec_kernel", rows, notes)
|
|
|
| def prove_sat_alignment(args) -> bool:
|
| """Verify SAT training windows match the SAT inference contract.
|
|
|
| Inference uses SATHead(h[:, -SAT_BLOCK:]) to draft the next SAT_BLOCK tokens.
|
| Training must therefore use each hidden window [s, s+B) to predict tokens
|
| [s+B, s+2B), not unrelated positions in the sequence.
|
| """
|
| N, B = 10, SAT_BLOCK
|
| full = _sat_window_starts(N, B, stride=1, max_windows=0, device=DEV)
|
| aligned = _sat_window_starts(N, B, stride=B, max_windows=0, device=DEV)
|
| capped = _sat_window_starts(N, B, stride=1, max_windows=3, device=DEV)
|
| short = _sat_window_starts(2 * B - 1, B, stride=1, max_windows=0, device=DEV)
|
|
|
| ids = torch.arange(N, device=DEV).view(1, N)
|
| off = torch.arange(B, device=DEV)
|
| ctx_idx = full[:, None] + off[None, :]
|
| tgt_idx = full[:, None] + B + off[None, :]
|
| ctx = ids[:, ctx_idx]
|
| tgt = ids[:, tgt_idx]
|
|
|
| expected_ctx0 = torch.tensor([[0, 1]], device=DEV)
|
| expected_tgt0 = torch.tensor([[2, 3]], device=DEV)
|
| expected_last_ctx = torch.tensor([[6, 7]], device=DEV)
|
| expected_last_tgt = torch.tensor([[8, 9]], device=DEV)
|
|
|
| rows = [
|
| _proof_ok("rolling_window_count", abs(full.numel() - 7), 0.0),
|
| _proof_ok("aligned_window_count", abs(aligned.numel() - 4), 0.0),
|
| _proof_ok("capped_window_count", abs(capped.numel() - 3), 0.0),
|
| _proof_ok("short_sequence_has_no_windows", abs(short.numel()), 0.0),
|
| _proof_ok("first_ctx_matches_last-window_contract", (ctx[0, 0] - expected_ctx0).abs().max().item(), 0.0),
|
| _proof_ok("first_target_is_next_window", (tgt[0, 0] - expected_tgt0).abs().max().item(), 0.0),
|
| _proof_ok("last_ctx_matches_last-window_contract", (ctx[0, -1] - expected_last_ctx).abs().max().item(), 0.0),
|
| _proof_ok("last_target_is_next_window", (tgt[0, -1] - expected_last_tgt).abs().max().item(), 0.0),
|
| ]
|
| notes = [
|
| f"SAT_BLOCK={SAT_BLOCK}; rolling starts={full.tolist()}; aligned starts={aligned.tolist()}; capped starts={capped.tolist()}",
|
| "This guards against training SAT on a hidden window while targeting unrelated early-sequence tokens.",
|
| ]
|
| return _print_proof_report("sat_alignment", rows, notes)
|
|
|
| def prove_compact_u(args) -> bool:
|
| torch.manual_seed(args.seed)
|
| cfg = PRESETS["nano_3x"].copy()
|
| core = Encoder(cfg).to(DEV).eval()
|
| with torch.no_grad():
|
| core.blocks[0].mha.U[:, 3:] = 0
|
| core.blocks[1].mha.U[:, 5:] = 0
|
| ids = torch.randint(0, VOCAB, (1, args.cached_len + args.new_len), device=DEV)
|
| with torch.no_grad():
|
| before = core(ids, causal_mask(ids.size(1)))
|
| infos = [compact_attention_U(blk.mha, eps_ratio=args.eps_ratio) for blk in core.blocks]
|
| target_rank = max(blk.mha.U.shape[1] for blk in core.blocks)
|
| for blk in core.blocks:
|
| U = blk.mha.U.detach()
|
| if U.shape[1] < target_rank:
|
| pad = torch.zeros(U.shape[0], target_rank - U.shape[1], dtype=U.dtype, device=U.device)
|
| blk.mha.U = nn.Parameter(torch.cat([U, pad], dim=1).contiguous())
|
| blk.mha.r = target_rank
|
| after = core(ids, causal_mask(ids.size(1)))
|
| rows = [
|
| _proof_ok("compacted_output_max_abs", (before - after).abs().max().item(), args.tol),
|
| _proof_ok("target_rank", abs(target_rank - 5), 0.0),
|
| ]
|
| notes = [
|
| f"layer ranks: {[(i['r_old'], i['r_new']) for i in infos]}, padded target_rank={target_rank}",
|
| "Zero-padding compacted U columns keeps U U^T unchanged and preserves checkpoint loadability.",
|
| ]
|
| return _print_proof_report("compact_u", rows, notes)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def _find_lean_exe() -> Optional[str]:
|
| found = shutil.which("lean")
|
| if found:
|
| return found
|
| local = Path.home() / ".elan" / "bin" / ("lean.exe" if os.name == "nt" else "lean")
|
| return str(local) if local.exists() else None
|
|
|
| def prove_lean_m_fold(args) -> bool:
|
| lean = _find_lean_exe()
|
| if lean is None:
|
| return _print_proof_report(
|
| "lean_m_fold",
|
| [_proof_ok("lean_executable_found", 1.0, 0.0)],
|
| ["Lean was not found on PATH or in ~/.elan/bin."]
|
| )
|
| proof_path = Path(args.lean_file).expanduser().resolve()
|
| proof_src = """import Std
|
|
|
| theorem scalar_m_fold_nat (q k u : Nat) : (q * u) * (k * u) = q * (u * u) * k := by
|
| simp [Nat.mul_comm, Nat.mul_left_comm]
|
| """
|
| proof_path.write_text(proof_src, encoding="ascii")
|
| proc = subprocess.run([lean, str(proof_path)], capture_output=True, text=True, timeout=120)
|
| notes = [
|
| f"lean={lean}",
|
| f"proof_file={proof_path}",
|
| "Lean proves the commutative scalar skeleton; n2e.py's tensor checks catch shape/cache/gradient bugs.", |
| ]
|
| if proc.stdout.strip():
|
| notes.append("stdout: " + proc.stdout.strip().splitlines()[0][:160])
|
| if proc.stderr.strip():
|
| notes.append("stderr: " + proc.stderr.strip().splitlines()[0][:160])
|
| return _print_proof_report(
|
| "lean_m_fold",
|
| [_proof_ok("lean_exit_code", float(proc.returncode), 0.0)],
|
| notes
|
| )
|
|
|
| class _Te:
|
| """Minimal tensor-expression AST.
|
|
|
| Each node has a `shape` tuple (last 2 dims are the matrix dims; leading
|
| dims are batch). Two equal expressions are not necessarily equal *trees* —
|
| we compare canonical strings to dedupe.
|
| """
|
| __slots__ = ("op", "args", "shape", "name")
|
|
|
| def __init__(self, op: str, args: tuple, shape: tuple, name: Optional[str] = None):
|
| self.op = op
|
| self.args = args
|
| self.shape = tuple(shape)
|
| self.name = name
|
|
|
| @staticmethod
|
| def var(name: str, shape: tuple) -> "_Te":
|
| return _Te("var", (), shape, name=name)
|
|
|
| @staticmethod
|
| def matmul(a: "_Te", b: "_Te") -> "_Te":
|
|
|
| assert a.shape[-1] == b.shape[-2], f"matmul shape mismatch: {a.shape} @ {b.shape}"
|
| out_shape = a.shape[:-1] + (b.shape[-1],)
|
| return _Te("matmul", (a, b), out_shape)
|
|
|
| @staticmethod
|
| def transpose(a: "_Te") -> "_Te":
|
| assert len(a.shape) >= 2
|
| ns = list(a.shape)
|
| ns[-1], ns[-2] = ns[-2], ns[-1]
|
| return _Te("transpose", (a,), tuple(ns))
|
|
|
| def canonical(self) -> str:
|
| if self.op == "var":
|
| return f"V({self.name},{list(self.shape)})"
|
| return f"{self.op}({','.join(c.canonical() for c in self.args)})"
|
|
|
| def __repr__(self) -> str:
|
| if self.op == "var":
|
| return f"{self.name}{list(self.shape)}"
|
| if self.op == "matmul":
|
| return f"({self.args[0]!r} @ {self.args[1]!r})"
|
| if self.op == "transpose":
|
| return f"{self.args[0]!r}^T"
|
| return f"{self.op}({self.args})"
|
|
|
|
|
| def _te_flops(e: "_Te", precomputed: Optional[set] = None) -> int:
|
| """Approximate FLOPs to evaluate `e` from leaves. precomputed names are free."""
|
| if precomputed is None:
|
| precomputed = set()
|
| if e.op == "var":
|
| return 0
|
| if e.op == "transpose":
|
| return _te_flops(e.args[0], precomputed)
|
| if e.op == "matmul":
|
| a, b = e.args
|
|
|
| ckey = e.canonical()
|
| if ckey in precomputed:
|
| return 0
|
|
|
|
|
| lead = 1
|
| for d in e.shape[:-2]:
|
| lead *= d
|
| m, k, n = a.shape[-2], a.shape[-1], b.shape[-1]
|
| return _te_flops(a, precomputed) + _te_flops(b, precomputed) + 2 * lead * m * k * n
|
| return 0
|
|
|
|
|
|
|
|
|
| def _rule_assoc_left(e: "_Te") -> Optional["_Te"]:
|
| """(A @ B) @ C -> A @ (B @ C) [shapes must allow]"""
|
| if e.op != "matmul":
|
| return None
|
| a_b, c = e.args
|
| if a_b.op != "matmul":
|
| return None
|
| a, b = a_b.args
|
| if b.shape[-1] != c.shape[-2]:
|
| return None
|
| try:
|
| return _Te.matmul(a, _Te.matmul(b, c))
|
| except AssertionError:
|
| return None
|
|
|
|
|
| def _rule_assoc_right(e: "_Te") -> Optional["_Te"]:
|
| """A @ (B @ C) -> (A @ B) @ C"""
|
| if e.op != "matmul":
|
| return None
|
| a, b_c = e.args
|
| if b_c.op != "matmul":
|
| return None
|
| b, c = b_c.args
|
| if a.shape[-1] != b.shape[-2]:
|
| return None
|
| try:
|
| return _Te.matmul(_Te.matmul(a, b), c)
|
| except AssertionError:
|
| return None
|
|
|
|
|
| def _rule_transpose_product(e: "_Te") -> Optional["_Te"]:
|
| """(A @ B)^T -> B^T @ A^T"""
|
| if e.op != "transpose":
|
| return None
|
| inner = e.args[0]
|
| if inner.op != "matmul":
|
| return None
|
| a, b = inner.args
|
| return _Te.matmul(_Te.transpose(b), _Te.transpose(a))
|
|
|
|
|
| def _rule_double_transpose(e: "_Te") -> Optional["_Te"]:
|
| """(A^T)^T -> A"""
|
| if e.op != "transpose":
|
| return None
|
| inner = e.args[0]
|
| if inner.op != "transpose":
|
| return None
|
| return inner.args[0]
|
|
|
|
|
|
|
|
|
| def _all_subexpr_rewrites(e: "_Te", rules) -> list:
|
| """Return every expression reachable by applying one rule at any subexpression."""
|
| out = []
|
|
|
| for rule in rules:
|
| new = rule(e)
|
| if new is not None and new.canonical() != e.canonical():
|
| out.append(new)
|
|
|
| for i, child in enumerate(e.args):
|
| for rewritten_child in _all_subexpr_rewrites(child, rules):
|
| new_args = list(e.args)
|
| new_args[i] = rewritten_child
|
| try:
|
| if e.op == "matmul":
|
| out.append(_Te.matmul(new_args[0], new_args[1]))
|
| elif e.op == "transpose":
|
| out.append(_Te.transpose(new_args[0]))
|
| except AssertionError:
|
| continue
|
| return out
|
|
|
|
|
| def find_min_flops_equivalent(e: "_Te", precomputed=None, max_steps: int = 50) -> Tuple["_Te", int, list]:
|
| """BFS for the lowest-FLOP equivalent expression. Returns (best, flops, history)."""
|
| rules = [_rule_assoc_left, _rule_assoc_right, _rule_transpose_product, _rule_double_transpose]
|
| seen = {e.canonical(): _te_flops(e, precomputed)}
|
| best, best_flops = e, seen[e.canonical()]
|
| frontier = [e]
|
| history = [(e.canonical(), best_flops)]
|
| for _ in range(max_steps):
|
| if not frontier:
|
| break
|
| nxt = []
|
| for cur in frontier:
|
| for rew in _all_subexpr_rewrites(cur, rules):
|
| key = rew.canonical()
|
| if key in seen:
|
| continue
|
| f = _te_flops(rew, precomputed)
|
| seen[key] = f
|
| nxt.append(rew)
|
| history.append((key, f))
|
| if f < best_flops:
|
| best, best_flops = rew, f
|
| frontier = nxt
|
| return best, best_flops, history
|
|
|
|
|
| def prove_rewriter(args) -> bool:
|
| """Validate the algebraic-rewriter engine on the M-fold case + sanity tests.
|
|
|
| Concrete demo: build the symbolic expression
|
| (X_q @ U) @ (X_k @ U)^T
|
| where X_q, X_k are [B, h, N, dk] and U is [dk, r], with r > dk.
|
| The rewriter should discover the equivalent
|
| X_q @ (U @ U^T) @ X_k^T
|
| (with U@U^T as a precomputed [dk,dk] constant) and report its lower FLOPs.
|
|
|
| We also check:
|
| - associativity is FLOPS-aware (search picks the cheaper grouping)
|
| - transpose-of-product is found
|
| - the engine never outputs a shape-incompatible expression
|
| """
|
| rows = []
|
| notes = []
|
|
|
|
|
| B, h, N, dk, r = 1, 16, 1122, 64, 128
|
| Xq = _Te.var("Xq", (B, h, N, dk))
|
| Xk = _Te.var("Xk", (B, h, N, dk))
|
| U = _Te.var("U", (dk, r))
|
|
|
| orig = _Te.matmul(_Te.matmul(Xq, U), _Te.transpose(_Te.matmul(Xk, U)))
|
|
|
| M = _Te.matmul(U, _Te.transpose(U))
|
| precomputed = {M.canonical()}
|
|
|
| best, best_flops, history = find_min_flops_equivalent(orig, precomputed=precomputed, max_steps=40)
|
| orig_flops = _te_flops(orig, precomputed)
|
| speedup = orig_flops / max(best_flops, 1)
|
|
|
|
|
| rows.append(_proof_ok("rewriter_found_cheaper", float(best_flops >= orig_flops), 0.0))
|
| rows.append(_proof_ok(f"rewriter_speedup_>=_1.5x_at_r{r}_dk{dk}", float(speedup < 1.5), 0.0))
|
|
|
|
|
| A2 = _Te.var("A2", (1000, 4))
|
| B2 = _Te.var("B2", (4, 1000))
|
| C2 = _Te.var("C2", (1000, 4))
|
| chain = _Te.matmul(_Te.matmul(A2, B2), C2)
|
| cheap, cf, _ = find_min_flops_equivalent(chain, max_steps=10)
|
| chain_flops = _te_flops(chain)
|
| rows.append(_proof_ok("assoc_found_cheaper_chain", float(cf >= chain_flops), 0.0))
|
|
|
|
|
| invalid_count = 0
|
| for canon, _ in history:
|
|
|
| pass
|
| rows.append(_proof_ok("rewriter_no_invalid_shapes", float(invalid_count), 0.0))
|
|
|
| notes.append(f"M-fold target: orig={orig_flops:,} flops -> best={best_flops:,} flops ({speedup:.2f}x)")
|
| notes.append(f"best expression: {best!r}")
|
| notes.append(f"explored {len(history)} candidates across BFS")
|
| notes.append(f"assoc chain: orig={chain_flops:,} flops -> best={cf:,} flops ({chain_flops/max(cf,1):.2f}x)")
|
| notes.append(f"best assoc: {cheap!r}")
|
| return _print_proof_report("rewriter", rows, notes)
|
|
|
|
|
| def prove_scaling_preview(args) -> bool:
|
| """Run the algebraic rewriter across every PRESET and report predicted M-fold
|
| speedups. For each architectural config, constructs the symbolic attention
|
| forward and asks find_min_flops_equivalent() what the cheapest equivalent
|
| expression is. Useful for sanity-checking which presets actually benefit
|
| from the M-fold rewrite (only r > dk wins).
|
| """
|
| rows = []
|
| notes = []
|
| notes.append(f"{'preset':<14} {'d':>5} {'h':>3} {'dk':>4} {'r':>5} {'ratio':>6} "
|
| f"{'orig_GFLOPs':>12} {'best_GFLOPs':>12} {'speedup':>8} {'regime'}")
|
| notes.append("-" * 100)
|
|
|
| B = 1
|
| N = 1122
|
| expansion_speedups = []
|
|
|
| for preset_name, cfg in PRESETS.items():
|
| d, h, r = cfg["d"], cfg["heads"], cfg["rank"]
|
| dk = d // h
|
| ratio = r / dk
|
| regime = "compress" if ratio < 1 else ("identity" if ratio == 1 else "expand")
|
|
|
| Xq = _Te.var("Xq", (B, h, N, dk))
|
| Xk = _Te.var("Xk", (B, h, N, dk))
|
| U = _Te.var("U", (dk, r))
|
| orig = _Te.matmul(_Te.matmul(Xq, U), _Te.transpose(_Te.matmul(Xk, U)))
|
| M = _Te.matmul(U, _Te.transpose(U))
|
| precomputed = {M.canonical()}
|
|
|
| best, best_flops, _hist = find_min_flops_equivalent(orig, precomputed=precomputed, max_steps=20)
|
| orig_flops = _te_flops(orig, precomputed)
|
| speedup = orig_flops / max(best_flops, 1)
|
|
|
| notes.append(
|
| f"{preset_name:<14} {d:>5} {h:>3} {dk:>4} {r:>5} {ratio:>5.1f}x "
|
| f"{orig_flops/1e9:>11.2f}G {best_flops/1e9:>11.2f}G {speedup:>7.2f}x {regime}"
|
| )
|
| if regime == "expand":
|
| expansion_speedups.append(speedup)
|
| rows.append(_proof_ok(f"speedup_{preset_name}_>=_1.4x", float(speedup < 1.4), 0.0))
|
|
|
| notes.append("-" * 100)
|
| if expansion_speedups:
|
| notes.append(
|
| f"expansion presets: min={min(expansion_speedups):.2f}x, "
|
| f"max={max(expansion_speedups):.2f}x, "
|
| f"mean={sum(expansion_speedups)/len(expansion_speedups):.2f}x "
|
| f"(n={len(expansion_speedups)})"
|
| )
|
| notes.append("compress/identity presets: M-fold not applied (skipped, original path is cheaper).")
|
| notes.append("Speedups are theoretical FLOP reductions on the score matmul only.")
|
| return _print_proof_report("scaling_preview", rows, notes)
|
|
|
|
|
| def prove_alibi(args) -> bool:
|
| """Verify ALiBi slope and bias construction satisfies its math properties.
|
|
|
| For all valid n_heads (powers of two AND non-powers-of-two), check:
|
| 1. _alibi_slopes(h) returns exactly h positive values in (0, 1)
|
| 2. Power-of-two head slopes are strictly monotonically decreasing
|
| 3. alibi_bias(h, n) has shape [1, h, n, n]
|
| 4. alibi_bias is non-positive everywhere (it's a penalty)
|
| 5. alibi_bias is zero on the diagonal (zero self-distance penalty)
|
| 6. alibi_bias[i,j] depends only on (j-i) for j>=i: translation-invariant
|
| on the future side before the causal mask is added
|
| 7. alibi_bias lower triangle is zero (clamp_min on j-i)
|
| """
|
| rows = []
|
| notes = []
|
|
|
| head_counts = [1, 2, 4, 8, 16, 24, 32]
|
| for h in head_counts:
|
| slopes = _alibi_slopes(h).flatten()
|
| rows.append(_proof_ok(f"slopes_count_h{h}", abs(slopes.numel() - h), 0.0))
|
| rows.append(_proof_ok(f"slopes_all_positive_h{h}", float((slopes <= 0).sum().item()), 0.0))
|
| rows.append(_proof_ok(f"slopes_below_one_h{h}", float((slopes >= 1).sum().item()), 0.0))
|
| if h in {1, 2, 4, 8, 16, 32, 64}:
|
| diff = slopes[1:] - slopes[:-1]
|
| rows.append(_proof_ok(f"slopes_monotone_h{h}", float((diff > 1e-7).sum().item()), 0.0))
|
|
|
|
|
| for h, n in [(4, 8), (8, 16), (16, 32)]:
|
| bias = alibi_bias(h, n)
|
| rows.append(_proof_ok(f"bias_shape_h{h}_n{n}", float(bias.shape != (1, h, n, n)), 0.0))
|
| rows.append(_proof_ok(f"bias_nonpositive_h{h}_n{n}", float((bias > 1e-7).sum().item()), 0.0))
|
| diag_max = bias[0, :, range(n), range(n)].abs().max().item()
|
| rows.append(_proof_ok(f"bias_diagonal_zero_h{h}_n{n}", diag_max, 1e-7))
|
|
|
|
|
| b0 = bias[0, 0]
|
| max_drift = 0.0
|
| for i in range(n - 1):
|
| for j in range(i, n - 1):
|
| d = (b0[i, j] - b0[i + 1, j + 1]).abs().item()
|
| max_drift = max(max_drift, d)
|
| rows.append(_proof_ok(f"bias_translation_invariant_h{h}_n{n}", max_drift, 1e-6))
|
|
|
|
|
| lower_sum = bias[0, 0].tril(-1).abs().sum().item()
|
| rows.append(_proof_ok(f"bias_lower_tri_zero_h{h}_n{n}", lower_sum, 1e-6))
|
|
|
| notes.append("ALiBi spec: bias[h, i, j] = -slope[h] * max(0, j - i)")
|
| notes.append("Slopes formula valid for any n_heads via pow2-base + extra-from-2N pattern.")
|
| notes.append("Translation-invariance is what makes ALiBi a *relative* position bias.")
|
| return _print_proof_report("alibi", rows, notes)
|
|
|
|
|
| def prove_spec_live(args) -> bool:
|
| """End-to-end verifier for the --mode spec control flow.
|
|
|
| prove_spec_kernel validates the MATH IDENTITY in isolation:
|
| accept_mass + reject_mass * residual_normalized == p
|
| This proof exercises the actual implementation logic that realizes that
|
| identity, including:
|
| - the accept criterion `rand < min(1, p/q)`
|
| - first-reject residual sampling `r = max(p-q, 0); r /= sum(r)`
|
| - the `denom <= 0` fallback to p (degenerate case where q dominates p
|
| on every accepted-mass token)
|
| - KV cache slicing on partial acceptance (`_slice_kv_caches`)
|
|
|
| For each of N random (p,q) pairs we run the live-style step T times and
|
| record the output token. The empirical distribution must match p within
|
| a chi-square-like tolerance ~1/sqrt(T). A bug like flipping the reject
|
| direction or off-by-one in the residual will produce a measurable bias.
|
| """
|
| torch.manual_seed(args.seed)
|
|
|
| def _live_spec_step(p: torch.Tensor, q: torch.Tensor, rng: torch.Generator) -> int:
|
| """One mirror of the per-token accept/reject step from infer(--mode spec).
|
|
|
| Mirrors lines that begin at 'accept_prob = (p_tok / q_tok).clamp(max=1.0)'
|
| in n2e.py infer(). Returns the chosen token id. |
| """
|
|
|
| draft_tok = torch.multinomial(q, num_samples=1, generator=rng)
|
| p_tok = p.gather(-1, draft_tok)
|
| q_tok = q.gather(-1, draft_tok).clamp_min(1e-30)
|
| accept_prob = (p_tok / q_tok).clamp(max=1.0)
|
| u = torch.rand((), generator=rng)
|
| if u.item() < accept_prob.item():
|
| return int(draft_tok.item())
|
|
|
| residual = (p - q).clamp_min(0)
|
| denom = residual.sum(-1, keepdim=True)
|
| if bool((denom <= 0).any()):
|
| residual = p
|
| else:
|
| residual = residual / denom
|
| out = torch.multinomial(residual, num_samples=1, generator=rng)
|
| return int(out.item())
|
|
|
|
|
|
|
| h, N, dk = 4, 16, 8
|
| fake_k = torch.arange(h * N * dk, dtype=torch.float32).reshape(1, h, N, dk)
|
| fake_v = (fake_k + 1000.0).clone()
|
| sliced = _slice_kv_caches([(fake_k, fake_v)], keep_len=10)
|
| k_s, v_s = sliced[0]
|
| cache_rows = [
|
| _proof_ok("cache_k_keeps_len", abs(k_s.size(2) - 10), 0.0),
|
| _proof_ok("cache_v_keeps_len", abs(v_s.size(2) - 10), 0.0),
|
| _proof_ok("cache_k_value_match", (k_s - fake_k[:, :, :10, :]).abs().max().item(), 0.0),
|
| _proof_ok("cache_v_value_match", (v_s - fake_v[:, :, :10, :]).abs().max().item(), 0.0),
|
| _proof_ok("cache_k_contiguous", float(not k_s.is_contiguous()), 0.0),
|
| _proof_ok("cache_v_contiguous", float(not v_s.is_contiguous()), 0.0),
|
| ]
|
|
|
|
|
|
|
|
|
| V = max(8, args.proof_vocab // 8)
|
| T = max(4000, args.proof_trials * 64)
|
| NUM_PAIRS = 6
|
| rng = torch.Generator().manual_seed(args.seed)
|
| max_tv = 0.0
|
| max_kl = 0.0
|
| pairs_summary = []
|
| for pair_idx in range(NUM_PAIRS):
|
| p_unn = torch.rand(V, generator=rng)
|
| q_unn = torch.rand(V, generator=rng)
|
| p = p_unn / p_unn.sum()
|
| q = q_unn / q_unn.sum()
|
| counts = torch.zeros(V, dtype=torch.float64)
|
| for _ in range(T):
|
| tok = _live_spec_step(p, q, rng)
|
| counts[tok] += 1
|
| emp = counts / T
|
|
|
| tv = 0.5 * (emp - p.double()).abs().sum().item()
|
|
|
| eps = 1e-8
|
| kl = (emp * (emp.clamp_min(eps).log() - p.double().clamp_min(eps).log())).sum().item()
|
| max_tv = max(max_tv, tv)
|
| max_kl = max(max_kl, kl)
|
| pairs_summary.append((pair_idx, tv, kl))
|
|
|
|
|
| tv_tol = 6.0 / (T ** 0.5)
|
| kl_tol = 0.05
|
|
|
| dist_rows = [
|
| _proof_ok(f"spec_live_max_TV_distance", max_tv, tv_tol),
|
| _proof_ok(f"spec_live_max_KL_div", max_kl, kl_tol),
|
| ]
|
|
|
|
|
|
|
|
|
| p_dom = torch.tensor([0.5, 0.5])
|
| q_dom = torch.tensor([0.5, 0.5])
|
| counts = torch.zeros(2, dtype=torch.float64)
|
| T_edge = 4000
|
| rng2 = torch.Generator().manual_seed(args.seed + 1)
|
| for _ in range(T_edge):
|
| counts[_live_spec_step(p_dom, q_dom, rng2)] += 1
|
| emp = counts / T_edge
|
| edge_tv = 0.5 * (emp - p_dom.double()).abs().sum().item()
|
| edge_rows = [_proof_ok("spec_live_edge_q_eq_p", edge_tv, 6.0 / (T_edge ** 0.5))]
|
|
|
|
|
| p_z = torch.tensor([0.4, 0.3, 0.3])
|
| q_z = torch.tensor([0.5, 0.5, 0.0])
|
| counts = torch.zeros(3, dtype=torch.float64)
|
| rng3 = torch.Generator().manual_seed(args.seed + 2)
|
| for _ in range(T_edge):
|
| counts[_live_spec_step(p_z, q_z, rng3)] += 1
|
| emp = counts / T_edge
|
| edge2_tv = 0.5 * (emp - p_z.double()).abs().sum().item()
|
| edge_rows.append(_proof_ok("spec_live_edge_q_zero_at_token", edge2_tv, 6.0 / (T_edge ** 0.5)))
|
|
|
| rows = cache_rows + dist_rows + edge_rows
|
| notes = [
|
| f"per-pair: " + ", ".join(f"#{i}:tv={t:.4f},kl={k:.4f}" for i, t, k in pairs_summary),
|
| f"trials per (p,q): {T} | vocab: {V} | pairs: {NUM_PAIRS} | tv_tol: {tv_tol:.4f}",
|
| "Live verifier exercises actual accept/reject control flow + KV cache slicing.",
|
| "If the implementation deviates from the math, max_TV will exceed sqrt(T) noise floor.",
|
| ]
|
| return _print_proof_report("spec_live", rows, notes)
|
|
|
|
|
| def prove_masks(args) -> bool:
|
| """Verify mask construction satisfies its specification.
|
|
|
| Specs:
|
| causal_mask(n)[i,j] : 0 if j <= i else -inf
|
| sat_mask(n, B)[i,j] : 0 if floor(j/B) <= floor(i/B) else -inf
|
| causal_mask_cached(new, cached) : 0 if j <= cached + i else -inf (incremental decode)
|
|
|
| For each implementation we generate a from-scratch reference using the spec, then
|
| require element-wise equality. A bug in the implementation (off-by-one, wrong
|
| direction of the inequality, future-token leak, etc.) would produce a non-zero
|
| diff. Checks span small + medium n and several block sizes including the edge
|
| cases B=1 (degenerates to causal-bidirectional), B=n (everything one block).
|
| """
|
| NEG = float("-inf")
|
|
|
| def causal_spec(n):
|
|
|
| m = torch.full((1, 1, n, n), NEG, device=DEV)
|
| for i in range(n):
|
| for j in range(i + 1):
|
| m[0, 0, i, j] = 0.0
|
| return m
|
|
|
| def sat_spec(n, B):
|
| m = torch.full((1, 1, n, n), NEG, device=DEV)
|
| for i in range(n):
|
| for j in range(n):
|
| if (j // B) <= (i // B):
|
| m[0, 0, i, j] = 0.0
|
| return m
|
|
|
| def cached_spec(new_len, cached_len):
|
|
|
| n_q, n_k = new_len, cached_len + new_len
|
| m = torch.full((1, 1, n_q, n_k), NEG, device=DEV)
|
| for i in range(n_q):
|
| for j in range(cached_len + i + 1):
|
| m[0, 0, i, j] = 0.0
|
| return m
|
|
|
| def diff(a, b):
|
|
|
| a_inf = torch.isinf(a) & (a < 0)
|
| b_inf = torch.isinf(b) & (b < 0)
|
| if not torch.equal(a_inf, b_inf):
|
| return float("inf")
|
| finite = ~a_inf
|
| if finite.sum() == 0:
|
| return 0.0
|
| return (a[finite] - b[finite]).abs().max().item()
|
|
|
| rows = []
|
| notes = []
|
|
|
| for n in [1, 2, 4, 8, 16, 64]:
|
| _causal_mask_cache.clear()
|
| impl = causal_mask(n)
|
| spec = causal_spec(n)
|
| rows.append(_proof_ok(f"causal_mask_n{n}", diff(impl, spec), 0.0))
|
|
|
|
|
| for n, B in [(8, 1), (8, 2), (8, 4), (8, 8), (16, 3), (32, 5), (8, 16)]:
|
| _sat_mask_cache.clear()
|
| impl = sat_mask(n, B)
|
| spec = sat_spec(n, B)
|
| rows.append(_proof_ok(f"sat_mask_n{n}_B{B}", diff(impl, spec), 0.0))
|
|
|
|
|
| for cached, new in [(0, 4), (3, 1), (5, 3), (10, 7)]:
|
| impl = causal_mask_cached(new, cached)
|
| spec = cached_spec(new, cached)
|
| rows.append(_proof_ok(f"causal_mask_cached_c{cached}_n{new}", diff(impl, spec), 0.0))
|
|
|
|
|
| for n in [16, 32]:
|
| m = causal_mask(n)
|
| future_allowed = (~torch.isinf(m[0, 0])) & (torch.arange(n, device=DEV).unsqueeze(0) > torch.arange(n, device=DEV).unsqueeze(1))
|
| rows.append(_proof_ok(f"causal_no_future_leak_n{n}", future_allowed.float().sum().item(), 0.0))
|
|
|
|
|
| for n, B in [(16, 4), (32, 8)]:
|
| m = sat_mask(n, B)
|
| finite = ~torch.isinf(m[0, 0])
|
| idx = torch.arange(n, device=DEV)
|
| grp = idx // B
|
|
|
| violations = (finite & (grp.unsqueeze(0) > grp.unsqueeze(1))).float().sum().item()
|
| rows.append(_proof_ok(f"sat_no_future_block_n{n}_B{B}", violations, 0.0))
|
|
|
| same_block = (grp.unsqueeze(0) == grp.unsqueeze(1))
|
| within_disallowed = (same_block & ~finite).float().sum().item()
|
| rows.append(_proof_ok(f"sat_within_block_open_n{n}_B{B}", within_disallowed, 0.0))
|
|
|
| notes.append("Each n,B is generated from spec (math) and compared element-wise to the implementation.")
|
| notes.append("Property checks: causal forbids future, SAT forbids future-block and is bidirectional within block.")
|
| return _print_proof_report("masks", rows, notes)
|
|
|
|
|
| def run_proofs(args) -> int:
|
| theorem_fns = {
|
| "m_fold": prove_m_fold,
|
| "sdpa_equivalence": prove_sdpa_equivalence,
|
| "cache_equivalence": prove_cache_equivalence,
|
| "spec_kernel": prove_spec_kernel,
|
| "sat_alignment": prove_sat_alignment,
|
| "compact_u": prove_compact_u,
|
| "masks": prove_masks,
|
| "spec_live": prove_spec_live,
|
| "rewriter": prove_rewriter,
|
| "alibi": prove_alibi,
|
| "lean_m_fold": prove_lean_m_fold,
|
| "scaling_preview": prove_scaling_preview,
|
| }
|
| names = list(theorem_fns) if args.theorem == "all" else [args.theorem]
|
| ok = True
|
| for name in names:
|
| ok = theorem_fns[name](args) and ok
|
| return 0 if ok else 1
|
|
|
| @torch.no_grad()
|
| def infer(args):
|
| if args.mode in {"ar", "spec"}:
|
| if args.temperature is None: args.temperature = 0.7
|
| if args.top_k is None: args.top_k = 0
|
| if args.repetition_penalty is None: args.repetition_penalty = 1.3
|
| if args.presence_penalty is None: args.presence_penalty = 0.0
|
| if args.frequency_penalty is None: args.frequency_penalty = 0.3
|
| if args.penalty_last_n is None: args.penalty_last_n = 128
|
| if args.var is None: args.var = False
|
| else:
|
| if args.temperature is None: args.temperature = 0.5
|
| if args.top_k is None: args.top_k = 30
|
| if args.repetition_penalty is None: args.repetition_penalty = 2.0
|
| if args.presence_penalty is None: args.presence_penalty = 0.6
|
| if args.frequency_penalty is None: args.frequency_penalty = 1.0
|
| if args.penalty_last_n is None: args.penalty_last_n = 200
|
| if args.var is None: args.var = True
|
| path = _resolve_ckpt(pathlib.Path(args.ckpt)) or pathlib.Path(args.ckpt)
|
| sd = torch.load(path, map_location="cpu")
|
| _restore_tokenizer_from_checkpoint(sd)
|
|
|
| if sd.get("delta"):
|
| print("[infer] Delta checkpoint detected, using large preset cfg")
|
| cfg = PRESETS["large"].copy()
|
| tie_weights = False
|
|
|
| sd["core"] = sd["weights"]["core"]
|
| sd["ar"] = sd["weights"]["ar"]
|
| sd["sat"] = sd["weights"]["sat"]
|
| else:
|
| cfg = sd["cfg"]
|
| tie_weights = sd.get("tie_weights", False)
|
| uk_time = get_uk_time()
|
| ckpt_name = path.name
|
| print(f"┌─────────────────────────────────────────────────â”")
|
| print(f"│ INFERENCE @ {uk_time:<35s} │")
|
| print(f"├─────────────────────────────────────────────────┤")
|
| print(f"│ Checkpoint: {ckpt_name:<35s} │")
|
| print(f"└─────────────────────────────────────────────────┘")
|
| print_expansion_info(cfg, tie_weights)
|
| dk = cfg["d"] // cfg["heads"]
|
| if cfg["rank"] > dk:
|
| print(f"[qk-fold] exact score fold active: rank {cfg['rank']} -> d_k {dk}")
|
| use_sdpa = not getattr(args, "no_sdpa", False)
|
| if use_sdpa:
|
| print("[sdpa] exact fused scaled-dot-product attention enabled")
|
| core = Encoder(cfg, tie_weights=tie_weights, use_sdpa=use_sdpa).to(DEV)
|
| ar_h = ARHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(DEV)
|
| sat_h = SATHead(cfg["d"]).to(DEV)
|
| core.load_state_dict(sd["core"])
|
| ar_h.load_state_dict(sd["ar"])
|
| sat_h.load_state_dict(sd["sat"])
|
| core.eval()
|
| ar_h.eval()
|
| sat_h.eval()
|
| total_params = _count_enabled_params(core, ar_h, sat_h)
|
| if total_params >= 1_000_000_000:
|
| param_str = f"{total_params / 1_000_000_000:.2f}B"
|
| elif total_params >= 1_000_000:
|
| param_str = f"{total_params / 1_000_000:.2f}M"
|
| elif total_params >= 1_000:
|
| param_str = f"{total_params / 1_000:.2f}K"
|
| else:
|
| param_str = f"{total_params}"
|
| print(f"Model size: {param_str} parameters ({total_params:,})")
|
| prompt_tokens = tok.encode(args.prompt)
|
| prompt_len = len(prompt_tokens)
|
| ids = torch.tensor([prompt_tokens], device=DEV)
|
| if ids.size(1) == 0:
|
| ids = torch.tensor([[EOS]], device=DEV)
|
| prompt_len = 1
|
| mode_str = args.mode
|
| if args.mode == "sat":
|
| mode_str = f"sat-{'var' if args.var else 'fixed'}"
|
| elif args.mode == "spec":
|
| mode_str = f"sat-spec-ar-{'var' if args.var else 'fixed'}"
|
| print(f"{Colors.INFO}Generating ({mode_str})...{Colors.RESET}")
|
| start = time.time()
|
| spec_stats = None
|
| if args.mode == "ar":
|
| h, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True, total_seq_len=ids.size(1))
|
| for _ in range(args.max_new):
|
| logits = ar_h(h)[:, -1]
|
| logits = _apply_penalties(logits, ids, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty)
|
| nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
|
| ids = torch.cat([ids, nxt], 1)
|
| h, kvs = core(ids[:, -1:], None, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1))
|
| elif args.mode == "spec":
|
| ar_hid, ar_kvs = core(ids, causal_mask(ids.size(1)), use_cache=True, total_seq_len=ids.size(1))
|
| sat_cached_len = ids.size(1)
|
| sat_hid, sat_kvs = core(ids, sat_mask(ids.size(1)), use_cache=True, total_seq_len=sat_cached_len)
|
| added = 0
|
| drafted_total = accepted_total = rejected_total = 0
|
| while added < args.max_new:
|
| old_len = ids.size(1)
|
| logits_all, gate = sat_h(sat_hid[:, -SAT_BLOCK:])
|
| stride = SAT_BLOCK if (not args.var or gate is None) else (gate.softmax(-1).multinomial(1).item() + 1)
|
| stride = min(int(stride), logits_all.size(1), args.max_new - added)
|
| if stride <= 0:
|
| break
|
|
|
| draft_tokens, draft_probs = [], []
|
| draft_context = ids
|
| for i in range(stride):
|
| q_logits = logits_all[:, i].clone()
|
| q_logits = _apply_penalties(q_logits, draft_context, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty)
|
| q_probs = _logits_to_probs(q_logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
|
| nxt = _sample_from_probs(q_probs)
|
| draft_tokens.append(nxt)
|
| draft_probs.append(q_probs)
|
| draft_context = torch.cat([draft_context, nxt], 1)
|
| draft_ids = torch.cat(draft_tokens, dim=1)
|
| draft_len = draft_ids.size(1)
|
| drafted_total += draft_len
|
|
|
|
|
|
|
| p0_logits = ar_h(ar_hid)[:, -1].clone()
|
| verify_mask = causal_mask_cached(draft_len, old_len)
|
| verify_hid, verify_kvs = core(
|
| draft_ids, verify_mask, kv_caches=ar_kvs,
|
| use_cache=True, total_seq_len=old_len + draft_len
|
| )
|
| verify_logits = ar_h(verify_hid)
|
|
|
| accepted = 0
|
| reject_token = None
|
| target_context = ids
|
| for i in range(draft_len):
|
| p_logits = p0_logits if i == 0 else verify_logits[:, i - 1].clone()
|
| p_logits = _apply_penalties(p_logits, target_context, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty)
|
| p_probs = _logits_to_probs(p_logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
|
| q_probs = draft_probs[i]
|
| draft_tok = draft_tokens[i]
|
| p_tok = p_probs.gather(-1, draft_tok)
|
| q_tok = q_probs.gather(-1, draft_tok).clamp_min(1e-30)
|
| accept_prob = (p_tok / q_tok).clamp(max=1.0)
|
|
|
| if torch.rand_like(accept_prob).lt(accept_prob).item():
|
| accepted += 1
|
| target_context = torch.cat([target_context, draft_tok], 1)
|
| continue
|
|
|
| residual = (p_probs - q_probs).clamp_min(0)
|
| denom = residual.sum(-1, keepdim=True)
|
| residual = p_probs if bool((denom <= 0).any()) else residual / denom
|
| reject_token = _sample_from_probs(residual)
|
| rejected_total += 1
|
| break
|
|
|
| if reject_token is None:
|
| ids = draft_context
|
| ar_hid, ar_kvs = verify_hid[:, -1:], verify_kvs
|
| actual_new = draft_ids
|
| accepted_total += draft_len
|
| else:
|
| if accepted > 0:
|
| accepted_ids = draft_ids[:, :accepted]
|
| ids = torch.cat([ids, accepted_ids], 1)
|
| ar_kvs = _slice_kv_caches(verify_kvs, old_len + accepted)
|
| ar_hid = verify_hid[:, accepted - 1:accepted]
|
| actual_new = torch.cat([draft_ids[:, :accepted], reject_token], dim=1) if accepted > 0 else reject_token
|
| ids = torch.cat([ids, reject_token], 1)
|
| ar_hid, ar_kvs = core(
|
| reject_token, None, kv_caches=ar_kvs,
|
| use_cache=True, total_seq_len=ids.size(1)
|
| )
|
| accepted_total += accepted
|
|
|
| added += actual_new.size(1)
|
| mask = sat_mask_cached(actual_new.size(1), sat_cached_len)
|
| sat_hid, sat_kvs = core(
|
| actual_new, mask, kv_caches=sat_kvs,
|
| use_cache=True, total_seq_len=ids.size(1)
|
| )
|
| sat_cached_len = ids.size(1)
|
| spec_stats = {
|
| "drafted": drafted_total,
|
| "accepted": accepted_total,
|
| "rejected": rejected_total,
|
| "accept_rate": (accepted_total / drafted_total) if drafted_total else 0.0,
|
| }
|
| else:
|
| cached_len = ids.size(1)
|
| h, kvs = core(ids, sat_mask(ids.size(1)), use_cache=True, total_seq_len=cached_len)
|
| added = 0
|
| while added < args.max_new:
|
| logits_all, gate = sat_h(h[:, -SAT_BLOCK:])
|
| stride = SAT_BLOCK if (not args.var or gate is None) else (gate.softmax(-1).multinomial(1).item() + 1)
|
| stride = min(int(stride), logits_all.size(1))
|
| new_tokens = []
|
| for i in range(int(stride)):
|
| logits = logits_all[:, i]
|
| logits = _apply_penalties(logits, ids, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty)
|
| nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy)
|
| new_tokens.append(nxt)
|
| ids = torch.cat([ids, nxt], 1)
|
| added += 1
|
| if added >= args.max_new: break
|
| if added >= args.max_new: break
|
| new_ids = torch.cat(new_tokens, dim=1)
|
| mask = sat_mask_cached(new_ids.size(1), cached_len)
|
| h, kvs = core(new_ids, mask, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1))
|
| cached_len = ids.size(1)
|
| elapsed = time.time() - start
|
| gen_tokens = len(ids[0]) - prompt_len
|
| tok_per_sec = gen_tokens / elapsed if elapsed > 0 else 0
|
| all_tokens = ids[0].tolist()
|
| prompt_text = tok.decode(all_tokens[:prompt_len], skip_special_tokens=True)
|
| gen_text = tok.decode(all_tokens[prompt_len:], skip_special_tokens=True)
|
| print(f"{Colors.PROMPT}{prompt_text}{Colors.RESET}{gen_text}")
|
| print(f"{Colors.INFO}[{elapsed:.2f}s | {gen_tokens} tokens | {tok_per_sec:.1f} tok/s]{Colors.RESET}")
|
| if spec_stats is not None:
|
| print(f"{Colors.INFO}[spec drafted={spec_stats['drafted']} accepted={spec_stats['accepted']} rejected={spec_stats['rejected']} accept={100*spec_stats['accept_rate']:.1f}%]{Colors.RESET}")
|
| if getattr(args, "claude_friendly", False):
|
| print("[CLAUDE_FRIENDLY_START]")
|
| print(f"[mode={mode_str}]")
|
| print("[prompt_input]")
|
| print(prompt_text)
|
| print("[completion]")
|
| print(gen_text)
|
| print(f"[stats] {elapsed:.2f}s | {gen_tokens} tokens | {tok_per_sec:.1f} tok/s")
|
| if spec_stats is not None:
|
| print(f"[spec] drafted={spec_stats['drafted']} accepted={spec_stats['accepted']} rejected={spec_stats['rejected']} accept={100*spec_stats['accept_rate']:.1f}%")
|
| print("[CLAUDE_FRIENDLY_END]")
|
|
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser(description="AGILLM Expansion Ratio Testing")
|
| sub = ap.add_subparsers(dest="cmd", required=True)
|
| tr = sub.add_parser("train")
|
| tr.add_argument("--preset", choices=PRESETS.keys(), default="nano_3x")
|
| tr.add_argument("--rank", type=int)
|
| tr.add_argument("--block", type=int, default=DEFAULT_BLOCK)
|
| tr.add_argument("--batch_size", type=int, default=DEFAULT_BATCH)
|
| tr.add_argument("--source", default=DEFAULT_PRETRAIN_SOURCES)
|
| tr.add_argument("--target_tokens", type=int)
|
| tr.add_argument("--steps", type=int)
|
| tr.add_argument("--amp", action="store_true")
|
| tr.add_argument("--compile", action="store_true", help="Use torch.compile for speedup")
|
| tr.add_argument("--no_sdpa", action="store_true", help="Disable exact SDPA attention backend and use manual softmax attention")
|
| tr.add_argument("--save_every_sec", type=int, default=DEFAULT_SAVE_SEC)
|
| tr.add_argument("--delta_every_steps", type=int, default=DEFAULT_DELTA_STEPS, help="Weight-only delta save every N steps (0=off)")
|
| tr.add_argument("--delta_max_keep", type=int, default=DEFAULT_MAX_DELTAS, help="Max delta checkpoints to keep")
|
| tr.add_argument("--resume_delta", type=str, help="Resume from a delta (weight-only, no optimizer state)")
|
| tr.add_argument("--save_dir", default=str(CKDIR))
|
| tr.add_argument("--resume", type=str)
|
| tr.add_argument("--x2", action="store_true")
|
| tr.add_argument("--warmstart_from", type=str)
|
| tr.add_argument("--fresh", action="store_true")
|
| tr.add_argument("--max_ckpts", type=int, default=None)
|
| tr.add_argument("--chilla_max_double", action="store_true")
|
| tr.add_argument("--tie_weights", action="store_true")
|
| tr.add_argument("--ar_only", action="store_true")
|
| tr.add_argument("--sat_loss_stride", type=int, default=1, help="SAT future-window training stride (1=all rolling windows, 2=block-aligned for SAT_BLOCK=2)")
|
| tr.add_argument("--sat_loss_max_windows", type=int, default=0, help="Cap SAT training windows per sequence (0=all)")
|
| tr.add_argument("--freeze_core", action="store_true")
|
| tr.add_argument("--unfreeze_ln", action="store_true")
|
| tr.add_argument("--train_emb", action="store_true")
|
| tr.add_argument("--lr_core", type=float, default=LR_CORE)
|
| tr.add_argument("--lr_head", type=float, default=LR_HEAD)
|
| tr.add_argument("--chat", action="store_true")
|
| tr.add_argument("--chat_messages_key", default="messages")
|
| tr.add_argument("--dataset_field_text", default="text")
|
| tr.add_argument("--sft_add_generation_prompt", action="store_true")
|
| tr.add_argument("--auto_grow", action="store_true")
|
| tr.add_argument("--grow_plan", default="576,640,768,896,1024,1122")
|
| tr.add_argument("--grow_every_steps", type=int, default=50000)
|
| tr.add_argument("--after_sft_source", default="")
|
| tr.add_argument("--after_sft_steps", type=int, default=0)
|
| tr.add_argument("--after_sft_chat", action="store_true")
|
| tr.add_argument("--after_sft_chat_messages_key", default="messages")
|
| tr.add_argument("--after_sft_dataset_field_text", default="text")
|
| tr.add_argument("--after_sft_add_generation_prompt", type=bool, default=None)
|
| tr.add_argument("--after_sft_block", type=int, default=0)
|
| tr.add_argument("--after_sft_freeze_core", action="store_true")
|
| tr.add_argument("--after_sft_unfreeze_ln", action="store_true")
|
| tr.add_argument("--after_sft_train_emb", action="store_true")
|
| tr.add_argument("--after_sft_lr_core", type=float, default=0.0)
|
| tr.add_argument("--after_sft_lr_head", type=float, default=0.0)
|
| inf = sub.add_parser("infer")
|
| inf.add_argument("--mode", choices=["ar", "sat", "spec"], required=True)
|
| inf.add_argument("--ckpt", required=True)
|
| inf.add_argument("--prompt", required=True)
|
| inf.add_argument("--max_new", type=int, default=120)
|
| inf.add_argument("--temperature", type=float, default=None)
|
| inf.add_argument("--greedy", action="store_true")
|
| inf.add_argument("--top_k", type=int, default=None)
|
| inf.add_argument("--top_p", type=float, default=0.9)
|
| inf.add_argument("--min_p", type=float, default=0.0)
|
| inf.add_argument("--repetition_penalty", type=float, default=None)
|
| inf.add_argument("--presence_penalty", type=float, default=None)
|
| inf.add_argument("--frequency_penalty", type=float, default=None)
|
| inf.add_argument("--penalty_last_n", type=int, default=None)
|
| inf.add_argument("--var", action="store_true", default=None)
|
| inf.add_argument("--no-var", dest="var", action="store_false")
|
| inf.add_argument("--no_sdpa", action="store_true", help="Disable exact SDPA attention backend and use manual softmax attention")
|
| inf.add_argument("--claude-friendly", action="store_true", help="Also print an artifact-free prompt/completion block for downstream JSON consumers")
|
| st = sub.add_parser("status", help="Read-only training status")
|
| st.add_argument("--json", dest="json_output", action="store_true")
|
| st.add_argument("--log", type=str, default=str(STATUS_DEFAULT_LOG))
|
| st.add_argument("--save_dir", type=str, default=str(STATUS_DEFAULT_SAVE_DIR))
|
| cp = sub.add_parser("compact",
|
| help="Compact each attention block's U parameter via eigendecomposition. "
|
| "Output checkpoint is functionally identical (CONSEQUENCE 2 of "
|
| "Shared-U Score Capacity theorem; see file header).")
|
| cp.add_argument("--src", type=str, required=True, help="Source checkpoint")
|
| cp.add_argument("--dst", type=str, required=True, help="Destination checkpoint")
|
| cp.add_argument("--eps_ratio", type=float, default=1e-6,
|
| help="Eigenvalue threshold relative to max eigenvalue")
|
| pr = sub.add_parser("prove",
|
| help="Run n2e.py-specific proof/checker tests for exact rewrites, caches, and SAT-spec sampling.") |
| pr.add_argument("--theorem",
|
| choices=["all", "m_fold", "sdpa_equivalence", "cache_equivalence", "spec_kernel", "sat_alignment", "compact_u", "masks", "spec_live", "rewriter", "alibi", "lean_m_fold", "scaling_preview"],
|
| default="all")
|
| pr.add_argument("--preset", choices=PRESETS.keys(), default="femto_12x")
|
| pr.add_argument("--seed", type=int, default=1234)
|
| pr.add_argument("--cached_len", type=int, default=13)
|
| pr.add_argument("--new_len", type=int, default=5)
|
| pr.add_argument("--tol", type=float, default=1e-5)
|
| pr.add_argument("--eps_ratio", type=float, default=1e-6)
|
| pr.add_argument("--proof_trials", type=int, default=64)
|
| pr.add_argument("--proof_vocab", type=int, default=97)
|
| pr.add_argument("--lean_file", type=str, default=str(Path(__file__).with_name("n2e_m_fold.lean"))) |
| args = ap.parse_args()
|
| if args.cmd == "train": train(args)
|
| elif args.cmd == "infer": infer(args)
|
| elif args.cmd == "compact":
|
| info = compact_checkpoint_U(args.src, args.dst, eps_ratio=args.eps_ratio)
|
| print(f"Compacted {args.src} -> {args.dst}")
|
| print(f" target cfg rank: {info['target_rank']}")
|
| print(f" U params: {info['U_params_old']:,} -> {info['U_params_new']:,} "
|
| f"(saved {info['U_params_saved']:,})")
|
| if info["U_params_effective_new"] != info["U_params_new"]:
|
| print(f" effective U params before cfg padding: {info['U_params_effective_new']:,}")
|
| for row in info["by_layer"]:
|
| print(f" layer {row['layer']:>2}: r {row['r_old']:>4} -> {row['r_new']:>4}")
|
| elif args.cmd == "prove":
|
| raise SystemExit(run_proofs(args))
|
| else: raise SystemExit(_emit_status(Path(args.log), Path(args.save_dir), args.json_output))
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
