| """0-shot lm_eval for LOOM checkpoints. |
| |
| Protocol: |
| official lm_eval.simple_evaluate, 0-shot, full-text continuation, |
| fully causal SUM log-likelihood, acc / acc_norm. |
| |
| Before scoring, the run's own eval/loss at the checkpoint step is recomputed (same |
| holdout batches as pretrain.run_eval) and must match to EVAL_HOLDOUT_MAX_GAP. A gap |
| there means the weights or code tree are wrong, not that the model is weak. |
| |
| Launch:: |
| |
| bash eval/run.sh /path/to/ckpt |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import shutil |
| import time |
| from contextlib import contextmanager |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
| import torch.nn.functional as F |
| from lm_eval.api.model import TemplateLM |
| from lm_eval.api.registry import register_model |
| from tqdm import tqdm |
|
|
| from datasets_cache import load_dataset, setup_hf_datasets_cache |
| from distributed_utils import ( |
| clear_prep_done_flag, |
| init_eval_distributed, |
| signal_prep_done, |
| wait_for_prep_done, |
| ) |
| from loom_loader import EvalCheckpoint, load_checkpoint_for_eval, provenance, set_param_dtype |
| from models.common import wrap_tensor |
| from models.prefixlm_attention import compute_aux_seq_tensors_scalars |
| from utils.device import get_device, set_device |
|
|
|
|
| DEFAULT_TASKS = ( |
| "openbookqa", |
| "winogrande", |
| "arc_challenge", |
| "arc_easy", |
| "hellaswag", |
| "social_iqa", |
| "piqa", |
| ) |
|
|
|
|
| class _DistAccel: |
| """Minimal Accelerator stand-in for lm_eval's gather / barrier calls.""" |
|
|
| def __init__(self, rank: int, world_size: int, local_rank: int, device: torch.device): |
| self.process_index = rank |
| self.local_process_index = local_rank |
| self.num_processes = world_size |
| self.device = device |
|
|
| @property |
| def is_local_main_process(self) -> bool: |
| return self.local_process_index == 0 |
|
|
| def gather(self, tensor: torch.Tensor) -> torch.Tensor: |
| if self.num_processes <= 1: |
| return tensor.unsqueeze(0) if tensor.ndim == 0 else tensor |
| src = tensor.detach().contiguous() |
| gathered = [torch.empty_like(src) for _ in range(self.num_processes)] |
| dist.all_gather(gathered, src) |
| return torch.stack(gathered) |
|
|
| def wait_for_everyone(self) -> None: |
| if self.num_processes > 1 and dist.is_initialized(): |
| dist.barrier() |
|
|
|
|
| def _build_fully_causal_ll_batch( |
| inp_toks: list[np.ndarray], |
| device: torch.device, |
| ) -> dict[str, torch.Tensor]: |
| """Pack sequences with prefix_lens=0 so attention is fully causal (lm_eval).""" |
| prefix_lens = np.zeros(len(inp_toks), dtype=np.int32) |
| causal_lens = np.asarray([int(x.size) for x in inp_toks], dtype=np.int32) |
| if causal_lens.size == 0 or int(causal_lens.max(initial=0)) <= 0: |
| raise ValueError("empty causal batch") |
| n_seq = int(causal_lens.shape[0]) |
| batch_max_tokens = max(int(causal_lens.max()), n_seq + 1, 1) |
| flat_inputs = np.concatenate(inp_toks) |
| flat_pos = np.concatenate([np.arange(int(x.size), dtype=np.int32) for x in inp_toks]) |
| tensors, scalars = compute_aux_seq_tensors_scalars( |
| prefix_lens, causal_lens, batch_max_tokens |
| ) |
| batch: dict[str, torch.Tensor] = { |
| "inputs": torch.from_numpy(flat_inputs).to(device, non_blocking=True), |
| "position_ids": torch.from_numpy(flat_pos).to(device, non_blocking=True), |
| } |
| for key, value in tensors.items(): |
| batch[key] = torch.from_numpy(value).to(device, non_blocking=True) |
| for key, value in scalars.items(): |
| batch[key] = wrap_tensor(torch.tensor(value, device="cpu")) |
| return batch |
|
|
|
|
| @register_model("loom") |
| class LoomLM(TemplateLM): |
| """TemplateLM adapter: official causal SUM log-likelihood over packed PrefixLM=0.""" |
|
|
| def __init__( |
| self, |
| ckpt: EvalCheckpoint, |
| *, |
| rank: int, |
| world_size: int, |
| local_rank: int, |
| batch_size: int = 32, |
| max_length: int = 1024, |
| ): |
| super().__init__() |
| self._ckpt = ckpt |
| self.tokenizer = ckpt.tokenizer |
| self._device = get_device(local_rank) |
| self._rank = rank |
| self._world_size = world_size |
| self.batch_size = int(batch_size) |
| self._max_length = int(max_length) |
| self.accelerator = _DistAccel(rank, world_size, local_rank, self._device) |
| self.backend = "causal" |
|
|
| @property |
| def eot_token_id(self) -> int: |
| eos = self.tokenizer.eos_token_id |
| if eos is None: |
| raise ValueError("tokenizer has no eos_token_id") |
| return int(eos) |
|
|
| @property |
| def prefix_token_id(self) -> int: |
| bos = self.tokenizer.bos_token_id |
| return int(bos) if bos is not None else self.eot_token_id |
|
|
| @property |
| def max_length(self) -> int: |
| return self._max_length |
|
|
| @property |
| def device(self) -> torch.device: |
| return self._device |
|
|
| def tok_encode( |
| self, |
| string: str, |
| add_special_tokens: bool | None = None, |
| **kwargs, |
| ) -> list[int]: |
| |
| del add_special_tokens |
| return list(self.tokenizer.encode(string, add_special_tokens=False)) |
|
|
| @torch.inference_mode() |
| def _loglikelihood_tokens( |
| self, |
| requests: list[tuple[tuple[str, str], list[int], list[int]]], |
| disable_tqdm: bool = False, |
| **kwargs, |
| ) -> list[tuple[float, bool]]: |
| results: list[tuple[float, bool]] = [(0.0, False)] * len(requests) |
| batch_size = max(int(self.batch_size), 1) |
| max_len = self.max_length |
| device = self.device |
| pbar = tqdm( |
| total=len(requests), |
| disable=disable_tqdm or (self.rank != 0), |
| desc="lm_eval loglikelihood", |
| ) |
| for start in range(0, len(requests), batch_size): |
| chunk = requests[start : start + batch_size] |
| inps: list[np.ndarray] = [] |
| conts: list[np.ndarray] = [] |
| valid_idx: list[int] = [] |
| for i, (_pair, ctx, cont) in enumerate(chunk): |
| if not ctx or not cont: |
| continue |
| if len(cont) > max_len: |
| continue |
| whole = list(ctx) + list(cont) |
| inp = whole[-(max_len + 1) :][:-1] |
| if not inp: |
| continue |
| inps.append(np.asarray(inp, dtype=np.int32)) |
| conts.append(np.asarray(cont, dtype=np.int32)) |
| valid_idx.append(i) |
| if not inps: |
| pbar.update(len(chunk)) |
| continue |
| batch = _build_fully_causal_ll_batch(inps, device) |
| _, logits = self._ckpt.model(carry=None, batch=batch) |
| cu = batch["cu_seqlens"].tolist() |
| for j, i in enumerate(valid_idx): |
| cl = int(conts[j].size) |
| end = cu[j] + int(inps[j].size) |
| rows = logits[end - cl : end].float() |
| cont = torch.from_numpy(conts[j]).to(device=device, dtype=torch.long) |
| logp = torch.log_softmax(rows, dim=-1).gather(-1, cont[:, None]).squeeze(-1) |
| greedy_ok = bool((rows.argmax(dim=-1) == cont).all().item()) |
| results[start + i] = (float(logp.double().sum().item()), greedy_ok) |
| pbar.update(len(chunk)) |
| pbar.close() |
| return results |
|
|
| def loglikelihood_rolling(self, requests, disable_tqdm: bool = False) -> list[float]: |
| raise NotImplementedError("LoomLM only implements multiple-choice loglikelihood") |
|
|
| def generate_until(self, requests, disable_tqdm: bool = False) -> list[str]: |
| raise NotImplementedError("LoomLM only implements multiple-choice loglikelihood") |
|
|
|
|
| _task_dir_env = os.environ.get("LMEVAL_TASK_DIR", "").strip() |
| _LOCAL_TASK_DIR = ( |
| Path(_task_dir_env) if _task_dir_env else Path(__file__).resolve().parent / "tasks" |
| ) |
|
|
|
|
| def _parse_tasks() -> list[str]: |
| raw = os.environ.get("LMEVAL_TASKS", ",".join(DEFAULT_TASKS)).strip() |
| return [t.strip() for t in raw.split(",") if t.strip()] |
|
|
|
|
| _DATASET_WARMUP = ( |
| ("allenai/openbookqa", "main", "train"), |
| ("allenai/openbookqa", "main", "test"), |
| ("allenai/winogrande", "winogrande_debiased", "train"), |
| ("allenai/winogrande", "winogrande_debiased", "validation"), |
| ("allenai/ai2_arc", "ARC-Challenge", "train"), |
| ("allenai/ai2_arc", "ARC-Challenge", "test"), |
| ("allenai/ai2_arc", "ARC-Easy", "train"), |
| ("allenai/ai2_arc", "ARC-Easy", "test"), |
| ("Rowan/hellaswag", None, "train"), |
| ("Rowan/hellaswag", None, "validation"), |
| ("jet-ai/social_i_qa", None, "train"), |
| ("jet-ai/social_i_qa", None, "validation"), |
| ("baber/piqa", None, "train"), |
| ("baber/piqa", None, "validation"), |
| ) |
|
|
|
|
| def _warmup_hf_datasets(tasks: list[str], task_manager, num_fewshot: int | None) -> None: |
| """Materialize raw + process_docs arrow caches on rank 0 (avoid N-way races on a shared filesystem).""" |
| for path, name, split in _DATASET_WARMUP: |
| print(f"[lm_eval_loom] warmup {path} {name} {split}", flush=True) |
| kwargs = {"split": split, "trust_remote_code": True} |
| ds = load_dataset(path, name, **kwargs) if name else load_dataset(path, **kwargs) |
| print(f"[lm_eval_loom] warmup ok n={len(ds)}", flush=True) |
| from lm_eval.tasks import get_task_dict |
|
|
| shots = 0 if num_fewshot is None else int(num_fewshot) |
| task_dict = get_task_dict(tasks, task_manager) |
| for name, task in task_dict.items(): |
| if shots: |
| task.set_config(key="num_fewshot", value=shots) |
| print(f"[lm_eval_loom] warmup build_all_requests {name} shots={shots}", flush=True) |
| task.build_all_requests(limit=None, rank=0, world_size=1) |
| print(f"[lm_eval_loom] warmup task {name} n={len(task.instances)}", flush=True) |
|
|
|
|
| _LOCAL_DS_CACHE = Path(os.environ.get("EVAL_LOCAL_DS_CACHE", "/tmp/loom_lmeval_hf_datasets")) |
| _SHARED_DS_NAMES = ( |
| "allenai___openbookqa", |
| "allenai___winogrande", |
| "allenai___ai2_arc", |
| "Rowan___hellaswag", |
| "jet-ai___social_i_qa", |
| "baber___piqa", |
| ) |
|
|
|
|
| def _stage_local_dataset_cache(local_rank: int) -> None: |
| """Copy warmed arrow caches onto node-local disk so local ranks do not write the shared filesystem.""" |
| src_root = Path(os.environ.get("HF_DATASETS_CACHE", str(_LOCAL_DS_CACHE))) |
| dest = _LOCAL_DS_CACHE |
| ready = dest / ".ready" |
| if local_rank == 0: |
| dest.mkdir(parents=True, exist_ok=True) |
| for name in _SHARED_DS_NAMES: |
| src = src_root / name |
| dst = dest / name |
| if src.is_dir() and src.resolve() != dst.resolve() and not dst.exists(): |
| print(f"[lm_eval_loom] stage {name} -> {dst}", flush=True) |
| |
| tmp = dest / f".{name}.{os.getpid()}" |
| shutil.copytree(src, tmp) |
| try: |
| os.rename(tmp, dst) |
| except OSError: |
| shutil.rmtree(tmp, ignore_errors=True) |
| ready.touch() |
| else: |
| while not ready.exists(): |
| time.sleep(0.5) |
| os.environ["HF_DATASETS_CACHE"] = str(dest) |
| import datasets.config as ds_config |
|
|
| ds_config.HF_DATASETS_CACHE = dest |
| print(f"[lm_eval_loom] HF_DATASETS_CACHE={dest}", flush=True) |
|
|
|
|
| def _parse_num_fewshot() -> int | None: |
| raw = os.environ.get("EVAL_NUM_FEWSHOT", "0").strip() |
| if raw in {"", "auto", "none", "None"}: |
| return None |
| return int(raw) |
|
|
|
|
| def _train_eval_loss_at(ckpt: EvalCheckpoint) -> float | None: |
| """The run's own eval/loss at the checkpoint step, from tensorboard_scalars.json |
| (in the ckpt dir, its config dir, or the run a snapshot / archive was taken from).""" |
| step = ckpt.info.get("step") |
| roots = [Path(ckpt.info["ckpt_path"]), Path(ckpt.info["config_dir"])] |
| for meta, key in (("snapshot_meta.json", "source_checkpoint_path"), ("final_meta.json", "source")): |
| meta_path = roots[0] / meta |
| if meta_path.is_file(): |
| src = json.loads(meta_path.read_text()).get(key) |
| if src: |
| roots.append(Path(src)) |
| for root in roots: |
| path = root / "tensorboard_scalars.json" |
| if step is None or not path.is_file(): |
| continue |
| points = json.loads(path.read_text()).get("eval/loss") or [] |
| hits = [v for s, v in points if int(s) == int(step)] |
| if hits: |
| return float(hits[-1]) |
| return None |
|
|
|
|
| def _train_world_size(weights_dir: str) -> int: |
| """Ranks that wrote the checkpoint (one __<rank>_0.distcp each) = the run's eval sharding. |
| A model-only export is rewritten as one file, so its meta records the source rank count.""" |
| for meta in ("snapshot_meta.json", "final_meta.json"): |
| meta_path = Path(weights_dir).parent / meta |
| if meta_path.is_file(): |
| n = json.loads(meta_path.read_text()).get("source_world_size") |
| if n: |
| return int(n) |
| n = len(list(Path(weights_dir).glob("__*_0.distcp"))) |
| if n <= 0: |
| raise FileNotFoundError(f"no __*_0.distcp under {weights_dir}") |
| return n |
|
|
|
|
| @contextmanager |
| def _training_numerics(ckpt: EvalCheckpoint): |
| """Run as pretrain.run_eval does: fwd_bwd_dtype parameters and MoE capacity dropping on. |
| The eval-time dtype and capacity setting are restored afterwards (fp32 masters are kept, |
| not re-derived from bf16).""" |
| fwd_dtype = getattr(torch, ckpt.config.fwd_bwd_dtype) |
| params = list(ckpt.model.parameters()) |
| saved = [p.data for p in params] if any(p.dtype != fwd_dtype for p in params) else None |
| saved_env = os.environ.get("MOE_SKIP_CAPACITY") |
| os.environ["MOE_SKIP_CAPACITY"] = "0" |
| if saved is not None: |
| set_param_dtype(ckpt.model, fwd_dtype) |
| try: |
| yield |
| finally: |
| if saved is not None: |
| for p, data in zip(params, saved): |
| p.data = data |
| if saved_env is None: |
| os.environ.pop("MOE_SKIP_CAPACITY", None) |
| else: |
| os.environ["MOE_SKIP_CAPACITY"] = saved_env |
|
|
|
|
| @torch.inference_mode() |
| def _holdout_ce(ckpt: EvalCheckpoint, device: torch.device, rank: int, world_size: int) -> dict | None: |
| """Recompute the run's own eval/loss: same holdout split, same per-rank multipack |
| batches (pads included, they share MoE capacity), same token budget and stop rule |
| as pretrain.run_eval, under the same numerics. Matching it to ~1e-4 pins weights + |
| tokenizer + code tree.""" |
| if os.environ.get("EVAL_HOLDOUT_CHECK", "1") in {"0", "false", "False"}: |
| return None |
| from models.common import IGNORE_LABEL_ID |
| from pretrain import _make_dataset |
|
|
| cfg = ckpt.config |
| train_world = int(os.environ.get("EVAL_HOLDOUT_TRAIN_WORLD", "0")) or _train_world_size(ckpt.info["weights"]) |
| val_path = cfg.data.val_path or cfg.data.path |
| epoch = 0 if cfg.data.val_path is not None else int(cfg.eval_epoch) |
| batch_max_length = int(cfg.micro_batch_samples) * int(ckpt.info["max_seq_len"]) |
| mine = list(range(rank, train_world, world_size)) |
| iters = { |
| r: iter(_make_dataset( |
| cfg, dataset_path=val_path, batch_max_length=batch_max_length, drop_last_batch=False, |
| rank=r, world_size=train_world, fixed_epoch=epoch, |
| )) |
| for r in mine |
| } |
| stats = torch.zeros(2, device=device, dtype=torch.float64) |
| steps = 0 |
| while True: |
| for r in mine: |
| batch, scalars = next(iters[r]) |
| labels = batch.pop("labels").to(device, torch.long) |
| model_batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} |
| for k, v in scalars.items(): |
| if not str(k).startswith("_resume_"): |
| model_batch[k] = wrap_tensor(torch.tensor(v, device="cpu")) |
| _, logits = ckpt.model(carry=None, batch=model_batch) |
| stats[0] += F.cross_entropy( |
| logits.float(), labels, ignore_index=IGNORE_LABEL_ID, reduction="sum" |
| ).double() |
| stats[1] += (labels != IGNORE_LABEL_ID).sum() |
| steps += 1 |
| seen = stats[1].clone() |
| if world_size > 1 and dist.is_initialized(): |
| dist.all_reduce(seen) |
| if seen.item() >= cfg.eval_token_budget: |
| break |
| if world_size > 1 and dist.is_initialized(): |
| dist.all_reduce(stats) |
| ce = float(stats[0] / stats[1]) |
| ref = _train_eval_loss_at(ckpt) |
| return { |
| "source": f"{val_path}/epoch_{epoch}", |
| "train_world": train_world, |
| "steps": steps, |
| "tokens": int(stats[1]), |
| "ce": ce, |
| "train_eval_loss_at_step": ref, |
| "gap": None if ref is None else ce - ref, |
| } |
|
|
|
|
| def _write_outputs( |
| results: dict, |
| out_dir: Path, |
| ckpt: EvalCheckpoint, |
| tasks: list[str], |
| num_fewshot: int | None, |
| extra: dict, |
| ) -> None: |
| out_dir.mkdir(parents=True, exist_ok=True) |
| slim = {k: v for k, v in results.items() if k != "samples"} |
| slim["ckpt_path"] = ckpt.info["ckpt_path"] |
| slim["checkpoint"] = ckpt.info |
| slim["harness"] = "lm_eval.simple_evaluate" |
| slim["protocol"] = "official_0shot_causal_sum_ll" |
| slim["tasks"] = tasks |
| slim["num_fewshot"] = num_fewshot |
| slim.update(extra) |
| (out_dir / "results.json").write_text( |
| json.dumps(slim, indent=2, default=str), |
| encoding="utf-8", |
| ) |
| for task, rows in (results.get("samples") or {}).items(): |
| with open(out_dir / f"samples_{task}.jsonl", "w", encoding="utf-8") as f: |
| for row in rows: |
| keep = {k: row.get(k) for k in ("doc_id", "target", "acc", "acc_norm")} |
| keep["ll"] = [float(r[0]) for r in row.get("filtered_resps", [])] |
| f.write(json.dumps(keep, default=str) + "\n") |
| try: |
| from lm_eval.utils import make_table |
|
|
| table = make_table(results) |
| except Exception as exc: |
| table = f"(make_table failed: {exc})\n" |
| prov = extra["provenance"] |
| hold = extra.get("holdout") |
| info = ckpt.info |
| lines = [ |
| f"ckpt={info['ckpt_path']} step={info['step']} H={info['num_loops']}", |
| f"code={prov['code_root']}@{prov['commit'][:12]} dirty={prov['model_code_dirty']}", |
| f"tokenizer={info['tokenizer_path']}", |
| f"harness=lm_eval.simple_evaluate scoring=causal_sum_ll shots={num_fewshot} " |
| f"bs={extra['batch_size']} world={extra['world_size']} " |
| f"param_dtype={info['param_dtype']} moe_skip_capacity={prov['env'].get('MOE_SKIP_CAPACITY')}", |
| "holdout: skipped" if hold is None else ( |
| f"holdout ({info['fwd_dtype']}, capacity on) CE={hold['ce']:.4f} on {hold['tokens']} tok " |
| f"train eval/loss@step={hold['train_eval_loss_at_step']} gap={hold['gap']}" |
| ), |
| f"tasks={','.join(tasks)}", |
| "", |
| table, |
| "", |
| ] |
| text = "\n".join(lines) |
| (out_dir / "summary.txt").write_text(text, encoding="utf-8") |
| print(text, flush=True) |
|
|
|
|
| def main() -> None: |
| setup_hf_datasets_cache() |
| rank, world_size, is_main = init_eval_distributed() |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| set_device(local_rank) |
| device = get_device(local_rank) |
| if world_size > 1 and dist.is_initialized(): |
| dist.barrier() |
|
|
| ckpt_path = os.environ.get("CKPT_PATH", "").strip() |
| if not ckpt_path: |
| raise SystemExit("CKPT_PATH is required") |
| out_dir = Path(os.environ.get("EVAL_OUTPUT_DIR", str(Path(ckpt_path) / "eval_results" / "official_lmeval"))) |
| batch_size = int(os.environ.get("EVAL_BATCH_SIZE_MCQ", "32")) |
| num_fewshot = _parse_num_fewshot() |
| limit_raw = os.environ.get("LMEVAL_LIMIT", "").strip() |
| limit = float(limit_raw) if limit_raw else None |
| if limit is not None and limit >= 1: |
| limit = int(limit) |
| log_samples = os.environ.get("LMEVAL_LOG_SAMPLES", "0") not in {"0", "false", "False"} |
| bootstrap = int(os.environ.get("LMEVAL_BOOTSTRAP_ITERS", "1000")) |
| max_length = int(os.environ.get("LMEVAL_MAX_LENGTH", "1024")) |
| tasks = _parse_tasks() |
|
|
| from lm_eval.evaluator import simple_evaluate |
| from lm_eval.tasks import TaskManager |
|
|
| task_manager = TaskManager( |
| verbosity="INFO" if is_main else "ERROR", |
| include_path=str(_LOCAL_TASK_DIR), |
| ) |
|
|
| if is_main: |
| print( |
| f"[lm_eval_loom] rank0/{world_size} ckpt={ckpt_path} tasks={tasks} " |
| f"bs={batch_size} shots={num_fewshot} task_dir={_LOCAL_TASK_DIR} out={out_dir}", |
| flush=True, |
| ) |
| clear_prep_done_flag() |
| _warmup_hf_datasets(tasks, task_manager, num_fewshot) |
| signal_prep_done() |
| else: |
| print(f"[lm_eval_loom] rank {rank}: waiting for dataset warmup", flush=True) |
| wait_for_prep_done(rank, world_size) |
| if world_size > 1 and dist.is_initialized(): |
| dist.barrier() |
| _stage_local_dataset_cache(local_rank) |
| if world_size > 1 and dist.is_initialized(): |
| dist.barrier() |
|
|
| ckpt_epoch = os.environ.get("CKPT_EPOCH", "").strip() |
| ckpt = load_checkpoint_for_eval(ckpt_path, int(ckpt_epoch) if ckpt_epoch else None) |
| prov = provenance(ckpt.config.arch.name) |
| if is_main: |
| print(f"[lm_eval_loom] checkpoint {json.dumps(ckpt.info, default=str)}", flush=True) |
| print(f"[lm_eval_loom] provenance {json.dumps(prov)}", flush=True) |
|
|
| t0 = time.time() |
| with _training_numerics(ckpt): |
| holdout = _holdout_ce(ckpt, device, rank, world_size) |
| if is_main and holdout is not None: |
| print(f"[lm_eval_loom] holdout {json.dumps(holdout)} ({time.time() - t0:.0f}s)", flush=True) |
| max_gap = float(os.environ.get("EVAL_HOLDOUT_MAX_GAP", "0.005")) |
| if holdout["gap"] is None: |
| |
| |
| if os.environ.get("EVAL_HOLDOUT_CHECK", "1") != "record": |
| raise SystemExit( |
| f"no eval/loss at step {ckpt.info.get('step')} in any tensorboard_scalars.json; " |
| "nothing to check the load against (EVAL_HOLDOUT_CHECK=0 skips the check, " |
| "EVAL_HOLDOUT_CHECK=record keeps the CE without checking)" |
| ) |
| print(f"[lm_eval_loom] holdout: no eval/loss at step {ckpt.info.get('step')}, CE recorded unchecked", flush=True) |
| elif abs(holdout["gap"]) > max_gap: |
| raise SystemExit( |
| f"holdout CE {holdout['ce']:.4f} is {holdout['gap']:+.4f} off the run's own " |
| f"eval/loss; weights, tokenizer or code tree do not match (EVAL_HOLDOUT_MAX_GAP={max_gap})" |
| ) |
|
|
| lm = LoomLM( |
| ckpt, |
| rank=rank, |
| world_size=world_size, |
| local_rank=local_rank, |
| batch_size=batch_size, |
| max_length=max_length, |
| ) |
|
|
| if world_size > 1 and dist.is_initialized(): |
| dist.barrier() |
|
|
| t1 = time.time() |
| results = simple_evaluate( |
| model=lm, |
| tasks=tasks, |
| num_fewshot=num_fewshot, |
| batch_size=batch_size, |
| limit=limit, |
| log_samples=log_samples, |
| bootstrap_iters=bootstrap, |
| fewshot_random_seed=int(os.environ.get("LMEVAL_FEWSHOT_SEED", "1234")), |
| verbosity="INFO" if is_main else "ERROR", |
| task_manager=task_manager, |
| ) |
|
|
| if is_main: |
| if results is None: |
| raise SystemExit("lm_eval.simple_evaluate returned None on rank 0") |
| extra = { |
| "provenance": prov, |
| "holdout": holdout, |
| "batch_size": batch_size, |
| "world_size": world_size, |
| "limit": limit, |
| "seconds": {"holdout": t1 - t0, "lm_eval": time.time() - t1}, |
| } |
| _write_outputs(results, out_dir, ckpt, tasks, num_fewshot, extra) |
|
|
| if world_size > 1 and dist.is_initialized(): |
| dist.barrier() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
