Text Generation
Transformers
Safetensors
OpenPeerLLM
English
Arabic
riemann_openpeer
amd
ntk
code
fsdp
rocm
coding
mi325x
decoder
code gen
directml
openpeer
ryzen-ai
reasoning
rewriting
multimodal
paraphrase
prerelease
ntk-trainer
growth-ladder
style-transfer
text-rewriting
riemann-computing
progressive-growth
custom_code
Eval Results (legacy)
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import logging | |
| import random | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| import torch | |
| from transformers import AutoTokenizer | |
| from riemann.growth_ladder import get_stage | |
| from riemann.modeling_riemann import RiemannForCausalLM | |
| from riemann.device import main_device, cpu_device | |
| from riemann.paths import DATA_DIR, CHECKPOINT_ROOT, CONTROLLER_ROOT | |
| logger = logging.getLogger("riemann.fit_ntk") | |
| TOKENIZER_NAME = "gpt2" | |
| def load_examples(path: Path, limit: int | None, seed: int) -> list: | |
| from ntkmirror import load_jsonl_examples | |
| examples = list(load_jsonl_examples(str(path))) | |
| if limit is not None and len(examples) > limit: | |
| rng = random.Random(seed) | |
| rng.shuffle(examples) | |
| examples = examples[:limit] | |
| return examples | |
| def split_validation(examples: list, fraction: float, seed: int) -> tuple[list, list]: | |
| if fraction <= 0 or len(examples) < 20: | |
| return examples, [] | |
| rng = random.Random(seed) | |
| idx = list(range(len(examples))) | |
| rng.shuffle(idx) | |
| n_val = max(8, int(len(examples) * fraction)) | |
| val_idx = set(idx[:n_val]) | |
| train = [e for i, e in enumerate(examples) if i not in val_idx] | |
| val = [e for i, e in enumerate(examples) if i in val_idx] | |
| return train, val | |
| def warm_start_from(tuner, prior_path: Path, device, examples: list, args) -> dict: | |
| from ntkmirror.controller import SignedLogMaskState, _SignedLogMaskModule | |
| prior = SignedLogMaskState.load(str(prior_path), map_location="cpu") | |
| n_layers = len(tuner.decoder_layers) | |
| hidden = tuner.hidden_size | |
| keep = (prior.layer_indices < n_layers) & (prior.channel_indices < hidden) | |
| layer_idx = prior.layer_indices[keep] | |
| chan_idx = prior.channel_indices[keep] | |
| raw = prior.raw[keep] | |
| fresh_layer_idx, fresh_chan_idx = tuner._score_gates( | |
| examples, | |
| score_batches=min(16, max(1, len(examples) // max(1, args.batch_size))), | |
| batch_size=args.batch_size, | |
| max_length=args.max_length, | |
| ) | |
| fresh_layer_idx = fresh_layer_idx.cpu() | |
| fresh_chan_idx = fresh_chan_idx.cpu() | |
| existing = set(zip(layer_idx.tolist(), chan_idx.tolist())) | |
| new_pairs = [ | |
| (l, c) | |
| for l, c in zip(fresh_layer_idx.tolist(), fresh_chan_idx.tolist()) | |
| if (l, c) not in existing and (l >= prior.n_layers or c >= prior.hidden_size) | |
| ] | |
| budget = max(0, args.gates - layer_idx.numel()) | |
| new_pairs = new_pairs[:budget] | |
| if new_pairs: | |
| nl = torch.tensor([p[0] for p in new_pairs], dtype=torch.long) | |
| nc = torch.tensor([p[1] for p in new_pairs], dtype=torch.long) | |
| layer_idx = torch.cat([layer_idx, nl]) | |
| chan_idx = torch.cat([chan_idx, nc]) | |
| raw = torch.cat([raw, torch.zeros(len(new_pairs), dtype=raw.dtype)]) | |
| if tuner.controller is not None: | |
| tuner.controller.remove() | |
| tuner.controller = _SignedLogMaskModule( | |
| tuner.decoder_layers, | |
| layer_idx.to(device), | |
| chan_idx.to(device), | |
| hidden_size=hidden, | |
| max_log_gate=tuner.max_log_gate, | |
| hook_site=tuner.hook_site, | |
| raw_init=raw.to(device), | |
| ).to(device) | |
| def _keep_warm_controller(*_a, **_k): | |
| return { | |
| "selected_gates": float(tuner.controller.raw.numel()), | |
| "select_seconds": 0.0, | |
| "hook_site": tuner.hook_site, | |
| } | |
| tuner.initialize_controller = _keep_warm_controller | |
| return { | |
| "prior_gates": int(prior.n_gates), | |
| "carried_gates": int(keep.sum()), | |
| "new_gates_for_grown_dims": len(new_pairs), | |
| "total_gates": int(layer_idx.numel()), | |
| "prior_layers": int(prior.n_layers), | |
| "prior_hidden": int(prior.hidden_size), | |
| } | |
| def main() -> None: | |
| logging.basicConfig( | |
| level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" | |
| ) | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--stage", default="stage0_124m") | |
| parser.add_argument("--gates", type=int, default=4096) | |
| parser.add_argument("--steps", type=int, default=240) | |
| parser.add_argument("--lr", type=float, default=5e-3) | |
| parser.add_argument("--batch-size", type=int, default=4) | |
| parser.add_argument("--max-length", type=int, default=512) | |
| parser.add_argument("--max-examples", type=int, default=20000) | |
| parser.add_argument("--val-fraction", type=float, default=0.05) | |
| parser.add_argument("--val-interval", type=int, default=20) | |
| parser.add_argument("--early-stop-patience", type=int, default=3) | |
| parser.add_argument("--kl-to-base", type=float, default=0.0) | |
| parser.add_argument("--seed", type=int, default=1234) | |
| parser.add_argument("--init-from", default=None) | |
| parser.add_argument("--cpu", action="store_true") | |
| parser.add_argument("--train-jsonl", default=str(DATA_DIR / "ntk_finetune.jsonl")) | |
| parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT)) | |
| parser.add_argument("--controller-root", default=str(CONTROLLER_ROOT)) | |
| parser.add_argument("--out", default=None) | |
| args = parser.parse_args() | |
| stage = get_stage(args.stage) | |
| checkpoint_dir = Path(args.checkpoint_root) / stage.name | |
| if not (checkpoint_dir / "config.json").exists(): | |
| raise FileNotFoundError( | |
| f"No trained checkpoint at {checkpoint_dir}. Run pretrain.py for {stage.name} first." | |
| ) | |
| controller_root = Path(args.controller_root) | |
| controller_root.mkdir(parents=True, exist_ok=True) | |
| out_path = Path(args.out) if args.out else controller_root / f"{stage.name}.pt" | |
| train_path = Path(args.train_jsonl) | |
| if not train_path.exists(): | |
| raise FileNotFoundError(f"{train_path} not found. Run riemann.data first.") | |
| try: | |
| from ntkmirror import ForwardFineTuner | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "ntkmirror is required: pip install git+https://github.com/leochlon/ntkmirror.git" | |
| ) from exc | |
| device = cpu_device() if args.cpu else main_device() | |
| torch.manual_seed(args.seed) | |
| tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| logger.info("Loading %s from %s onto %s", stage.name, checkpoint_dir, device) | |
| model = RiemannForCausalLM.from_pretrained(str(checkpoint_dir)) | |
| model.config.use_cache = False | |
| model.to(device) | |
| model.eval() | |
| examples = load_examples(train_path, args.max_examples, args.seed) | |
| train_examples, val_examples = split_validation( | |
| examples, args.val_fraction, args.seed | |
| ) | |
| logger.info( | |
| "NTK fit data: %d train / %d validation examples (from %s)", | |
| len(train_examples), | |
| len(val_examples), | |
| train_path, | |
| ) | |
| tuner = ForwardFineTuner( | |
| model, | |
| tokenizer, | |
| gates=args.gates, | |
| model_name=f"my-name-is-riemann/{stage.name}", | |
| tokenizer_name=TOKENIZER_NAME, | |
| ) | |
| warm = None | |
| if args.init_from: | |
| prior_path = Path(args.init_from) | |
| if not prior_path.exists(): | |
| prior_path = controller_root / f"{args.init_from}.pt" | |
| if not prior_path.exists(): | |
| raise FileNotFoundError(f"No prior controller at {prior_path}") | |
| warm = warm_start_from(tuner, prior_path, device, train_examples, args) | |
| logger.info( | |
| "Warm start from %s: carried %d of %d gates, added %d for grown dims, %d total (%d layers, hidden %d -> %d layers, hidden %d)", | |
| prior_path.name, | |
| warm["carried_gates"], | |
| warm["prior_gates"], | |
| warm["new_gates_for_grown_dims"], | |
| warm["total_gates"], | |
| warm["prior_layers"], | |
| warm["prior_hidden"], | |
| len(tuner.decoder_layers), | |
| tuner.hidden_size, | |
| ) | |
| base_nll = None | |
| if val_examples: | |
| base_eval = tuner.evaluate_nll( | |
| val_examples, batch_size=args.batch_size, max_length=args.max_length | |
| ) | |
| base_nll = float(base_eval["nll"]) | |
| logger.info( | |
| "Validation before fitting: nll %.4f, token acc %.4f", | |
| base_nll, | |
| base_eval["token_acc"], | |
| ) | |
| t0 = time.time() | |
| stats = tuner.fit( | |
| train_examples, | |
| steps=args.steps, | |
| lr=args.lr, | |
| batch_size=args.batch_size, | |
| max_length=args.max_length, | |
| validation_examples=val_examples or None, | |
| validation_interval=args.val_interval if val_examples else None, | |
| early_stop_patience=args.early_stop_patience if val_examples else 0, | |
| select_best_on_validation=bool(val_examples), | |
| retain_examples=train_examples if args.kl_to_base > 0 else None, | |
| kl_to_base=args.kl_to_base, | |
| verbose=True, | |
| ) | |
| elapsed = time.time() - t0 | |
| final_nll = None | |
| if val_examples: | |
| final_eval = tuner.evaluate_nll( | |
| val_examples, batch_size=args.batch_size, max_length=args.max_length | |
| ) | |
| final_nll = float(final_eval["nll"]) | |
| logger.info( | |
| "Validation after fitting: nll %.4f (base %.4f, delta %+.4f), token acc %.4f", | |
| final_nll, | |
| base_nll, | |
| final_nll - base_nll, | |
| final_eval["token_acc"], | |
| ) | |
| metadata = { | |
| "stage": stage.name, | |
| "checkpoint": str(checkpoint_dir), | |
| "train_examples": len(train_examples), | |
| "validation_examples": len(val_examples), | |
| "steps": args.steps, | |
| "lr": args.lr, | |
| "gates_requested": args.gates, | |
| "fit_seconds": elapsed, | |
| "validation_nll_base": base_nll, | |
| "validation_nll_final": final_nll, | |
| "warm_start": warm, | |
| } | |
| tuner.save(str(out_path), metadata=metadata) | |
| summary_path = out_path.with_suffix(".json") | |
| summary = { | |
| k: (v if not isinstance(v, torch.Tensor) else v.tolist()) | |
| for k, v in stats.items() | |
| if not isinstance(v, list) | |
| } | |
| summary.update(metadata) | |
| summary_path.write_text( | |
| json.dumps(summary, indent=2, default=str), encoding="utf-8" | |
| ) | |
| logger.info("Saved gate controller to %s (summary at %s)", out_path, summary_path) | |
| if __name__ == "__main__": | |
| main() | |