#!/usr/bin/env python3 """ SymbolicLight 0.8B evaluation script. Features: 1. Load a step_*.pt checkpoint 2. Compute held-out CE / PPL with val.bin or a train-tail fallback 3. Run generation quality tests on English and code prompts 4. Report sparsity statistics Examples: python eval_08.py --generate_only python eval_08.py --data_bin /path/to/data_bin --generate python eval_08.py --checkpoint_path ../weights/pytorch/latest.pt --generate_only python eval_08.py --data_bin /path/to/data_bin --json """ import argparse import io import json import math import os import pickle import struct import sys import time import warnings import zipfile from collections import OrderedDict, namedtuple from pathlib import Path from typing import Optional, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, str(Path(__file__).parent)) from model import SymbolicLightConfig, SymbolicLightModel PACKAGE_ROOT = Path(__file__).resolve().parent.parent _SL_TOKENIZER_PATH = str(PACKAGE_ROOT / "tokenizer" / "sl_tokenizer.model") _StorageRef = namedtuple("_StorageRef", "dtype key location size") _TensorRef = namedtuple("_TensorRef", "storage offset size stride requires_grad hooks") def _rebuild_tensor_ref(storage, storage_offset, size, stride, requires_grad=False, backward_hooks=None, metadata=None): return _TensorRef( storage=storage, offset=storage_offset, size=tuple(size), stride=tuple(stride), requires_grad=requires_grad, hooks=backward_hooks, ) def _rebuild_parameter_ref(data, requires_grad, backward_hooks): return data class _CheckpointMetadataUnpickler(pickle.Unpickler): """Read PyTorch checkpoint metadata without materializing storages.""" def persistent_load(self, pid): kind, storage_type, key, location, size = pid if kind != "storage": raise pickle.UnpicklingError(f"Unsupported persistent id: {pid}") return _StorageRef(storage_type.__name__, key, location, size) def find_class(self, module, name): if module == "torch._utils" and name in {"_rebuild_tensor", "_rebuild_tensor_v2"}: return _rebuild_tensor_ref if module == "torch._utils" and name == "_rebuild_parameter": return _rebuild_parameter_ref if module == "collections" and name == "OrderedDict": return OrderedDict return super().find_class(module, name) DEFAULT_08B_CONFIG = dict( vocab_size=57344, embed_dim=1536, n_layers=22, n_heads=24, head_dim=64, intermediate_dim=6144, max_seq_len=512, ) DEFAULT_CHECKPOINT = os.path.join( PACKAGE_ROOT, "weights", "pytorch", "latest.pt" ) def parse_args(): p = argparse.ArgumentParser(description="SymbolicLight 0.8B evaluation / PPL / generation test") p.add_argument("--checkpoint_path", type=str, default=DEFAULT_CHECKPOINT, help="Path to .pt checkpoint") p.add_argument("--tokenizer_path", type=str, default=_SL_TOKENIZER_PATH, help="Path to SL tokenizer model") p.add_argument("--device", type=str, default="cpu", help="cuda / cpu") p.add_argument("--allow_windows_cuda", action="store_true", help="Allow experimental CUDA transfer on Windows") p.add_argument("--load_dtype", type=str, default="auto", choices=["auto", "fp32", "fp16"], help="Weight dtype during loading; auto uses fp16 on Windows") p.add_argument("--data_bin", type=str, default=None, help="Directory containing val.bin/val.meta.json or train.bin/train.meta.json") p.add_argument("--batch_size", type=int, default=4) p.add_argument("--max_batches", type=int, default=128, help="Maximum number of eval batches") p.add_argument("--seq_len", type=int, default=None, help="Override sequence length; default reads from checkpoint config") p.add_argument("--allow_train_tail_fallback", action="store_true", help="If val.bin is absent, evaluate on the tail of train.bin") p.add_argument("--tail_ratio", type=float, default=0.01, help="When falling back to train.bin, reserve this tail fraction") p.add_argument("--tail_tokens_min", type=int, default=2_000_000, help="Minimum number of tokens for pseudo-val region") p.add_argument("--generate", action="store_true", help="Run generation quality tests after PPL evaluation") p.add_argument("--generate_only", action="store_true", help="Skip PPL evaluation and run generation only") p.add_argument("--prompts", type=str, nargs="*", default=None, help="Custom generation prompts appended after the built-in set") p.add_argument("--max_new_tokens", type=int, default=128, help="Maximum number of generated tokens per prompt") p.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature") p.add_argument("--top_k", type=int, default=50, help="Top-k sampling cutoff") p.add_argument("--top_p", type=float, default=0.9, help="Top-p nucleus sampling threshold") p.add_argument("--repetition_penalty", type=float, default=1.1, help="Repetition penalty (1.0 disables the penalty)") p.add_argument("--json", action="store_true", help="Print final eval summary as JSON") return p.parse_args() def load_checkpoint(checkpoint_path: str, device: str = "cpu") -> dict: """Load a checkpoint while supporting both older and newer formats.""" path = Path(checkpoint_path).resolve() if not path.exists(): raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}") print(f"[Load] Loading checkpoint: {path}") ckpt = torch.load(str(path), map_location=device, weights_only=True, mmap=True) return ckpt def load_checkpoint_metadata(checkpoint_path: str) -> dict: """Load checkpoint metadata only, keeping tensor storages as lightweight refs.""" path = Path(checkpoint_path).resolve() if not path.exists(): raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}") with zipfile.ZipFile(path) as archive: data_pkl = next( name for name in archive.namelist() if name.endswith("/data.pkl") ) payload = archive.read(data_pkl) return _CheckpointMetadataUnpickler(io.BytesIO(payload)).load() def build_config_from_checkpoint(ckpt: dict, args) -> Tuple[SymbolicLightConfig, int]: """Restore the model configuration from a checkpoint.""" global_step = int(ckpt.get("global_step", ckpt.get("step", 0))) if "config" in ckpt: cfg_dict = ckpt["config"] if isinstance(cfg_dict, dict): valid_keys = {f.name for f in SymbolicLightConfig.__dataclass_fields__.values()} filtered = {k: v for k, v in cfg_dict.items() if k in valid_keys} config = SymbolicLightConfig(**filtered) else: config = cfg_dict else: print("[Load] [WARN] checkpoint has no config field; using the default 0.8B configuration") config = SymbolicLightConfig(**DEFAULT_08B_CONFIG) if args.seq_len is not None: config.max_seq_len = int(args.seq_len) return config, global_step def prepare_model_state(ckpt: dict) -> dict: """Prepare a checkpoint state_dict for loading into SymbolicLightModel.""" raw_state = ckpt.get("model", ckpt.get("model_state_dict", None)) if raw_state is None: raw_state = {k: v for k, v in ckpt.items() if isinstance(v, torch.Tensor)} cleaned = {} for k, v in raw_state.items(): new_k = k.replace("module.", "").replace("_orig_mod.", "") cleaned[new_k] = v cleaned = {k: v for k, v in cleaned.items() if "v_mem" not in k} return cleaned def load_model_state(model: SymbolicLightModel, ckpt: dict, assign: bool = False) -> dict: """Load model weights with DDP and torch.compile key cleanup.""" cleaned = prepare_model_state(ckpt) incompatible = model.load_state_dict(cleaned, strict=False, assign=assign) return { "missing_keys": list(incompatible.missing_keys), "unexpected_keys": list(incompatible.unexpected_keys), } def build_model_from_checkpoint(config: SymbolicLightConfig, ckpt: dict) -> Tuple[SymbolicLightModel, dict]: """Build the 0.8B model without duplicating parameter storage during load.""" with torch.device("meta"): model = SymbolicLightModel(config) load_info = load_model_state(model, ckpt, assign=True) return model, load_info def _dtype_from_storage(storage_dtype: str) -> torch.dtype: mapping = { "FloatStorage": torch.float32, "HalfStorage": torch.float16, "BFloat16Storage": torch.bfloat16, "DoubleStorage": torch.float64, "BoolStorage": torch.bool, "ByteStorage": torch.uint8, "CharStorage": torch.int8, "ShortStorage": torch.int16, "IntStorage": torch.int32, "LongStorage": torch.int64, } if storage_dtype not in mapping: raise TypeError(f"Unsupported checkpoint storage dtype: {storage_dtype}") return mapping[storage_dtype] def _numpy_dtype_from_storage(storage_dtype: str): mapping = { "FloatStorage": np.float32, "HalfStorage": np.float16, "DoubleStorage": np.float64, "BoolStorage": np.bool_, "ByteStorage": np.uint8, "CharStorage": np.int8, "ShortStorage": np.int16, "IntStorage": np.int32, "LongStorage": np.int64, } if storage_dtype not in mapping: raise TypeError(f"Unsupported checkpoint storage dtype: {storage_dtype}") return mapping[storage_dtype] def _checkpoint_zip_prefix(archive: zipfile.ZipFile) -> str: data_pkl = next(name for name in archive.namelist() if name.endswith("/data.pkl")) return data_pkl.rsplit("/", 1)[0] def _zip_stored_data_offset(archive: zipfile.ZipFile, info: zipfile.ZipInfo) -> int: archive.fp.seek(info.header_offset) header = archive.fp.read(30) if len(header) != 30 or header[:4] != b"PK\x03\x04": raise zipfile.BadZipFile(f"Invalid local header for {info.filename}") name_len, extra_len = struct.unpack(" Optional[torch.dtype]: if load_dtype == "auto": return torch.float16 if os.name == "nt" else None if load_dtype == "fp16": return torch.float16 return None def _tensor_from_zip_ref( archive: zipfile.ZipFile, prefix: str, ref: _TensorRef, target_float_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: storage = ref.storage np_dtype = _numpy_dtype_from_storage(storage.dtype) info = archive.getinfo(f"{prefix}/data/{storage.key}") if info.compress_type == zipfile.ZIP_STORED: data_offset = _zip_stored_data_offset(archive, info) array = np.memmap( archive.filename, mode="r", dtype=np_dtype, offset=data_offset, shape=(storage.size,), ) else: raw = archive.read(info.filename) array = np.frombuffer(raw, dtype=np_dtype).copy() with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="The given NumPy array is not writable") tensor = torch.from_numpy(array) if ref.offset: tensor = tensor[ref.offset:] tensor = tensor.as_strided(ref.size, ref.stride) if target_float_dtype is not None and tensor.is_floating_point(): tensor = tensor.to(target_float_dtype) else: tensor = tensor.clone() return tensor def _set_module_tensor(model: nn.Module, name: str, tensor: torch.Tensor) -> bool: module_name, _, leaf = name.rpartition(".") module = model.get_submodule(module_name) if module_name else model if leaf in module._parameters: old_param = module._parameters[leaf] requires_grad = True if old_param is None else old_param.requires_grad module._parameters[leaf] = nn.Parameter(tensor, requires_grad=requires_grad) return True if leaf in module._buffers: module._buffers[leaf] = tensor return True return False def build_model_from_checkpoint_zip( config: SymbolicLightConfig, checkpoint_path: str, ckpt_metadata: dict, target_float_dtype: Optional[torch.dtype] = None, ) -> Tuple[SymbolicLightModel, dict]: """Load model weights directly from the checkpoint zip storages. This avoids a Windows PyTorch mmap failure observed with large CUDA-origin checkpoints where metadata loads correctly but tensor access segfaults. """ raw_state = ckpt_metadata.get("model", ckpt_metadata.get("model_state_dict", None)) if raw_state is None: raise KeyError("Checkpoint does not contain a model state_dict") cleaned = { k.replace("module.", "").replace("_orig_mod.", ""): v for k, v in raw_state.items() if "v_mem" not in k } with torch.device("meta"): model = SymbolicLightModel(config) expected_keys = set(model.state_dict().keys()) loaded_keys = set() unexpected_keys = [] path = Path(checkpoint_path).resolve() with zipfile.ZipFile(path) as archive: prefix = _checkpoint_zip_prefix(archive) for name, ref in cleaned.items(): if name not in expected_keys: unexpected_keys.append(name) continue tensor = _tensor_from_zip_ref( archive, prefix, ref, target_float_dtype=target_float_dtype, ) if _set_module_tensor(model, name, tensor): loaded_keys.add(name) else: unexpected_keys.append(name) return model, { "missing_keys": sorted(expected_keys - loaded_keys), "unexpected_keys": sorted(unexpected_keys), } def resolve_eval_region(data_bin_dir: Path, allow_train_tail_fallback: bool, tail_ratio: float, tail_tokens_min: int): """Locate the evaluation region in val.bin or a fallback train tail.""" val_bin = data_bin_dir / "val.bin" val_meta = data_bin_dir / "val.meta.json" if val_bin.exists() and val_meta.exists(): with open(val_meta, encoding="utf-8") as f: meta = json.load(f) return val_bin, meta, 0, int(meta["total_tokens"]), "held-out val.bin" if not allow_train_tail_fallback: raise FileNotFoundError( "val.bin / val.meta.json not found. " "Provide a separate validation split or pass --allow_train_tail_fallback." ) train_bin = data_bin_dir / "train.bin" train_meta = data_bin_dir / "train.meta.json" if not (train_bin.exists() and train_meta.exists()): raise FileNotFoundError(f"train.bin / train.meta.json does not exist: {data_bin_dir}") with open(train_meta, encoding="utf-8") as f: meta = json.load(f) total_tokens = int(meta["total_tokens"]) tail_tokens = max(int(total_tokens * tail_ratio), int(tail_tokens_min)) tail_tokens = min(tail_tokens, total_tokens) start_idx = max(0, total_tokens - tail_tokens) return train_bin, meta, start_idx, total_tokens, "pseudo-val tail of train.bin" def run_ppl_eval(model, config, args, device): """Compute CE / PPL on a memmap dataset.""" bin_path, meta, region_start, region_end, split_desc = resolve_eval_region( Path(args.data_bin), allow_train_tail_fallback=args.allow_train_tail_fallback, tail_ratio=args.tail_ratio, tail_tokens_min=args.tail_tokens_min, ) dtype_str = meta.get("dtype", "uint16") if dtype_str != "uint16": raise ValueError(f"This evaluation path only supports uint16 memmap data, got {dtype_str}") seq_len = int(config.max_seq_len) usable_tokens = max(0, region_end - region_start) available_windows = max(0, (usable_tokens - 1) // seq_len) eval_batches = min(args.max_batches, available_windows // max(1, args.batch_size)) if eval_batches <= 0: raise ValueError( f"Not enough usable validation windows: usable_tokens={usable_tokens}, " f"seq_len={seq_len}, batch_size={args.batch_size}" ) data = np.memmap(bin_path, dtype=np.uint16, mode="r") autocast_enabled = device.type == "cuda" if not args.json: print(f"Eval split: {split_desc}") print(f"Eval region: [{region_start:,} - {region_end:,}] ({usable_tokens:,} tokens)") if "pseudo-val" in split_desc: print("[WARNING] Using the tail of train.bin as pseudo-validation; this is not a strict held-out split.") print(f"Seq len: {seq_len}") print(f"Batches: {eval_batches} x batch_size={args.batch_size}") total_loss = 0.0 total_tokens = 0 t_start = time.time() with torch.no_grad(): for batch_idx in range(eval_batches): batch_tokens = [] base_offset = region_start + batch_idx * args.batch_size * seq_len for item_idx in range(args.batch_size): offset = base_offset + item_idx * seq_len if offset + seq_len + 1 > region_end: break tokens = data[offset:offset + seq_len + 1].astype(np.int64) batch_tokens.append(tokens) if not batch_tokens: break batch = torch.tensor(np.array(batch_tokens), dtype=torch.long, device=device) batch = batch.clamp(0, config.vocab_size - 1) inputs = batch[:, :-1] targets = batch[:, 1:] if autocast_enabled: with torch.autocast(device_type="cuda", dtype=torch.bfloat16): logits = model(inputs) else: logits = model(inputs) loss_sum = F.cross_entropy( logits.float().reshape(-1, config.vocab_size), targets.reshape(-1), reduction="sum", ) total_loss += float(loss_sum.item()) total_tokens += int(targets.numel()) if not args.json and (batch_idx == 0 or (batch_idx + 1) % 20 == 0): mean_ce = total_loss / max(total_tokens, 1) ppl = math.exp(min(mean_ce, 20.0)) print(f"[Eval] batch {batch_idx+1}/{eval_batches} | " f"CE={mean_ce:.4f} | PPL={ppl:.1f}") elapsed = time.time() - t_start mean_ce = total_loss / max(total_tokens, 1) ppl = math.exp(min(mean_ce, 20.0)) sparsity_stats = {} try: sparsity_stats = model.get_sparsity_stats() except Exception: pass return { "split": split_desc, "mean_ce": float(mean_ce), "ppl": float(ppl), "tokens": int(total_tokens), "batches": int(eval_batches), "elapsed_sec": round(elapsed, 2), "sparsity": sparsity_stats, } DEFAULT_GENERATION_PROMPTS = [ "Artificial intelligence can improve society by", "A good scientific explanation should", "Once upon a time, in a small village,", "def fibonacci(n):\n ", "The Earth completes one orbit around the Sun in about", "Photosynthesis is the process by which", "In a world shaped by rapid scientific progress,", "The capital of France is", ] def run_generation_test(model, tokenizer, args, device, global_step, config): """Run generation quality tests with cached incremental decoding.""" prompts = list(DEFAULT_GENERATION_PROMPTS) if args.prompts: prompts.extend(args.prompts) print("\n" + "=" * 60) print(f" Generation Quality Test (step {global_step})") print(f" temperature={args.temperature} | top_k={args.top_k} | " f"rep_penalty={args.repetition_penalty} | " f"max_tokens={args.max_new_tokens}") print("=" * 60) results = [] for i, prompt in enumerate(prompts): ids = tokenizer.encode(prompt) input_ids = torch.tensor([ids], dtype=torch.long, device=device) prompt_len = input_ids.size(1) t_start = time.time() output_ids = model.generate( input_ids, max_new_tokens=args.max_new_tokens, temperature=args.temperature, top_k=args.top_k, ) t_elapsed = time.time() - t_start generated_ids = output_ids[0].tolist() generated_text = tokenizer.decode(generated_ids) new_tokens = len(generated_ids) - prompt_len tok_per_sec = new_tokens / max(t_elapsed, 1e-6) print(f"\n{'─'*60}") print(f"[{i+1}/{len(prompts)}] Prompt: {prompt}") print(f"{'─'*60}") print(generated_text) print(f" ({new_tokens} new tokens | {tok_per_sec:.1f} tok/s | {t_elapsed:.2f}s)") results.append({ "prompt": prompt, "generated": generated_text, "new_tokens": new_tokens, "tok_per_sec": round(tok_per_sec, 1), }) print(f"\n{'='*60}") print(f" Generation test complete: {len(results)} prompts") print(f"{'='*60}") return results class TokenizerWrapper: """Small compatibility wrapper around the available tokenizer implementations.""" def __init__(self, tokenizer_path: str): self.path = tokenizer_path try: from train_tokenizer import SLTokenizer self.tok = SLTokenizer(tokenizer_path) self._type = "sl" print(f"[Tokenizer] SLTokenizer loaded: {tokenizer_path}") except Exception: import sentencepiece as spm self.tok = spm.SentencePieceProcessor(model_file=tokenizer_path) self._type = "sp" print(f"[Tokenizer] SentencePiece loaded: {tokenizer_path}") def encode(self, text: str) -> list: if self._type == "sl": return self.tok.encode(text, add_bos=True, add_eos=False) return self.tok.encode(text) def decode(self, ids: list) -> str: if self._type == "sl": return self.tok.decode(ids) return self.tok.decode(ids) def eos_id(self) -> int: if self._type == "sl": return self.tok.sp.eos_id() return self.tok.eos_id() def main(): args = parse_args() device = torch.device( "cuda" if args.device == "cuda" and torch.cuda.is_available() else "cpu" ) if os.name == "nt" and device.type == "cuda" and not args.allow_windows_cuda: print("[Device] Windows CUDA path is experimental for this checkpoint; using CPU for safe smoke tests.") print("[Device] Re-run with --allow_windows_cuda only if you want to test the CUDA transfer path.") device = torch.device("cpu") if os.name == "nt": print("[Load] Windows detected: using zip-storage checkpoint loader") target_float_dtype = _resolve_load_dtype(args.load_dtype) if target_float_dtype is not None: print(f"[Load] Loading floating-point weights as {target_float_dtype}") ckpt = load_checkpoint_metadata(args.checkpoint_path) config, global_step = build_config_from_checkpoint(ckpt, args) model, load_info = build_model_from_checkpoint_zip( config, args.checkpoint_path, ckpt, target_float_dtype=target_float_dtype, ) else: ckpt = load_checkpoint(args.checkpoint_path, device="cpu") config, global_step = build_config_from_checkpoint(ckpt, args) model, load_info = build_model_from_checkpoint(config, ckpt) model.to(device) model.eval() del ckpt n_params = sum(p.numel() for p in model.parameters()) if not args.json: print("=" * 60) print(" SymbolicLight 0.8B Evaluation") print("=" * 60) print(f"Checkpoint: {args.checkpoint_path}") print(f"Global step: {global_step}") print(f"Parameters: {n_params/1e6:.1f}M ({n_params/1e9:.3f}B)") print(f"Config: embed_dim={config.embed_dim}, n_layers={config.n_layers}, " f"n_heads={config.n_heads}, seq_len={config.max_seq_len}") print(f"Device: {device}") if load_info["missing_keys"]: print(f"[WARN] Missing keys: {load_info['missing_keys'][:5]}...") if load_info["unexpected_keys"]: print(f"[WARN] Unexpected keys: {load_info['unexpected_keys'][:5]}...") if args.generate_only: tokenizer = TokenizerWrapper(args.tokenizer_path) gen_results = run_generation_test( model, tokenizer, args, device, global_step, config ) if args.json: summary = { "checkpoint": args.checkpoint_path, "global_step": global_step, "parameters": n_params, "generation_results": gen_results, } print(json.dumps(summary, ensure_ascii=False, indent=2)) return if not args.data_bin: raise ValueError( "PPL evaluation requires --data_bin. " "If you only want generation tests, pass --generate_only." ) eval_result = run_ppl_eval(model, config, args, device) if not args.json: print("-" * 60) print(f"[Eval Done] mean_CE={eval_result['mean_ce']:.4f} | " f"PPL={eval_result['ppl']:.2f} | " f"tokens={eval_result['tokens']:,} | " f"time={eval_result['elapsed_sec']}s") if eval_result["sparsity"]: print("[Sparsity]") for k, v in eval_result["sparsity"].items(): print(f" {k}: {v*100:.1f}% silent") gen_results = None if args.generate: tokenizer = TokenizerWrapper(args.tokenizer_path) gen_results = run_generation_test( model, tokenizer, args, device, global_step, config ) if args.json: summary = { "checkpoint": args.checkpoint_path, "global_step": global_step, "parameters": n_params, **eval_result, } if gen_results: summary["generation_results"] = gen_results print(json.dumps(summary, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()