Coverage for transformer_lens/tools/model_registry/verify_models.py: 29%
731 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Batch model verification tool for the TransformerLens model registry.
3Iterates through supported models, estimates memory requirements, runs benchmarks
4phase-by-phase, and updates the registry with status, phase scores, and notes.
6Usage:
7 python -m transformer_lens.tools.model_registry.verify_models [options]
9Examples:
10 # Dry run to see what would be tested
11 python -m transformer_lens.tools.model_registry.verify_models --dry-run
13 # Verify top 10 models per architecture on CPU
14 python -m transformer_lens.tools.model_registry.verify_models --device cpu
16 # Verify only GPT2 models, limit to 3
17 python -m transformer_lens.tools.model_registry.verify_models --architectures GPT2LMHeadModel --limit 3
19 # Resume from a previous interrupted run
20 python -m transformer_lens.tools.model_registry.verify_models --resume
22 # Re-verify already-tested models for a specific architecture
23 python -m transformer_lens.tools.model_registry.verify_models --reverify --architectures Olmo2ForCausalLM
24"""
26import argparse
27import gc
28import json
29import logging
30import re
31import signal
32import time
33from dataclasses import dataclass, field
34from datetime import datetime
35from pathlib import Path
36from typing import Optional
38from transformer_lens.benchmarks.text_quality_profiles import (
39 P4_SCORING_VERSION,
40 p4_pass_threshold,
41)
42from transformer_lens.utilities.heterogeneous_config import het_safe_view
44# Exit code used for graceful interrupts (Ctrl+C). The wrapper script
45# recognises this and stops without marking the in-flight model as failed.
46_EXIT_GRACEFUL_INTERRUPT = 42
48# Module-level flag set by the SIGINT handler so the main loop can stop
49# between models without corrupting state.
50_interrupt_requested = False
52from . import REMOTE_CODE_MODEL_PREFIXES
53from .registry_io import (
54 MODALITY_PHASES,
55 QUANTIZED_NOTE,
56 STATUS_FAILED,
57 STATUS_PROVISIONAL,
58 STATUS_SKIPPED,
59 STATUS_UNVERIFIED,
60 STATUS_VERIFIED,
61 TEXT_PHASES,
62 add_verification_record,
63 extract_phase_scores,
64 is_incompatible_quantized,
65 load_supported_models_raw,
66 pass_status,
67 required_quant_library_for_model,
68 update_model_status,
69)
71logger = logging.getLogger(__name__)
73# Data directory for registry files
74_DATA_DIR = Path(__file__).parent / "data"
75_CHECKPOINT_PATH = _DATA_DIR / "verification_checkpoint.json"
78def _handle_sigint(signum, frame): # noqa: ARG001
79 """Handle Ctrl+C by setting a flag instead of raising immediately.
81 The main verification loop checks this flag between models so it can
82 save the checkpoint cleanly and exit without marking the current model
83 as failed.
84 """
85 global _interrupt_requested # noqa: PLW0603
86 if _interrupt_requested:
87 # Second Ctrl+C — force exit immediately
88 print("\nForce quit.")
89 raise SystemExit(1)
90 _interrupt_requested = True
91 print("\n\nInterrupt received — finishing current model before stopping.")
92 print("(Press Ctrl+C again to force quit immediately.)\n")
95# Pattern matching HuggingFace API tokens (hf_ followed by 20+ alphanumeric chars)
96_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}")
99def _sanitize_note(note: Optional[str]) -> Optional[str]:
100 """Sanitize a note string to remove sensitive information.
102 Strips HuggingFace tokens and replaces verbose gated-repo error messages
103 with a concise summary.
104 """
105 if not note:
106 return note
107 # Replace any HF tokens that leaked into the message
108 note = _HF_TOKEN_RE.sub("HF_TOKEN", note)
109 # Replace verbose gated-repo 401 errors with a clean summary
110 if "gated repo" in note:
111 url_match = re.search(r"https://huggingface\.co/([^\s.]+)", note)
112 model_ref = url_match.group(1) if url_match else "unknown"
113 return f"Config unavailable: Gated repo ({model_ref})"
114 return note
117def _phases_to_run(arch: str, phases: list[int]) -> list[int]:
118 """Restrict requested phases to those the adapter supports.
120 An adapter's ``applicable_phases`` declares which text phases (1-4) it covers. Phases
121 7/8/9 are gated separately by ``is_multimodal``/``is_audio``/``is_visual_model`` in
122 ``main_benchmark._phase_enabled`` (the mirror of this gate), so they are never filtered
123 out here. An empty result means none of the requested phases apply to this architecture
124 (SSM / recurrent families run all four).
125 """
126 from transformer_lens.factories.architecture_adapter_factory import (
127 SUPPORTED_ARCHITECTURES,
128 )
130 applicable = getattr(SUPPORTED_ARCHITECTURES.get(arch), "applicable_phases", list(TEXT_PHASES))
131 return [p for p in phases if p in applicable or p in MODALITY_PHASES]
134def _full_and_core_phases(arch: str) -> tuple[set[int], set[int]]:
135 """``(full verification set, core subset)`` for this architecture -- the single
136 source for both default phase selection and the status-writing decision."""
137 from transformer_lens.utilities.architectures import AUDIO_TEXT_ARCHITECTURES
139 kind = classify_architecture(arch)
140 if kind == "audio":
141 return {1, 8}, {1, 8}
142 if kind == "vision":
143 # Vision encoders have no tokenizer and no text tower: Phases 2/3 compare
144 # text logits/loss, Phase 4 needs text generation, and Phase 7 covers
145 # vision+text multimodal models, not these. Phase 1 (HF parity) plus
146 # Phase 9 (pixel forward/cache/stability) are the whole story.
147 return {1, 9}, {1, 9}
148 if kind == "multimodal":
149 return {1, 2, 3, 4, 7}, {1, 4, 7}
150 if arch in AUDIO_TEXT_ARCHITECTURES:
151 # Phase 8 (audio-conditioned forward) out of the core set so a partial {1,4}
152 # run still verifies; a full run records and gates it via _check_phase_scores.
153 return {1, 2, 3, 4, 8}, {1, 4}
154 return {1, 2, 3, 4}, {1, 4}
157def _default_phases_for_architecture(arch: str) -> list[int]:
158 """Phases to run when the caller names none — a full verification."""
159 return sorted(_full_and_core_phases(arch)[0])
162def _get_current_model_status(model_id: str, arch_id: str) -> int:
163 """Look up a model's current status in the registry.
165 Returns STATUS_UNVERIFIED (0) if the model is not found.
166 """
167 data = load_supported_models_raw()
168 for entry in data.get("models", []):
169 if not isinstance(entry, dict):
170 continue
171 if entry.get("model_id") == model_id and entry.get("architecture_id") == arch_id:
172 return entry.get("status", STATUS_UNVERIFIED)
173 return STATUS_UNVERIFIED
176@dataclass
177class ModelCandidate:
178 """A model selected for verification."""
180 model_id: str
181 architecture_id: str
182 estimated_params: Optional[int] = None
183 estimated_memory_gb: Optional[float] = None
186@dataclass
187class VerificationProgress:
188 """Tracks progress across a verification run."""
190 tested: list[str] = field(default_factory=list)
191 skipped: list[str] = field(default_factory=list)
192 failed: list[str] = field(default_factory=list)
193 verified: list[str] = field(default_factory=list)
194 # Structural-only (--no-hf-reference) passes; kept out of the verified tally.
195 provisional: list[str] = field(default_factory=list)
196 start_time: Optional[str] = None
198 def to_dict(self) -> dict:
199 return {
200 "tested": self.tested,
201 "skipped": self.skipped,
202 "failed": self.failed,
203 "verified": self.verified,
204 "provisional": self.provisional,
205 "start_time": self.start_time,
206 }
208 @classmethod
209 def from_dict(cls, data: dict) -> "VerificationProgress":
210 return cls(
211 tested=data.get("tested", []),
212 skipped=data.get("skipped", []),
213 failed=data.get("failed", []),
214 verified=data.get("verified", []),
215 provisional=data.get("provisional", []),
216 start_time=data.get("start_time"),
217 )
220def published_param_count(model_id: str) -> Optional[int]:
221 """Exact parameter count from the hub's safetensors metadata, or None.
223 Metadata only -- no weights are fetched. Preferred over the config formula
224 below, which assumes every layer carries full attention and an MLP and so
225 over-counts a hybrid Mamba/attention stack roughly fourfold
226 (NVIDIA-Nemotron-Nano-9B-v2: 36.6B estimated against 8.89B published, enough
227 to skip the model as too large for memory it does not need).
228 """
229 from transformer_lens.utilities.hf_utils import get_hf_token
231 try:
232 from huggingface_hub import HfApi
234 info = HfApi().model_info(model_id, expand=["safetensors"], token=get_hf_token())
235 except Exception:
236 # Unpublished metadata, a gated repo or a network blip: fall back rather
237 # than fail, since the config formula needs no hub metadata.
238 return None
239 total = getattr(info.safetensors, "total", None) if info.safetensors else None
240 return int(total) if total else None
243def estimate_model_params(model_id: str) -> int:
244 """Parameter count for this model: published metadata first, else a config estimate.
246 Fetches only metadata and the config JSON (~KB), never weights.
248 Args:
249 model_id: HuggingFace model ID
251 Returns:
252 Number of parameters — exact when the hub publishes it, else estimated
253 from config dimensions using the standard TransformerLens
254 parameter-count formula.
256 Raises:
257 Exception: If config cannot be fetched or parsed
258 """
259 published = published_param_count(model_id)
260 if published:
261 return published
263 trust_remote_code = any(model_id.startswith(prefix) for prefix in REMOTE_CODE_MODEL_PREFIXES)
264 from transformer_lens.utilities.hf_utils import (
265 autoconfig_with_remote_post_init_compat,
266 get_hf_token,
267 )
269 config = autoconfig_with_remote_post_init_compat(
270 model_id, trust_remote_code=trust_remote_code, token=get_hf_token()
271 )
273 # For multimodal models (LLaVA, Gemma3 multimodal), the language model config
274 # is nested under text_config. Fall through to the top-level config otherwise.
275 lang_config = getattr(config, "text_config", config)
277 # Encoder-decoder models (e.g. T5Gemma) nest dimensions under decoder/encoder
278 # subconfigs rather than the top level; prefer the decoder for the estimate.
279 if not (hasattr(lang_config, "hidden_size") or hasattr(lang_config, "d_model")): 279 ↛ 280line 279 didn't jump to line 280 because the condition on line 279 was never true
280 for _sub in ("decoder", "encoder"):
281 _subcfg = getattr(config, _sub, None)
282 if _subcfg is not None and (
283 hasattr(_subcfg, "hidden_size") or hasattr(_subcfg, "d_model")
284 ):
285 lang_config = _subcfg
286 break
288 # Heterogeneous configs (transformers>=5.15 Gemma 4) raise on global reads of
289 # per-layer fields like head_dim; the view resolves them to majority values.
290 lang_config = het_safe_view(lang_config)
292 # Extract dimensions from config (different models use different attribute names)
293 d_model = (
294 getattr(lang_config, "hidden_size", None)
295 or getattr(lang_config, "d_model", None)
296 or getattr(lang_config, "model_dim", None) # OpenELM
297 or 0
298 )
299 n_heads_raw = (
300 getattr(lang_config, "num_attention_heads", None)
301 or getattr(lang_config, "n_head", None)
302 or getattr(lang_config, "num_query_heads", None) # OpenELM (may be per-layer list)
303 or getattr(lang_config, "num_heads", None) # Mamba-2 SSM heads
304 or 0
305 )
306 # OpenELM uses per-layer lists for heads; take the max for estimation
307 n_heads = max(n_heads_raw) if isinstance(n_heads_raw, (list, tuple)) else n_heads_raw
308 n_layers = (
309 getattr(lang_config, "num_hidden_layers", None)
310 or getattr(lang_config, "n_layer", None)
311 or getattr(lang_config, "num_transformer_layers", None) # OpenELM
312 or 0
313 )
314 d_mlp = (
315 getattr(lang_config, "intermediate_size", None)
316 or getattr(lang_config, "d_inner", None)
317 or getattr(lang_config, "n_inner", None)
318 or getattr(lang_config, "ffn_dim", None) # OPT
319 or getattr(lang_config, "d_ff", None) # T5
320 )
321 # Gemma 3n exposes a per-layer intermediate_size list (uniform in all released
322 # checkpoints); collapse to max for the scalar param estimate.
323 if isinstance(d_mlp, (list, tuple)): 323 ↛ 324line 323 didn't jump to line 324 because the condition on line 323 was never true
324 d_mlp = max(d_mlp) if d_mlp else None
325 # OpenELM uses per-layer ffn_multipliers instead of a fixed intermediate_size
326 if not d_mlp and d_model: 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true
327 ffn_multipliers = getattr(lang_config, "ffn_multipliers", None)
328 if isinstance(ffn_multipliers, (list, tuple)):
329 d_mlp = int(max(ffn_multipliers) * d_model)
330 else:
331 # Many architectures (GPT-2, Bloom, GPT-Neo, GPT-J) leave d_mlp/n_inner
332 # as None and default to 4 * hidden_size internally.
333 d_mlp = 4 * d_model
334 d_vocab = getattr(lang_config, "vocab_size", None) or 0
336 if d_model == 0 or n_layers == 0: 336 ↛ 337line 336 didn't jump to line 337 because the condition on line 336 was never true
337 raise ValueError(f"Could not extract model dimensions from config for {model_id}")
339 # Attention-less architectures (Mamba SSMs) have no heads. Use nominal
340 # values so the estimate doesn't attribute phantom attention params.
341 is_attention_less = n_heads == 0
342 if is_attention_less: 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true
343 n_heads = 1
344 d_head = d_model
345 else:
346 d_head = getattr(lang_config, "head_dim", None) or (d_model // n_heads)
348 # Attention parameters: W_Q, W_K, W_V, W_O per layer (skipped for SSMs)
349 if is_attention_less: 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true
350 n_params = 0
351 else:
352 n_params = n_layers * (d_model * d_head * n_heads * 4)
354 # MLP parameters (if present)
355 if d_mlp is not None and d_mlp > 0: 355 ↛ 398line 355 didn't jump to line 398 because the condition on line 355 was always true
356 # Check for gated MLP (LLaMA, Gemma, Mistral, Qwen, T5 gated-gelu, etc.)
357 has_gate = getattr(lang_config, "is_gated_act", False) or (
358 hasattr(lang_config, "intermediate_size")
359 and (
360 getattr(lang_config, "hidden_act", None) in ("silu", "gelu", "swiglu")
361 or getattr(lang_config, "model_type", None)
362 in (
363 "llama",
364 "gemma",
365 "gemma2",
366 "gemma3",
367 "mistral",
368 "mixtral",
369 "qwen2",
370 "qwen3",
371 "qwen3_moe",
372 "phi3",
373 "stablelm",
374 )
375 )
376 )
377 mlp_multiplier = 3 if has_gate else 2
378 n_params += n_layers * (d_model * d_mlp * mlp_multiplier)
380 # MoE expert scaling
381 num_experts = (
382 getattr(lang_config, "num_local_experts", None)
383 or getattr(lang_config, "num_experts", None)
384 or getattr(lang_config, "n_routed_experts", None) # DeepSeek-V2/V3
385 )
386 if num_experts and num_experts > 1: 386 ↛ 389line 386 didn't jump to line 389 because the condition on line 386 was never true
387 # Qwen3MoE and similar store per-expert hidden size in moe_intermediate_size;
388 # intermediate_size refers to a dense fallback MLP that we don't use here.
389 moe_d_mlp = getattr(lang_config, "moe_intermediate_size", None) or d_mlp
390 # MLP params scale with num_experts; add gate params per expert
391 mlp_per_layer = d_model * moe_d_mlp * mlp_multiplier
392 moe_per_layer = (mlp_per_layer + d_model) * num_experts
393 # Replace the non-MoE MLP contribution
394 n_params -= n_layers * (d_model * d_mlp * mlp_multiplier)
395 n_params += n_layers * moe_per_layer
397 # Embedding parameters (not in the block-param formula but relevant for memory)
398 n_params += d_vocab * d_model
400 return n_params
403def estimate_benchmark_memory_gb(
404 n_params: int,
405 dtype: str = "float32",
406 phases: Optional[list[int]] = None,
407 use_hf_reference: bool = True,
408 device: str = "cpu",
409) -> float:
410 """Estimate peak memory needed for benchmark suite.
412 Phases run sequentially, so peak memory is the maximum of any single phase,
413 not the sum. The multiplier represents how many model copies exist at peak:
415 Phase 1 (HF ref on): HF ref + Bridge → 2.0x peak
416 Phase 1 (HF ref off): Bridge only → 1.0x peak
417 Phase 2: Bridge only (runtime self-checks) → 1.0x model + overhead
418 Phase 3: Bridge + weight-processing state-dict transient → 2.0x model + overhead
419 Phase 4: Bridge + GPT-2 scorer (~500MB) → ~1.0x model + 0.5 GB
421 Args:
422 n_params: Number of model parameters
423 dtype: Data type for memory calculation
424 phases: Which phases will be run (None = all phases)
425 use_hf_reference: Whether Phase 1 loads an HF reference alongside the
426 Bridge. Mirrors the ``--no-hf-reference`` CLI flag.
428 Returns:
429 Estimated peak memory in GB
430 """
431 bytes_per_param = {"float32": 4, "float16": 2, "bfloat16": 2}
432 bpp = bytes_per_param.get(dtype, 4)
433 model_size_gb = n_params * bpp / (1024**3)
435 # Phase-4 judge overhead: measured 2.33 GB RSS loading Qwen2.5-0.5B fp32
436 # on CPU (494M params). Kept slightly above the measurement; over-counting
437 # is the safe direction.
438 # The CPU-pinned judge never occupies accelerator memory; charging it to
439 # a cuda budget produces spurious VRAM skips.
440 judge_overhead_gb = 2.5 if device == "cpu" else 0.0
442 # Activation/framework overhead as a fraction of model size
443 overhead_fraction = 0.2
445 # Determine peak memory across all requested phases
446 phase_peaks = []
448 if phases is None: 448 ↛ 449line 448 didn't jump to line 449 because the condition on line 448 was never true
449 phases = list(TEXT_PHASES)
451 for p in phases:
452 if p == 1:
453 # HF ref + Bridge (2 copies) or Bridge alone
454 multiplier = 2.0 if use_hf_reference else 1.0
455 phase_peaks.append(model_size_gb * multiplier * (1 + overhead_fraction))
456 elif p == 2: 456 ↛ 458line 456 didn't jump to line 458 because the condition on line 456 was never true
457 # Bridge only (runtime self-checks + saved-HF equivalence)
458 phase_peaks.append(model_size_gb * 1.0 * (1 + overhead_fraction))
459 elif p == 3: 459 ↛ 461line 459 didn't jump to line 461 because the condition on line 459 was never true
460 # Bridge + the full state-dict copy materialized during weight processing
461 phase_peaks.append(model_size_gb * 2.0 * (1 + overhead_fraction))
462 elif p == 4: 462 ↛ 451line 462 didn't jump to line 451 because the condition on line 462 was always true
463 # Bridge + judge
464 phase_peaks.append(model_size_gb * (1 + overhead_fraction) + judge_overhead_gb)
466 return max(phase_peaks) if phase_peaks else model_size_gb
469def get_available_memory_gb(device: str) -> float:
470 """Detect available memory on the target device.
472 Args:
473 device: "cpu" or "cuda"
475 Returns:
476 Available memory in GB
477 """
478 if device.startswith("cuda"):
479 try:
480 import torch
482 if torch.cuda.is_available():
483 device_idx = 0
484 if ":" in device:
485 device_idx = int(device.split(":")[1])
486 props = torch.cuda.get_device_properties(device_idx)
487 return props.total_memory / (1024**3)
488 except Exception:
489 pass
490 return 8.0 # Conservative default for GPU
492 # CPU: use psutil if available, else conservative default
493 try:
494 import psutil
496 return psutil.virtual_memory().available / (1024**3)
497 except ImportError:
498 return 16.0 # Conservative default for CPU
501def select_models_for_verification(
502 per_arch: int = 10,
503 architectures: Optional[list[str]] = None,
504 limit: Optional[int] = None,
505 resume_progress: Optional[VerificationProgress] = None,
506 retry_failed: bool = False,
507 reverify: bool = False,
508) -> list[ModelCandidate]:
509 """Select models for verification from the registry.
511 Loads supported_models.json (already sorted by downloads).
512 Takes the top N unverified models per architecture.
514 Args:
515 per_arch: Maximum models to verify per architecture
516 architectures: Filter to specific architectures (None = all)
517 limit: Total model cap (None = no cap)
518 resume_progress: If resuming, skip already-tested models
519 retry_failed: If True, include previously failed models for re-testing
520 reverify: If True, ignore previous status and re-test all matching models
522 Returns:
523 List of ModelCandidate objects to verify
524 """
525 already_tested: set[str] = set()
526 if resume_progress and not reverify:
527 already_tested = set(resume_progress.tested)
528 if retry_failed:
529 # Remove failed models from already_tested so they get re-selected
530 failed_set = set(resume_progress.failed)
531 already_tested -= failed_set
533 data = load_supported_models_raw()
534 models = data.get("models", [])
536 # Group by architecture
537 by_arch: dict[str, list[dict]] = {}
538 for model in models:
539 arch = model["architecture_id"]
540 by_arch.setdefault(arch, []).append(model)
542 # Determine which architectures to scan
543 if architectures:
544 arch_ids = architectures
545 else:
546 arch_ids = sorted(by_arch.keys())
548 candidates: list[ModelCandidate] = []
550 for arch in arch_ids:
551 arch_models = by_arch.get(arch, [])
552 count = 0
554 for model in arch_models:
555 model_id = model["model_id"]
557 # Skip already-verified or already-tested models
558 if not reverify:
559 model_status = model.get("status", 0)
560 if model_status == STATUS_VERIFIED or model_status == STATUS_SKIPPED:
561 continue
562 if model_status == STATUS_FAILED and not retry_failed:
563 continue
564 if model_id in already_tested:
565 continue
567 # Check per-arch limit
568 if count >= per_arch:
569 break
571 count += 1
572 candidates.append(ModelCandidate(model_id=model_id, architecture_id=arch))
574 # Check total limit
575 if limit and len(candidates) >= limit:
576 return candidates
578 return candidates
581# Phase-score extraction and pass-status logic live in registry_io
582# (extract_phase_scores / pass_status) — the shared home for both
583# registry-writing paths, this module and main_benchmark.update_model_registry.
586def _extract_prompt_profile(results: list) -> Optional[str]:
587 """Effective Phase-4 prompt profile from the benchmark details, or None
588 when no Phase-4 result exists. The default "continuation" is reported so
589 the registry write can clear a stale non-default key."""
590 for result in results:
591 if result.phase == 4 and result.details:
592 profile = result.details.get("prompt_profile")
593 if isinstance(profile, str): 593 ↛ 590line 593 didn't jump to line 590 because the condition on line 593 was always true
594 return profile
595 return None
598# Per-phase minimum score thresholds (0-100).
599# Phase 1: Core correctness (bridge vs HF) — must pass everything.
600# Phase 2: Hook/cache/gradient tests — most should pass.
601# Phase 3: Weight processing tests — most should pass.
602# Phase 4: Text quality — inherently fuzzy, keep lenient.
603_MIN_PHASE_SCORES: dict[int, float] = {
604 1: 100.0,
605 2: 75.0,
606 3: 75.0,
607 # Phase 4 floor == the benchmark pass line; a gap between them lets a
608 # failing score carry a clean "completed" note.
609 4: p4_pass_threshold(),
610 7: 75.0,
611 8: 75.0,
612 9: 75.0,
613}
614_DEFAULT_MIN_PHASE_SCORE = 50.0
616from transformer_lens.utilities.architectures import classify_architecture
618# Tests that MUST pass for a phase to be considered passing, regardless of
619# the overall percentage score. If any required test fails, the phase fails
620# even if the score is above the minimum threshold.
621_REQUIRED_PHASE_TESTS: dict[int, list[str]] = {
622 2: ["logits_equivalence", "loss_equivalence"],
623 3: ["logits_equivalence", "loss_equivalence"],
624 7: ["multimodal_forward"],
625 8: ["audio_forward", "audio_text_forward"],
626 9: ["vision_forward", "vision_cache"],
627}
629# Failure text for a modality phase that produced no score — either absent from
630# phase_scores (all tests skipped) or explicitly NULL.
631_MODALITY_NULL_MESSAGES: dict[int, str] = {
632 7: "P7=NULL (multimodal tests skipped — processor unavailable)",
633 8: "P8=NULL (audio tests skipped — no results)",
634 9: "P9=NULL (vision tests skipped — no results)",
635}
638def _measured_nothing(phase_scores: dict) -> bool:
639 """True when no phase produced a score, so the run verified nothing.
641 An adapter's ``applicable_phases`` can prune the requested phases to empty.
642 That leaves ``required_phases`` empty as well, so no score check fires and a
643 run that measured nothing would otherwise be recorded as VERIFIED.
645 ``is not None`` rather than truthiness: 0.0 is a real score.
646 """
647 return not any(score is not None for score in phase_scores.values())
650def _check_phase_scores(
651 phase_scores: dict[int, Optional[float]],
652 all_results: list,
653 required_phases: Optional[set[int]] = None,
654) -> Optional[str]:
655 """Check phase scores against per-phase minimum thresholds and required tests.
657 A phase fails if:
658 1. Its overall score is below the minimum threshold, OR
659 2. Any of its required tests (per _REQUIRED_PHASE_TESTS) failed.
661 Phase 4 (text quality) is excluded — it is a quality metric, not a
662 correctness check. Low text quality is surfaced in the verification
663 note via _build_verified_note() but never causes a model to fail.
665 Args:
666 phase_scores: Per-phase scores; a phase whose tests all skipped is
667 absent entirely (``extract_phase_scores`` omits empty phases).
668 all_results: Benchmark results, used to name the failing tests.
669 required_phases: Phases this architecture must produce a score for —
670 the core set from ``_full_and_core_phases``. A required modality
671 phase that is absent counts as NULL, not as a pass.
673 Returns an error message if any phase fails, or None if all phases pass.
674 The message includes the names of failed tests.
675 """
676 from transformer_lens.benchmarks.utils import BenchmarkSeverity
678 failing_phases: list[str] = []
680 # A required modality phase whose tests all skipped never reaches phase_scores.
681 for phase in sorted((required_phases or set()) - set(phase_scores)):
682 if phase in _MODALITY_NULL_MESSAGES: 682 ↛ 681line 682 didn't jump to line 681 because the condition on line 682 was always true
683 failing_phases.append(_MODALITY_NULL_MESSAGES[phase])
685 for phase, score in sorted(phase_scores.items()):
686 if score is None: 686 ↛ 690line 686 didn't jump to line 690 because the condition on line 686 was never true
687 # Phase 7 (multimodal), 8 (audio), or 9 (vision) with a NULL score
688 # means the modality tests never ran. This is a verification
689 # failure, not something to silently skip.
690 if phase in _MODALITY_NULL_MESSAGES:
691 failing_phases.append(_MODALITY_NULL_MESSAGES[phase])
692 continue
694 # Phase 4 is a quality metric, not a pass/fail check — skip it here.
695 # Low text quality is reported in the note by _build_verified_note().
696 if phase == 4:
697 continue
699 # Check 1: overall score threshold
700 threshold = _MIN_PHASE_SCORES.get(phase, _DEFAULT_MIN_PHASE_SCORE)
701 if score < threshold:
702 failed_tests = [
703 r.name
704 for r in all_results
705 if r.phase == phase and not r.passed and r.severity != BenchmarkSeverity.SKIPPED
706 ]
707 tests_str = ", ".join(failed_tests) if failed_tests else "unknown"
708 failing_phases.append(f"P{phase}={score}% < {threshold}% (failed: {tests_str})")
709 continue # Already failing; no need to also check required tests
711 # Check 2: required tests must pass
712 required_tests = _REQUIRED_PHASE_TESTS.get(phase, [])
713 if required_tests:
714 failed_required = [
715 r.name
716 for r in all_results
717 if r.phase == phase
718 and r.name in required_tests
719 and not r.passed
720 and r.severity != BenchmarkSeverity.SKIPPED
721 ]
722 if failed_required: 722 ↛ 723line 722 didn't jump to line 723 because the condition on line 722 was never true
723 tests_str = ", ".join(failed_required)
724 failing_phases.append(f"P{phase}={score}% but required tests failed: {tests_str}")
726 if failing_phases:
727 return f"Below threshold: {'; '.join(failing_phases)}"
728 return None
731def _build_verified_note(
732 phase_scores: dict[int, Optional[float]],
733 all_results: list,
734) -> str:
735 """Build a verification note summarizing phase scores.
737 Phase 4 (text quality) is excluded from the score summary since it's a
738 quality metric, not a pass/fail comparison. It only contributes a "low
739 text quality" flag when below threshold.
740 """
741 from transformer_lens.benchmarks.utils import BenchmarkSeverity
743 issue_parts: list[str] = []
744 low_text_quality = False
746 for phase in sorted(phase_scores):
747 score = phase_scores[phase]
748 if score is None: 748 ↛ 749line 748 didn't jump to line 749 because the condition on line 748 was never true
749 continue
750 # Phase 4 is a quality score, not a pass/fail comparison — don't
751 # include it in the normal score summary.
752 if phase == 4:
753 threshold = _MIN_PHASE_SCORES.get(4, _DEFAULT_MIN_PHASE_SCORE)
754 if score < threshold: 754 ↛ 755line 754 didn't jump to line 755 because the condition on line 754 was never true
755 low_text_quality = True
756 continue
758 if score < 100.0: 758 ↛ 759line 758 didn't jump to line 759 because the condition on line 758 was never true
759 failed_tests = [
760 r.name
761 for r in all_results
762 if r.phase == phase and not r.passed and r.severity != BenchmarkSeverity.SKIPPED
763 ]
764 if failed_tests:
765 issue_parts.append(f"P{phase}={score}% (failed: {', '.join(failed_tests)})")
766 else:
767 issue_parts.append(f"P{phase}={score}%")
769 p4_uncovered = next(
770 (
771 r.message
772 for r in all_results
773 if r.phase == 4
774 and r.severity == BenchmarkSeverity.SKIPPED
775 and r.message.startswith("P4 skipped:")
776 ),
777 None,
778 )
779 suffix = ""
780 if p4_uncovered: 780 ↛ 782line 780 didn't jump to line 782 because the condition on line 780 was never true
781 # Keep the gap visible in the registry until prompt coverage is added.
782 reason = p4_uncovered.split("—")[0].replace("P4 skipped:", "").strip()
783 suffix = f"; P4 skipped (uncovered: {reason} — file a coverage issue)"
785 if issue_parts and low_text_quality: 785 ↛ 786line 785 didn't jump to line 786 because the condition on line 785 was never true
786 return (
787 f"Full verification completed with issues, low text quality: {'; '.join(issue_parts)}"
788 + suffix
789 )
790 if issue_parts: 790 ↛ 791line 790 didn't jump to line 791 because the condition on line 790 was never true
791 return f"Full verification completed with issues: {'; '.join(issue_parts)}" + suffix
792 if low_text_quality: 792 ↛ 793line 792 didn't jump to line 793 because the condition on line 792 was never true
793 return "Full verification completed with issues, low text quality" + suffix
794 return "Full verification completed" + suffix
797def _preserved_issue_suffix(model_id: str, eff_phases) -> str:
798 """Sub-100 scores from phases not re-run this pass stay visible in the
799 note; a partial pass must not overwrite tracked residue."""
800 from transformer_lens.tools.model_registry.registry_io import (
801 load_supported_models_raw,
802 )
804 try:
805 entry = next(
806 (
807 m
808 for m in load_supported_models_raw().get("models", [])
809 if m.get("model_id") == model_id
810 ),
811 None,
812 )
813 except OSError:
814 return ""
815 if entry is None: 815 ↛ 816line 815 didn't jump to line 816 because the condition on line 815 was never true
816 return ""
817 residue = []
818 for phase in (2, 3, 7, 8, 9):
819 if phase in (eff_phases or []):
820 continue
821 score = entry.get(f"phase{phase}_score")
822 if score is not None and score < 100.0:
823 residue.append(f"P{phase}={score}%")
824 if not residue:
825 return ""
826 return f" (prior issues retained: {', '.join(residue)})"
829def _p1_only_core_note(p4_score, all_results: list) -> str:
830 """Note for a core run where P1 passed but P4 did not contribute a pass.
832 A skipped P4 is a coverage gap, not a quality failure — the stale
833 (possibly old-scale) score must not be relabeled "poor"."""
834 from transformer_lens.benchmarks.utils import BenchmarkSeverity
836 p4_skip_msg = next(
837 (
838 r.message
839 for r in all_results
840 if r.phase == 4
841 and r.severity == BenchmarkSeverity.SKIPPED
842 and r.message.startswith("P4 skipped:")
843 ),
844 None,
845 )
846 if p4_skip_msg is not None:
847 reason = p4_skip_msg.split("—")[0].replace("P4 skipped:", "").strip()
848 return f"Core verification passed; P4 skipped ({reason})"
849 if p4_score is None:
850 return "Core verification passed, but text quality benchmark errored. Needs review"
851 return f"Core verification passed, but text quality poor (P4={p4_score}). Needs review"
854def _clear_hf_cache(quiet: bool = False) -> None:
855 """Remove downloaded model weights from the HuggingFace cache to free disk."""
856 from pathlib import Path
858 cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
859 if not cache_dir.exists(): 859 ↛ 860line 859 didn't jump to line 860 because the condition on line 859 was never true
860 return
862 from transformer_lens.benchmarks.text_quality import JUDGE_MODEL_ID
864 # The pinned Phase-4 judge is needed by every run; deleting it here would
865 # force a re-download per family.
866 judge_dir = "models--" + JUDGE_MODEL_ID.replace("/", "--")
868 freed = 0
869 for blobs_dir in cache_dir.glob("models--*/blobs"):
870 if blobs_dir.parent.name == judge_dir:
871 continue
872 for blob in blobs_dir.iterdir():
873 try:
874 size = blob.stat().st_size
875 blob.unlink()
876 freed += size
877 except OSError:
878 pass
880 if not quiet and freed > 0: 880 ↛ 881line 880 didn't jump to line 881 because the condition on line 880 was never true
881 print(f" Cleared {freed / (1024**3):.1f} GB from HuggingFace cache")
884def _save_checkpoint(progress: VerificationProgress) -> None:
885 """Save verification progress to checkpoint file."""
886 with open(_CHECKPOINT_PATH, "w") as f:
887 json.dump(progress.to_dict(), f, indent=2)
888 f.write("\n")
891def _skip_model(
892 model_id: str, arch: str, note: str, progress: VerificationProgress, quiet: bool
893) -> None:
894 """Record a model as skipped with ``note``, preserving an existing verified or provisional
895 status, and checkpoint. Callers ``continue`` the loop afterwards.
896 """
897 if not quiet:
898 print(f" SKIP: {note}")
899 if _get_current_model_status(model_id, arch) not in (STATUS_VERIFIED, STATUS_PROVISIONAL):
900 update_model_status(model_id, arch, STATUS_SKIPPED, note=note, sanitize_fn=_sanitize_note)
901 elif not quiet:
902 print(" (preserving existing verified/provisional status)")
903 progress.skipped.append(model_id)
904 _save_checkpoint(progress)
907def _load_checkpoint() -> Optional[VerificationProgress]:
908 """Load verification progress from checkpoint file."""
909 if not _CHECKPOINT_PATH.exists():
910 return None
911 try:
912 with open(_CHECKPOINT_PATH) as f:
913 data = json.load(f)
914 return VerificationProgress.from_dict(data)
915 except (json.JSONDecodeError, KeyError):
916 return None
919def verify_models(
920 candidates: list[ModelCandidate],
921 device: str = "cpu",
922 max_memory_gb: Optional[float] = None,
923 dtype: str = "float32",
924 use_hf_reference: bool = True,
925 phases: Optional[list[int]] = None,
926 quiet: bool = False,
927 progress: Optional[VerificationProgress] = None,
928) -> VerificationProgress:
929 """Run verification benchmarks on a list of model candidates.
931 Args:
932 candidates: Models to verify
933 device: Device for benchmarks
934 max_memory_gb: Memory limit (auto-detected if None)
935 dtype: Dtype for memory estimation
936 use_hf_reference: Whether to compare against HuggingFace model
937 phases: Which benchmark phases to run
938 quiet: Suppress verbose output
939 progress: Existing progress for resume
941 Returns:
942 VerificationProgress with results
943 """
944 from transformer_lens.benchmarks.main_benchmark import run_benchmark_suite
946 if progress is None:
947 progress = VerificationProgress(start_time=datetime.now().isoformat())
949 if max_memory_gb is None:
950 max_memory_gb = get_available_memory_gb(device)
951 if not quiet:
952 print(f"Auto-detected available memory: {max_memory_gb:.1f} GB")
954 # phases stays None = full verification for the model.
956 # Pre-load the Phase-4 judge so it persists across all models in the batch
957 # instead of being loaded and destroyed for each one.
958 _judge_model = None
959 _judge_tokenizer = None
960 if phases is None or 4 in phases:
961 try:
962 from transformer_lens.benchmarks.text_quality import (
963 JUDGE_MODEL_ID,
964 JUDGE_REVISION,
965 load_judge,
966 )
968 _judge_model, _judge_tokenizer = load_judge()
969 if not quiet:
970 print(f"Pre-loaded Phase 4 judge {JUDGE_MODEL_ID}@{JUDGE_REVISION[:8]}")
971 except Exception as e:
972 if not quiet:
973 print(f"Warning: Could not pre-load Phase 4 judge: {e}")
974 print(" Phase 4 will load its own judge per model.")
976 total = len(candidates)
977 for i, candidate in enumerate(candidates, 1):
978 # Check for graceful interrupt between models
979 if _interrupt_requested:
980 if not quiet:
981 print(
982 f"\nStopping gracefully. Progress saved "
983 f"({len(progress.verified)} verified, "
984 f"{len(progress.provisional)} provisional)."
985 )
986 _save_checkpoint(progress)
987 raise SystemExit(_EXIT_GRACEFUL_INTERRUPT)
989 model_id = candidate.model_id
990 arch = candidate.architecture_id
992 if not quiet:
993 print(f"\n{'='*70}")
994 print(f"[{i}/{total}] {model_id} ({arch})")
995 print(f"{'='*70}")
997 progress.tested.append(model_id)
999 # Step 0: Skip formats with no HF loader path (GGUF / MLX / FP4 / FP8).
1000 if is_incompatible_quantized(model_id):
1001 _skip_model(model_id, arch, QUANTIZED_NOTE, progress, quiet)
1002 continue
1004 # Step 0a: skip HF-loadable quantized models when their loader lib is missing.
1005 required_lib = required_quant_library_for_model(model_id)
1006 if required_lib is not None:
1007 import importlib.util
1009 if importlib.util.find_spec(required_lib) is None:
1010 note = f"Skipped: {required_lib} not installed (required to load this quantized format)"
1011 _skip_model(model_id, arch, note, progress, quiet)
1012 continue
1014 # Step 0b: Check adapter-level phase applicability. applicable_phases=[]
1015 # means no phase applies (a genuinely unsupported architecture — rare);
1016 # SSM / recurrent families run the full [1, 2, 3, 4].
1017 from transformer_lens.factories.architecture_adapter_factory import (
1018 SUPPORTED_ARCHITECTURES,
1019 )
1021 adapter_cls = SUPPORTED_ARCHITECTURES.get(arch)
1022 eff_phases = phases if phases is not None else _default_phases_for_architecture(arch)
1023 phases_to_run = _phases_to_run(arch, eff_phases)
1024 if adapter_cls is not None and not phases_to_run:
1025 applicable = getattr(adapter_cls, "applicable_phases", list(TEXT_PHASES))
1026 note = (
1027 f"Architecture {arch} has applicable_phases={applicable}; "
1028 f"verify_models coverage is deferred. Verification lives "
1029 f"in integration tests."
1030 )
1031 _skip_model(model_id, arch, note, progress, quiet)
1032 continue
1034 # Step 1: Estimate parameters
1035 try:
1036 n_params = estimate_model_params(model_id)
1037 candidate.estimated_params = n_params
1038 if not quiet:
1039 print(f" Estimated parameters: {n_params:,}")
1040 except Exception as e:
1041 _skip_model(model_id, arch, f"Config unavailable: {str(e)[:200]}", progress, quiet)
1042 continue
1044 # Step 2: Check memory
1045 estimated_mem = estimate_benchmark_memory_gb(
1046 n_params, dtype, phases=phases_to_run, use_hf_reference=use_hf_reference, device=device
1047 )
1048 candidate.estimated_memory_gb = estimated_mem
1049 if not quiet:
1050 print(
1051 f" Estimated benchmark memory: {estimated_mem:.1f} GB (limit: {max_memory_gb:.1f} GB)"
1052 )
1054 if estimated_mem > max_memory_gb:
1055 note = f"Estimated {estimated_mem:.1f} GB exceeds {max_memory_gb:.1f} GB limit"
1056 _skip_model(model_id, arch, note, progress, quiet)
1057 continue
1059 # Step 3: Run benchmarks (all phases in a single call to share models)
1060 all_results: list = []
1061 error_msg: Optional[str] = None
1063 needs_remote_code = any(
1064 model_id.startswith(prefix) for prefix in REMOTE_CODE_MODEL_PREFIXES
1065 )
1067 # Convert string dtype to torch.dtype for benchmark suite
1068 import torch
1070 _dtype_map = {
1071 "float32": torch.float32,
1072 "float16": torch.float16,
1073 "bfloat16": torch.bfloat16,
1074 }
1075 torch_dtype = _dtype_map[dtype]
1077 from transformer_lens.benchmarks.text_quality_profiles import resolve_profile
1078 from transformer_lens.tools.model_registry.registry_io import (
1079 registry_prompt_profile,
1080 )
1082 resolved_profile = str(resolve_profile(model_id, arch, registry_prompt_profile(model_id)))
1083 if not quiet:
1084 print(f" Prompt profile: {resolved_profile}")
1085 print(f" Running phases {phases} in a single benchmark call...")
1086 try:
1087 all_results = run_benchmark_suite(
1088 model_id,
1089 device=device,
1090 dtype=torch_dtype,
1091 use_hf_reference=use_hf_reference,
1092 verbose=not quiet,
1093 phases=phases_to_run,
1094 trust_remote_code=needs_remote_code,
1095 judge_model=_judge_model,
1096 judge_tokenizer=_judge_tokenizer,
1097 prompt_profile=resolved_profile,
1098 )
1099 except Exception as e:
1100 error_msg = str(e)
1101 if not quiet:
1102 print(f" Benchmark failed: {error_msg[:200]}")
1104 phase_scores = extract_phase_scores(all_results)
1106 if not error_msg:
1107 # Only require the core phases this run actually requested, so a
1108 # partial run (e.g. --phases 1 2) isn't failed for a missing P7.
1109 _, core_for_arch = _full_and_core_phases(arch)
1110 score_error = _check_phase_scores(
1111 phase_scores, all_results, required_phases=core_for_arch & set(eff_phases)
1112 )
1113 if score_error:
1114 error_msg = score_error
1116 if error_msg:
1117 is_oom = "out of memory" in error_msg.lower() or "oom" in error_msg.lower()
1118 if is_oom:
1119 note = "OOM during benchmark"
1120 else:
1121 # Include the specific error from failed results (e.g., tokenizer
1122 # errors, load failures) so the note explains WHY it failed.
1123 root_errors = [r.message for r in all_results if not r.passed and r.message]
1124 if root_errors:
1125 # Deduplicate and use first unique error as the detail
1126 unique_errors = list(dict.fromkeys(root_errors))
1127 detail = unique_errors[0][:150]
1128 note = f"{error_msg[:100]} — {detail}"
1129 else:
1130 note = error_msg[:200]
1131 final_status = STATUS_FAILED
1132 else:
1133 note = _build_verified_note(phase_scores, all_results)
1134 final_status = STATUS_VERIFIED
1136 # When running a partial phase set (e.g., --phases 4 for backfill),
1137 # only update the phase scores that were run. Don't change the
1138 # model's overall status or note — those reflect the full
1139 # verification and should only be set by a complete run.
1140 kind = classify_architecture(arch)
1141 is_multimodal = kind == "multimodal"
1142 is_audio = kind == "audio"
1143 is_vision = kind == "vision"
1144 full_phases, core_required = _full_and_core_phases(arch)
1145 is_partial_run = set(eff_phases) != full_phases
1147 if is_partial_run and phase_scores:
1148 # Only write scores for phases that were actually requested.
1149 # Bridge load failures can produce Phase 1-tagged error results
1150 # even during Phase 4-only runs — don't let those corrupt
1151 # existing scores for unrequested phases.
1152 filtered_scores = {p: s for p, s in phase_scores.items() if p in eff_phases}
1153 if filtered_scores:
1154 if not quiet:
1155 score_parts = [f"P{p}={s}%" for p, s in sorted(filtered_scores.items())]
1156 print(f" Partial phase update: {', '.join(score_parts)}")
1158 # Core verification: P1+P4 for text-only, P1+P4+P7 for
1159 # multimodal, P9 for vision.
1160 is_core_verification = set(eff_phases) >= core_required
1161 partial_status = None
1162 partial_note = None
1164 if is_core_verification:
1165 p1 = filtered_scores.get(1)
1166 p4 = filtered_scores.get(4)
1167 p1_pass = p1 is not None and p1 >= _MIN_PHASE_SCORES.get(
1168 1, _DEFAULT_MIN_PHASE_SCORE
1169 )
1170 p4_pass = p4 is not None and p4 >= _MIN_PHASE_SCORES.get(
1171 4, _DEFAULT_MIN_PHASE_SCORE
1172 )
1174 # For multimodal, Phase 7 is required. A score below 75%
1175 # or a missing score (NULL — processor unavailable) both
1176 # count as failures.
1177 p7_pass = True
1178 if is_multimodal:
1179 p7 = filtered_scores.get(7)
1180 if p7 is not None:
1181 p7_pass = p7 >= _MIN_PHASE_SCORES.get(7, _DEFAULT_MIN_PHASE_SCORE)
1182 else:
1183 p7_pass = False
1185 # For audio models, Phase 8 is required; Phase 4 is not applicable
1186 p8_pass = True
1187 if is_audio:
1188 p4_pass = True # Audio models skip text quality
1189 p8 = filtered_scores.get(8)
1190 if p8 is not None:
1191 p8_pass = p8 >= _MIN_PHASE_SCORES.get(8, _DEFAULT_MIN_PHASE_SCORE)
1192 else:
1193 p8_pass = False
1195 # For vision models, Phase 9 is required; text phases 1/4
1196 # are not applicable (no text tower).
1197 p9_pass = True
1198 if is_vision:
1199 p1_pass = True
1200 p4_pass = True
1201 p9 = filtered_scores.get(9)
1202 if p9 is not None:
1203 p9_pass = p9 >= _MIN_PHASE_SCORES.get(9, _DEFAULT_MIN_PHASE_SCORE)
1204 else:
1205 p9_pass = False
1207 if p1_pass and p4_pass and p7_pass and p8_pass and p9_pass:
1208 partial_status = STATUS_VERIFIED
1209 partial_note = "Core verification completed" + _preserved_issue_suffix(
1210 model_id, eff_phases
1211 )
1212 elif p1_pass and p4_pass and not p7_pass:
1213 p7_score = filtered_scores.get(7)
1214 if p7_score is None:
1215 partial_status = STATUS_FAILED
1216 partial_note = (
1217 "Core verification failed: multimodal tests skipped "
1218 "(processor unavailable)"
1219 )
1220 else:
1221 partial_status = STATUS_FAILED
1222 partial_note = (
1223 f"Core verification failed: multimodal tests "
1224 f"scored {p7_score}% (requires >= "
1225 f"{_MIN_PHASE_SCORES.get(7, _DEFAULT_MIN_PHASE_SCORE):g}%)"
1226 )
1227 elif p1_pass and p4_pass and not p9_pass:
1228 p9_score = filtered_scores.get(9)
1229 partial_status = STATUS_FAILED
1230 if p9_score is None:
1231 partial_note = (
1232 "Core verification failed: vision tests skipped (no results)"
1233 )
1234 else:
1235 partial_note = (
1236 f"Core verification failed: vision tests "
1237 f"scored {p9_score}% (requires >= "
1238 f"{_MIN_PHASE_SCORES.get(9, _DEFAULT_MIN_PHASE_SCORE):g}%)"
1239 )
1240 elif p1_pass:
1241 partial_status = STATUS_VERIFIED
1242 partial_note = _p1_only_core_note(p4, all_results)
1243 else:
1244 # P1 failed — build a descriptive failure note
1245 partial_status = STATUS_FAILED
1246 if error_msg:
1247 partial_note = f"CORE FAILED: {error_msg[:200]}"
1248 else:
1249 # Score-based failure — include details
1250 from transformer_lens.benchmarks.utils import (
1251 BenchmarkSeverity,
1252 )
1254 failed_tests = [
1255 r.name
1256 for r in all_results
1257 if r.phase == 1
1258 and not r.passed
1259 and r.severity != BenchmarkSeverity.SKIPPED
1260 ]
1261 tests_str = ", ".join(failed_tests) if failed_tests else "unknown"
1262 partial_note = f"CORE FAILED: P1={p1}% (failed: {tests_str})"
1264 # A structural-only core pass (--no-hf-reference) is provisional,
1265 # never numerically compared to HF, so it must not count as verified.
1266 if partial_status == STATUS_VERIFIED and not use_hf_reference:
1267 partial_status = STATUS_PROVISIONAL
1268 partial_note = f"Structural only (no HF reference): {partial_note}"
1270 if not quiet:
1271 print(f" {partial_note}")
1273 update_model_status(
1274 model_id,
1275 arch,
1276 status=partial_status,
1277 phase_scores=filtered_scores,
1278 note=partial_note,
1279 prompt_profile=_extract_prompt_profile(all_results),
1280 )
1281 # A provisional run was not numerically verified; do not write a
1282 # verification-history record (VerificationHistory.is_verified()
1283 # treats any record as verified — a second "counts as verified" path).
1284 if partial_status != STATUS_PROVISIONAL:
1285 add_verification_record(
1286 model_id,
1287 arch,
1288 notes=partial_note,
1289 sanitize_fn=_sanitize_note,
1290 prompt_profile=_extract_prompt_profile(all_results),
1291 p4_scoring_version=(P4_SCORING_VERSION if 4 in filtered_scores else None),
1292 )
1293 if partial_status == STATUS_FAILED:
1294 progress.failed.append(model_id)
1295 elif partial_status == STATUS_PROVISIONAL:
1296 progress.provisional.append(model_id)
1297 elif partial_status is None:
1298 # Scores were written but the status deliberately was not:
1299 # this phase set cannot establish core verification for
1300 # this architecture class. Reporting it as verified is how
1301 # a model ends up with all-pass scores and status 0.
1302 if not quiet:
1303 print(
1304 f" Scores updated, status unchanged — core verification for "
1305 f"{arch} requires phases {sorted(core_required)}, "
1306 f"this run had {sorted(eff_phases)}."
1307 )
1308 progress.skipped.append(model_id)
1309 else:
1310 progress.verified.append(model_id)
1311 else:
1312 if not quiet:
1313 print(f" No results for requested phases {eff_phases} — skipping update")
1314 progress.skipped.append(model_id)
1315 elif final_status == STATUS_VERIFIED and _measured_nothing(phase_scores):
1316 if not quiet:
1317 print(
1318 f" No phase produced a score (requested {eff_phases}) — "
1319 f"status left unchanged"
1320 )
1321 progress.skipped.append(model_id)
1322 elif final_status == STATUS_VERIFIED:
1323 # A passing run is VERIFIED only if it was numerically compared to an
1324 # HF reference; a --no-hf-reference (structural-only) pass is PROVISIONAL.
1325 written_status = pass_status(use_hf_reference)
1326 is_provisional = written_status == STATUS_PROVISIONAL
1327 if is_provisional:
1328 note = f"Structural only (no HF reference): {note}"
1329 if not quiet:
1330 label = "PROVISIONAL" if is_provisional else "VERIFIED"
1331 print(
1332 f" {label}: P1={phase_scores.get(1)}%, "
1333 f"P2={phase_scores.get(2)}%, P3={phase_scores.get(3)}%, "
1334 f"P4={phase_scores.get(4)}%, P7={phase_scores.get(7)}%, "
1335 f"P8={phase_scores.get(8)}%, P9={phase_scores.get(9)}%"
1336 )
1337 update_model_status(
1338 model_id,
1339 arch,
1340 written_status,
1341 phase_scores=phase_scores,
1342 note=note,
1343 prompt_profile=_extract_prompt_profile(all_results),
1344 )
1345 # Provisional runs are not numerically verified — no history record
1346 # (is_verified() would otherwise report them as verified).
1347 if not is_provisional:
1348 add_verification_record(
1349 model_id,
1350 arch,
1351 notes=note,
1352 prompt_profile=_extract_prompt_profile(all_results),
1353 p4_scoring_version=(P4_SCORING_VERSION if 4 in phase_scores else None),
1354 )
1355 if is_provisional:
1356 progress.provisional.append(model_id)
1357 else:
1358 progress.verified.append(model_id)
1359 else:
1360 if not quiet:
1361 print(f" FAILED: {note}")
1362 if any(v is not None for v in phase_scores.values()):
1363 print(
1364 f" Partial scores saved: P1={phase_scores.get(1)}%, "
1365 f"P2={phase_scores.get(2)}%, P3={phase_scores.get(3)}%, "
1366 f"P4={phase_scores.get(4)}%, P7={phase_scores.get(7)}%, "
1367 f"P8={phase_scores.get(8)}%, P9={phase_scores.get(9)}%"
1368 )
1369 update_model_status(
1370 model_id,
1371 arch,
1372 STATUS_FAILED,
1373 note=note,
1374 phase_scores=phase_scores,
1375 sanitize_fn=_sanitize_note,
1376 prompt_profile=_extract_prompt_profile(all_results),
1377 )
1378 add_verification_record(
1379 model_id,
1380 arch,
1381 notes=note,
1382 sanitize_fn=_sanitize_note,
1383 prompt_profile=_extract_prompt_profile(all_results),
1384 p4_scoring_version=(P4_SCORING_VERSION if 4 in phase_scores else None),
1385 )
1386 progress.failed.append(model_id)
1388 # Post-model cleanup
1389 gc.collect()
1390 try:
1391 import torch
1393 if torch.cuda.is_available():
1394 torch.cuda.empty_cache()
1395 torch.cuda.synchronize()
1396 if device == "mps" and hasattr(torch, "mps") and torch.backends.mps.is_available():
1397 torch.mps.synchronize()
1398 torch.mps.empty_cache()
1400 # Log MPS memory state for debugging long runs
1401 if device == "mps" and not quiet and hasattr(torch.mps, "current_allocated_memory"):
1402 alloc_mb = torch.mps.current_allocated_memory() / (1024 * 1024)
1403 driver_mb = torch.mps.driver_allocated_memory() / (1024 * 1024)
1404 print(f" MPS memory: {alloc_mb:.0f} MB allocated, " f"{driver_mb:.0f} MB driver")
1405 except ImportError:
1406 pass
1408 # Brief pause to let the OS and MPS reclaim memory between models
1409 if device in ("mps", "cuda"):
1410 time.sleep(3)
1412 # Periodically clear the HuggingFace cache to prevent disk exhaustion
1413 if i % 50 == 0:
1414 _clear_hf_cache(quiet)
1416 _save_checkpoint(progress)
1418 # Clean up pre-loaded scoring model
1419 if _judge_model is not None:
1420 del _judge_model
1421 del _judge_tokenizer
1422 gc.collect()
1424 return progress
1427def _print_dry_run(
1428 candidates: list[ModelCandidate],
1429 dtype: str,
1430 max_memory_gb: float,
1431 phases: Optional[list[int]] = None,
1432 use_hf_reference: bool = True,
1433 device: str = "cpu",
1434) -> None:
1435 """Print what would be tested in a dry run."""
1436 print(f"\nDry run: {len(candidates)} models would be tested")
1437 print(f"Memory limit: {max_memory_gb:.1f} GB | Dtype: {dtype}")
1438 print()
1440 # Group by architecture
1441 by_arch: dict[str, list[ModelCandidate]] = {}
1442 for c in candidates:
1443 by_arch.setdefault(c.architecture_id, []).append(c)
1445 skippable = 0
1446 testable = 0
1448 for arch in sorted(by_arch.keys()):
1449 models = by_arch[arch]
1450 # Same per-architecture default as the real run, so the dry run's
1451 # memory estimates cover the phases that will actually execute.
1452 eff_phases = phases if phases is not None else _default_phases_for_architecture(arch)
1453 phases_to_run = _phases_to_run(arch, eff_phases)
1454 print(f" {arch} ({len(models)} models):")
1455 for c in models:
1456 try:
1457 n_params = estimate_model_params(c.model_id)
1458 mem = estimate_benchmark_memory_gb(
1459 n_params,
1460 dtype,
1461 phases=phases_to_run,
1462 use_hf_reference=use_hf_reference,
1463 device=device,
1464 )
1465 status = "OK" if mem <= max_memory_gb else "SKIP (too large)"
1466 if mem > max_memory_gb:
1467 skippable += 1
1468 else:
1469 testable += 1
1470 print(f" {c.model_id}: ~{n_params/1e6:.0f}M params, ~{mem:.1f} GB [{status}]")
1471 except Exception as e:
1472 skippable += 1
1473 print(f" {c.model_id}: config error ({e})")
1474 print()
1476 print(f"Summary: {testable} testable, {skippable} would be skipped")
1479def _print_summary(progress: VerificationProgress) -> None:
1480 """Print a summary of the verification run."""
1481 total = len(progress.tested)
1482 print(f"\n{'='*70}")
1483 print("Verification Summary")
1484 print(f"{'='*70}")
1485 print(f" Total tested: {total}")
1486 print(f" Verified: {len(progress.verified)}")
1487 print(f" Provisional: {len(progress.provisional)}")
1488 print(f" Skipped: {len(progress.skipped)}")
1489 print(f" Failed: {len(progress.failed)}")
1491 if progress.verified:
1492 print(f"\n Verified models:")
1493 for m in progress.verified:
1494 print(f" - {m}")
1496 if progress.provisional:
1497 print(f"\n Provisional models (structural only, no HF reference):")
1498 for m in progress.provisional:
1499 print(f" - {m}")
1501 if progress.failed:
1502 print(f"\n Failed models:")
1503 for m in progress.failed:
1504 print(f" - {m}")
1506 if progress.skipped:
1507 print(f"\n Skipped models:")
1508 for m in progress.skipped[:20]:
1509 print(f" - {m}")
1510 if len(progress.skipped) > 20:
1511 print(f" ... and {len(progress.skipped) - 20} more")
1514def main() -> None:
1515 """CLI entry point for batch model verification."""
1516 parser = argparse.ArgumentParser(
1517 description="Batch verify models in the TransformerLens registry",
1518 formatter_class=argparse.RawDescriptionHelpFormatter,
1519 epilog="""
1520Examples:
1521 %(prog)s --dry-run Show what would be tested
1522 %(prog)s --limit 3 Test 3 models total
1523 %(prog)s --architectures GPT2LMHeadModel --per-arch 5
1524 %(prog)s --device cuda --max-memory 24
1525 %(prog)s --resume Resume from checkpoint
1526 %(prog)s --reverify --architectures Olmo2ForCausalLM Re-verify already-tested models
1527 %(prog)s --model google/gemma-2b Verify a single model by ID
1528 """,
1529 )
1530 parser.add_argument(
1531 "--per-arch",
1532 type=int,
1533 default=10,
1534 help="Max models to verify per architecture (default: 10)",
1535 )
1536 parser.add_argument(
1537 "--device",
1538 type=str,
1539 default="cpu",
1540 help="Device for benchmarks (default: cpu)",
1541 )
1542 parser.add_argument(
1543 "--max-memory",
1544 type=float,
1545 default=None,
1546 help="Memory limit in GB (default: auto-detect)",
1547 )
1548 parser.add_argument(
1549 "--architectures",
1550 nargs="+",
1551 default=None,
1552 help="Filter to specific architectures",
1553 )
1554 parser.add_argument(
1555 "--limit",
1556 type=int,
1557 default=None,
1558 help="Total model cap",
1559 )
1560 parser.add_argument(
1561 "--resume",
1562 action="store_true",
1563 help="Resume from checkpoint",
1564 )
1565 parser.add_argument(
1566 "--dry-run",
1567 action="store_true",
1568 help="Show what would be tested without running benchmarks",
1569 )
1570 parser.add_argument(
1571 "--no-hf-reference",
1572 action="store_true",
1573 help=(
1574 "Skip HuggingFace reference comparison (Phase 1 is structural-only). "
1575 "A passing run is recorded as PROVISIONAL, not verified — re-run without "
1576 "this flag for a real HF-compared verification."
1577 ),
1578 )
1579 parser.add_argument(
1580 "--phases",
1581 nargs="+",
1582 type=int,
1583 default=None,
1584 help=(
1585 "Which benchmark phases to run (default: a full verification for each "
1586 "model's architecture — 1 2 3 4 for text, 1 2 3 4 7 for multimodal, "
1587 "1 8 for audio, 1 9 for vision)"
1588 ),
1589 )
1590 parser.add_argument(
1591 "--dtype",
1592 type=str,
1593 default="float32",
1594 choices=["float32", "float16", "bfloat16"],
1595 help="Dtype for memory estimation (default: float32)",
1596 )
1597 parser.add_argument(
1598 "--quiet",
1599 action="store_true",
1600 help="Suppress verbose output",
1601 )
1602 parser.add_argument(
1603 "--retry-failed",
1604 action="store_true",
1605 help="Re-run previously failed models instead of skipping them",
1606 )
1607 parser.add_argument(
1608 "--reverify",
1609 action="store_true",
1610 help="Re-run verification for already-verified/skipped/failed models. "
1611 "Ignores previous status and re-tests matching models from scratch.",
1612 )
1613 parser.add_argument(
1614 "--model",
1615 type=str,
1616 nargs="+",
1617 default=None,
1618 help="Verify one or more models by HuggingFace model ID. "
1619 "Looks up architecture from the registry automatically.",
1620 )
1622 args = parser.parse_args()
1624 # Setup logging
1625 logging.basicConfig(
1626 level=logging.WARNING if args.quiet else logging.INFO,
1627 format="%(asctime)s [%(levelname)s] %(message)s",
1628 )
1630 # Auto-detect memory
1631 max_memory_gb = args.max_memory
1632 if max_memory_gb is None:
1633 max_memory_gb = get_available_memory_gb(args.device)
1635 # Load checkpoint if resuming
1636 progress = None
1637 if args.resume:
1638 progress = _load_checkpoint()
1639 if progress:
1640 print(f"Resuming from checkpoint: {len(progress.tested)} models already tested")
1641 else:
1642 print("No checkpoint found, starting fresh")
1644 # If retrying failed, clean them from checkpoint and reset status in registry
1645 if args.retry_failed and progress and not args.dry_run:
1646 failed_set = set(progress.failed)
1647 if failed_set:
1648 # Reset status in supported_models.json
1649 registry_data = load_supported_models_raw()
1650 for entry in registry_data.get("models", []):
1651 if entry["model_id"] in failed_set and entry.get("status") == STATUS_FAILED:
1652 update_model_status(
1653 entry["model_id"],
1654 entry["architecture_id"],
1655 STATUS_UNVERIFIED,
1656 )
1657 # Clean checkpoint
1658 progress.tested = [m for m in progress.tested if m not in failed_set]
1659 progress.failed = []
1660 _save_checkpoint(progress)
1661 print(f" Cleared {len(failed_set)} failed models for retry")
1663 # Select models — either --model list or the normal batch selection
1664 if args.model:
1665 # Look up architecture for each model from the registry
1666 registry_data = load_supported_models_raw()
1667 candidates = []
1668 for model_id in args.model:
1669 arch_id = None
1670 for entry in registry_data.get("models", []):
1671 if entry["model_id"] == model_id:
1672 arch_id = entry["architecture_id"]
1673 break
1674 if arch_id is None:
1675 print(f"Model '{model_id}' not found in supported_models.json, skipping")
1676 continue
1677 candidates.append(ModelCandidate(model_id=model_id, architecture_id=arch_id))
1678 if not candidates:
1679 print("No valid models found in registry")
1680 return
1681 print(f"Model list mode: {len(candidates)} model(s)")
1682 else:
1683 candidates = select_models_for_verification(
1684 per_arch=args.per_arch,
1685 architectures=args.architectures,
1686 limit=args.limit,
1687 resume_progress=progress,
1688 retry_failed=args.retry_failed,
1689 reverify=args.reverify,
1690 )
1692 if not candidates:
1693 print("No models to verify (all matching models already tested)")
1694 return
1696 print(f"Selected {len(candidates)} models for verification")
1698 # Dry run
1699 if args.dry_run:
1700 _print_dry_run(
1701 candidates,
1702 args.dtype,
1703 max_memory_gb,
1704 phases=args.phases,
1705 use_hf_reference=not args.no_hf_reference,
1706 device=args.device,
1707 )
1708 return
1710 # Install graceful interrupt handler (Ctrl+C stops between models)
1711 signal.signal(signal.SIGINT, _handle_sigint)
1713 # Run verification
1714 start = time.time()
1715 progress = verify_models(
1716 candidates,
1717 device=args.device,
1718 max_memory_gb=max_memory_gb,
1719 dtype=args.dtype,
1720 use_hf_reference=not args.no_hf_reference,
1721 phases=args.phases,
1722 quiet=args.quiet,
1723 progress=progress,
1724 )
1725 elapsed = time.time() - start
1727 _print_summary(progress)
1728 print(f"\nTotal time: {elapsed:.1f}s")
1730 # Clean up checkpoint on successful completion
1731 if _CHECKPOINT_PATH.exists():
1732 _CHECKPOINT_PATH.unlink()
1733 print("Checkpoint cleared (run complete)")
1736if __name__ == "__main__": 1736 ↛ 1737line 1736 didn't jump to line 1737 because the condition on line 1736 was never true
1737 main()